liquid_fixpoint/
suggestion_solver.rs

1//! Validity check + Quantifier Elimination. Used to make suggestions.
2
3use std::{
4    collections::{HashMap, HashSet},
5    fmt::Write as _,
6    io,
7};
8
9use crate::{
10    BinOp, BinRel, ConstDecl, Constant, DataDecl, Expr, FlatConstraint, FunDef, Identifier, Pred,
11    Sort, Types,
12    sexp::{Atom, Sexp},
13    smt_fmt::{SmtEncoder, SmtOptions},
14    z3_subprocess::{SatStatus, Z3Error, Z3Subprocess},
15};
16
17#[derive(Debug)]
18pub enum SuggestionSolverError {
19    Z3Process(io::Error),
20    QETimeout,
21    ContainsQuantifier,
22    NoResults,
23    FailedSanityCheck,
24}
25
26impl From<io::Error> for SuggestionSolverError {
27    fn from(err: io::Error) -> Self {
28        Self::Z3Process(err)
29    }
30}
31
32impl From<Z3Error> for SuggestionSolverError {
33    fn from(err: Z3Error) -> Self {
34        match err {
35            Z3Error::Io(err) => Self::Z3Process(err),
36            Z3Error::Timeout => Self::QETimeout,
37        }
38    }
39}
40
41pub struct SuggestionSolver<T: Types> {
42    session: Z3Subprocess,
43    env: ConstraintGlobalEnv<T>,
44}
45
46impl<T: Types> SuggestionSolver<T> {
47    pub fn new(
48        global_consts: &[ConstDecl<T>],
49        funs: &[FunDef<T>],
50        datatype_decls: Vec<DataDecl<T>>,
51    ) -> Result<Self, SuggestionSolverError> {
52        let datatype_decls =
53            crate::constraint_with_env::topo_sort_data_declarations(datatype_decls);
54        let (env, declarations) = ConstraintGlobalEnv::new(&datatype_decls, global_consts, funs);
55        let mut session = Z3Subprocess::new()?;
56        session.reset(&declarations)?;
57        Ok(Self { session, env })
58    }
59
60    pub(crate) fn check_validity(
61        &mut self,
62        constraint: &FlatConstraint<T>,
63    ) -> Result<bool, SuggestionSolverError> {
64        let locals = ConstraintLocalEnv::new(&self.env, &[], &constraint.binders);
65        let mut assertions = constraint
66            .preconditions()
67            .iter()
68            .map(|pred| self.env.pred(pred))
69            .collect::<Vec<_>>();
70        assertions.push(format!("(not {})", self.env.pred(&constraint.head)));
71        let mut setup = locals.declarations;
72        for assertion in &assertions {
73            writeln!(setup, "(assert {assertion})").unwrap();
74        }
75        Ok(matches!(self.session.check_sat(&setup)?, SatStatus::Unsat))
76    }
77
78    pub(crate) fn qe_and_simplify(
79        &mut self,
80        constraint: &FlatConstraint<T>,
81        binder_consts: &[ConstDecl<T>],
82    ) -> Result<Expr<T>, SuggestionSolverError> {
83        let locals = ConstraintLocalEnv::new(&self.env, binder_consts, &constraint.binders);
84        let assumptions = constraint
85            .preconditions()
86            .iter()
87            .map(|pred| self.env.pred(pred))
88            .collect::<Vec<_>>();
89        let body =
90            format!("(=> {} {})", conjunction(&assumptions), self.env.pred(&constraint.head));
91        let implication = self.env.quantified_implication(constraint, &body);
92        self.session.scoped(&locals.declarations, |session| {
93            Self::qe_and_simplify_in_context(
94                session,
95                &self.env,
96                &locals,
97                &assumptions,
98                &body,
99                &implication,
100            )
101        })
102    }
103
104    fn qe_and_simplify_in_context(
105        session: &mut Z3Subprocess,
106        env: &ConstraintGlobalEnv<T>,
107        locals: &ConstraintLocalEnv<'_, T>,
108        assumptions: &[String],
109        body: &str,
110        implication: &str,
111    ) -> Result<Expr<T>, SuggestionSolverError> {
112        let goals =
113            session.apply(&format!("(assert {implication})"), "(try-for (then qe nnf) 10000)")?;
114        let goal = goals.last().ok_or(SuggestionSolverError::NoResults)?;
115        let mut candidate = goal
116            .iter()
117            .map(|formula| locals.decode(formula))
118            .collect::<Result<Vec<_>, _>>()?
119            .into_iter()
120            .min_by_key(Expr::total_num_disjuncts)
121            .ok_or(SuggestionSolverError::NoResults)?;
122
123        let mut setup = String::new();
124        for assumption in assumptions {
125            writeln!(setup, "(assert {assumption})").unwrap();
126        }
127        if session.scoped(&setup, |session| prune_vacuous(&mut candidate, env, session))? {
128            return Ok(Expr::FALSE);
129        }
130
131        let sanity = format!("(not (=> {} {body}))", env.expr(&candidate));
132        if matches!(session.check_sat(&format!("(assert {sanity})"))?, SatStatus::Unsat) {
133            Ok(candidate)
134        } else {
135            Err(SuggestionSolverError::FailedSanityCheck)
136        }
137    }
138}
139
140struct ConstraintGlobalEnv<T: Types> {
141    symbols: HashMap<String, T::Var>,
142    declared_names: HashSet<String>,
143    encoder: SmtEncoder<T>,
144}
145
146impl<T: Types> ConstraintGlobalEnv<T> {
147    fn new(
148        datatype_decls: &[DataDecl<T>],
149        global_consts: &[ConstDecl<T>],
150        funs: &[FunDef<T>],
151    ) -> (Self, String) {
152        let mut symbols = HashMap::new();
153        for decl in datatype_decls {
154            for ctor in &decl.ctors {
155                let name = ctor.name.display().to_string();
156                symbols.insert(name, ctor.name.clone());
157                for field in &ctor.fields {
158                    symbols.insert(field.name.display().to_string(), field.name.clone());
159                }
160            }
161        }
162        for decl in global_consts {
163            symbols.insert(decl.name.display().to_string(), decl.name.clone());
164        }
165        for fun in funs {
166            symbols.insert(fun.name.display().to_string(), fun.name.clone());
167        }
168        let encoder = SmtEncoder::new(datatype_decls, SmtOptions::SUGGESTIONS);
169        let mut out = String::new();
170        for decl in datatype_decls {
171            encoder.write_datatype_decl(&mut out, decl).unwrap();
172        }
173
174        let mut declared = HashSet::new();
175        for decl in global_consts {
176            let name = decl.name.display().to_string();
177            if declared.insert(name.clone()) {
178                encoder
179                    .write_const_decl(&mut out, &name, &decl.sort)
180                    .unwrap();
181            }
182        }
183        for fun in funs {
184            let name = fun.name.display().to_string();
185            if declared.insert(name.clone()) {
186                encoder.write_fun_def(&mut out, fun).unwrap();
187            }
188        }
189        (Self { symbols, declared_names: declared, encoder }, out)
190    }
191
192    fn pred(&self, pred: &Pred<T>) -> String {
193        self.encoder.pred(pred).unwrap()
194    }
195
196    fn expr(&self, expr: &Expr<T>) -> String {
197        self.encoder.expr(expr).unwrap()
198    }
199
200    fn quantified_implication(&self, cstr: &FlatConstraint<T>, body: &str) -> String {
201        let mut quantified = body.to_string();
202        for (name, sort) in &cstr.binders {
203            let mut binder_sort = String::new();
204            self.encoder.write_sort(&mut binder_sort, sort).unwrap();
205            quantified = format!("(forall (({} {binder_sort})) {quantified})", name.display());
206        }
207        quantified
208    }
209}
210
211/// Declarations active in a single constraint's solver scope (as a string to
212/// send to Z3), and maps to decode its response back to T.
213struct ConstraintLocalEnv<'a, T: Types> {
214    declarations: String,
215    symbols: HashMap<String, T::Var>,
216    global_symbols: &'a HashMap<String, T::Var>,
217}
218
219impl<'a, T: Types> ConstraintLocalEnv<'a, T> {
220    fn new(
221        globals: &'a ConstraintGlobalEnv<T>,
222        binder_consts: &[ConstDecl<T>],
223        binders: &[(T::Var, Sort<T>)],
224    ) -> Self {
225        let mut declarations = String::new();
226        let mut declared = HashSet::new();
227        let mut symbols = HashMap::new();
228        for (var, sort) in binder_consts
229            .iter()
230            .map(|decl| (&decl.name, &decl.sort))
231            .chain(binders.iter().map(|(var, sort)| (var, sort)))
232        {
233            let name = var.display().to_string();
234            if !globals.declared_names.contains(&name) && declared.insert(name.clone()) {
235                globals
236                    .encoder
237                    .write_const_decl(&mut declarations, &name, sort)
238                    .unwrap();
239            }
240            if !globals.symbols.contains_key(&name) {
241                symbols.insert(name, var.clone());
242            }
243        }
244        for (var, _) in binders {
245            symbols.insert(var.display().to_string(), var.clone());
246        }
247        Self { declarations, symbols, global_symbols: &globals.symbols }
248    }
249
250    fn symbol(&self, name: &str) -> Option<&T::Var> {
251        self.symbols
252            .get(name)
253            .or_else(|| self.global_symbols.get(name))
254    }
255
256    fn decode(&self, sexp: &Sexp) -> Result<Expr<T>, SuggestionSolverError> {
257        match sexp {
258            Sexp::Atom(Atom::B(value)) => Ok(Expr::Constant(Constant::Boolean(*value))),
259            Sexp::Atom(Atom::I(value)) if *value >= 0 => {
260                Ok(Expr::Constant(Constant::Numeral(*value as u128)))
261            }
262            Sexp::Atom(Atom::I(value)) => {
263                Ok(Expr::Neg(Box::new(Expr::Constant(Constant::Numeral(value.unsigned_abs())))))
264            }
265            Sexp::Atom(Atom::S(name)) => self.decode_symbol(name),
266            Sexp::Atom(atom) => panic!("unsupported atom in Z3 output: `{atom:?}`"),
267            Sexp::List(items) => self.decode_application(items),
268        }
269    }
270
271    fn decode_symbol(&self, name: &str) -> Result<Expr<T>, SuggestionSolverError> {
272        match name {
273            "true" => Ok(Expr::TRUE),
274            "false" => Ok(Expr::FALSE),
275            _ => {
276                Ok(Expr::Var(
277                    self.symbol(name)
278                        .cloned()
279                        .unwrap_or_else(|| panic!("unknown symbol in Z3 output: `{name}`")),
280                ))
281            }
282        }
283    }
284
285    fn decode_application(&self, items: &[Sexp]) -> Result<Expr<T>, SuggestionSolverError> {
286        let Some((head, args)) = items.split_first() else {
287            panic!("empty application in Z3 output");
288        };
289        let head = application_head(head);
290        if matches!(head, "forall" | "exists") {
291            return Err(SuggestionSolverError::ContainsQuantifier);
292        }
293        let decoded = || {
294            args.iter()
295                .map(|arg| self.decode(arg))
296                .collect::<Result<Vec<_>, _>>()
297        };
298        match head {
299            "let" => {
300                // FIXME: substitution below is not scope-aware. For example,
301                // (let ((x 1)) (let ((x 2)) x)) should decode to 2, but replacing
302                // the outer x also rewrites the inner binding to (let ((1 2)) 1).
303                // A fix must also preserve simultaneous let bindings: in
304                // (let ((x 1)) (let ((x 2) (y x)) y)), y refers to the outer x,
305                // so the result is 1, not 2. Decode all RHSs in the outer scope
306                // before extending the environment for the body.
307                let args = exact_arg(head, args, 2);
308                let Sexp::List(bindings) = &args[0] else {
309                    panic!("malformed let bindings in Z3 output: {:?}", args[0]);
310                };
311                let mut substitutions = HashMap::new();
312                for binding in bindings {
313                    let Sexp::List(binding) = binding else {
314                        panic!("malformed let binding in Z3 output: {binding:?}");
315                    };
316                    let [Sexp::Atom(Atom::S(name)), value] = binding.as_slice() else {
317                        panic!("malformed let binding in Z3 output: {binding:?}");
318                    };
319                    substitutions.insert(name.as_str(), value);
320                }
321                self.decode(&substitute_sexp(&args[1], &substitutions))
322            }
323            "and" => Ok(Expr::And(decoded()?)),
324            "or" => Ok(Expr::Or(decoded()?)),
325            "not" => Ok(Expr::Not(Box::new(self.decode(&exact_arg(head, args, 1)[0])?))),
326            "=>" => {
327                let args = exact_arg(head, args, 2);
328                Ok(Expr::Imp(Box::new([self.decode(&args[0])?, self.decode(&args[1])?])))
329            }
330            "=" => {
331                let args = exact_arg(head, args, 2);
332                Ok(Expr::Atom(
333                    BinRel::Eq,
334                    Box::new([self.decode(&args[0])?, self.decode(&args[1])?]),
335                ))
336            }
337            "<" | "<=" | ">" | ">=" => {
338                let rel = match head {
339                    "<" => BinRel::Lt,
340                    "<=" => BinRel::Le,
341                    ">" => BinRel::Gt,
342                    ">=" => BinRel::Ge,
343                    _ => unreachable!(),
344                };
345                let args = exact_arg(head, args, 2);
346                Ok(Expr::Atom(rel, Box::new([self.decode(&args[0])?, self.decode(&args[1])?])))
347            }
348            "ite" => {
349                let args = exact_arg(head, args, 3);
350                Ok(Expr::IfThenElse(Box::new([
351                    self.decode(&args[0])?,
352                    self.decode(&args[1])?,
353                    self.decode(&args[2])?,
354                ])))
355            }
356            "-" if args.len() == 1 => Ok(Expr::Neg(Box::new(self.decode(&args[0])?))),
357            "+" | "-" | "*" | "div" | "mod" => {
358                if args.len() < 2 {
359                    panic!(
360                        "expected at least two arguments for `{head}` in Z3 output, got {}",
361                        args.len()
362                    );
363                }
364                let op = match head {
365                    "+" => BinOp::Add,
366                    "-" => BinOp::Sub,
367                    "*" => BinOp::Mul,
368                    "div" => BinOp::Div,
369                    "mod" => BinOp::Mod,
370                    _ => unreachable!(),
371                };
372                let mut args = decoded()?.into_iter();
373                let lhs = args
374                    .next()
375                    .expect("binary operator has at least two arguments");
376                Ok(args.fold(lhs, |lhs, rhs| Expr::BinaryOp(op, Box::new([lhs, rhs]))))
377            }
378            name => {
379                let var = self
380                    .symbol(name)
381                    .unwrap_or_else(|| panic!("unknown application in Z3 output: `{name}`"));
382                Ok(Expr::App(Box::new(Expr::Var(var.clone())), None, decoded()?, None))
383            }
384        }
385    }
386}
387
388fn substitute_sexp(sexp: &Sexp, substitutions: &HashMap<&str, &Sexp>) -> Sexp {
389    match sexp {
390        Sexp::Atom(Atom::S(name)) => {
391            substitutions
392                .get(name.as_str())
393                .copied()
394                .unwrap_or(sexp)
395                .clone()
396        }
397        Sexp::Atom(_) => sexp.clone(),
398        Sexp::List(items) => {
399            Sexp::List(
400                items
401                    .iter()
402                    .map(|item| substitute_sexp(item, substitutions))
403                    .collect(),
404            )
405        }
406    }
407}
408
409fn prune_vacuous<T: Types>(
410    expr: &mut Expr<T>,
411    env: &ConstraintGlobalEnv<T>,
412    session: &mut Z3Subprocess,
413) -> Result<bool, SuggestionSolverError> {
414    match expr {
415        Expr::And(conjuncts) => {
416            let copy = conjuncts.clone();
417            for (i, conjunct) in conjuncts.iter_mut().enumerate() {
418                let mut setup = String::new();
419                for (j, sibling) in copy.iter().enumerate() {
420                    if i != j {
421                        writeln!(setup, "(assert {})", env.expr(sibling)).unwrap();
422                    }
423                }
424                if session.scoped(&setup, |session| prune_vacuous(conjunct, env, session))? {
425                    *expr = Expr::FALSE;
426                    return Ok(true);
427                }
428            }
429            Ok(false)
430        }
431        Expr::Or(disjuncts) => {
432            let mut kept = Vec::with_capacity(disjuncts.len());
433            for mut disjunct in std::mem::take(disjuncts) {
434                if !prune_vacuous(&mut disjunct, env, session)? {
435                    kept.push(disjunct);
436                }
437            }
438            match kept.len() {
439                0 => {
440                    *expr = Expr::FALSE;
441                    Ok(true)
442                }
443                1 => {
444                    *expr = kept.pop().unwrap();
445                    Ok(false)
446                }
447                _ => {
448                    *disjuncts = kept;
449                    Ok(false)
450                }
451            }
452        }
453        _ => {
454            Ok(matches!(
455                session.check_sat(&format!("(assert {})", env.expr(expr)))?,
456                SatStatus::Unsat
457            ))
458        }
459    }
460}
461
462fn application_head(head: &Sexp) -> &str {
463    if let Sexp::Atom(Atom::S(name)) = head {
464        return name;
465    }
466    if let Sexp::List(items) = head
467        && atom_symbol(items.first()) == Some("as")
468        && let Some(Sexp::Atom(Atom::S(name))) = items.get(1)
469    {
470        return name;
471    }
472    panic!("unsupported application head in Z3 output: `{head:?}`")
473}
474
475fn atom_symbol(sexp: Option<&Sexp>) -> Option<&str> {
476    match sexp {
477        Some(Sexp::Atom(Atom::S(symbol))) => Some(symbol),
478        _ => None,
479    }
480}
481
482fn exact_arg<'a>(operator: &str, args: &'a [Sexp], expected: usize) -> &'a [Sexp] {
483    if args.len() == expected {
484        args
485    } else {
486        panic!(
487            "unexpected argument count for `{operator}` in Z3 output: expected {expected}, got {}",
488            args.len()
489        )
490    }
491}
492
493fn conjunction(assertions: &[String]) -> String {
494    match assertions {
495        [] => "true".to_string(),
496        [only] => only.clone(),
497        _ => format!("(and {})", assertions.join(" ")),
498    }
499}
liquid_fixpoint/suggestion_solver.rs

liquid_fixpoint/
suggestion_solver.rs
