From d9c354f8092a45eacaf2b68226a49c58e44947b5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 21 Oct 2025 20:19:29 +0000 Subject: [PATCH] Add Python bindings for finite sets Co-authored-by: NikolajBjorner <3085284+NikolajBjorner@users.noreply.github.com> --- src/api/python/z3/z3.py | 219 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 218 insertions(+), 1 deletion(-) diff --git a/src/api/python/z3/z3.py b/src/api/python/z3/z3.py index 051265a78..104647284 100644 --- a/src/api/python/z3/z3.py +++ b/src/api/python/z3/z3.py @@ -678,6 +678,8 @@ def is_sort(s : Any) -> bool: def _to_sort_ref(s, ctx): if z3_debug(): _z3_assert(isinstance(s, Sort), "Z3 Sort expected") + if Z3_is_finite_set_sort(ctx.ref(), s): + return FiniteSetSortRef(s, ctx) k = _sort_kind(ctx, s) if k == Z3_BOOL_SORT: return BoolSortRef(s, ctx) @@ -1184,7 +1186,11 @@ def _to_expr_ref(a, ctx): k = Z3_get_ast_kind(ctx_ref, a) if k == Z3_QUANTIFIER_AST: return QuantifierRef(a, ctx) - sk = Z3_get_sort_kind(ctx_ref, Z3_get_sort(ctx_ref, a)) + # Check for finite set sort before checking sort kind + s = Z3_get_sort(ctx_ref, a) + if Z3_is_finite_set_sort(ctx_ref, s): + return FiniteSetRef(a, ctx) + sk = Z3_get_sort_kind(ctx_ref, s) if sk == Z3_BOOL_SORT: return BoolRef(a, ctx) if sk == Z3_INT_SORT: @@ -5147,6 +5153,217 @@ def IsSubset(a, b): return BoolRef(Z3_mk_set_subset(ctx.ref(), a.as_ast(), b.as_ast()), ctx) +######################################### +# +# Finite Sets +# +######################################### + + +class FiniteSetSortRef(SortRef): + """Finite set sort.""" + + def element_sort(self): + """Return the element sort of this finite set sort.""" + return _to_sort_ref(Z3_get_finite_set_sort_basis(self.ctx_ref(), self.ast), self.ctx) + + def cast(self, val): + """Try to cast val as a finite set expression.""" + if is_expr(val): + if self.eq(val.sort()): + return val + else: + _z3_assert(False, "Cannot cast to finite set sort") + if isinstance(val, set): + elem_sort = self.element_sort() + result = FiniteSetEmpty(self) + for e in val: + result = FiniteSetUnion(result, FiniteSetSingleton(_py2expr(e, self.ctx, elem_sort))) + return result + _z3_assert(False, "Cannot cast to finite set sort") + + def subsort(self, other): + return False + + def is_int(self): + return False + + def is_bool(self): + return False + + def is_datatype(self): + return False + + def is_array(self): + return False + + def is_bv(self): + return False + + +def is_finite_set(a): + """Return True if a is a Z3 finite set expression. + >>> s = FiniteSetSort(IntSort()) + >>> is_finite_set(FiniteSetEmpty(s)) + True + >>> is_finite_set(IntVal(1)) + False + """ + return isinstance(a, FiniteSetRef) + + +def is_finite_set_sort(s): + """Return True if s is a Z3 finite set sort. + >>> is_finite_set_sort(FiniteSetSort(IntSort())) + True + >>> is_finite_set_sort(IntSort()) + False + """ + return isinstance(s, FiniteSetSortRef) + + +class FiniteSetRef(ExprRef): + """Finite set expression.""" + + def sort(self): + return FiniteSetSortRef(Z3_get_sort(self.ctx_ref(), self.as_ast()), self.ctx) + + def __or__(self, other): + """Return the union of self and other.""" + return FiniteSetUnion(self, other) + + def __and__(self, other): + """Return the intersection of self and other.""" + return FiniteSetIntersect(self, other) + + def __sub__(self, other): + """Return the set difference of self and other.""" + return FiniteSetDifference(self, other) + + +def FiniteSetSort(elem_sort): + """Create a finite set sort over element sort elem_sort. + >>> s = FiniteSetSort(IntSort()) + >>> s + FiniteSet(Int) + """ + return FiniteSetSortRef(Z3_mk_finite_set_sort(elem_sort.ctx_ref(), elem_sort.ast), elem_sort.ctx) + + +def FiniteSetEmpty(set_sort): + """Create an empty finite set of the given sort. + >>> s = FiniteSetSort(IntSort()) + >>> FiniteSetEmpty(s) + set.empty + """ + ctx = set_sort.ctx + return FiniteSetRef(Z3_mk_finite_set_empty(ctx.ref(), set_sort.ast), ctx) + + +def FiniteSetSingleton(elem): + """Create a singleton finite set containing elem. + >>> FiniteSetSingleton(IntVal(1)) + set.singleton(1) + """ + ctx = elem.ctx + return FiniteSetRef(Z3_mk_finite_set_singleton(ctx.ref(), elem.as_ast()), ctx) + + +def FiniteSetUnion(s1, s2): + """Create the union of two finite sets. + >>> a = Const('a', FiniteSetSort(IntSort())) + >>> b = Const('b', FiniteSetSort(IntSort())) + >>> FiniteSetUnion(a, b) + set.union(a, b) + """ + ctx = _ctx_from_ast_arg_list([s1, s2]) + return FiniteSetRef(Z3_mk_finite_set_union(ctx.ref(), s1.as_ast(), s2.as_ast()), ctx) + + +def FiniteSetIntersect(s1, s2): + """Create the intersection of two finite sets. + >>> a = Const('a', FiniteSetSort(IntSort())) + >>> b = Const('b', FiniteSetSort(IntSort())) + >>> FiniteSetIntersect(a, b) + set.intersect(a, b) + """ + ctx = _ctx_from_ast_arg_list([s1, s2]) + return FiniteSetRef(Z3_mk_finite_set_intersect(ctx.ref(), s1.as_ast(), s2.as_ast()), ctx) + + +def FiniteSetDifference(s1, s2): + """Create the set difference of two finite sets. + >>> a = Const('a', FiniteSetSort(IntSort())) + >>> b = Const('b', FiniteSetSort(IntSort())) + >>> FiniteSetDifference(a, b) + set.difference(a, b) + """ + ctx = _ctx_from_ast_arg_list([s1, s2]) + return FiniteSetRef(Z3_mk_finite_set_difference(ctx.ref(), s1.as_ast(), s2.as_ast()), ctx) + + +def FiniteSetMember(elem, set): + """Check if elem is a member of the finite set. + >>> a = Const('a', FiniteSetSort(IntSort())) + >>> FiniteSetMember(IntVal(1), a) + set.in(1, a) + """ + ctx = _ctx_from_ast_arg_list([elem, set]) + return BoolRef(Z3_mk_finite_set_member(ctx.ref(), elem.as_ast(), set.as_ast()), ctx) + + +def FiniteSetSize(set): + """Get the size (cardinality) of a finite set. + >>> a = Const('a', FiniteSetSort(IntSort())) + >>> FiniteSetSize(a) + set.size(a) + """ + ctx = set.ctx + return ArithRef(Z3_mk_finite_set_size(ctx.ref(), set.as_ast()), ctx) + + +def FiniteSetSubset(s1, s2): + """Check if s1 is a subset of s2. + >>> a = Const('a', FiniteSetSort(IntSort())) + >>> b = Const('b', FiniteSetSort(IntSort())) + >>> FiniteSetSubset(a, b) + set.subset(a, b) + """ + ctx = _ctx_from_ast_arg_list([s1, s2]) + return BoolRef(Z3_mk_finite_set_subset(ctx.ref(), s1.as_ast(), s2.as_ast()), ctx) + + +def FiniteSetMap(f, set): + """Apply function f to all elements of the finite set. + >>> f = Function('f', IntSort(), IntSort()) + >>> a = Const('a', FiniteSetSort(IntSort())) + >>> FiniteSetMap(f, a) + set.map(f, a) + """ + ctx = _ctx_from_ast_arg_list([f, set]) + return FiniteSetRef(Z3_mk_finite_set_map(ctx.ref(), f.as_ast(), set.as_ast()), ctx) + + +def FiniteSetFilter(f, set): + """Filter a finite set using predicate f. + >>> f = Function('f', IntSort(), BoolSort()) + >>> a = Const('a', FiniteSetSort(IntSort())) + >>> FiniteSetFilter(f, a) + set.filter(f, a) + """ + ctx = _ctx_from_ast_arg_list([f, set]) + return FiniteSetRef(Z3_mk_finite_set_filter(ctx.ref(), f.as_ast(), set.as_ast()), ctx) + + +def FiniteSetRange(low, high): + """Create a finite set of integers in the range [low, high). + >>> FiniteSetRange(IntVal(0), IntVal(5)) + set.range(0, 5) + """ + ctx = _ctx_from_ast_arg_list([low, high]) + return FiniteSetRef(Z3_mk_finite_set_range(ctx.ref(), low.as_ast(), high.as_ast()), ctx) + + ######################################### # # Datatypes