diff --git a/z3/src/ast.rs b/z3/src/ast.rs index ea3b2838..09bef91d 100644 --- a/z3/src/ast.rs +++ b/z3/src/ast.rs @@ -760,6 +760,48 @@ impl<'ctx> Bool<'ctx> { } } +pub fn atmost<'a, 'ctx, I: IntoIterator>>( + ctx: &'ctx Context, + args: I, + k: u32, +) -> Bool<'ctx> +where + 'ctx: 'a, +{ + let args: Vec<_> = args.into_iter().map(|f| f.z3_ast).collect(); + _atmost(ctx, args.as_ref(), k) +} + +fn _atmost<'ctx>(ctx: &'ctx Context, args: &[Z3_ast], k: u32) -> Bool<'ctx> { + unsafe { + Bool::wrap( + ctx, + Z3_mk_atmost(ctx.z3_ctx, args.len().try_into().unwrap(), args.as_ptr(), k), + ) + } +} + +pub fn atleast<'a, 'ctx, I: IntoIterator>>( + ctx: &'ctx Context, + args: I, + k: u32, +) -> Bool<'ctx> +where + 'ctx: 'a, +{ + let args: Vec<_> = args.into_iter().map(|f| f.z3_ast).collect(); + _atleast(ctx, args.as_ref(), k) +} + +fn _atleast<'ctx>(ctx: &'ctx Context, args: &[Z3_ast], k: u32) -> Bool<'ctx> { + unsafe { + Bool::wrap( + ctx, + Z3_mk_atleast(ctx.z3_ctx, args.len().try_into().unwrap(), args.as_ptr(), k), + ) + } +} + impl<'ctx> Int<'ctx> { pub fn new_const>(ctx: &'ctx Context, name: S) -> Int<'ctx> { let sort = Sort::int(ctx); diff --git a/z3/tests/lib.rs b/z3/tests/lib.rs index 4d6fc7b6..0f206684 100644 --- a/z3/tests/lib.rs +++ b/z3/tests/lib.rs @@ -2,7 +2,7 @@ use log::info; use std::convert::TryInto; use std::ops::Add; use std::time::Duration; -use z3::ast::{Array, Ast, Bool, Int, BV}; +use z3::ast::{atleast, atmost, Array, Ast, Bool, Int, BV}; use z3::*; use num::{bigint::BigInt, rational::BigRational}; @@ -1921,3 +1921,55 @@ fn test_consequences() { assert!(cons.pop().unwrap().to_string() == "(=> (not c) (not c))"); assert!(cons.pop().unwrap().to_string() == "(=> d d)"); } + +#[test] +fn test_atmost() { + let cfg = Config::new(); + let ctx = &Context::new(&cfg); + let solver = Solver::new(ctx); + let a = Bool::new_const(ctx, "a"); + let b = Bool::new_const(ctx, "b"); + let c = Bool::new_const(ctx, "c"); + let d = Bool::new_const(ctx, "d"); + solver.assert(&a.implies(&b)); + solver.assert(&b.implies(&c)); + + solver.push(); + let am = atmost(ctx, [&a, &b, &c, &d], 2); + solver.assert(&am); + assert!(matches!(solver.check(), SatResult::Sat)); + solver.pop(1); + + solver.push(); + solver.assert(&a); + let am = atmost(ctx, [&a, &b, &c, &d], 0); + solver.assert(&am); + assert!(matches!(solver.check(), SatResult::Unsat)); + solver.pop(1); +} + +#[test] +fn test_atleast() { + let cfg = Config::new(); + let ctx = &Context::new(&cfg); + let solver = Solver::new(ctx); + let a = Bool::new_const(ctx, "a"); + let b = Bool::new_const(ctx, "b"); + let c = Bool::new_const(ctx, "c"); + let d = Bool::new_const(ctx, "d"); + solver.assert(&a.implies(&b)); + solver.assert(&b.implies(&c)); + + solver.push(); + let am = atleast(ctx, [&a, &b, &c, &d], 4); + solver.assert(&am); + assert!(matches!(solver.check(), SatResult::Sat)); + solver.pop(1); + + solver.push(); + solver.assert(&a.not()); + let am = atleast(ctx, [&a, &b, &c, &d], 4); + solver.assert(&am); + assert!(matches!(solver.check(), SatResult::Unsat)); + solver.pop(1); +}