flux_infer/
suggestions.rs

1use std::collections::HashSet;
2
3use flux_middle::rty;
4use itertools::Itertools;
5use liquid_fixpoint::{
6    SuggestionSolver,
7    SuggestionSolverError::{self, Z3Process},
8    check_validity, qe_and_simplify,
9};
10use rustc_data_structures::fx::FxIndexMap;
11
12use crate::{
13    fixpoint_encoding::{
14        ConstKey, FixpointCtxt, FixpointSolution, PossibleSolutions, SuggestionCtxt, TagIdx,
15        fixpoint,
16    },
17    wkvars::WKVarInstantiator,
18};
19
20pub type TagToFlatConstraint = FxIndexMap<TagIdx, fixpoint::FlatConstraint>;
21
22pub(crate) fn make_flat_constraint_map(constraint: &fixpoint::Constraint) -> TagToFlatConstraint {
23    constraint
24        .flatten(|var| matches!(var, fixpoint::Var::Underscore))
25        .into_iter()
26        .flat_map(|flat_constraint| {
27            // We can't send a kvar to the SMT. If there's a kvar on the LHS we
28            // can underapproximate it with TRUE, but if it's in head position
29            // we don't know what to do.
30            if let Some(tag) = flat_constraint.tag
31                && !matches!(flat_constraint.head, fixpoint::Pred::KVar(..))
32            {
33                Some((tag, flat_constraint))
34            } else {
35                None
36            }
37        })
38        .collect()
39}
40
41pub(crate) fn subst_fixpoint_solutions(
42    flat_constraint: &mut fixpoint::FlatConstraint,
43    fixpoint_solution: &FxIndexMap<rty::KVid, FixpointSolution>,
44) {
45    flat_constraint.assumptions = flat_constraint
46        .assumptions
47        .iter()
48        .flat_map(|pred| {
49            // Remove all trivially true assumptions
50            if pred.is_trivially_true() {
51                vec![].into_iter()
52            // Substitute the kvar solutions in
53            } else if let fixpoint::Pred::KVar(kvid, args) = pred
54                && let Some((var_sorts, solution)) = &fixpoint_solution.get(kvid)
55            {
56                assert!(var_sorts.len() == args.len());
57                let subst = args
58                    .iter()
59                    .zip(var_sorts.iter())
60                    .map(|(arg, (subst_var, _))| (*subst_var, arg.clone()))
61                    .collect();
62                let subst_solution = solution.substitute(&subst);
63                subst_solution
64                    .as_conjunction()
65                    .into_iter()
66                    .map(fixpoint::Pred::Expr)
67                    .collect_vec()
68                    .into_iter()
69            } else {
70                vec![pred.clone()].into_iter()
71            }
72        })
73        .collect();
74}
75
76pub(crate) fn find_possible_solutions<'genv, 'tcx, Tag>(
77    fxctx: &mut FixpointCtxt<'genv, 'tcx, Tag>,
78    tag_idx: TagIdx,
79    suggestion_ctx: &SuggestionCtxt,
80    solver: &mut Option<SuggestionSolver<fixpoint::FixpointTypes>>,
81) -> Result<PossibleSolutions, liquid_fixpoint::SuggestionSolverError>
82where
83    Tag: std::hash::Hash + Eq + Copy,
84{
85    let Some(flat_constraint) = suggestion_ctx.flat_constraints.get(&tag_idx) else {
86        return Ok(Default::default());
87    };
88    let head_expr = match &flat_constraint.head {
89        fixpoint::Pred::Expr(e) => Some(e.clone()),
90        _ => None,
91    };
92    let mut possible_solutions: PossibleSolutions = Default::default();
93    if solver.is_none() {
94        *solver = Some(SuggestionSolver::new(
95            &suggestion_ctx.const_decls,
96            &suggestion_ctx.fun_defs,
97            suggestion_ctx.data_decls.clone(),
98        )?);
99    }
100    let solver = solver.as_mut().unwrap();
101    let wkvars_and_constraints = flat_constraint.wkvars_and_constrs();
102    for (wkvar, flat_constraint, other_constrs) in wkvars_and_constraints {
103        let mut valid = true;
104        for other_constr in other_constrs {
105            let validity = match check_validity(solver, &other_constr) {
106                Ok(b) => b,
107                Err(Z3Process(e)) => panic!("Z3 process returned an error {}", e),
108                Err(_) => false,
109            };
110            if !validity {
111                valid = false;
112                break;
113            }
114        }
115        if !valid {
116            continue;
117        }
118        let ConstKey::WKVar(wkvid, self_args) = fxctx
119            .ecx
120            .const_env
121            .wkvar_map_rev
122            .get(&wkvar.wkvid)
123            .cloned()
124            .unwrap()
125        else {
126            panic!()
127        };
128        let fvars: HashSet<fixpoint::Var> = wkvar
129            .args
130            .iter()
131            .flat_map(|arg| {
132                arg.free_vars().into_iter().filter(|fvar| {
133                    matches!(fvar, fixpoint::Var::Local(_) | fixpoint::Var::Param(_))
134                })
135            })
136            .collect();
137        let rty_args: Vec<rty::Expr> = wkvar
138            .args
139            .iter()
140            .map(|arg| fxctx.fixpoint_to_expr(arg))
141            .try_collect()
142            .unwrap();
143        let (binder_consts, mut new_flat_constraint) = flat_constraint.remove_binders(&fvars);
144        // Remove any wkvars and drop assumptions that are just wkvars or true
145        new_flat_constraint.assumptions = new_flat_constraint
146            .assumptions
147            .into_iter()
148            .filter_map(|assumption| {
149                let assumption = assumption.strip_wkvars();
150                if !assumption.is_trivially_true() { Some(assumption) } else { None }
151            })
152            .collect();
153        let result = qe_and_simplify(solver, &new_flat_constraint, &binder_consts);
154        let fallback = |fxctx: &mut FixpointCtxt<'genv, 'tcx, Tag>| {
155            head_expr
156                .as_ref()
157                .and_then(|head| fxctx.fixpoint_to_expr(head).ok())
158                .and_then(|head| {
159                    WKVarInstantiator::try_instantiate_wkvar_args(self_args, &rty_args, &head)
160                })
161        };
162        let solution = match result {
163            Ok(fe) => {
164                match fxctx.fixpoint_to_expr(&fe) {
165                    Ok(expr) if expr.is_trivially_false() || expr.is_trivially_true() => None,
166                    Ok(_) if fe.total_num_disjuncts() > 3 => fallback(fxctx),
167                    Ok(expr) => {
168                        WKVarInstantiator::try_instantiate_wkvar_args(self_args, &rty_args, &expr)
169                            .or_else(|| fallback(fxctx))
170                    }
171                    Err(_) => None,
172                }
173            }
174            Err(SuggestionSolverError::Z3Process(e)) => panic!("Z3 returned an error: {}", e),
175            Err(SuggestionSolverError::FailedSanityCheck) => panic!("Failed sanity check!"),
176            Err(_) => fallback(fxctx),
177        };
178        if let Some(solution) = solution {
179            possible_solutions
180                .entry(wkvid.clone())
181                .or_default()
182                .push(solution);
183        }
184    }
185    Ok(possible_solutions)
186}
flux_infer/suggestions.rs

flux_infer/
suggestions.rs
