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

charon_lib/ast/type_level/
substitute.rs
