Add support for Algebraic Datatypes in JavaScript/TypeScript bindings (#7734)

* Initial plan

* Add datatype type definitions to types.ts (work in progress)

Co-authored-by: NikolajBjorner <[email protected]>

* Complete datatype type definitions with working TypeScript compilation

Co-authored-by: NikolajBjorner <[email protected]>

* Implement core datatype functionality with TypeScript compilation success

Co-authored-by: NikolajBjorner <[email protected]>

* Complete datatype implementation with full Context integration and tests

Co-authored-by: NikolajBjorner <[email protected]>

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: NikolajBjorner <[email protected]>

Author: Copilot

src/api/js/src/high-level/high-level.test.ts

@@ -890,4 +890,74 @@ describe('high-level', () => {
       expect(model.eval(z).eqIdentity(Int.val(5))).toBeTruthy();
     });
   });
+
+  describe('datatypes', () => {
+    it('should create simple enum datatype', async () => {
+      const { Datatype, Int, Bool, Solver } = api.Context('main');
+
+      // Create a simple Color enum datatype
+      const Color = Datatype('Color');
+      Color.declare('red');
+      Color.declare('green');
+      Color.declare('blue');
+      
+      const ColorSort = Color.create();
+      
+      // Test that we can access the constructors
+      expect(typeof (ColorSort as any).red).not.toBe('undefined');
+      expect(typeof (ColorSort as any).green).not.toBe('undefined');
+      expect(typeof (ColorSort as any).blue).not.toBe('undefined');
+      
+      // Test that we can access the recognizers
+      expect(typeof (ColorSort as any).is_red).not.toBe('undefined');
+      expect(typeof (ColorSort as any).is_green).not.toBe('undefined');
+      expect(typeof (ColorSort as any).is_blue).not.toBe('undefined');
+    });
+
+    it('should create recursive list datatype', async () => {
+      const { Datatype, Int, Solver } = api.Context('main');
+
+      // Create a recursive List datatype like in the Python example
+      const List = Datatype('List');
+      List.declare('cons', ['car', Int.sort()], ['cdr', List]);
+      List.declare('nil');
+      
+      const ListSort = List.create();
+      
+      // Test that constructors and accessors exist
+      expect(typeof (ListSort as any).cons).not.toBe('undefined');
+      expect(typeof (ListSort as any).nil).not.toBe('undefined');
+      expect(typeof (ListSort as any).is_cons).not.toBe('undefined');
+      expect(typeof (ListSort as any).is_nil).not.toBe('undefined');
+      expect(typeof (ListSort as any).car).not.toBe('undefined');
+      expect(typeof (ListSort as any).cdr).not.toBe('undefined');
+    });
+
+    it('should create mutually recursive tree datatypes', async () => {
+      const { Datatype, Int } = api.Context('main');
+
+      // Create mutually recursive Tree and TreeList datatypes
+      const Tree = Datatype('Tree');
+      const TreeList = Datatype('TreeList');
+      
+      Tree.declare('leaf', ['value', Int.sort()]);
+      Tree.declare('node', ['children', TreeList]);
+      TreeList.declare('nil');
+      TreeList.declare('cons', ['car', Tree], ['cdr', TreeList]);
+      
+      const [TreeSort, TreeListSort] = Datatype.createDatatypes(Tree, TreeList);
+      
+      // Test that both datatypes have their constructors
+      expect(typeof (TreeSort as any).leaf).not.toBe('undefined');
+      expect(typeof (TreeSort as any).node).not.toBe('undefined');
+      expect(typeof (TreeListSort as any).nil).not.toBe('undefined');
+      expect(typeof (TreeListSort as any).cons).not.toBe('undefined');
+      
+      // Test accessors exist
+      expect(typeof (TreeSort as any).value).not.toBe('undefined');
+      expect(typeof (TreeSort as any).children).not.toBe('undefined');
+      expect(typeof (TreeListSort as any).car).not.toBe('undefined');
+      expect(typeof (TreeListSort as any).cdr).not.toBe('undefined');
+    });
+  });
 });

src/api/js/src/high-level/high-level.ts

@@ -17,6 +17,8 @@ import {
   Z3_ast_print_mode,
   Z3_ast_vector,
   Z3_context,
+  Z3_constructor,
+  Z3_constructor_list,
   Z3_decl_kind,
   Z3_error_code,
   Z3_func_decl,
@@ -88,6 +90,10 @@ import {
   FuncEntry,
   SMTSetSort,
   SMTSet,
+  Datatype,
+  DatatypeSort,
+  DatatypeExpr,
+  DatatypeCreation,
 } from './types';
 import { allSatisfy, assert, assertExhaustive } from './utils';
 
@@ -825,6 +831,17 @@ export function createApi(Z3: Z3Core): Z3HighLevel {
       }
     }
 
+    const Datatype = Object.assign(
+      (name: string): DatatypeImpl => {
+        return new DatatypeImpl(ctx, name);
+      },
+      {
+        createDatatypes(...datatypes: DatatypeImpl[]): DatatypeSortImpl[] {
+          return createDatatypes(...datatypes);
+        }
+      }
+    );
+
     ////////////////
     // Operations //
     ////////////////
@@ -2647,6 +2664,185 @@ export function createApi(Z3: Z3Core): Z3HighLevel {
       }
     }
 
+    ////////////////////////////
+    // Datatypes
+    ////////////////////////////
+
+    class DatatypeImpl implements Datatype<Name> {
+      readonly ctx: Context<Name>;
+      readonly name: string;
+      public constructors: Array<[string, Array<[string, Sort<Name> | Datatype<Name>]>]> = [];
+
+      constructor(ctx: Context<Name>, name: string) {
+        this.ctx = ctx;
+        this.name = name;
+      }
+
+      declare(name: string, ...fields: Array<[string, Sort<Name> | Datatype<Name>]>): this {
+        this.constructors.push([name, fields]);
+        return this;
+      }
+
+      create(): DatatypeSort<Name> {
+        const datatypes = createDatatypes(this);
+        return datatypes[0];
+      }
+    }
+
+    class DatatypeSortImpl extends SortImpl implements DatatypeSort<Name> {
+      declare readonly __typename: DatatypeSort['__typename'];
+
+      numConstructors(): number {
+        return Z3.get_datatype_sort_num_constructors(contextPtr, this.ptr);
+      }
+
+      constructorDecl(idx: number): FuncDecl<Name> {
+        const ptr = Z3.get_datatype_sort_constructor(contextPtr, this.ptr, idx);
+        return new FuncDeclImpl(ptr);
+      }
+
+      recognizer(idx: number): FuncDecl<Name> {
+        const ptr = Z3.get_datatype_sort_recognizer(contextPtr, this.ptr, idx);
+        return new FuncDeclImpl(ptr);
+      }
+
+      accessor(constructorIdx: number, accessorIdx: number): FuncDecl<Name> {
+        const ptr = Z3.get_datatype_sort_constructor_accessor(contextPtr, this.ptr, constructorIdx, accessorIdx);
+        return new FuncDeclImpl(ptr);
+      }
+
+      cast(other: CoercibleToExpr<Name>): DatatypeExpr<Name>;
+      cast(other: DatatypeExpr<Name>): DatatypeExpr<Name>;
+      cast(other: CoercibleToExpr<Name> | DatatypeExpr<Name>): DatatypeExpr<Name> {
+        if (isExpr(other)) {
+          assert(this.eqIdentity(other.sort), 'Value cannot be converted to this datatype');
+          return other as DatatypeExpr<Name>;
+        }
+        throw new Error('Cannot coerce value to datatype expression');
+      }
+
+      subsort(other: Sort<Name>) {
+        _assertContext(other.ctx);
+        return this.eqIdentity(other);
+      }
+    }
+
+    class DatatypeExprImpl extends ExprImpl<Z3_ast, DatatypeSortImpl> implements DatatypeExpr<Name> {
+      declare readonly __typename: DatatypeExpr['__typename'];
+    }
+
+    function createDatatypes(...datatypes: DatatypeImpl[]): DatatypeSortImpl[] {
+      if (datatypes.length === 0) {
+        throw new Error('At least one datatype must be provided');
+      }
+
+      // All datatypes must be from the same context
+      const dtCtx = datatypes[0].ctx;
+      for (const dt of datatypes) {
+        if (dt.ctx !== dtCtx) {
+          throw new Error('All datatypes must be from the same context');
+        }
+      }
+
+      const sortNames = datatypes.map(dt => dt.name);
+      const constructorLists: Z3_constructor_list[] = [];
+      const scopedConstructors: Z3_constructor[] = [];
+
+      try {
+        // Create constructor lists for each datatype
+        for (const dt of datatypes) {
+          const constructors: Z3_constructor[] = [];
+
+          for (const [constructorName, fields] of dt.constructors) {
+            const fieldNames: string[] = [];
+            const fieldSorts: Z3_sort[] = [];
+            const fieldRefs: number[] = [];
+
+            for (const [fieldName, fieldSort] of fields) {
+              fieldNames.push(fieldName);
+              
+              if (fieldSort instanceof DatatypeImpl) {
+                // Reference to another datatype being defined
+                const refIndex = datatypes.indexOf(fieldSort);
+                if (refIndex === -1) {
+                  throw new Error(`Referenced datatype "${fieldSort.name}" not found in datatypes being created`);
+                }
+                // For recursive references, we pass null and the ref index
+                fieldSorts.push(null as any); // null will be handled by the Z3 API
+                fieldRefs.push(refIndex);
+              } else {
+                // Regular sort
+                fieldSorts.push((fieldSort as Sort<Name>).ptr);
+                fieldRefs.push(0);
+              }
+            }
+
+            const constructor = Z3.mk_constructor(
+              contextPtr,
+              Z3.mk_string_symbol(contextPtr, constructorName),
+              Z3.mk_string_symbol(contextPtr, `is_${constructorName}`),
+              fieldNames.map(name => Z3.mk_string_symbol(contextPtr, name)),
+              fieldSorts,
+              fieldRefs
+            );
+            constructors.push(constructor);
+            scopedConstructors.push(constructor);
+          }
+
+          const constructorList = Z3.mk_constructor_list(contextPtr, constructors);
+          constructorLists.push(constructorList);
+        }
+
+        // Create the datatypes
+        const sortSymbols = sortNames.map(name => Z3.mk_string_symbol(contextPtr, name));
+        const resultSorts = Z3.mk_datatypes(contextPtr, sortSymbols, constructorLists);
+
+        // Create DatatypeSortImpl instances
+        const results: DatatypeSortImpl[] = [];
+        for (let i = 0; i < resultSorts.length; i++) {
+          const sortImpl = new DatatypeSortImpl(resultSorts[i]);
+          
+          // Attach constructor, recognizer, and accessor functions dynamically
+          const numConstructors = sortImpl.numConstructors();
+          for (let j = 0; j < numConstructors; j++) {
+            const constructor = sortImpl.constructorDecl(j);
+            const recognizer = sortImpl.recognizer(j);
+            const constructorName = constructor.name().toString();
+            
+            // Attach constructor function
+            if (constructor.arity() === 0) {
+              // Nullary constructor (constant)
+              (sortImpl as any)[constructorName] = constructor.call();
+            } else {
+              (sortImpl as any)[constructorName] = constructor;
+            }
+            
+            // Attach recognizer function
+            (sortImpl as any)[`is_${constructorName}`] = recognizer;
+            
+            // Attach accessor functions
+            for (let k = 0; k < constructor.arity(); k++) {
+              const accessor = sortImpl.accessor(j, k);
+              const accessorName = accessor.name().toString();
+              (sortImpl as any)[accessorName] = accessor;
+            }
+          }
+          
+          results.push(sortImpl);
+        }
+
+        return results;
+      } finally {
+        // Clean up resources
+        for (const constructor of scopedConstructors) {
+          Z3.del_constructor(contextPtr, constructor);
+        }
+        for (const constructorList of constructorLists) {
+          Z3.del_constructor_list(contextPtr, constructorList);
+        }
+      }
+    }
+
     class QuantifierImpl<
         QVarSorts extends NonEmptySortArray<Name>,
         QSort extends BoolSort<Name> | SMTArraySort<Name, QVarSorts>,
@@ -3029,6 +3225,7 @@ export function createApi(Z3: Z3Core): Z3HighLevel {
       BitVec,
       Array,
       Set,
+      Datatype,
 
       ////////////////
       // Operations //