Implement Iterator over a Model

z3/src/model.rs

@@ -126,6 +126,17 @@ impl<'ctx> Model<'ctx> {
             None
         }
     }
+
+    fn len(&self) -> u32 {
+        unsafe {
+            Z3_model_get_num_consts(self.ctx.z3_ctx, self.z3_mdl)
+                + Z3_model_get_num_funcs(self.ctx.z3_ctx, self.z3_mdl)
+        }
+    }
+
+    pub fn iter(&'ctx self) -> ModelIter<'ctx> {
+        self.into_iter()
+    }
 }
 
 impl<'ctx> fmt::Display for Model<'ctx> {
@@ -153,6 +164,62 @@ impl<'ctx> Drop for Model<'ctx> {
     }
 }
 
+#[derive(Debug)]
+/// https://z3prover.github.io/api/html/classz3py_1_1_model_ref.html#a7890b7c9bc70cf2a26a343c22d2c8367
+pub struct ModelIter<'ctx> {
+    model: &'ctx Model<'ctx>,
+    idx: u32,
+    len: u32,
+}
+
+impl<'ctx> IntoIterator for &'ctx Model<'ctx> {
+    type Item = FuncDecl<'ctx>;
+    type IntoIter = ModelIter<'ctx>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        ModelIter {
+            model: self,
+            idx: 0,
+            len: self.len(),
+        }
+    }
+}
+
+impl<'ctx> Iterator for ModelIter<'ctx> {
+    type Item = FuncDecl<'ctx>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.idx >= self.len {
+            None
+        } else {
+            let num_consts =
+                unsafe { Z3_model_get_num_consts(self.model.ctx.z3_ctx, self.model.z3_mdl) };
+            if self.idx < num_consts {
+                let const_decl = unsafe {
+                    Z3_model_get_const_decl(self.model.ctx.z3_ctx, self.model.z3_mdl, self.idx)
+                };
+                self.idx += 1;
+                Some(unsafe { FuncDecl::wrap(self.model.ctx, const_decl) })
+            } else {
+                let func_decl = unsafe {
+                    Z3_model_get_func_decl(
+                        self.model.ctx.z3_ctx,
+                        self.model.z3_mdl,
+                        self.idx - num_consts,
+                    )
+                };
+                self.idx += 1;
+                Some(unsafe { FuncDecl::wrap(self.model.ctx, func_decl) })
+            }
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = (self.len - self.idx) as usize;
+        (len, Some(len))
+    }
+}
+
 #[test]
 fn test_unsat() {
     use crate::{ast, Config, SatResult};

z3/tests/lib.rs

@@ -1654,3 +1654,60 @@ fn return_number_args_in_given_entry() {
 
     assert!(model.to_string() == "f -> {\n  1 2 -> 20\n  else -> 10\n}\n");
 }
+
+#[test]
+/// https://stackoverflow.com/questions/13395391/z3-finding-all-satisfying-models
+fn iterate_all_solutions() {
+    let cfg = Config::new();
+    let ctx = &Context::new(&cfg);
+    let solver = Solver::new(ctx);
+    let a = &Int::new_const(ctx, "a");
+    let b = &Int::new_const(ctx, "b");
+    let one = Int::from_u64(ctx, 1);
+    let two = Int::from_u64(ctx, 2);
+    let five = &Int::from_u64(ctx, 5);
+
+    solver.assert(&one.le(a));
+    solver.assert(&a.le(five));
+    solver.assert(&one.le(b));
+    solver.assert(&b.le(five));
+    solver.assert(&a.ge(&(two * b)));
+
+    let mut solutions = std::collections::HashSet::new();
+    while solver.check() == SatResult::Sat {
+        let model = solver.get_model().unwrap();
+        let mut modifications = Vec::new();
+        let this_solution = model
+            .iter()
+            .map(|fd| {
+                modifications.push(
+                    fd.apply(&[])
+                        ._eq(&model.get_const_interp(&fd.apply(&[])).unwrap())
+                        .not(),
+                );
+                format!(
+                    "{} = {}",
+                    fd.name(),
+                    model.get_const_interp(&fd.apply(&[])).unwrap()
+                )
+            })
+            .collect::<Vec<_>>()
+            .join(", ");
+        solutions.insert(format!("[{this_solution}]"));
+        solver.assert(&Bool::or(ctx, &modifications.iter().collect::<Vec<_>>()));
+    }
+
+    assert!(
+        solutions
+            == vec![
+                "[b = 1, a = 2]".to_string(),
+                "[b = 2, a = 4]".to_string(),
+                "[b = 1, a = 3]".to_string(),
+                "[b = 2, a = 5]".to_string(),
+                "[b = 1, a = 4]".to_string(),
+                "[b = 1, a = 5]".to_string()
+            ]
+            .into_iter()
+            .collect()
+    )
+}