flux_middle/rty/
subst.rs

1use std::{cmp::Ordering, marker::PhantomData};
2
3use flux_common::{bug, tracked_span_bug};
4use rustc_type_ir::DebruijnIndex;
5
6use self::fold::FallibleTypeFolder;
7use super::fold::{TypeFolder, TypeSuperFoldable};
8use crate::rty::*;
9
10/// Substitution for late bound variables
11pub(super) struct BoundVarReplacer<D> {
12    current_index: DebruijnIndex,
13    delegate: D,
14}
15
16pub trait BoundVarReplacerDelegate {
17    type Error = !;
18
19    fn replace_expr(&mut self, var: BoundReft) -> Result<Expr, Self::Error>;
20    fn replace_region(&mut self, br: BoundRegion) -> Result<Region, Self::Error>;
21}
22
23pub(crate) struct FnMutDelegate<F1, F2, E = !> {
24    pub exprs: F1,
25    pub regions: F2,
26    _error: PhantomData<E>,
27}
28
29impl<F1, F2, E> FnMutDelegate<F1, F2, E>
30where
31    F1: FnMut(BoundReft) -> Result<Expr, E>,
32    F2: FnMut(BoundRegion) -> Result<Region, E>,
33{
34    pub(crate) fn new(exprs: F1, regions: F2) -> Self {
35        Self { exprs, regions, _error: PhantomData }
36    }
37}
38
39impl<F1, F2, E> BoundVarReplacerDelegate for FnMutDelegate<F1, F2, E>
40where
41    F1: FnMut(BoundReft) -> Result<Expr, E>,
42    F2: FnMut(BoundRegion) -> Result<Region, E>,
43{
44    type Error = E;
45
46    fn replace_expr(&mut self, var: BoundReft) -> Result<Expr, E> {
47        (self.exprs)(var)
48    }
49
50    fn replace_region(&mut self, br: BoundRegion) -> Result<Region, E> {
51        (self.regions)(br)
52    }
53}
54
55impl<D> BoundVarReplacer<D> {
56    pub(super) fn new(delegate: D) -> BoundVarReplacer<D> {
57        BoundVarReplacer { delegate, current_index: INNERMOST }
58    }
59}
60
61impl<D> FallibleTypeFolder for BoundVarReplacer<D>
62where
63    D: BoundVarReplacerDelegate,
64{
65    type Error = D::Error;
66
67    fn try_enter_binder(&mut self, _: &BoundVariableKinds) {
68        self.current_index.shift_in(1);
69    }
70
71    fn try_exit_binder(&mut self) {
72        self.current_index.shift_out(1);
73    }
74
75    fn try_fold_expr(&mut self, e: &Expr) -> Result<Expr, D::Error> {
76        if let ExprKind::Var(Var::Bound(debruijn, breft)) = e.kind() {
77            match debruijn.cmp(&self.current_index) {
78                Ordering::Less => Ok(Expr::bvar(*debruijn, breft.var, breft.kind)),
79                Ordering::Equal => {
80                    Ok(self
81                        .delegate
82                        .replace_expr(*breft)?
83                        .shift_in_escaping(self.current_index.as_u32()))
84                }
85                Ordering::Greater => Ok(Expr::bvar(debruijn.shifted_out(1), breft.var, breft.kind)),
86            }
87        } else {
88            e.try_super_fold_with(self)
89        }
90    }
91
92    fn try_fold_region(&mut self, re: &Region) -> Result<Region, D::Error> {
93        if let ReBound(debruijn, br) = *re {
94            match debruijn.cmp(&self.current_index) {
95                Ordering::Less => Ok(*re),
96                Ordering::Equal => {
97                    let region = self.delegate.replace_region(br)?;
98                    if let ReBound(debruijn1, br) = region {
99                        // If the callback returns a late-bound region,
100                        // that region should always use the INNERMOST
101                        // debruijn index. Then we adjust it to the
102                        // correct depth.
103                        tracked_span_assert_eq!(debruijn1, INNERMOST);
104                        Ok(Region::ReBound(debruijn, br))
105                    } else {
106                        Ok(region)
107                    }
108                }
109                Ordering::Greater => Ok(ReBound(debruijn.shifted_out(1), br)),
110            }
111        } else {
112            Ok(*re)
113        }
114    }
115}
116
117/// Substitution for generics, i.e., early bound types, lifetimes, const generics, and refinements.
118/// Note that a substitution for refinement parameters (a list of expressions) must always be
119/// specified, while the behavior of other generics parameters (types, lifetimes and consts) can be
120/// configured with [`GenericsSubstDelegate`].
121pub struct GenericsSubstFolder<'a, D> {
122    current_index: DebruijnIndex,
123    delegate: D,
124    refinement_args: &'a [Expr],
125}
126
127pub trait GenericsSubstDelegate {
128    type Error = !;
129
130    fn sort_for_param(&mut self, param_ty: ParamTy) -> Result<Sort, Self::Error>;
131    fn ty_for_param(&mut self, param_ty: ParamTy) -> Result<Ty, Self::Error>;
132    fn ctor_for_param(&mut self, param_ty: ParamTy) -> Result<SubsetTyCtor, Self::Error>;
133    fn region_for_param(&mut self, ebr: EarlyParamRegion) -> Result<Region, Self::Error>;
134    fn expr_for_param_const(&self, param_const: ParamConst) -> Expr;
135    fn const_for_param(&mut self, param: &Const) -> Const;
136}
137
138/// A substitution with an explicit list of generic arguments.
139pub(crate) struct GenericArgsDelegate<'a, 'tcx>(
140    pub(crate) &'a [GenericArg],
141    pub(crate) TyCtxt<'tcx>,
142);
143
144impl GenericsSubstDelegate for GenericArgsDelegate<'_, '_> {
145    fn sort_for_param(&mut self, param_ty: ParamTy) -> Result<Sort, !> {
146        match self.0.get(param_ty.index as usize) {
147            Some(GenericArg::Base(ctor)) => Ok(ctor.sort()),
148            Some(arg) => {
149                tracked_span_bug!("expected base type for generic parameter, found `{arg:?}`")
150            }
151            None => tracked_span_bug!("type parameter out of range {param_ty:?}"),
152        }
153    }
154
155    fn ty_for_param(&mut self, param_ty: ParamTy) -> Result<Ty, !> {
156        match self.0.get(param_ty.index as usize) {
157            Some(GenericArg::Ty(ty)) => Ok(ty.clone()),
158            Some(arg) => tracked_span_bug!("expected type for generic parameter, found `{arg:?}`"),
159            None => tracked_span_bug!("type parameter out of range {param_ty:?}"),
160        }
161    }
162
163    fn ctor_for_param(&mut self, param_ty: ParamTy) -> Result<SubsetTyCtor, !> {
164        match self.0.get(param_ty.index as usize) {
165            Some(GenericArg::Base(ctor)) => Ok(ctor.clone()),
166            Some(arg) => {
167                tracked_span_bug!("expected base type for generic parameter, found `{arg:?}`")
168            }
169            None => tracked_span_bug!("type parameter out of range"),
170        }
171    }
172
173    fn region_for_param(&mut self, ebr: EarlyParamRegion) -> Result<Region, !> {
174        match self.0.get(ebr.index as usize) {
175            Some(GenericArg::Lifetime(re)) => Ok(*re),
176            Some(arg) => bug!("expected region for generic parameter, found `{arg:?}`"),
177            None => bug!("region parameter out of range"),
178        }
179    }
180
181    fn const_for_param(&mut self, param: &Const) -> Const {
182        if let ConstKind::Param(param_const) = &param.kind {
183            match self.0.get(param_const.index as usize) {
184                Some(GenericArg::Const(cst)) => cst.clone(),
185                Some(arg) => bug!("expected const for generic parameter, found `{arg:?}`"),
186                None => bug!("generic parameter out of range"),
187            }
188        } else {
189            param.clone()
190        }
191    }
192
193    fn expr_for_param_const(&self, param_const: ParamConst) -> Expr {
194        match self.0.get(param_const.index as usize) {
195            Some(GenericArg::Const(cst)) => Expr::from_const(self.1, cst),
196            Some(arg) => bug!("expected const for generic parameter, found `{arg:?}`"),
197            None => bug!("generic parameter out of range"),
198        }
199    }
200}
201
202/// A substitution meant to be used only for sorts. It'll panic if used on a type. This is used to
203/// break cycles during wf checking. During wf-checking we use [`rty::Sort`], but we can't yet
204/// generate (in general) an [`rty::GenericArg`] because conversion from [`fhir`] into [`rty`]
205/// requires the results of wf checking. Perhaps, we could also solve this problem by doing
206/// wf-checking with a different "IR" for sorts that sits in between [`fhir`] and [`rty`].
207///
208/// [`rty::Sort`]: crate::rty::Sort
209/// [`rty::GenericArg`]: crate::rty::GenericArg
210/// [`fhir`]: crate::fhir
211/// [`rty`]: crate::rty
212pub(crate) struct GenericsSubstForSort<F, E>
213where
214    F: FnMut(ParamTy) -> Result<Sort, E>,
215{
216    /// Implementation of [`GenericsSubstDelegate::sort_for_param`]
217    pub(crate) sort_for_param: F,
218}
219
220impl<F, E> GenericsSubstDelegate for GenericsSubstForSort<F, E>
221where
222    F: FnMut(ParamTy) -> Result<Sort, E>,
223{
224    type Error = E;
225
226    fn sort_for_param(&mut self, param_ty: ParamTy) -> Result<Sort, E> {
227        (self.sort_for_param)(param_ty)
228    }
229
230    fn ty_for_param(&mut self, param_ty: ParamTy) -> Result<Ty, E> {
231        bug!("unexpected type param {param_ty:?}");
232    }
233
234    fn ctor_for_param(&mut self, param_ty: ParamTy) -> Result<SubsetTyCtor, E> {
235        bug!("unexpected base type param {param_ty:?}");
236    }
237
238    fn region_for_param(&mut self, ebr: EarlyParamRegion) -> Result<Region, E> {
239        bug!("unexpected region param {ebr:?}");
240    }
241
242    fn const_for_param(&mut self, param: &Const) -> Const {
243        bug!("unexpected const param {param:?}");
244    }
245
246    fn expr_for_param_const(&self, param_const: ParamConst) -> Expr {
247        bug!("unexpected param_const {param_const:?}");
248    }
249}
250
251impl<'a, D> GenericsSubstFolder<'a, D> {
252    pub fn new(delegate: D, refine: &'a [Expr]) -> Self {
253        Self { current_index: INNERMOST, delegate, refinement_args: refine }
254    }
255}
256
257impl<D: GenericsSubstDelegate> FallibleTypeFolder for GenericsSubstFolder<'_, D> {
258    type Error = D::Error;
259
260    fn try_enter_binder(&mut self, _: &BoundVariableKinds) {
261        self.current_index.shift_in(1);
262    }
263
264    fn try_exit_binder(&mut self) {
265        self.current_index.shift_out(1);
266    }
267
268    fn try_fold_sort(&mut self, sort: &Sort) -> Result<Sort, D::Error> {
269        if let Sort::Param(param_ty) = sort {
270            self.delegate.sort_for_param(*param_ty)
271        } else {
272            sort.try_super_fold_with(self)
273        }
274    }
275
276    fn try_fold_ty(&mut self, ty: &Ty) -> Result<Ty, D::Error> {
277        match ty.kind() {
278            TyKind::Param(param_ty) => self.delegate.ty_for_param(*param_ty),
279            TyKind::Indexed(BaseTy::Param(param_ty), idx) => {
280                let idx = idx.try_fold_with(self)?;
281                Ok(self
282                    .delegate
283                    .ctor_for_param(*param_ty)?
284                    .replace_bound_reft(&idx)
285                    .to_ty())
286            }
287            _ => ty.try_super_fold_with(self),
288        }
289    }
290
291    fn try_fold_subset_ty(&mut self, sty: &SubsetTy) -> Result<SubsetTy, D::Error> {
292        if let BaseTy::Param(param_ty) = &sty.bty {
293            let idx = sty.idx.try_fold_with(self)?;
294            let pred = sty.pred.try_fold_with(self)?;
295            Ok(self
296                .delegate
297                .ctor_for_param(*param_ty)?
298                .replace_bound_reft(&idx)
299                .strengthen(pred))
300        } else {
301            sty.try_super_fold_with(self)
302        }
303    }
304
305    fn try_fold_region(&mut self, re: &Region) -> Result<Region, D::Error> {
306        if let ReEarlyParam(ebr) = *re { self.delegate.region_for_param(ebr) } else { Ok(*re) }
307    }
308
309    fn try_fold_expr(&mut self, expr: &Expr) -> Result<Expr, D::Error> {
310        match expr.kind() {
311            ExprKind::Var(Var::EarlyParam(var)) => Ok(self.expr_for_param(var.index)),
312            ExprKind::Var(Var::ConstGeneric(param_const)) => {
313                Ok(self.delegate.expr_for_param_const(*param_const))
314            }
315            _ => expr.try_super_fold_with(self),
316        }
317    }
318
319    fn try_fold_const(&mut self, c: &Const) -> Result<Const, D::Error> {
320        Ok(self.delegate.const_for_param(c))
321    }
322}
323
324impl<D> GenericsSubstFolder<'_, D> {
325    fn expr_for_param(&self, idx: u32) -> Expr {
326        self.refinement_args[idx as usize].shift_in_escaping(self.current_index.as_u32())
327    }
328}
329
330pub(crate) struct SortSubst<D> {
331    delegate: D,
332}
333
334impl<D> SortSubst<D> {
335    pub(crate) fn new(delegate: D) -> Self {
336        Self { delegate }
337    }
338}
339
340impl<D: SortSubstDelegate> TypeFolder for SortSubst<D> {
341    fn fold_sort(&mut self, sort: &Sort) -> Sort {
342        match sort {
343            Sort::Var(var) => self.delegate.sort_for_param(*var),
344            Sort::BitVec(BvSize::Param(var)) => Sort::BitVec(self.delegate.bv_size_for_param(*var)),
345            _ => sort.super_fold_with(self),
346        }
347    }
348}
349
350trait SortSubstDelegate {
351    fn sort_for_param(&self, var: ParamSort) -> Sort;
352    fn bv_size_for_param(&self, var: ParamSort) -> BvSize;
353}
354
355impl SortSubstDelegate for &[SortArg] {
356    fn sort_for_param(&self, var: ParamSort) -> Sort {
357        match &self[var.index()] {
358            SortArg::Sort(sort) => sort.clone(),
359            SortArg::BvSize(_) => tracked_span_bug!("unexpected bv size for sort param"),
360        }
361    }
362
363    fn bv_size_for_param(&self, var: ParamSort) -> BvSize {
364        match self[var.index()] {
365            SortArg::BvSize(size) => size,
366            SortArg::Sort(_) => tracked_span_bug!("unexpected sort for bv size param"),
367        }
368    }
369}
370
371impl SortSubstDelegate for &[Sort] {
372    fn sort_for_param(&self, var: ParamSort) -> Sort {
373        self[var.index()].clone()
374    }
375
376    fn bv_size_for_param(&self, _var: ParamSort) -> BvSize {
377        tracked_span_bug!("unexpected bv size parameter")
378    }
379}
flux_middle/rty/subst.rs

flux_middle/rty/
subst.rs
