sirius/solver/
t1.rs

1//! Linearized systems of constraints and a pipe to [Z3](https://github.com/z3prover/z3) for QF_LIA solving.
2
3use js_sys::Function;
4use std::collections::HashMap;
5use std::io::Write as _;
6use std::path::PathBuf;
7use std::process::{Command, Stdio};
8use wasm_bindgen::prelude::*;
9
10use super::Constraint;
11use super::Verdict;
12use super::poly::{
13    Poly,
14    coef::Coef,
15    mono::{Mono, Var},
16};
17
18/// A linearized representation of a system of constraints.
19/// Non-linear atoms get fresh [`Var`]s and names like `X*Y` or `Z^2`.
20pub struct Linearized {
21    pub facts: Vec<Constraint>,
22    pub goal: Constraint,
23    pub nonneg: Vec<Var>,
24    pub pure_linear: bool,
25    pub names: Vec<String>,
26}
27
28impl Linearized {
29    pub fn new(
30        facts: &[Constraint],
31        goal: &Constraint,
32        vars: &[String],
33        nonneg: &dyn Fn(Var) -> bool,
34    ) -> Self {
35        let base = vars.len() as u32;
36        let mut atoms: Vec<Mono> = Vec::new();
37        let mut linearize_poly = |p: &Poly| -> Poly {
38            let terms = p.terms().into_iter().map(|(c, m)| {
39                if m.total_degree() <= 1 {
40                    (c, m.clone())
41                } else {
42                    let idx = match atoms.iter().position(|a| *a == m) {
43                        Some(i) => i,
44                        None => {
45                            atoms.push(m.clone());
46                            atoms.len() - 1
47                        }
48                    };
49                    (c, Mono::new([(base + idx as u32, 1)]))
50                }
51            });
52            Poly::from_terms(terms)
53        };
54
55        let mut linearize_constraint = |c: &Constraint| Constraint {
56            lhs: linearize_poly(&c.lhs),
57            cmp: c.cmp,
58            rhs: linearize_poly(&c.rhs),
59        };
60
61        let facts: Vec<Constraint> = facts.iter().map(&mut linearize_constraint).collect();
62        let goal = linearize_constraint(goal);
63
64        let mut nonneg_vars: Vec<Var> = (0..base).filter(|&v| nonneg(v)).collect();
65        let mut names: Vec<String> = vars.to_vec();
66        for (i, m) in atoms.iter().enumerate() {
67            // An atom over nonneg variables is nonneg
68            if m.exps().iter().all(|&(v, _)| nonneg(v)) {
69                nonneg_vars.push(base + i as u32);
70            }
71            let rendered: Vec<String> = m
72                .exps()
73                .iter()
74                .map(|&(v, e)| {
75                    let name = &vars[v as usize];
76                    if e > 1 {
77                        format!("{name}^{e}")
78                    } else {
79                        name.to_string()
80                    }
81                })
82                .collect();
83            names.push(rendered.join("*"));
84        }
85
86        Self {
87            pure_linear: atoms.is_empty(),
88            facts,
89            goal,
90            nonneg: nonneg_vars,
91            names,
92        }
93    }
94}
95
96enum Z3Kind {
97    Callback(Box<dyn Fn(&str) -> String>),
98    Cli,
99}
100
101/// A cached Z3 pipe, using either the command line or a JavaScript callback.
102pub struct Z3 {
103    /// Obligations submitted.
104    pub queries: usize,
105    /// Obligations passed to z3.
106    pub z3_calls: usize,
107    cache: HashMap<u64, Verdict>,
108    kind: Z3Kind,
109    // optional filesystem cache
110    cache_dir: Option<PathBuf>,
111}
112
113impl Z3 {
114    /// Expect the `z3` command to be available in the environment.
115    /// Optionally provide a filesystem cache for model output.
116    pub fn new_cli(cache_dir: Option<PathBuf>) -> Option<Self> {
117        let z3_available = Command::new("z3")
118            .arg("-version")
119            .stdout(Stdio::null())
120            .stderr(Stdio::null())
121            .status()
122            .is_ok_and(|s| s.success());
123
124        z3_available.then_some(Self {
125            cache: HashMap::new(),
126            queries: 0,
127            z3_calls: 0,
128            kind: Z3Kind::Cli,
129            cache_dir,
130        })
131    }
132
133    /// Pass a javascript callback `string->string` for smt2 input.
134    pub fn new_wasm(callback: Function) -> Self {
135        let wrapped_callback = move |s: &str| {
136            let res = callback
137                .call1(&JsValue::NULL, &JsValue::from_str(s))
138                .and_then(|val| val.as_string().ok_or("callback returned non-string".into()));
139            match res {
140                Ok(model) => model,
141                Err(_) => "".into(),
142            }
143        };
144
145        Self {
146            cache: HashMap::new(),
147            queries: 0,
148            z3_calls: 0,
149            kind: Z3Kind::Callback(Box::new(wrapped_callback)),
150            cache_dir: None,
151        }
152    }
153
154    /// Try to prove a `goal` given `facts` by linearizing the whole system.
155    pub fn entails_lia(
156        &mut self,
157        facts: &[Constraint],
158        goal: &Constraint,
159        nonneg: &[Var],
160        names: &[&str],
161    ) -> Verdict {
162        self.queries += 1;
163        let mut vars: Vec<Var> = Vec::new();
164        for c in facts.iter().chain(std::iter::once(goal)) {
165            for p in [&c.lhs, &c.rhs] {
166                debug_assert!(p.total_degree() <= 1);
167                for v in p.vars() {
168                    if !vars.contains(&v) {
169                        vars.push(v);
170                    }
171                }
172            }
173        }
174        vars.sort_unstable();
175
176        let script = build_script(facts, goal, nonneg, &vars);
177        let key = fnv1a_hash(script.as_bytes());
178        if let Some(v) = self.cache.get(&key) {
179            return v.clone();
180        }
181
182        self.z3_calls += 1;
183        let verdict = self.run_z3(&script, &vars);
184        if let Some(dir) = &self.cache_dir {
185            let _ = std::fs::create_dir_all(dir);
186            let mut file = String::new();
187            file.push_str(&format!("; obligation: {}\n", render(goal, names)));
188            for f in facts {
189                file.push_str(&format!("; given:      {}\n", render(f, names)));
190            }
191            file.push_str(&format!("; verdict:    {}\n", verdict_tag(&verdict)));
192            file.push_str(&script);
193            let _ = std::fs::write(dir.join(format!("{key:016x}.smt2")), file);
194        }
195        self.cache.insert(key, verdict.clone());
196        verdict
197    }
198
199    fn run_z3(&mut self, script: &str, vars: &[Var]) -> Verdict {
200        let stdout = match &mut self.kind {
201            Z3Kind::Cli => {
202                let Ok(mut child) = Command::new("z3")
203                    .arg("-smt2")
204                    .arg("-in")
205                    .stdin(Stdio::piped())
206                    .stdout(Stdio::piped())
207                    .stderr(Stdio::null())
208                    .spawn()
209                else {
210                    return Verdict::Unknown;
211                };
212                if let Some(stdin) = child.stdin.take() {
213                    let mut stdin = stdin;
214                    if stdin.write_all(script.as_bytes()).is_err() {
215                        let _ = child.kill();
216                        return Verdict::Unknown;
217                    }
218                }
219                let Ok(out) = child.wait_with_output() else {
220                    return Verdict::Unknown;
221                };
222                let stdout = String::from_utf8_lossy(&out.stdout);
223                stdout.to_string()
224            }
225            Z3Kind::Callback(callback) => callback(script),
226        };
227
228        let mut stdout_lines = stdout.lines();
229
230        match stdout_lines.next().map(str::trim) {
231            Some("unsat") => Verdict::Proved,
232            Some("sat") => {
233                let rest: String = stdout_lines.collect::<Vec<_>>().join(" ");
234                Verdict::RefutedAt(parse_model(&rest, vars))
235            }
236            _ => Verdict::Unknown,
237        }
238    }
239}
240
241fn verdict_tag(v: &Verdict) -> &'static str {
242    match v {
243        Verdict::Proved => "proved (unsat)",
244        Verdict::RefutedBy(_) => "refuted (sat)",
245        Verdict::RefutedAt { .. } => "refuted (sat)",
246        Verdict::Unknown => "unknown",
247    }
248}
249
250fn render(c: &Constraint, names: &[&str]) -> String {
251    format!(
252        "{} {} {}",
253        c.lhs.display_with(Some(names)),
254        c.cmp,
255        c.rhs.display_with(Some(names))
256    )
257}
258
259fn build_script(facts: &[Constraint], goal: &Constraint, nonneg: &[Var], vars: &[Var]) -> String {
260    let mut s = String::new();
261    s.push_str("(set-option :timeout 2000)\n");
262    s.push_str("(set-logic QF_LIA)\n");
263    for v in vars {
264        s.push_str(&format!("(declare-const v{} Int)\n", v));
265    }
266    for v in vars {
267        if nonneg.contains(v) {
268            s.push_str(&format!("(assert (>= v{} 0))\n", v));
269        }
270    }
271    for f in facts {
272        s.push_str(&format!("(assert {})\n", smt_cmp(f)));
273    }
274    s.push_str(&format!("(assert (not {}))\n", smt_cmp(goal)));
275    s.push_str("(check-sat)\n");
276    if !vars.is_empty() {
277        let list: Vec<String> = vars.iter().map(|v| format!("v{}", v)).collect();
278        s.push_str(&format!("(get-value ({}))\n", list.join(" ")));
279    }
280    s
281}
282
283fn smt_cmp(c: &Constraint) -> String {
284    format!(
285        "({} {} {})",
286        c.cmp.smt(),
287        smt_poly(&c.lhs),
288        smt_poly(&c.rhs)
289    )
290}
291
292fn smt_poly(p: &Poly) -> String {
293    if p.is_zero() {
294        return "0".to_string();
295    }
296    let terms: Vec<String> = p
297        .terms()
298        .iter()
299        .map(|(c, m)| {
300            let coef = smt_int(*c);
301            match m.exps() {
302                [] => coef,
303                [(v, 1)] => {
304                    if *c == 1.into() {
305                        format!("v{}", v)
306                    } else {
307                        format!("(* {coef} v{})", v)
308                    }
309                }
310                _ => unreachable!("not linearized"),
311            }
312        })
313        .collect();
314    if terms.len() == 1 {
315        terms.into_iter().next().unwrap()
316    } else {
317        format!("(+ {})", terms.join(" "))
318    }
319}
320
321fn smt_int(c: Coef) -> String {
322    if !c.is_zero() && !c.is_positive() {
323        format!("(- {:?})", c.abs())
324    } else {
325        format!("{:?}", c.abs())
326    }
327}
328
329fn parse_model(text: &str, vars: &[Var]) -> Vec<(Var, i128)> {
330    let mut model = Vec::new();
331    let cleaned = text.replace(['(', ')'], " ");
332    let tokens: Vec<&str> = cleaned.split_whitespace().collect();
333    let mut i = 0;
334    while i < tokens.len() {
335        if let Some(idx) = tokens[i]
336            .strip_prefix('v')
337            .and_then(|s| s.parse::<u32>().ok())
338        {
339            let (value, used) = match tokens.get(i + 1) {
340                Some(&"-") => (
341                    tokens
342                        .get(i + 2)
343                        .and_then(|t| t.parse::<i128>().ok())
344                        .map(|k| -k),
345                    3,
346                ),
347                Some(t) => (t.parse::<i128>().ok(), 2),
348                None => (None, 1),
349            };
350            if let Some(value) = value
351                && vars.contains(&(idx as Var))
352            {
353                model.push(((idx as Var), value));
354            }
355            i += used;
356        } else {
357            i += 1;
358        }
359    }
360    model
361}
362
363fn fnv1a_hash(bytes: &[u8]) -> u64 {
364    let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
365    for &b in bytes {
366        hash ^= u64::from(b);
367        hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
368    }
369    hash
370}
371
372#[cfg(test)]
373mod tests {
374    use crate::solver::poly::{Poly, poly};
375    use crate::solver::{Cmp, Constraint};
376
377    use super::*;
378
379    #[test]
380    fn model_script() {
381        let facts = [Constraint::new(poly!(n), Cmp::Eq, poly!(m))];
382        let goal = Constraint::new(poly!(m), Cmp::Eq, poly!(n));
383        let script = "(set-option :timeout 2000)
384(set-logic QF_LIA)
385(declare-const v109 Int)
386(declare-const v110 Int)
387(assert (= v110 v109))
388(assert (not (= v109 v110)))
389(check-sat)
390(get-value (v109 v110))
391";
392        assert_eq!(script, build_script(&facts, &goal, &[], &[109, 110]))
393    }
394
395    #[test]
396    fn entailment_with_context() {
397        let mut s = Z3::new_cli(None).unwrap();
398        // N >= 1 => 0 < N
399        let facts = [Constraint::new(poly!(n), Cmp::Ge, poly!(1))];
400        let goal = Constraint::new(poly!(0), Cmp::Lt, poly!(n));
401        assert_eq!(s.entails_lia(&facts, &goal, &[], &["N"]), Verdict::Proved);
402        // N >= 1 =/> 1 < N
403        let goal = Constraint::new(poly!(1), Cmp::Lt, poly!(n));
404        match s.entails_lia(&facts, &goal, &['n' as Var], &["N"]) {
405            Verdict::RefutedAt(point) => assert_eq!(point, vec![('n' as Var, 1)]),
406            v => panic!("expected refutation, got {v:?}"),
407        }
408    }
409
410    #[test]
411    fn integer_reasoning_not_real() {
412        let mut s = Z3::new_cli(None).unwrap();
413        // 2N >= 1 => N >= 1
414        let facts = [Constraint::new(poly!(2 * n), Cmp::Ge, poly!(1))];
415        let goal = Constraint::new(poly!(n), Cmp::Ge, poly!(1));
416        assert_eq!(s.entails_lia(&facts, &goal, &[], &["N"]), Verdict::Proved);
417    }
418
419    #[test]
420    fn commute() {
421        let mut s = Z3::new_cli(None).unwrap();
422        // N == M => M == N
423        let facts = [Constraint::new(poly!(n), Cmp::Eq, poly!(m))];
424        let goal = Constraint::new(poly!(m), Cmp::Eq, poly!(n));
425        assert_eq!(
426            s.entails_lia(&facts, &goal, &[], &["N", "M"]),
427            Verdict::Proved
428        );
429        // cached
430        assert_eq!(
431            s.entails_lia(&facts, &goal, &[], &["N", "M"]),
432            Verdict::Proved
433        );
434    }
435
436    #[test]
437    fn linearized() {
438        let vars = ["N".into(), "M".into()];
439        let nm = Poly::var(0u32, 1).mul(&Poly::var(1u32, 1));
440
441        let fact = Constraint {
442            lhs: nm.clone(),
443            cmp: Cmp::Ge,
444            rhs: poly!(4),
445        };
446        let goal = Constraint {
447            lhs: poly!(4),
448            cmp: Cmp::Le,
449            rhs: nm,
450        };
451
452        let lin = Linearized::new(&[fact], &goal, &vars, &|_| true);
453        assert!(!lin.pure_linear);
454        assert_eq!(lin.facts[0].lhs, lin.goal.rhs);
455        assert_eq!(lin.facts[0].lhs, Poly::var(2u32, 1));
456
457        assert!(lin.nonneg.contains(&2));
458        assert_eq!(lin.names[2], "N*M");
459    }
460}
sirius/solver/t1.rs

sirius/solver/
t1.rs
