charon_lib/ast/type_level/
substitute.rs

1use crate::ast::*;
2use derive_generic_visitor::*;
3use std::borrow::Cow;
4use std::convert::Infallible;
5use std::fmt::Debug;
6use std::iter::Iterator;
7
8/// Visitor for type-level variables. Used to visit the variables contained in a value, as seen
9/// from the outside of the value. This means that any variable bound inside the value will be
10/// skipped, and all the seen De Bruijn indices will count from the outside of the value. The
11/// returned value, if any, will be put in place of the variable.
12pub trait VarsVisitor {
13    fn visit_erased_region(&mut self) -> Option<Region> {
14        None
15    }
16    fn visit_region_var(&mut self, _v: RegionDbVar) -> Option<Region> {
17        None
18    }
19    fn visit_type_var(&mut self, _v: TypeDbVar) -> Option<Ty> {
20        None
21    }
22    fn visit_const_generic_var(&mut self, _v: ConstGenericDbVar) -> Option<ConstantExprKind> {
23        None
24    }
25    fn visit_clause_var(&mut self, _v: ClauseDbVar) -> Option<TraitRefKind> {
26        None
27    }
28    fn visit_self_clause(&mut self) -> Option<TraitRefKind> {
29        None
30    }
31}
32
33/// Visitor for the [TyVisitable::substitute] function.
34/// This substitutes variables bound at the level where we start to substitute (level 0).
35#[derive(Visitor)]
36pub(crate) struct SubstVisitor<'a> {
37    generics: &'a GenericArgs,
38    self_ref: Option<&'a TraitRefKind>,
39    /// Whether to substitute explicit variables only (types, regions, const generics).
40    explicits_only: bool,
41    had_error: bool,
42}
43impl<'a> SubstVisitor<'a> {
44    pub(crate) fn new(
45        generics: &'a GenericArgs,
46        self_ref: Option<&'a TraitRefKind>,
47        explicits_only: bool,
48    ) -> Self {
49        Self {
50            generics,
51            self_ref,
52            explicits_only,
53            had_error: false,
54        }
55    }
56
57    pub fn visit<T: TyVisitable>(mut self, mut x: T) -> Result<T, GenericsMismatch> {
58        x.visit_vars(&mut self);
59        if self.had_error {
60            Err(GenericsMismatch)
61        } else {
62            Ok(x)
63        }
64    }
65
66    /// Returns the value for this variable, if any.
67    fn process_var<Id, T>(
68        &mut self,
69        var: DeBruijnVar<Id>,
70        get: impl Fn(Id) -> Option<&'a T>,
71    ) -> Option<T>
72    where
73        Id: Copy,
74        T: Clone + TyVisitable,
75        DeBruijnVar<Id>: Into<T>,
76    {
77        match var {
78            DeBruijnVar::Bound(dbid, varid) => {
79                Some(if let Some(dbid) = dbid.sub(DeBruijnId::one()) {
80                    // This is bound outside the binder we're substituting for.
81                    DeBruijnVar::Bound(dbid, varid).into()
82                } else {
83                    match get(varid) {
84                        Some(v) => v.clone(),
85                        None => {
86                            self.had_error = true;
87                            return None;
88                        }
89                    }
90                })
91            }
92            DeBruijnVar::Free(..) => None,
93        }
94    }
95}
96impl VarsVisitor for SubstVisitor<'_> {
97    fn visit_region_var(&mut self, v: RegionDbVar) -> Option<Region> {
98        self.process_var(v, |id| self.generics.regions.get(id))
99    }
100    fn visit_type_var(&mut self, v: TypeDbVar) -> Option<Ty> {
101        self.process_var(v, |id| self.generics.types.get(id))
102    }
103    fn visit_const_generic_var(&mut self, v: ConstGenericDbVar) -> Option<ConstantExprKind> {
104        self.process_var(v, |id| {
105            self.generics.const_generics.get(id).map(|c| &c.kind)
106        })
107    }
108    fn visit_clause_var(&mut self, v: ClauseDbVar) -> Option<TraitRefKind> {
109        if self.explicits_only {
110            None
111        } else {
112            self.process_var(v, |id| Some(&self.generics.trait_refs.get(id)?.kind))
113        }
114    }
115    fn visit_self_clause(&mut self) -> Option<TraitRefKind> {
116        Some(self.self_ref.cloned().expect(
117            "used `substitute` on an item coming from a trait; \
118            use `substitute_with_self` or `substitute_inner_binder` instead.",
119        ))
120    }
121}
122
123#[derive(Debug)]
124pub struct GenericsMismatch;
125
126/// Types that are involved at the type-level and may be substituted around.
127pub trait TyVisitable: Sized + AstVisitable {
128    /// Visit the variables contained in `self`, as seen from the outside of `self`. This means
129    /// that any variable bound inside `self` will be skipped, and all the seen De Bruijn indices
130    /// will count from the outside of `self`.
131    fn visit_vars(&mut self, v: &mut impl VarsVisitor) {
132        #[derive(Visitor)]
133        struct Wrap<'v, V> {
134            v: &'v mut V,
135            depth: DeBruijnId,
136        }
137        impl<V> VisitorWithBinderDepth for Wrap<'_, V> {
138            fn binder_depth_mut(&mut self) -> &mut DeBruijnId {
139                &mut self.depth
140            }
141        }
142        impl<V: VarsVisitor> VisitAstMut for Wrap<'_, V> {
143            fn visit<T: AstVisitable>(&mut self, x: &mut T) -> ControlFlow<Self::Break> {
144                VisitWithBinderDepth::new(self).visit(x)
145            }
146
147            fn exit_region(&mut self, r: &mut Region) {
148                match r {
149                    Region::Var(var)
150                        if let Some(var) = var.move_out_from_depth(self.depth)
151                            && let Some(new_r) = self.v.visit_region_var(var) =>
152                    {
153                        *r = new_r.move_under_binders(self.depth);
154                    }
155                    Region::Erased | Region::Body(..)
156                        if let Some(new_r) = self.v.visit_erased_region() =>
157                    {
158                        *r = new_r.move_under_binders(self.depth);
159                    }
160                    _ => (),
161                }
162            }
163            fn exit_ty(&mut self, ty: &mut Ty) {
164                if let TyKind::TypeVar(var) = ty.kind()
165                    && let Some(var) = var.move_out_from_depth(self.depth)
166                    && let Some(new_ty) = self.v.visit_type_var(var)
167                {
168                    *ty = new_ty.move_under_binders(self.depth);
169                }
170            }
171            fn exit_constant_expr(&mut self, ce: &mut ConstantExpr) {
172                if let ConstantExprKind::Var(var) = &mut ce.kind
173                    && let Some(var) = var.move_out_from_depth(self.depth)
174                    && let Some(new_cg) = self.v.visit_const_generic_var(var)
175                {
176                    ce.kind = new_cg.move_under_binders(self.depth);
177                }
178            }
179            fn exit_trait_ref_kind(&mut self, kind: &mut TraitRefKind) {
180                match kind {
181                    TraitRefKind::SelfId => {
182                        if let Some(new_kind) = self.v.visit_self_clause() {
183                            *kind = new_kind.move_under_binders(self.depth);
184                        }
185                    }
186                    TraitRefKind::Clause(var) => {
187                        if let Some(var) = var.move_out_from_depth(self.depth)
188                            && let Some(new_kind) = self.v.visit_clause_var(var)
189                        {
190                            *kind = new_kind.move_under_binders(self.depth);
191                        }
192                    }
193                    _ => {}
194                }
195            }
196        }
197        Wrap {
198            v,
199            depth: DeBruijnId::zero(),
200        }
201        .visit(self);
202    }
203
204    /// Substitute the generic variables inside `self` by replacing them with the provided values.
205    /// Note: if `self` is an item that comes from a `TraitDecl`, you must use
206    /// `substitute_with_self` or `substitute_inner_binder`, otherwise you'll get panics.
207    fn substitute(self, generics: &GenericArgs) -> Self {
208        SubstVisitor::new(generics, None, false)
209            .visit(self)
210            .unwrap()
211    }
212    /// Substitute the generic variables inside `self` by replacing them with the provided values.
213    /// This is appropriate when substituting an inner binder.
214    fn substitute_inner_binder(self, generics: &GenericArgs) -> Self {
215        self.substitute_with_self(generics, &TraitRefKind::SelfId)
216    }
217    /// Substitute only the type, region and const generic args.
218    fn substitute_explicits(self, generics: &GenericArgs) -> Self {
219        SubstVisitor::new(generics, None, true).visit(self).unwrap()
220    }
221    /// Substitute the generic variables as well as the `TraitRefKind::SelfId` trait ref.
222    fn substitute_with_self(self, generics: &GenericArgs, self_ref: &TraitRefKind) -> Self {
223        self.try_substitute_with_self(generics, self_ref).unwrap()
224    }
225    /// Substitute the generic variables as well as the `TraitRefKind::SelfId` trait ref.
226    fn substitute_with_tref(self, tref: &TraitRef) -> Self {
227        let pred = tref.trait_decl_ref.clone().erase();
228        self.substitute_with_self(&pred.generics, &tref.kind)
229    }
230    /// Substitute the generic variables as well as the `TraitRefKind::SelfId` trait ref.
231    fn try_substitute_with_tref(self, tref: &TraitRef) -> Result<Self, GenericsMismatch> {
232        let pred = tref.trait_decl_ref.clone().erase();
233        self.try_substitute_with_self(&pred.generics, &tref.kind)
234    }
235
236    fn try_substitute(self, generics: &GenericArgs) -> Result<Self, GenericsMismatch> {
237        SubstVisitor::new(generics, None, false).visit(self)
238    }
239    fn try_substitute_with_self(
240        self,
241        generics: &GenericArgs,
242        self_ref: &TraitRefKind,
243    ) -> Result<Self, GenericsMismatch> {
244        SubstVisitor::new(generics, Some(self_ref), false).visit(self)
245    }
246
247    /// Move under one binder.
248    fn move_under_binder(self) -> Self {
249        self.move_under_binders(DeBruijnId::one())
250    }
251
252    /// Move under `depth` binders.
253    fn move_under_binders(mut self, depth: DeBruijnId) -> Self {
254        if !depth.is_zero() {
255            let Continue(()) = self.visit_db_id::<Infallible>(|id| {
256                *id = id.plus(depth);
257                Continue(())
258            });
259        }
260        self
261    }
262
263    /// Move from under one binder.
264    fn move_from_under_binder(self) -> Option<Self> {
265        self.move_from_under_binders(DeBruijnId::one())
266    }
267
268    /// Move the value out of `depth` binders. Returns `None` if it contains a variable bound in
269    /// one of these `depth` binders.
270    fn move_from_under_binders(mut self, depth: DeBruijnId) -> Option<Self> {
271        self.visit_db_id::<()>(|id| match id.sub(depth) {
272            Some(sub) => {
273                *id = sub;
274                Continue(())
275            }
276            None => Break(()),
277        })
278        .is_continue()
279        .then_some(self)
280    }
281
282    /// Visit the de Bruijn ids contained in `self`, as seen from the outside of `self`. This means
283    /// that any variable bound inside `self` will be skipped, and all the seen indices will count
284    /// from the outside of self.
285    fn visit_db_id<B>(
286        &mut self,
287        f: impl FnMut(&mut DeBruijnId) -> ControlFlow<B>,
288    ) -> ControlFlow<B> {
289        struct Wrap<F> {
290            f: F,
291            depth: DeBruijnId,
292        }
293        impl<B, F> Visitor for Wrap<F>
294        where
295            F: FnMut(&mut DeBruijnId) -> ControlFlow<B>,
296        {
297            type Break = B;
298        }
299        impl<B, F> VisitAstMut for Wrap<F>
300        where
301            F: FnMut(&mut DeBruijnId) -> ControlFlow<B>,
302        {
303            fn enter_region_binder<T: AstVisitable>(&mut self, _: &mut RegionBinder<T>) {
304                self.depth = self.depth.incr()
305            }
306            fn exit_region_binder<T: AstVisitable>(&mut self, _: &mut RegionBinder<T>) {
307                self.depth = self.depth.decr()
308            }
309            fn enter_binder<T: AstVisitable>(&mut self, _: &mut Binder<T>) {
310                self.depth = self.depth.incr()
311            }
312            fn exit_binder<T: AstVisitable>(&mut self, _: &mut Binder<T>) {
313                self.depth = self.depth.decr()
314            }
315
316            fn visit_de_bruijn_id(&mut self, x: &mut DeBruijnId) -> ControlFlow<Self::Break> {
317                if let Some(mut shifted) = x.sub(self.depth) {
318                    (self.f)(&mut shifted)?;
319                    *x = shifted.plus(self.depth)
320                }
321                Continue(())
322            }
323        }
324        self.drive_mut(&mut Wrap {
325            f,
326            depth: DeBruijnId::zero(),
327        })
328    }
329
330    /// Replace all the erased regions by the output of the provided function. Binders levels are
331    /// handled automatically.
332    fn replace_erased_regions(mut self, f: impl FnMut() -> Region) -> Self {
333        #[derive(Visitor)]
334        struct RefreshErasedRegions<F>(F);
335        impl<F: FnMut() -> Region> VarsVisitor for RefreshErasedRegions<F> {
336            fn visit_erased_region(&mut self) -> Option<Region> {
337                Some((self.0)())
338            }
339        }
340        self.visit_vars(&mut RefreshErasedRegions(f));
341        self
342    }
343}
344
345impl<T: AstVisitable> TyVisitable for T {}
346
347/// A value of type `T` applied to some `GenericArgs`, except we havent applied them yet to avoid a
348/// deep clone.
349#[derive(Debug, Clone)]
350pub struct Substituted<'a, T> {
351    pub val: &'a T,
352    pub generics: Cow<'a, GenericArgs>,
353    pub trait_self: Option<&'a TraitRefKind>,
354}
355
356impl<'a, T> Substituted<'a, T> {
357    pub fn new(val: &'a T, generics: &'a GenericArgs) -> Self {
358        Self {
359            val,
360            generics: Cow::Borrowed(generics),
361            trait_self: None,
362        }
363    }
364    pub fn new_for_trait(
365        val: &'a T,
366        generics: &'a GenericArgs,
367        trait_self: &'a TraitRefKind,
368    ) -> Self {
369        Self {
370            val,
371            generics: Cow::Borrowed(generics),
372            trait_self: Some(trait_self),
373        }
374    }
375    pub fn new_for_trait_ref(val: &'a T, tref: &'a TraitRef) -> Self {
376        Self {
377            val,
378            generics: Cow::Owned(*tref.trait_decl_ref.clone().erase().generics),
379            trait_self: Some(&tref.kind),
380        }
381    }
382
383    pub fn rebind<U>(&self, val: &'a U) -> Substituted<'a, U> {
384        Substituted {
385            val,
386            generics: self.generics.clone(),
387            trait_self: self.trait_self,
388        }
389    }
390
391    pub fn substitute(&self) -> T
392    where
393        T: TyVisitable + Clone,
394    {
395        self.try_substitute().unwrap()
396    }
397    pub fn try_substitute(&self) -> Result<T, GenericsMismatch>
398    where
399        T: TyVisitable + Clone,
400    {
401        match self.trait_self {
402            None => self.val.clone().try_substitute(&self.generics),
403            Some(trait_self) => self
404                .val
405                .clone()
406                .try_substitute_with_self(&self.generics, trait_self),
407        }
408    }
409
410    pub fn iter<Item: 'a>(&self) -> impl Iterator<Item = Substituted<'a, Item>>
411    where
412        &'a T: IntoIterator<Item = &'a Item>,
413    {
414        self.val.into_iter().map(move |x| self.rebind(x))
415    }
416}
417
418/// A value of type `T` bound by the generic parameters of item
419/// `item`. Used when dealing with multiple items at a time, to
420/// ensure we don't mix up generics.
421///
422/// To get the value, use `under_binder_of` or `subst_for`.
423#[derive(Debug, Clone, Copy)]
424pub struct ItemBinder<ItemId, T> {
425    pub item_id: ItemId,
426    val: T,
427}
428
429impl<ItemId, T> ItemBinder<ItemId, T>
430where
431    ItemId: Debug + Copy + PartialEq,
432{
433    pub fn new(item_id: ItemId, val: T) -> Self {
434        Self { item_id, val }
435    }
436
437    pub fn as_ref(&self) -> ItemBinder<ItemId, &T> {
438        ItemBinder {
439            item_id: self.item_id,
440            val: &self.val,
441        }
442    }
443
444    pub fn map_bound<U>(self, f: impl FnOnce(T) -> U) -> ItemBinder<ItemId, U> {
445        ItemBinder {
446            item_id: self.item_id,
447            val: f(self.val),
448        }
449    }
450
451    fn assert_item_id(&self, item_id: ItemId) {
452        assert_eq!(
453            self.item_id, item_id,
454            "Trying to use item bound for {:?} as if it belonged to {:?}",
455            self.item_id, item_id
456        );
457    }
458
459    /// Assert that the value is bound for item `item_id`, and returns it. This is used when we
460    /// plan to store the returned value inside that item.
461    pub fn under_binder_of(self, item_id: ItemId) -> T {
462        self.assert_item_id(item_id);
463        self.val
464    }
465
466    /// Given generic args for `item_id`, assert that the value is bound for `item_id` and
467    /// substitute it with the provided generic arguments. Because the arguments are bound in the
468    /// context of another item, so it the resulting substituted value.
469    pub fn substitute<OtherItem: Debug + Copy + PartialEq>(
470        self,
471        args: ItemBinder<OtherItem, &GenericArgs>,
472    ) -> ItemBinder<OtherItem, T>
473    where
474        ItemId: Into<ItemId>,
475        T: TyVisitable,
476    {
477        args.map_bound(|args| self.val.substitute(args))
478    }
479}
480
481/// Dummy item identifier that represents the current item when not ambiguous.
482#[derive(Debug, Clone, Copy, PartialEq, Eq)]
483pub struct CurrentItem;
484
485impl<T> ItemBinder<CurrentItem, T> {
486    pub fn under_current_binder(self) -> T {
487        self.val
488    }
489}
charon_lib/ast/type_level/substitute.rs

charon_lib/ast/type_level/
substitute.rs
