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    /// A pipe with nothing behind it: every query comes back [`Unknown`](Verdict::Unknown), so
134    /// only what the cheaper tiers settle gets settled.
135    pub fn new_offline() -> Self {
136        Self {
137            cache: HashMap::new(),
138            queries: 0,
139            z3_calls: 0,
140            kind: Z3Kind::Callback(Box::new(|_| String::new())),
141            cache_dir: None,
142        }
143    }
144
145    /// Pass a javascript callback `string->string` for smt2 input.
146    pub fn new_wasm(callback: Function) -> Self {
147        let wrapped_callback = move |s: &str| {
148            let res = callback
149                .call1(&JsValue::NULL, &JsValue::from_str(s))
150                .and_then(|val| val.as_string().ok_or("callback returned non-string".into()));
151            match res {
152                Ok(model) => model,
153                Err(_) => "".into(),
154            }
155        };
156
157        Self {
158            cache: HashMap::new(),
159            queries: 0,
160            z3_calls: 0,
161            kind: Z3Kind::Callback(Box::new(wrapped_callback)),
162            cache_dir: None,
163        }
164    }
165
166    /// Try to prove a `goal` given `facts` by linearizing the whole system.
167    pub fn entails_lia(
168        &mut self,
169        facts: &[Constraint],
170        goal: &Constraint,
171        nonneg: &[Var],
172        names: &[&str],
173    ) -> Verdict {
174        self.queries += 1;
175        let mut vars: Vec<Var> = Vec::new();
176        for c in facts.iter().chain(std::iter::once(goal)) {
177            for p in [&c.lhs, &c.rhs] {
178                debug_assert!(p.total_degree() <= 1);
179                for v in p.vars() {
180                    if !vars.contains(&v) {
181                        vars.push(v);
182                    }
183                }
184            }
185        }
186        vars.sort_unstable();
187
188        let script = build_script(facts, goal, nonneg, &vars);
189        let key = fnv1a_hash(script.as_bytes());
190        if let Some(v) = self.cache.get(&key) {
191            return v.clone();
192        }
193
194        self.z3_calls += 1;
195        let verdict = self.run_z3(&script, &vars);
196        if let Some(dir) = &self.cache_dir {
197            let _ = std::fs::create_dir_all(dir);
198            let mut file = String::new();
199            file.push_str(&format!("; obligation: {}\n", render(goal, names)));
200            for f in facts {
201                file.push_str(&format!("; given:      {}\n", render(f, names)));
202            }
203            file.push_str(&format!("; verdict:    {}\n", verdict_tag(&verdict)));
204            file.push_str(&script);
205            let _ = std::fs::write(dir.join(format!("{key:016x}.smt2")), file);
206        }
207        self.cache.insert(key, verdict.clone());
208        verdict
209    }
210
211    fn run_z3(&mut self, script: &str, vars: &[Var]) -> Verdict {
212        let stdout = match &mut self.kind {
213            Z3Kind::Cli => {
214                let Ok(mut child) = Command::new("z3")
215                    .arg("-smt2")
216                    .arg("-in")
217                    .stdin(Stdio::piped())
218                    .stdout(Stdio::piped())
219                    .stderr(Stdio::null())
220                    .spawn()
221                else {
222                    return Verdict::Unknown;
223                };
224                if let Some(stdin) = child.stdin.take() {
225                    let mut stdin = stdin;
226                    if stdin.write_all(script.as_bytes()).is_err() {
227                        let _ = child.kill();
228                        return Verdict::Unknown;
229                    }
230                }
231                let Ok(out) = child.wait_with_output() else {
232                    return Verdict::Unknown;
233                };
234                let stdout = String::from_utf8_lossy(&out.stdout);
235                stdout.to_string()
236            }
237            Z3Kind::Callback(callback) => callback(script),
238        };
239
240        let mut stdout_lines = stdout.lines();
241
242        match stdout_lines.next().map(str::trim) {
243            Some("unsat") => Verdict::Proved,
244            Some("sat") => {
245                let rest: String = stdout_lines.collect::<Vec<_>>().join(" ");
246                Verdict::RefutedAt(parse_model(&rest, vars))
247            }
248            _ => Verdict::Unknown,
249        }
250    }
251}
252
253fn verdict_tag(v: &Verdict) -> &'static str {
254    match v {
255        Verdict::Proved => "proved (unsat)",
256        Verdict::RefutedBy(_) => "refuted (sat)",
257        Verdict::RefutedAt { .. } => "refuted (sat)",
258        Verdict::Unknown => "unknown",
259    }
260}
261
262fn render(c: &Constraint, names: &[&str]) -> String {
263    format!(
264        "{} {} {}",
265        c.lhs.display_with(Some(names)),
266        c.cmp,
267        c.rhs.display_with(Some(names))
268    )
269}
270
271fn build_script(facts: &[Constraint], goal: &Constraint, nonneg: &[Var], vars: &[Var]) -> String {
272    let mut s = String::new();
273    s.push_str("(set-option :timeout 2000)\n");
274    s.push_str("(set-logic QF_LIA)\n");
275    for v in vars {
276        s.push_str(&format!("(declare-const v{} Int)\n", v));
277    }
278    for v in vars {
279        if nonneg.contains(v) {
280            s.push_str(&format!("(assert (>= v{} 0))\n", v));
281        }
282    }
283    for f in facts {
284        s.push_str(&format!("(assert {})\n", smt_cmp(f)));
285    }
286    s.push_str(&format!("(assert (not {}))\n", smt_cmp(goal)));
287    s.push_str("(check-sat)\n");
288    if !vars.is_empty() {
289        let list: Vec<String> = vars.iter().map(|v| format!("v{}", v)).collect();
290        s.push_str(&format!("(get-value ({}))\n", list.join(" ")));
291    }
292    s
293}
294
295fn smt_cmp(c: &Constraint) -> String {
296    format!(
297        "({} {} {})",
298        c.cmp.smt(),
299        smt_poly(&c.lhs),
300        smt_poly(&c.rhs)
301    )
302}
303
304fn smt_poly(p: &Poly) -> String {
305    if p.is_zero() {
306        return "0".to_string();
307    }
308    let terms: Vec<String> = p
309        .terms()
310        .iter()
311        .map(|(c, m)| {
312            let coef = smt_int(*c);
313            match m.exps() {
314                [] => coef,
315                [(v, 1)] => {
316                    if *c == 1.into() {
317                        format!("v{}", v)
318                    } else {
319                        format!("(* {coef} v{})", v)
320                    }
321                }
322                _ => unreachable!("not linearized"),
323            }
324        })
325        .collect();
326    if terms.len() == 1 {
327        terms.into_iter().next().unwrap()
328    } else {
329        format!("(+ {})", terms.join(" "))
330    }
331}
332
333fn smt_int(c: Coef) -> String {
334    if !c.is_zero() && !c.is_positive() {
335        format!("(- {:?})", c.abs())
336    } else {
337        format!("{:?}", c.abs())
338    }
339}
340
341fn parse_model(text: &str, vars: &[Var]) -> Vec<(Var, i128)> {
342    let mut model = Vec::new();
343    let cleaned = text.replace(['(', ')'], " ");
344    let tokens: Vec<&str> = cleaned.split_whitespace().collect();
345    let mut i = 0;
346    while i < tokens.len() {
347        if let Some(idx) = tokens[i]
348            .strip_prefix('v')
349            .and_then(|s| s.parse::<u32>().ok())
350        {
351            let (value, used) = match tokens.get(i + 1) {
352                Some(&"-") => (
353                    tokens
354                        .get(i + 2)
355                        .and_then(|t| t.parse::<i128>().ok())
356                        .map(|k| -k),
357                    3,
358                ),
359                Some(t) => (t.parse::<i128>().ok(), 2),
360                None => (None, 1),
361            };
362            if let Some(value) = value
363                && vars.contains(&(idx as Var))
364            {
365                model.push(((idx as Var), value));
366            }
367            i += used;
368        } else {
369            i += 1;
370        }
371    }
372    model
373}
374
375fn fnv1a_hash(bytes: &[u8]) -> u64 {
376    let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
377    for &b in bytes {
378        hash ^= u64::from(b);
379        hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
380    }
381    hash
382}
383
384#[cfg(test)]
385mod tests {
386    use crate::solver::poly::{Poly, poly};
387    use crate::solver::{Cmp, Constraint};
388
389    use super::*;
390
391    #[test]
392    fn model_script() {
393        let facts = [Constraint::new(poly!(n), Cmp::Eq, poly!(m))];
394        let goal = Constraint::new(poly!(m), Cmp::Eq, poly!(n));
395        let script = "(set-option :timeout 2000)
396(set-logic QF_LIA)
397(declare-const v109 Int)
398(declare-const v110 Int)
399(assert (= v110 v109))
400(assert (not (= v109 v110)))
401(check-sat)
402(get-value (v109 v110))
403";
404        assert_eq!(script, build_script(&facts, &goal, &[], &[109, 110]))
405    }
406
407    #[test]
408    fn entailment_with_context() {
409        let mut s = Z3::new_cli(None).unwrap();
410        // N >= 1 => 0 < N
411        let facts = [Constraint::new(poly!(n), Cmp::Ge, poly!(1))];
412        let goal = Constraint::new(poly!(0), Cmp::Lt, poly!(n));
413        assert_eq!(s.entails_lia(&facts, &goal, &[], &["N"]), Verdict::Proved);
414        // N >= 1 =/> 1 < N
415        let goal = Constraint::new(poly!(1), Cmp::Lt, poly!(n));
416        match s.entails_lia(&facts, &goal, &['n' as Var], &["N"]) {
417            Verdict::RefutedAt(point) => assert_eq!(point, vec![('n' as Var, 1)]),
418            v => panic!("expected refutation, got {v:?}"),
419        }
420    }
421
422    #[test]
423    fn integer_reasoning_not_real() {
424        let mut s = Z3::new_cli(None).unwrap();
425        // 2N >= 1 => N >= 1
426        let facts = [Constraint::new(poly!(2 * n), Cmp::Ge, poly!(1))];
427        let goal = Constraint::new(poly!(n), Cmp::Ge, poly!(1));
428        assert_eq!(s.entails_lia(&facts, &goal, &[], &["N"]), Verdict::Proved);
429    }
430
431    #[test]
432    fn commute() {
433        let mut s = Z3::new_cli(None).unwrap();
434        // N == M => M == N
435        let facts = [Constraint::new(poly!(n), Cmp::Eq, poly!(m))];
436        let goal = Constraint::new(poly!(m), Cmp::Eq, poly!(n));
437        assert_eq!(
438            s.entails_lia(&facts, &goal, &[], &["N", "M"]),
439            Verdict::Proved
440        );
441        // cached
442        assert_eq!(
443            s.entails_lia(&facts, &goal, &[], &["N", "M"]),
444            Verdict::Proved
445        );
446    }
447
448    #[test]
449    fn linearized() {
450        let vars = ["N".into(), "M".into()];
451        let nm = Poly::var(0u32, 1).mul(&Poly::var(1u32, 1));
452
453        let fact = Constraint {
454            lhs: nm.clone(),
455            cmp: Cmp::Ge,
456            rhs: poly!(4),
457        };
458        let goal = Constraint {
459            lhs: poly!(4),
460            cmp: Cmp::Le,
461            rhs: nm,
462        };
463
464        let lin = Linearized::new(&[&fact], &goal, &vars, &|_| true);
465        assert!(!lin.pure_linear);
466        assert_eq!(lin.facts[0].lhs, lin.goal.rhs);
467        assert_eq!(lin.facts[0].lhs, Poly::var(2u32, 1));
468
469        assert!(lin.nonneg.contains(&2));
470        assert_eq!(lin.names[2], "N*M");
471    }
472}
sirius/solver/t1.rs

sirius/solver/
t1.rs
