From 7f5495b134eebac0a06da12481dda1a045350f3a Mon Sep 17 00:00:00 2001 From: Nikolaj Bjorner Date: Wed, 28 Oct 2015 17:09:42 -0700 Subject: [PATCH] adding PB operators to Python API Signed-off-by: Nikolaj Bjorner --- src/api/python/z3.py | 38 +++++++++++++++++++++++++++++++++++++ src/api/python/z3printer.py | 22 ++++++++++++++++++++- 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/src/api/python/z3.py b/src/api/python/z3.py index 408b88e65..c918533ca 100644 --- a/src/api/python/z3.py +++ b/src/api/python/z3.py @@ -7303,6 +7303,44 @@ def Product(*args): _args, sz = _to_ast_array(args) return ArithRef(Z3_mk_mul(ctx.ref(), sz, _args), ctx) +def AtMost(*args): + """Create an at-most Pseudo-Boolean k constraint. + + >>> a, b, c = Bools('a b c') + >>> f = AtMost(a, b, c, 2) + """ + args = _get_args(args) + if __debug__: + _z3_assert(len(args) > 1, "Non empty list of arguments expected") + ctx = _ctx_from_ast_arg_list(args) + if __debug__: + _z3_assert(ctx != None, "At least one of the arguments must be a Z3 expression") + args1 = _coerce_expr_list(args[:-1], ctx) + k = args[-1] + _args, sz = _to_ast_array(args1) + return BoolRef(Z3_mk_atmost(ctx.ref(), sz, _args, k), ctx) + +def PbLe(args, k): + """Create a Pseudo-Boolean inequality k constraint. + + >>> a, b, c = Bools('a b c') + >>> f = PbLe(((a,1),(b,3),(c,2)), 3) + """ + args = _get_args(args) + args, coeffs = zip(*args) + if __debug__: + _z3_assert(len(args) > 0, "Non empty list of arguments expected") + ctx = _ctx_from_ast_arg_list(args) + if __debug__: + _z3_assert(ctx != None, "At least one of the arguments must be a Z3 expression") + args = _coerce_expr_list(args, ctx) + _args, sz = _to_ast_array(args) + _coeffs = (ctypes.c_int * len(coeffs))() + for i in range(len(coeffs)): + _coeffs[i] = coeffs[i] + return BoolRef(Z3_mk_pble(ctx.ref(), sz, _args, _coeffs, k), ctx) + + def solve(*args, **keywords): """Solve the constraints `*args`. diff --git a/src/api/python/z3printer.py b/src/api/python/z3printer.py index fe6378357..b4b83dc2b 100644 --- a/src/api/python/z3printer.py +++ b/src/api/python/z3printer.py @@ -31,7 +31,8 @@ _z3_op_to_str = { Z3_OP_BASHR : '>>', Z3_OP_BSHL : '<<', Z3_OP_BLSHR : 'LShR', Z3_OP_CONCAT : 'Concat', Z3_OP_EXTRACT : 'Extract', Z3_OP_BV2INT : 'BV2Int', Z3_OP_ARRAY_MAP : 'Map', Z3_OP_SELECT : 'Select', Z3_OP_STORE : 'Store', - Z3_OP_CONST_ARRAY : 'K' + Z3_OP_CONST_ARRAY : 'K', + Z3_OP_PB_AT_MOST : 'AtMost', Z3_OP_PB_LE : 'PbLe', Z3_OP_PB_GE : 'PbGe' } # List of infix operators @@ -841,6 +842,19 @@ class Formatter: def pp_K(self, a, d, xs): return seq1(self.pp_name(a), [ self.pp_sort(a.domain()), self.pp_expr(a.arg(0), d+1, xs) ]) + def pp_atmost(self, a, d, f, xs): + k = Z3_get_decl_int_parameter(a.ctx_ref(), a.decl().ast, 0) + return seq1(self.pp_name(a), [seq3([ self.pp_expr(ch, d+1, xs) for ch in a.children()]), to_format(k)]) + + def pp_pbcmp(self, a, d, f, xs): + chs = a.children() + rchs = range(len(chs)) + k = Z3_get_decl_int_parameter(a.ctx_ref(), a.decl().ast, 0) + ks = [Z3_get_decl_int_parameter(a.ctx_ref(), a.decl().ast, i+1) for i in rchs] + ls = [ seq3([self.pp_expr(chs[i], d+1,xs), to_format(ks[i])]) for i in rchs] + return seq1(self.pp_name(a), [seq3(ls), to_format(k)]) + + def pp_app(self, a, d, xs): if z3.is_int_value(a): return self.pp_int(a) @@ -875,6 +889,12 @@ class Formatter: return self.pp_map(a, d, xs) elif k == Z3_OP_CONST_ARRAY: return self.pp_K(a, d, xs) + elif k == Z3_OP_PB_AT_MOST: + return self.pp_atmost(a, d, f, xs) + elif k == Z3_OP_PB_LE: + return self.pp_pbcmp(a, d, f, xs) + elif k == Z3_OP_PB_GE: + return self.pp_pbcmp(a, d, f, xs) elif z3.is_pattern(a): return self.pp_pattern(a, d, xs) elif self.is_infix(k):