feat: atmost and atleast

prove-rs · Nov 5, 2024 · 088b4c5 · 088b4c5
1 parent 1a1e734
commit 088b4c5
Show file tree

Hide file tree

Showing 2 changed files with 105 additions and 1 deletion.
diff --git a/z3/src/ast.rs b/z3/src/ast.rs
@@ -760,6 +760,58 @@ impl<'ctx> Bool<'ctx> {
     }
 }
 
+pub fn atmost<'a, 'ctx, I: IntoIterator<Item = &'a Bool<'ctx>>>(
+    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<Item = &'a Bool<'ctx>>>(
+    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<S: Into<Symbol>>(ctx: &'ctx Context, name: S) -> Int<'ctx> {
         let sort = Sort::int(ctx);

diff --git 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);
+}