flux_infer/
wkvars.rs

1use std::ops::ControlFlow;
2
3use flux_middle::rty::{
4    self,
5    fold::{
6        FallibleTypeFolder, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitable, TypeVisitor,
7    },
8};
9use itertools::Itertools;
10use rustc_data_structures::unord::{UnordMap, UnordSet};
11use rustc_type_ir::{DebruijnIndex, INNERMOST};
12
13pub struct WKVarInstantiator<'a> {
14    /// Map from the actuals passed to this Weak KVar to its params
15    ///
16    /// In theory this could be a Vec<rty::Expr>, but the instantiator
17    /// is configured right now to only return a single solution.
18    args_to_param: &'a UnordMap<rty::Expr, rty::Expr>,
19    /// Set of self args
20    self_args: &'a UnordSet<rty::Expr>,
21    /// Were any of the self args used in the expr?
22    any_self_args_used: bool,
23    /// In theory, this could (and probably should) map to multiple
24    /// solutions, i.e. a Vec<rty::Expr>.
25    memo: &'a mut UnordMap<rty::Expr, rty::Expr>,
26    current_index: DebruijnIndex,
27}
28
29impl FallibleTypeFolder for WKVarInstantiator<'_> {
30    /// We fail instantiation if we can't replace all free variables;
31    /// return the name of the first unreplaceable free variable found.
32    type Error = rty::Var;
33
34    fn try_enter_binder(&mut self, _vars: &rty::BoundVariableKinds) {
35        self.current_index.shift_in(1);
36    }
37
38    fn try_exit_binder(&mut self) {
39        self.current_index.shift_out(1);
40    }
41
42    fn try_fold_expr(&mut self, e: &rty::Expr) -> Result<rty::Expr, rty::Var> {
43        if let Some(instantiated_e) = self.memo.get(e) {
44            return Ok(instantiated_e.clone());
45        }
46
47        // NOTE: In theory there is a choice here: either we substitute the
48        // current expression for the parameter or we ignore it and continue
49        // going. We'll choose to be greedy and always substitute if possible,
50        // which I think will guarantee a solution if it exists, but I am not
51        // sure.
52        if let Some(param) = self.args_to_param.get(e) {
53            // Check if this expression is a self_arg (if we haven't seen a
54            // self_arg) yet.
55            self.any_self_args_used |= self.self_args.contains(e);
56            let param_expr = param.shift_in_escaping(self.current_index.as_u32());
57            self.memo.insert(e.clone(), param_expr.clone());
58            return Ok(param_expr);
59        }
60
61        if let rty::ExprKind::Var(var) = e.kind() {
62            // This is an escaping free var
63            Err(*var)
64        } else {
65            let instantiated_expr = e.try_super_fold_with(self)?;
66            self.memo.insert(e.clone(), instantiated_expr.clone());
67            Ok(instantiated_expr)
68        }
69    }
70}
71
72impl WKVarInstantiator<'_> {
73    /// If it succeeds: creates an expression that can replace the weak kvar,
74    /// which when used as a refinement will produce the original expr in this
75    /// branch after all substitutions have happened.
76    ///
77    /// It requires that the expression uses at least one of the first
78    /// `self_args` number of `wkvar_args`.
79    ///
80    /// There are a lot of patches we put in this algorithm to get around the fact
81    /// that it is purely syntactic. We currently try to eagerly eta reduce any
82    /// Ctor/Tuples + eta expand any args which aren't of that form
83    ///
84    /// FIXME(ck): This does not properly deal with expressions that have bound variables:
85    /// if the expression has a bound variable, we might fail the instantiation
86    /// when it should succeed.
87    pub fn try_instantiate_wkvar_args(
88        self_args: usize,
89        wkvar_args: &[rty::Expr],
90        expr: &rty::Expr,
91    ) -> Option<rty::Binder<rty::Expr>> {
92        // println!("trying to instantiate {:?} using args {:?}", expr, wkvar_args);
93        let expr_without_metadata = expr.erase_metadata();
94        let expr_eta_expanded_rel = expr_without_metadata.expand_bin_rels();
95        let mut args_to_param = UnordMap::default();
96        std::iter::zip(
97            // Eta reduce and erase metadata.
98            wkvar_args
99                .iter()
100                .map(|arg| arg.eta_reduce_projs().erase_metadata()),
101            // We'll make an instantiator that is generic because this instatiation
102            // may (and probably will) be used in multiple places
103            (0..wkvar_args.len()).map(|var_num| {
104                rty::Expr::bvar(INNERMOST, var_num.into(), rty::BoundReftKind::Anon)
105            }),
106        )
107        .for_each(|(arg, param)| {
108            ProductEtaExpander::eta_expand_products(param, arg, &mut args_to_param);
109        });
110        let self_args = wkvar_args[..self_args].iter().cloned().collect();
111        let mut instantiator = WKVarInstantiator {
112            args_to_param: &args_to_param,
113            self_args: &self_args,
114            any_self_args_used: false,
115            memo: &mut UnordMap::default(),
116            // This remains 0 because we use it to track how to shift our params,
117            // so the scope is the same.
118            current_index: INNERMOST,
119        };
120        // println!("eta expanded args: {:?}", args_to_param);
121        instantiator
122            .try_fold_expr(&expr_eta_expanded_rel)
123            // .map_err(|e| println!("Couldn't unify with {:?}", e))
124            .ok()
125            .and_then(|instantiated_e| {
126                if !instantiator.any_self_args_used && !instantiator.self_args.is_empty() {
127                    // println!("Dropping instantiation {:?} because no self args were used", instantiated_e);
128                    None
129                } else {
130                    Some(rty::Binder::bind_with_sorts(
131                        instantiated_e,
132                        &std::iter::repeat_n(rty::Sort::Err, wkvar_args.len()).collect_vec(),
133                    ))
134                }
135            })
136    }
137}
138
139pub struct WKVarSubst {
140    /// A map from wkvid to each of the substitutions we made for it
141    /// in whatever we folded over.
142    ///
143    /// In theory there should only ever be one value in this map.
144    pub subst_instantiations: UnordMap<rty::WKVid, Vec<rty::Expr>>,
145    pub wkvar_instantiations: UnordMap<rty::WKVid, rty::Binder<rty::Expr>>,
146    /// Keep wkvars after substituting
147    pub keep_wkvars: bool,
148}
149
150impl WKVarSubst {
151    pub fn new(
152        wkvar_instantiations: UnordMap<rty::WKVid, rty::Binder<rty::Expr>>,
153        keep_wkvars: bool,
154    ) -> Self {
155        WKVarSubst { subst_instantiations: Default::default(), wkvar_instantiations, keep_wkvars }
156    }
157}
158
159impl TypeFolder for WKVarSubst {
160    fn fold_expr(&mut self, e: &rty::Expr) -> rty::Expr {
161        match e.kind() {
162            rty::ExprKind::WKVar(wkvar @ rty::WKVar { wkvid, .. })
163                if let Some(bound_e) = self.wkvar_instantiations.get(wkvid) =>
164            {
165                // The bound expression has bound vars that need to be replaced
166                // by the args given to the wkvar (in order).
167                let instantiated_e = bound_e.replace_bound_refts(&wkvar.args);
168                self.subst_instantiations
169                    .entry(wkvid.clone())
170                    .and_modify(|insts| insts.push(instantiated_e.clone()))
171                    .or_insert_with(|| vec![instantiated_e.clone()]);
172                if self.keep_wkvars {
173                    rty::Expr::and(instantiated_e, e.clone())
174                } else {
175                    instantiated_e
176                }
177            }
178            // Replace wkvars with true (i.e. eliminate them) if we aren't keeping them
179            rty::ExprKind::WKVar(_) if !self.keep_wkvars => rty::Expr::tt(),
180            _ => e.super_fold_with(self),
181        }
182    }
183}
184
185struct ProductEtaExpander<'a> {
186    // An expression that evalutes to the current_expr
187    current_path: rty::Expr,
188    // Maps an interior part of the product to its eta expansion.
189    expr_to_eta_expansion: &'a mut UnordMap<rty::Expr, rty::Expr>,
190}
191
192impl TypeVisitor for ProductEtaExpander<'_> {
193    fn visit_expr(&mut self, expr: &rty::Expr) -> ControlFlow<Self::BreakTy> {
194        match expr.kind() {
195            rty::ExprKind::Tuple(subexprs) | rty::ExprKind::Ctor(_, subexprs) => {
196                let current_path = self.current_path.clone();
197                let mk_proj = |field| {
198                    if let rty::ExprKind::Ctor(ctor, _) = expr.kind() {
199                        let def_id = match ctor {
200                            rty::Ctor::Struct(def_id) | rty::Ctor::Enum(def_id, _) => *def_id,
201                            rty::Ctor::RawPtr => panic!("RawPtr ctor in wkvar path"),
202                        };
203                        rty::FieldProj::Adt { def_id, field }
204                    } else {
205                        rty::FieldProj::Tuple { arity: subexprs.len(), field }
206                    }
207                };
208                for (i, subexpr) in subexprs.iter().enumerate() {
209                    self.current_path = rty::Expr::field_proj(&current_path, mk_proj(i as u32));
210                    subexpr.visit_with(self)?;
211                }
212                ControlFlow::Continue(())
213            }
214            _ => {
215                // NOTE: in theory this should be appending to a vec, not
216                // clobbering whatever lives at expr currently. But we don't
217                // currently support making multiple solutions in the weak kvar
218                // instantiation so I'm not bothering.
219                self.expr_to_eta_expansion
220                    .insert(expr.clone(), self.current_path.clone());
221                ControlFlow::Continue(())
222            }
223        }
224    }
225}
226
227/// Recursively "unpacks" a product by eta-expanding each part of it.
228/// Maps each part of the product to its eta-expanded path.
229///
230/// e.g.
231///
232///     in = {
233///       current_path: self,
234///       expr: TwoFields { 0: (a0.0, a0.1), 1: 42 })
235///     }
236///
237///     out =
238///       a0.0 => self.0.0
239///       a0.1 => self.0.1
240///       42 => self.1
241impl<'a> ProductEtaExpander<'a> {
242    fn eta_expand_products(
243        current_path: rty::Expr,
244        expr: rty::Expr,
245        expr_to_eta_expansion: &'a mut UnordMap<rty::Expr, rty::Expr>,
246    ) {
247        let mut expander = ProductEtaExpander { current_path, expr_to_eta_expansion };
248        let _ = expr.visit_with(&mut expander);
249    }
250}
flux_infer/wkvars.rs

flux_infer/
wkvars.rs
