python: smt solver interface

Author: ekiwi

python/Cargo.toml

@@ -15,3 +15,4 @@ pyo3 = { version = "0.26.0" , features = ["num-bigint"]}
 patronus = { path = "../patronus" }
 baa = { version = "0.17.1", features = ["rand1", "bigint"] }
 num-bigint = "0.4.6"
+rustc-hash.workspace = true

python/src/expr.rs

@@ -5,10 +5,9 @@ use num_bigint::BigInt;
 use patronus::expr::{TypeCheck, WidthInt};
 use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
-use std::fmt::format;
 
 #[pyclass]
-#[derive(Clone)]
+#[derive(Clone, Copy)]
 pub struct ExprRef(pub(crate) patronus::expr::ExprRef);
 
 /// Helper for binary ops that require a and b to be bitvectors of the same width

python/src/lib.rs

@@ -1,11 +1,13 @@
 mod ctx;
 mod expr;
 mod sim;
+mod smt;
 
 pub use ctx::Context;
 use ctx::{ContextGuardRead, ContextGuardWrite};
 pub use expr::*;
 pub use sim::{Simulator, interpreter};
+pub use smt::*;
 use std::path::PathBuf;
 
 use ::patronus::btor2;
@@ -259,5 +261,7 @@ fn patronus(_py: Python<'_>, m: &pyo3::Bound<'_, pyo3::types::PyModule>) -> PyRe
     m.add_function(wrap_pyfunction!(bit_vec, m)?)?;
     m.add_function(wrap_pyfunction!(bit_vec_val, m)?)?;
     m.add_function(wrap_pyfunction!(if_expr, m)?)?;
+    // smt
+    m.add_function(wrap_pyfunction!(solver, m)?)?;
     Ok(())
 }

python/src/smt.rs

@@ -0,0 +1,185 @@
+use crate::ExprRef;
+use crate::ctx::{ContextGuardRead, ContextGuardWrite};
+use baa::{BitVecOps, Value};
+use num_bigint::BigInt;
+use patronus::expr::{Context, ForEachChild, TypeCheck};
+use patronus::mc::get_smt_value;
+use patronus::smt::*;
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+use rustc_hash::{FxHashMap, FxHashSet};
+use std::fs::File;
+
+#[pyclass]
+pub struct SolverCtx {
+    underlying: SmtLibSolverCtx<File>,
+    declared_symbols: Vec<FxHashMap<String, patronus::expr::ExprRef>>,
+}
+
+impl SolverCtx {
+    fn new(underlying: SmtLibSolverCtx<File>) -> Self {
+        Self {
+            underlying,
+            declared_symbols: vec![FxHashMap::default()],
+        }
+    }
+
+    fn symbol_by_name(&self, name: &str) -> Option<patronus::expr::ExprRef> {
+        // from inner to outer
+        for map in self.declared_symbols.iter().rev() {
+            if let Some(symbol) = map.get(name) {
+                return Some(*symbol);
+            }
+        }
+        None
+    }
+}
+
+fn find_symbols(ctx: &Context, e: patronus::expr::ExprRef) -> FxHashSet<patronus::expr::ExprRef> {
+    let mut out = FxHashSet::default();
+    patronus::expr::traversal::bottom_up(ctx, e, |ctx, e, _| {
+        if ctx[e].is_symbol() {
+            out.insert(e);
+        }
+    });
+    out
+}
+
+#[pymethods]
+impl SolverCtx {
+    #[pyo3(signature = (*assertions))]
+    fn check(&mut self, assertions: Vec<ExprRef>) -> PyResult<CheckSatResult> {
+        if !assertions.is_empty() {
+            self.push()?;
+            for a in assertions.iter() {
+                self.add(*a)?;
+            }
+        }
+        let r = self
+            .underlying
+            .check_sat()
+            .map(CheckSatResult)
+            .map_err(convert_smt_err)?;
+        if !assertions.is_empty() {
+            self.pop()?;
+        }
+        Ok(r)
+    }
+
+    fn push(&mut self) -> PyResult<()> {
+        self.underlying.push().map_err(convert_smt_err)?;
+        self.declared_symbols.push(FxHashMap::default());
+        Ok(())
+    }
+
+    fn pop(&mut self) -> PyResult<()> {
+        self.underlying.pop().map_err(convert_smt_err)?;
+        self.declared_symbols.pop();
+        Ok(())
+    }
+
+    fn add(&mut self, assertion: ExprRef) -> PyResult<()> {
+        let ctx_guard = ContextGuardRead::default();
+        let ctx = ctx_guard.deref();
+        let a = assertion.0;
+        // scan the expression for any unknown symbols and declare them
+        let symbols = find_symbols(ctx, a);
+        for symbol in symbols.into_iter() {
+            let tpe = ctx[symbol].get_type(ctx);
+            let name = ctx[symbol].get_symbol_name(ctx).unwrap();
+            if let Some(existing) = self.symbol_by_name(name) {
+                // check for compatible type for existing symbols
+                let existing_tpe = ctx[existing].get_type(ctx);
+                if existing_tpe != tpe {
+                    return Err(PyRuntimeError::new_err(format!(
+                        "There is already a symbol `{name}` with incompatible type {existing_tpe} != {tpe}"
+                    )));
+                }
+            } else {
+                // declare if symbol does not exist
+                self.underlying
+                    .declare_const(ctx, symbol)
+                    .map_err(convert_smt_err)?;
+                self.declared_symbols
+                    .last_mut()
+                    .unwrap()
+                    .insert(name.to_string(), symbol);
+            }
+        }
+
+        self.underlying.assert(ctx, a).map_err(convert_smt_err)?;
+        Ok(())
+    }
+
+    fn model(&mut self) -> PyResult<Model> {
+        let mut ctx_guard = ContextGuardWrite::default();
+        let ctx = ctx_guard.deref_mut();
+        let mut entries = vec![];
+        for s in self.declared_symbols.iter().flat_map(|m| m.values()) {
+            let value = get_smt_value(ctx, &mut self.underlying, *s).map_err(convert_smt_err)?;
+            entries.push((*s, value));
+        }
+        Ok(Model(entries))
+    }
+}
+
+#[pyclass]
+pub struct Model(Vec<(patronus::expr::ExprRef, Value)>);
+
+#[pymethods]
+impl Model {
+    fn __str__(&self) -> String {
+        "TODO".to_string()
+    }
+
+    fn __len__(&self) -> usize {
+        self.0.len()
+    }
+
+    fn __getitem__(&self, symbol: ExprRef) -> Option<BigInt> {
+        self.0
+            .iter()
+            .find(|(e, _)| *e == symbol.0)
+            .map(|(_, value)| match value {
+                Value::Array(_) => {
+                    todo!("Array support!")
+                }
+                Value::BitVec(bv) => bv.to_big_int(),
+            })
+    }
+}
+
+#[pyclass]
+pub struct CheckSatResult(CheckSatResponse);
+
+#[pymethods]
+impl CheckSatResult {
+    fn __str__(&self) -> String {
+        match self.0 {
+            CheckSatResponse::Sat => "sat".to_string(),
+            CheckSatResponse::Unsat => "unsat".to_string(),
+            CheckSatResponse::Unknown => "unknonw".to_string(),
+        }
+    }
+}
+
+#[pyfunction]
+#[pyo3(name = "Solver")]
+pub fn solver(name: &str) -> PyResult<SolverCtx> {
+    match name.to_ascii_lowercase().as_str() {
+        "z3" => Ok(SolverCtx::new(Z3.start(None).map_err(convert_smt_err)?)),
+        "bitwuzla" => Ok(SolverCtx::new(
+            BITWUZLA.start(None).map_err(convert_smt_err)?,
+        )),
+        "yices" | "yices2" | "yices2-smt" => {
+            Ok(SolverCtx::new(YICES2.start(None).map_err(convert_smt_err)?))
+        }
+        _ => Err(PyRuntimeError::new_err(format!(
+            "Unknonw or unsupported solver: {name}"
+        ))),
+    }
+}
+
+fn convert_smt_err(e: Error) -> PyErr {
+    PyRuntimeError::new_err(format!("smt: {e}"))
+}

python/tests/test_btor2.py python/tests/test_all.pypython/tests/test_btor2.py renamed to python/tests/test_all.py

@@ -88,3 +88,23 @@ def test_transition_system_builder():
   [next] ite(en, add(count, 8'b00000001), count)
     """
     assert str(sys).strip() == expected_system.strip()
+
+
+def test_call_smt_solver():
+    a = BitVec('a', 3)
+    b = BitVec('b', 3)
+    s = Solver('z3')
+    r = s.check(a < b)
+    assert str(r) == "sat"
+
+    r = s.check(a < b, a > b)
+    assert str(r) == "unsat"
+
+    # to generate a model, we need to actually add the assertion!
+    s.add(a < b)
+    s.check()
+    m = s.model()
+    assert len(m) == 2
+    assert isinstance(m[a], int)
+    assert isinstance(m[b], int)
+    assert m[a] < m[b]