Add Z3_mk_polymorphic_datatype to Python, .NET, Go, and TypeScript bindings (#9181)

* Add Z3_mk_polymorphic_datatype to Python, .NET, Go, and TypeScript bindings

Agent-Logs-Url: https://github.com/Z3Prover/z3/sessions/13ef481d-61f5-47e1-8659-59cd91692fdd

Co-authored-by: NikolajBjorner <3085284+NikolajBjorner@users.noreply.github.com>

* Improve Python error message for polymorphic datatype self-reference check

Agent-Logs-Url: https://github.com/Z3Prover/z3/sessions/13ef481d-61f5-47e1-8659-59cd91692fdd

Co-authored-by: NikolajBjorner <3085284+NikolajBjorner@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: NikolajBjorner <3085284+NikolajBjorner@users.noreply.github.com>

src/api/dotnet/Context.cs

@@ -562,6 +562,63 @@ namespace Microsoft.Z3
             }
         }
 
+        /// <summary>
+        /// Create a type variable sort for use as a parameter in polymorphic datatypes.
+        /// </summary>
+        /// <param name="name">name of the type variable</param>
+        public Sort MkTypeVariable(Symbol name)
+        {
+            Debug.Assert(name != null);
+            CheckContextMatch(name);
+            return new Sort(this, Native.Z3_mk_type_variable(nCtx, name.NativeObject));
+        }
+
+        /// <summary>
+        /// Create a type variable sort for use as a parameter in polymorphic datatypes.
+        /// </summary>
+        /// <param name="name">name of the type variable</param>
+        public Sort MkTypeVariable(string name)
+        {
+            using var symbol = MkSymbol(name);
+            return MkTypeVariable(symbol);
+        }
+
+        /// <summary>
+        /// Create a polymorphic datatype sort with explicit type parameters.
+        /// Type parameters should be sorts created with <see cref="MkTypeVariable"/>.
+        /// </summary>
+        /// <param name="name">name of the datatype sort</param>
+        /// <param name="typeParams">array of type variable sorts</param>
+        /// <param name="constructors">array of constructors</param>
+        public DatatypeSort MkPolymorphicDatatypeSort(Symbol name, Sort[] typeParams, Constructor[] constructors)
+        {
+            Debug.Assert(name != null);
+            Debug.Assert(typeParams != null);
+            Debug.Assert(constructors != null);
+            Debug.Assert(constructors.All(c => c != null));
+
+            CheckContextMatch(name);
+            CheckContextMatch<Sort>(typeParams);
+            CheckContextMatch<Constructor>(constructors);
+            return new DatatypeSort(this,
+                Native.Z3_mk_polymorphic_datatype(nCtx, name.NativeObject,
+                    (uint)typeParams.Length, AST.ArrayToNative(typeParams),
+                    (uint)constructors.Length, Z3Object.ArrayToNative(constructors)));
+        }
+
+        /// <summary>
+        /// Create a polymorphic datatype sort with explicit type parameters.
+        /// Type parameters should be sorts created with <see cref="MkTypeVariable"/>.
+        /// </summary>
+        /// <param name="name">name of the datatype sort</param>
+        /// <param name="typeParams">array of type variable sorts</param>
+        /// <param name="constructors">array of constructors</param>
+        public DatatypeSort MkPolymorphicDatatypeSort(string name, Sort[] typeParams, Constructor[] constructors)
+        {
+            using var symbol = MkSymbol(name);
+            return MkPolymorphicDatatypeSort(symbol, typeParams, constructors);
+        }
+
         /// <summary>
         /// Update a datatype field at expression t with value v.
         /// The function performs a record update at t. The field

src/api/go/datatype.go

@@ -127,6 +127,41 @@ func (c *Context) MkDatatypeSort(name string, constructors []*Constructor) *Sort
 	return newSort(c, C.Z3_mk_datatype(c.ptr, sym.ptr, C.uint(numCons), &cons[0]))
 }
 
+// MkPolymorphicDatatypeSort creates a polymorphic datatype sort with explicit type parameters.
+// typeParams should be sorts created with MkTypeVariable.
+// Self-recursive field sorts should be passed as nil; use the fieldSortRefs parameter in
+// MkConstructor to indicate the recursive reference by index.
+func (c *Context) MkPolymorphicDatatypeSort(name string, typeParams []*Sort, constructors []*Constructor) *Sort {
+	sym := c.MkStringSymbol(name)
+
+	numParams := len(typeParams)
+	numCons := len(constructors)
+
+	var paramPtr *C.Z3_sort
+	if numParams > 0 {
+		paramPtrs := make([]C.Z3_sort, numParams)
+		for i, p := range typeParams {
+			paramPtrs[i] = p.ptr
+		}
+		paramPtr = &paramPtrs[0]
+	}
+
+	var consPtr *C.Z3_constructor
+	if numCons > 0 {
+		consPtrs := make([]C.Z3_constructor, numCons)
+		for i, cons := range constructors {
+			consPtrs[i] = cons.ptr
+		}
+		consPtr = &consPtrs[0]
+	}
+
+	return newSort(c, C.Z3_mk_polymorphic_datatype(
+		c.ptr, sym.ptr,
+		C.uint(numParams), paramPtr,
+		C.uint(numCons), consPtr,
+	))
+}
+
 // MkDatatypeSorts creates multiple mutually recursive datatype sorts.
 func (c *Context) MkDatatypeSorts(names []string, constructorLists [][]*Constructor) []*Sort {
 	numTypes := uint(len(names))

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

@@ -1172,9 +1172,16 @@ export function createApi(Z3: Z3Core, em?: any): Z3HighLevel {
         createDatatypes(...datatypes: DatatypeImpl[]): DatatypeSortImpl[] {
           return createDatatypes(...datatypes);
         },
+        createPolymorphicDatatype(typeParams: Sort<Name>[], datatype: DatatypeImpl): DatatypeSortImpl {
+          return createPolymorphicDatatype(typeParams, datatype);
+        },
       },
     );
 
+    function TypeVariable(name: string): Sort<Name> {
+      return new SortImpl(check(Z3.mk_type_variable(contextPtr, Z3.mk_string_symbol(contextPtr, name))));
+    }
+
     ////////////////
     // Operations //
     ////////////////
@@ -4689,6 +4696,10 @@ export function createApi(Z3: Z3Core, em?: any): Z3HighLevel {
         const datatypes = createDatatypes(this);
         return datatypes[0];
       }
+
+      createPolymorphic(typeParams: Sort<Name>[]): DatatypeSort<Name> {
+        return createPolymorphicDatatype(typeParams, this);
+      }
     }
 
     class DatatypeSortImpl extends SortImpl implements DatatypeSort<Name> {
@@ -4845,6 +4856,84 @@ export function createApi(Z3: Z3Core, em?: any): Z3HighLevel {
       }
     }
 
+    function createPolymorphicDatatype(typeParams: Sort<Name>[], datatype: DatatypeImpl): DatatypeSortImpl {
+      if (!(datatype instanceof DatatypeImpl)) {
+        throw new Error('Datatype instance expected');
+      }
+
+      const constructors: Z3_constructor[] = [];
+
+      try {
+        for (const [constructorName, fields] of datatype.constructors) {
+          const fieldNames: string[] = [];
+          const fieldSorts: Z3_sort[] = [];
+          const fieldRefs: number[] = [];
+
+          for (const [fieldName, fieldSort] of fields) {
+            fieldNames.push(fieldName);
+
+            if (fieldSort instanceof DatatypeImpl) {
+              // Self-recursive reference
+              if (fieldSort !== datatype) {
+                throw new Error(
+                  `Referenced datatype "${fieldSort.name}" is not the polymorphic datatype being created; mutual recursion is not supported in createPolymorphicDatatype`,
+                );
+              }
+              fieldSorts.push(null as any);
+              fieldRefs.push(0);
+            } else {
+              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);
+        }
+
+        const nameSymbol = Z3.mk_string_symbol(contextPtr, datatype.name);
+        const paramPtrs = typeParams.map(p => p.ptr);
+        const resultSort = Z3.mk_polymorphic_datatype(contextPtr, nameSymbol, paramPtrs, constructors);
+
+        const sortImpl = new DatatypeSortImpl(resultSort);
+
+        // 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();
+
+          if (constructor.arity() === 0) {
+            (sortImpl as any)[constructorName] = constructor.call();
+          } else {
+            (sortImpl as any)[constructorName] = constructor;
+          }
+
+          (sortImpl as any)[`is_${constructorName}`] = recognizer;
+
+          for (let k = 0; k < constructor.arity(); k++) {
+            const accessor = sortImpl.accessor(j, k);
+            const accessorName = accessor.name().toString();
+            (sortImpl as any)[accessorName] = accessor;
+          }
+        }
+
+        return sortImpl;
+      } finally {
+        for (const constructor of constructors) {
+          Z3.del_constructor(contextPtr, constructor);
+        }
+      }
+    }
+
     class QuantifierImpl<
         QVarSorts extends NonEmptySortArray<Name>,
         QSort extends BoolSort<Name> | SMTArraySort<Name, QVarSorts>,
@@ -5292,6 +5381,7 @@ export function createApi(Z3: Z3Core, em?: any): Z3HighLevel {
       Set,
       FiniteSet,
       Datatype,
+      TypeVariable,
 
       ////////////////
       // Operations //

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

@@ -479,6 +479,12 @@ export interface Context<Name extends string = 'main'> {
   /** @category Expressions */
   readonly Datatype: DatatypeCreation<Name>;
 
+  /**
+   * Create a type variable sort for use as a parameter in polymorphic datatypes.
+   * @category Sorts
+   */
+  TypeVariable(name: string): Sort<Name>;
+
   ////////////////
   // Operations //
   ////////////////
@@ -3136,6 +3142,15 @@ export interface Datatype<Name extends string = 'main'> {
    * For mutually recursive datatypes, use Context.createDatatypes instead.
    */
   create(): DatatypeSort<Name>;
+
+  /**
+   * Create a polymorphic datatype sort with explicit type parameters.
+   * Type parameters should be sorts created with Context.TypeVariable.
+   * Self-recursive fields may reference this Datatype object directly.
+   *
+   * @param typeParams Array of type variable sorts
+   */
+  createPolymorphic(typeParams: AnySort<Name>[]): DatatypeSort<Name>;
 }
 
 /**
@@ -3154,6 +3169,17 @@ export interface DatatypeCreation<Name extends string> {
    * @returns Array of created DatatypeSort instances
    */
   createDatatypes(...datatypes: Datatype<Name>[]): DatatypeSort<Name>[];
+
+  /**
+   * Create a single polymorphic datatype sort with explicit type parameters.
+   * Type parameters should be sorts created with Context.TypeVariable.
+   * Self-recursive fields in constructors may reference the Datatype object directly.
+   *
+   * @param typeParams Array of type variable sorts
+   * @param datatype Datatype declaration with constructors
+   * @returns Created DatatypeSort instance
+   */
+  createPolymorphicDatatype(typeParams: AnySort<Name>[], datatype: Datatype<Name>): DatatypeSort<Name>;
 }
 
 /**

src/api/python/z3/z3.py

@@ -5612,6 +5612,20 @@ class Datatype:
         """
         return CreateDatatypes([self])[0]
 
+    def create_polymorphic(self, type_params):
+        """Create a polymorphic Z3 datatype with explicit type variables.
+
+        `type_params` is a list of type variables created with `DeclareTypeVar`.
+        Constructor field sorts may reference these type variables.
+        Self-recursive fields may reference this datatype directly.
+
+        >>> A = DeclareTypeVar('A')
+        >>> Pair = Datatype('Pair')
+        >>> Pair.declare('pair', ('fst', A), ('snd', A))
+        >>> Pair = Pair.create_polymorphic([A])
+        """
+        return CreatePolymorphicDatatype(self, type_params)
+
 
 class ScopedConstructor:
     """Auxiliary object used to create Z3 datatypes."""
@@ -5733,6 +5747,76 @@ def CreateDatatypes(*ds):
     return tuple(result)
 
 
+def CreatePolymorphicDatatype(d, type_params):
+    """Create a single polymorphic Z3 datatype with explicit type parameters.
+
+    `d` is a `Datatype` helper object whose constructors have been declared.
+    `type_params` is a list of type variables created with `DeclareTypeVar`.
+    Constructor field sorts may reference these type variables, and self-recursive
+    fields may reference `d` directly.
+
+    >>> A = DeclareTypeVar('A')
+    >>> Pair = Datatype('Pair')
+    >>> Pair.declare('pair', ('fst', A), ('snd', A))
+    >>> Pair = CreatePolymorphicDatatype(Pair, [A])
+    """
+    if z3_debug():
+        _z3_assert(isinstance(d, Datatype), "Datatype expected")
+        _z3_assert(d.constructors != [], "Non-empty Datatype expected")
+    ctx = d.ctx
+    name = to_symbol(d.name, ctx)
+    num_params = len(type_params)
+    params_arr = (Sort * num_params)()
+    for i, p in enumerate(type_params):
+        if z3_debug():
+            _z3_assert(is_sort(p), "Z3 sort expected for type parameter")
+        params_arr[i] = p.ast
+    num_cs = len(d.constructors)
+    cs = (Constructor * num_cs)()
+    to_delete = []
+    for j in range(num_cs):
+        c = d.constructors[j]
+        cname = to_symbol(c[0], ctx)
+        rname = to_symbol(c[1], ctx)
+        fs = c[2]
+        num_fs = len(fs)
+        fnames = (Symbol * num_fs)()
+        sorts = (Sort * num_fs)()
+        refs = (ctypes.c_uint * num_fs)()
+        for k in range(num_fs):
+            fname = fs[k][0]
+            ftype = fs[k][1]
+            fnames[k] = to_symbol(fname, ctx)
+            if isinstance(ftype, Datatype):
+                if z3_debug():
+                    _z3_assert(ftype is d, "Only self-recursive references are supported in polymorphic datatypes. Use CreateDatatypes for mutually recursive datatypes.")
+                sorts[k] = None
+                refs[k] = 0
+            else:
+                if z3_debug():
+                    _z3_assert(is_sort(ftype), "Z3 sort expected")
+                sorts[k] = ftype.ast
+                refs[k] = 0
+        cs[j] = Z3_mk_constructor(ctx.ref(), cname, rname, num_fs, fnames, sorts, refs)
+        to_delete.append(ScopedConstructor(cs[j], ctx))
+    out = Z3_mk_polymorphic_datatype(ctx.ref(), name, num_params, params_arr, num_cs, cs)
+    dref = DatatypeSortRef(out, ctx)
+    num_cs_actual = dref.num_constructors()
+    for j in range(num_cs_actual):
+        cref = dref.constructor(j)
+        cref_name = cref.name()
+        cref_arity = cref.arity()
+        if cref_arity == 0:
+            cref = cref()
+        setattr(dref, cref_name, cref)
+        rref = dref.recognizer(j)
+        setattr(dref, "is_" + cref_name, rref)
+        for k in range(cref_arity):
+            aref = dref.accessor(j, k)
+            setattr(dref, aref.name(), aref)
+    return dref
+
+
 class DatatypeSortRef(SortRef):
     """Datatype sorts."""