Merge 12d4ec7eda into ad2934f8cf

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 //

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

@@ -3,6 +3,8 @@ import {
   Z3_ast_map,
   Z3_ast_vector,
   Z3_context,
+  Z3_constructor,
+  Z3_constructor_list,
   Z3_decl_kind,
   Z3_func_decl,
   Z3_func_entry,
@@ -362,6 +364,8 @@ export interface Context<Name extends string = 'main'> {
   readonly Array: SMTArrayCreation<Name>;
   /** @category Expressions */
   readonly Set: SMTSetCreation<Name>;
+  /** @category Expressions */
+  readonly Datatype: DatatypeCreation<Name>;
 
   ////////////////
   // Operations //
@@ -842,7 +846,8 @@ export interface Sort<Name extends string = 'main'> extends Ast<Name, Z3_sort> {
     | BoolSort['__typename']
     | ArithSort['__typename']
     | BitVecSort['__typename']
-    | SMTArraySort['__typename'];
+    | SMTArraySort['__typename']
+    | DatatypeSort['__typename'];
 
   kind(): Z3_sort_kind;
 
@@ -966,7 +971,8 @@ export interface Expr<Name extends string = 'main', S extends Sort<Name> = AnySo
     | Bool['__typename']
     | Arith['__typename']
     | BitVec['__typename']
-    | SMTArray['__typename'];
+    | SMTArray['__typename']
+    | DatatypeExpr['__typename'];
 
   get sort(): S;
 
@@ -1653,6 +1659,111 @@ export interface SMTSet<Name extends string = 'main', ElemSort extends AnySort<N
   subsetOf(b: SMTSet<Name, ElemSort>): Bool<Name>;
 
 }
+//////////////////////////////////////////
+//
+// Datatypes
+//
+//////////////////////////////////////////
+
+/**
+ * Helper class for declaring Z3 datatypes.
+ * 
+ * Follows the same pattern as Python Z3 API for declaring constructors
+ * before creating the actual datatype sort.
+ * 
+ * @example
+ * ```typescript
+ * const List = new ctx.Datatype('List');
+ * List.declare('cons', ['car', ctx.Int.sort()], ['cdr', List]);
+ * List.declare('nil');
+ * const ListSort = List.create();
+ * ```
+ * 
+ * @category Datatypes
+ */
+export interface Datatype<Name extends string = 'main'> {
+  readonly ctx: Context<Name>;
+  readonly name: string;
+
+  /**
+   * Declare a constructor for this datatype.
+   * 
+   * @param name Constructor name
+   * @param fields Array of [field_name, field_sort] pairs
+   */
+  declare(name: string, ...fields: Array<[string, AnySort<Name> | Datatype<Name>]>): this;
+
+  /**
+   * Create the actual datatype sort from the declared constructors.
+   * For mutually recursive datatypes, use Context.createDatatypes instead.
+   */
+  create(): DatatypeSort<Name>;
+}
+
+/**
+ * @category Datatypes
+ */
+export interface DatatypeCreation<Name extends string> {
+  /**
+   * Create a new datatype declaration helper.
+   */
+  (name: string): Datatype<Name>;
+
+  /**
+   * Create mutually recursive datatypes.
+   * 
+   * @param datatypes Array of Datatype declarations
+   * @returns Array of created DatatypeSort instances
+   */
+  createDatatypes(...datatypes: Datatype<Name>[]): DatatypeSort<Name>[];
+}
+
+/**
+ * A Sort representing an algebraic datatype.
+ * 
+ * After creation, this sort will have constructor, recognizer, and accessor
+ * functions dynamically attached based on the declared constructors.
+ * 
+ * @category Datatypes
+ */
+export interface DatatypeSort<Name extends string = 'main'> extends Sort<Name> {
+  /** @hidden */
+  readonly __typename: 'DatatypeSort';
+
+  /**
+   * Number of constructors in this datatype
+   */
+  numConstructors(): number;
+
+  /**
+   * Get the idx'th constructor function declaration
+   */
+  constructorDecl(idx: number): FuncDecl<Name>;
+
+  /**
+   * Get the idx'th recognizer function declaration
+   */
+  recognizer(idx: number): FuncDecl<Name>;
+
+  /**
+   * Get the accessor function declaration for the idx_a'th field of the idx_c'th constructor
+   */
+  accessor(constructorIdx: number, accessorIdx: number): FuncDecl<Name>;
+
+  cast(other: CoercibleToExpr<Name>): DatatypeExpr<Name>;
+
+  cast(other: DatatypeExpr<Name>): DatatypeExpr<Name>;
+}
+
+/**
+ * Represents expressions of datatype sorts.
+ * 
+ * @category Datatypes  
+ */
+export interface DatatypeExpr<Name extends string = 'main'> extends Expr<Name, DatatypeSort<Name>, Z3_ast> {
+  /** @hidden */
+  readonly __typename: 'DatatypeExpr';
+}
 
 /**
  * Defines the expression type of the body of a quantifier expression