charon_lib/ast/type_level/
substitute.rs

1use crate::ast::*;
2use derive_generic_visitor::*;
3use std::convert::Infallible;
4use std::fmt::Debug;
5use std::iter::Iterator;
6
7/// Visitor for type-level variables. Used to visit the variables contained in a value, as seen
8/// from the outside of the value. This means that any variable bound inside the value will be
9/// skipped, and all the seen De Bruijn indices will count from the outside of the value. The
10/// returned value, if any, will be put in place of the variable.
11pub trait VarsVisitor {
12    fn visit_erased_region(&mut self) -> Option<Region> {
13        None
14    }
15    fn visit_region_var(&mut self, _v: RegionDbVar) -> Option<Region> {
16        None
17    }
18    fn visit_type_var(&mut self, _v: TypeDbVar) -> Option<Ty> {
19        None
20    }
21    fn visit_const_generic_var(&mut self, _v: ConstGenericDbVar) -> Option<ConstantExprKind> {
22        None
23    }
24    fn visit_clause_var(&mut self, _v: ClauseDbVar) -> Option<ClauseReplacement> {
25        None
26    }
27    fn visit_self_clause(&mut self) -> Option<TraitRefKind> {
28        None
29    }
30    fn visit_metadata_value(&mut self, _value: &MetadataValue) {}
31}
32
33/// A replacement for a clause variable, expressed as seen from the outside of the value being
34/// visited.
35pub enum ClauseReplacement {
36    /// Replace this clause variable with another, which must be for the same predicate. This is
37    /// distinct from `Proof` to avoid infinite recursion when calling `apply`.
38    Var(ClauseDbVar),
39    /// Replace this clause variable with the given proof kind. The proof kind will be instantiated
40    /// with the clause regions then used as replacement.
41    Proof(RegionBinder<TraitRefKind>),
42}
43
44impl From<ClauseDbVar> for ClauseReplacement {
45    fn from(var: ClauseDbVar) -> Self {
46        Self::Var(var)
47    }
48}
49
50/// Visitor for the [TyVisitable::substitute] function.
51/// This substitutes variables bound at the level where we start to substitute (level 0).
52#[derive(Visitor)]
53pub(crate) struct SubstVisitor<'a> {
54    generics: &'a GenericArgs,
55    self_ref: Option<&'a TraitRefKind>,
56    /// Whether to substitute explicit variables only (types, regions, const generics).
57    explicits_only: bool,
58    had_error: bool,
59}
60impl<'a> SubstVisitor<'a> {
61    pub(crate) fn new(
62        generics: &'a GenericArgs,
63        self_ref: Option<&'a TraitRefKind>,
64        explicits_only: bool,
65    ) -> Self {
66        Self {
67            generics,
68            self_ref,
69            explicits_only,
70            had_error: false,
71        }
72    }
73
74    pub fn visit<T: TyVisitable>(mut self, mut x: T) -> Result<T, GenericsMismatch> {
75        if x.type_info().is_closed() {
76            return Ok(x);
77        }
78        x.visit_vars(&mut self);
79        if self.had_error {
80            Err(GenericsMismatch)
81        } else {
82            Ok(x)
83        }
84    }
85
86    /// Returns the value for this variable, if any.
87    fn process_var<Id, T>(
88        &mut self,
89        var: DeBruijnVar<Id>,
90        get: impl Fn(Id) -> Option<T>,
91    ) -> Option<T>
92    where
93        Id: Copy,
94        DeBruijnVar<Id>: Into<T>,
95    {
96        match var {
97            DeBruijnVar::Bound(dbid, varid) => {
98                Some(if let Some(dbid) = dbid.sub(DeBruijnId::one()) {
99                    // This is bound outside the binder we're substituting for.
100                    DeBruijnVar::Bound(dbid, varid).into()
101                } else {
102                    match get(varid) {
103                        Some(v) => v,
104                        None => {
105                            self.had_error = true;
106                            return None;
107                        }
108                    }
109                })
110            }
111            DeBruijnVar::Free(..) => None,
112        }
113    }
114}
115impl VarsVisitor for SubstVisitor<'_> {
116    fn visit_region_var(&mut self, v: RegionDbVar) -> Option<Region> {
117        self.process_var(v, |id| self.generics.regions.get(id).cloned())
118    }
119    fn visit_type_var(&mut self, v: TypeDbVar) -> Option<Ty> {
120        self.process_var(v, |id| self.generics.types.get(id).cloned())
121    }
122    fn visit_const_generic_var(&mut self, v: ConstGenericDbVar) -> Option<ConstantExprKind> {
123        self.process_var(v, |id| {
124            self.generics
125                .const_generics
126                .get(id)
127                .map(|c| c.kind().clone())
128        })
129    }
130    fn visit_clause_var(&mut self, v: ClauseDbVar) -> Option<ClauseReplacement> {
131        if self.explicits_only {
132            None
133        } else {
134            self.process_var(v, |id| {
135                let tref = self.generics.trait_refs.get(id)?;
136                let bound_kind = tref.0.map_ref(|tref| tref.kind.clone());
137                Some(ClauseReplacement::Proof(bound_kind))
138            })
139        }
140    }
141    fn visit_self_clause(&mut self) -> Option<TraitRefKind> {
142        Some(self.self_ref.cloned().expect(
143            "used `substitute` on an item coming from a trait; \
144            use `substitute_with_self` or `substitute_inner_binder` instead.",
145        ))
146    }
147    fn visit_metadata_value(&mut self, _value: &MetadataValue) {
148        self.had_error = true;
149    }
150}
151
152#[derive(Debug)]
153pub struct GenericsMismatch;
154
155/// Types that are involved at the type-level and may be substituted around.
156pub trait TyVisitable: Sized + AstVisitable {
157    /// Compute various bits of information about the contents of this value. See methods on
158    /// [`TypeInfo`].
159    fn type_info(&self) -> TypeInfo {
160        TypeInfo::compute(self)
161    }
162
163    /// Visit the variables contained in `self`, as seen from the outside of `self`. This means
164    /// that any variable bound inside `self` will be skipped, and all the seen De Bruijn indices
165    /// will count from the outside of `self`.
166    fn visit_vars(&mut self, v: &mut impl VarsVisitor) {
167        #[derive(Visitor)]
168        struct Wrap<'v, V> {
169            v: &'v mut V,
170            depth: DeBruijnId,
171        }
172        impl<V> VisitorWithBinderDepth for Wrap<'_, V> {
173            fn binder_depth_mut(&mut self) -> &mut DeBruijnId {
174                &mut self.depth
175            }
176        }
177        impl<V: VarsVisitor> VisitAstMut for Wrap<'_, V> {
178            fn visit<T: AstVisitable>(&mut self, x: &mut T) -> ControlFlow<Self::Break> {
179                VisitWithBinderDepth::new(self).visit(x)
180            }
181
182            fn exit_region(&mut self, r: &mut Region) {
183                match r {
184                    Region::Var(var)
185                        if let Some(var) = var.move_out_from_depth(self.depth)
186                            && let Some(new_r) = self.v.visit_region_var(var) =>
187                    {
188                        *r = new_r.move_under_binders(self.depth);
189                    }
190                    Region::Erased | Region::Body(..)
191                        if let Some(new_r) = self.v.visit_erased_region() =>
192                    {
193                        *r = new_r.move_under_binders(self.depth);
194                    }
195                    _ => (),
196                }
197            }
198            fn exit_ty(&mut self, ty: &mut Ty) {
199                if let TyKind::TypeVar(var) = ty.kind()
200                    && let Some(var) = var.move_out_from_depth(self.depth)
201                    && let Some(new_ty) = self.v.visit_type_var(var)
202                {
203                    *ty = new_ty.move_under_binders(self.depth);
204                }
205            }
206            fn exit_constant_expr_kind(&mut self, kind: &mut ConstantExprKind) {
207                if let ConstantExprKind::Var(var) = kind
208                    && let Some(var) = var.move_out_from_depth(self.depth)
209                    && let Some(new_cg) = self.v.visit_const_generic_var(var)
210                {
211                    *kind = new_cg.move_under_binders(self.depth);
212                }
213            }
214            fn exit_trait_ref_kind(&mut self, kind: &mut TraitRefKind) {
215                match kind {
216                    TraitRefKind::SelfId => {
217                        if let Some(new_kind) = self.v.visit_self_clause() {
218                            *kind = new_kind.move_under_binders(self.depth);
219                        }
220                    }
221                    TraitRefKind::Clause(var, regions) => {
222                        if let Some(var) = var.move_out_from_depth(self.depth)
223                            && let Some(replacement) = self.v.visit_clause_var(var)
224                        {
225                            *kind = match replacement {
226                                ClauseReplacement::Var(new_var) => TraitRefKind::Clause(
227                                    new_var.move_under_binders(self.depth),
228                                    regions.clone(),
229                                ),
230                                ClauseReplacement::Proof(proof) => {
231                                    proof.move_under_binders(self.depth).apply(regions.clone())
232                                }
233                            };
234                        }
235                    }
236                    _ => {}
237                }
238            }
239            fn enter_metadata_value(&mut self, value: &mut MetadataValue) {
240                self.v.visit_metadata_value(value);
241            }
242        }
243        Wrap {
244            v,
245            depth: DeBruijnId::zero(),
246        }
247        .visit(self);
248    }
249
250    /// Substitute the generic variables inside `self` by replacing them with the provided values.
251    /// Note: if `self` is an item that comes from a `TraitDecl`, you must use
252    /// `substitute_with_self` or `substitute_inner_binder`, otherwise you'll get panics.
253    fn substitute(self, generics: &GenericArgs) -> Self {
254        SubstVisitor::new(generics, None, false)
255            .visit(self)
256            .unwrap()
257    }
258    /// Substitute the generic variables inside `self` by replacing them with the provided values.
259    /// This is appropriate when substituting an inner binder.
260    fn substitute_inner_binder(self, generics: &GenericArgs) -> Self {
261        self.substitute_with_self(generics, &TraitRefKind::SelfId)
262    }
263    /// Substitute only the type, region and const generic args.
264    fn substitute_explicits(self, generics: &GenericArgs) -> Self {
265        SubstVisitor::new(generics, None, true).visit(self).unwrap()
266    }
267    /// Substitute the generic variables as well as the `TraitRefKind::SelfId` trait ref.
268    fn substitute_with_self(self, generics: &GenericArgs, self_ref: &TraitRefKind) -> Self {
269        self.try_substitute_with_self(generics, self_ref).unwrap()
270    }
271    /// Substitute the generic variables as well as the `TraitRefKind::SelfId` trait ref.
272    fn substitute_with_tref(self, tref: &TraitRef) -> Self {
273        self.substitute_with_self(&tref.trait_decl_ref.generics, &tref.kind)
274    }
275    /// Substitute the generic variables as well as the `TraitRefKind::SelfId` trait ref.
276    fn try_substitute_with_tref(self, tref: &TraitRef) -> Result<Self, GenericsMismatch> {
277        self.try_substitute_with_self(&tref.trait_decl_ref.generics, &tref.kind)
278    }
279
280    fn try_substitute(self, generics: &GenericArgs) -> Result<Self, GenericsMismatch> {
281        SubstVisitor::new(generics, None, false).visit(self)
282    }
283    fn try_substitute_with_self(
284        self,
285        generics: &GenericArgs,
286        self_ref: &TraitRefKind,
287    ) -> Result<Self, GenericsMismatch> {
288        SubstVisitor::new(generics, Some(self_ref), false).visit(self)
289    }
290
291    /// Move under one binder.
292    fn move_under_binder(self) -> Self {
293        self.move_under_binders(DeBruijnId::one())
294    }
295
296    /// Move under `depth` binders.
297    fn move_under_binders(mut self, depth: DeBruijnId) -> Self {
298        if !depth.is_zero() {
299            let Continue(()) = self.visit_db_id::<Infallible>(|id| {
300                *id = id.plus(depth);
301                Continue(())
302            });
303        }
304        self
305    }
306
307    /// Move from under one binder.
308    fn move_from_under_binder(self) -> Option<Self> {
309        self.move_from_under_binders(DeBruijnId::one())
310    }
311
312    /// Move the value out of `depth` binders. Returns `None` if it contains a variable bound in
313    /// one of these `depth` binders.
314    fn move_from_under_binders(mut self, depth: DeBruijnId) -> Option<Self> {
315        match self.type_info().max_de_bruijn_id() {
316            None => return Some(self),
317            Some(max) if max < depth => return None,
318            Some(_) => {}
319        }
320        self.visit_db_id::<()>(|id| match id.sub(depth) {
321            Some(sub) => {
322                *id = sub;
323                Continue(())
324            }
325            None => Break(()),
326        })
327        .is_continue()
328        .then_some(self)
329    }
330
331    /// Visit the de Bruijn ids contained in `self`, as seen from the outside of `self`. This means
332    /// that any variable bound inside `self` will be skipped, and all the seen indices will count
333    /// from the outside of self.
334    fn visit_db_id<B>(
335        &mut self,
336        f: impl FnMut(&mut DeBruijnId) -> ControlFlow<B>,
337    ) -> ControlFlow<B> {
338        if self.type_info().max_de_bruijn_id().is_none() {
339            return Continue(());
340        }
341
342        struct Wrap<F> {
343            f: F,
344            depth: DeBruijnId,
345        }
346        impl<B, F> Visitor for Wrap<F>
347        where
348            F: FnMut(&mut DeBruijnId) -> ControlFlow<B>,
349        {
350            type Break = B;
351        }
352        impl<F> VisitorWithBinderDepth for Wrap<F> {
353            fn binder_depth_mut(&mut self) -> &mut DeBruijnId {
354                &mut self.depth
355            }
356        }
357        impl<B, F> VisitAstMut for Wrap<F>
358        where
359            F: FnMut(&mut DeBruijnId) -> ControlFlow<B>,
360        {
361            fn visit<T: AstVisitable>(&mut self, x: &mut T) -> ControlFlow<Self::Break> {
362                VisitWithBinderDepth::new(self).visit(x)
363            }
364
365            fn visit_with_cached_type_info<T: AstVisitable>(
366                &mut self,
367                value: &mut WithCachedTypeInfo<T>,
368            ) -> ControlFlow<Self::Break> {
369                if value
370                    .type_info()
371                    .max_de_bruijn_id()
372                    .is_none_or(|max| max < self.depth)
373                {
374                    Continue(())
375                } else {
376                    self.visit_inner(value)
377                }
378            }
379
380            fn visit_de_bruijn_id(&mut self, x: &mut DeBruijnId) -> ControlFlow<Self::Break> {
381                if let Some(mut shifted) = x.sub(self.depth) {
382                    (self.f)(&mut shifted)?;
383                    *x = shifted.plus(self.depth)
384                }
385                Continue(())
386            }
387        }
388        Wrap {
389            f,
390            depth: DeBruijnId::zero(),
391        }
392        .visit(self)
393    }
394
395    /// Whether this value mentions any region var of its environment.
396    fn has_escaping_region_var(&self) -> bool {
397        struct HasEscapingRegionVar {
398            binder_depth: DeBruijnId,
399        }
400
401        impl Visitor for HasEscapingRegionVar {
402            type Break = ();
403        }
404
405        impl VisitorWithBinderDepth for HasEscapingRegionVar {
406            fn binder_depth_mut(&mut self) -> &mut DeBruijnId {
407                &mut self.binder_depth
408            }
409        }
410
411        impl VisitAst for HasEscapingRegionVar {
412            fn visit<T: AstVisitable>(&mut self, x: &T) -> ControlFlow<Self::Break> {
413                VisitWithBinderDepth::new(self).visit(x)
414            }
415
416            fn visit_region(&mut self, region: &Region) -> ControlFlow<Self::Break> {
417                match region {
418                    Region::Var(DeBruijnVar::Bound(dbid, _)) if *dbid >= self.binder_depth => {
419                        ControlFlow::Break(())
420                    }
421                    _ => ControlFlow::Continue(()),
422                }
423            }
424        }
425
426        HasEscapingRegionVar {
427            binder_depth: DeBruijnId::ZERO,
428        }
429        .visit(self)
430        .is_break()
431    }
432
433    /// Collect the regions contained in `self`.
434    fn collect_regions(&self) -> SeqHashSet<Region> {
435        #[derive(Default, Visitor)]
436        struct CollectRegions {
437            depth: DeBruijnId,
438            regions: SeqHashSet<Region>,
439        }
440        impl VisitorWithBinderDepth for CollectRegions {
441            fn binder_depth_mut(&mut self) -> &mut DeBruijnId {
442                &mut self.depth
443            }
444        }
445        impl VisitAst for CollectRegions {
446            fn visit<U: AstVisitable>(&mut self, value: &U) -> ControlFlow<Self::Break> {
447                VisitWithBinderDepth::new(self).visit(value)
448            }
449            fn enter_region(&mut self, region: &Region) {
450                if let Some(region) = region.move_from_under_binders(self.depth)
451                    && !region.is_erased()
452                {
453                    self.regions.insert(region);
454                }
455            }
456        }
457        CollectRegions::default()
458            .visit_by_val_infallible(self)
459            .regions
460    }
461
462    /// Replace all the regions by the output of the provided function. Binders levels are handled
463    /// automatically.
464    fn replace_regions(mut self, f: impl FnMut() -> Region) -> Self {
465        if !self.type_info().has_region() {
466            return self;
467        }
468
469        #[derive(Visitor)]
470        struct RefreshRegions<F>(F);
471        impl<F: FnMut() -> Region> VarsVisitor for RefreshRegions<F> {
472            fn visit_region_var(&mut self, _v: RegionDbVar) -> Option<Region> {
473                Some((self.0)())
474            }
475            fn visit_erased_region(&mut self) -> Option<Region> {
476                Some((self.0)())
477            }
478        }
479        self.visit_vars(&mut RefreshRegions(f));
480        self
481    }
482
483    /// Replace all the erased regions by the output of the provided function. Binders levels are
484    /// handled automatically.
485    fn replace_erased_regions(mut self, f: impl FnMut() -> Region) -> Self {
486        if !self.type_info().has_erased_or_body_regions() {
487            return self;
488        }
489
490        #[derive(Visitor)]
491        struct RefreshErasedRegions<F>(F);
492        impl<F: FnMut() -> Region> VarsVisitor for RefreshErasedRegions<F> {
493            fn visit_erased_region(&mut self) -> Option<Region> {
494                Some((self.0)())
495            }
496        }
497        self.visit_vars(&mut RefreshErasedRegions(f));
498        self
499    }
500}
501
502impl<T: AstVisitable> TyVisitable for T {}
503
504/// A value of type `T` applied to some `GenericArgs`, except we havent applied them yet to avoid a
505/// deep clone.
506#[derive(Debug, Clone)]
507pub struct Substituted<'a, T> {
508    pub val: &'a T,
509    pub generics: &'a GenericArgs,
510    pub trait_self: Option<&'a TraitRefKind>,
511}
512
513impl<'a, T> Substituted<'a, T> {
514    pub fn new(val: &'a T, generics: &'a GenericArgs) -> Self {
515        Self {
516            val,
517            generics,
518            trait_self: None,
519        }
520    }
521    pub fn new_for_trait(
522        val: &'a T,
523        generics: &'a GenericArgs,
524        trait_self: &'a TraitRefKind,
525    ) -> Self {
526        Self {
527            val,
528            generics,
529            trait_self: Some(trait_self),
530        }
531    }
532    pub fn new_for_trait_ref(val: &'a T, tref: &'a TraitRef) -> Self {
533        Self {
534            val,
535            generics: &tref.trait_decl_ref.generics,
536            trait_self: Some(&tref.kind),
537        }
538    }
539
540    pub fn rebind<U>(&self, val: &'a U) -> Substituted<'a, U> {
541        Substituted {
542            val,
543            generics: self.generics,
544            trait_self: self.trait_self,
545        }
546    }
547
548    pub fn substitute(&self) -> T
549    where
550        T: TyVisitable + Clone,
551    {
552        self.try_substitute().unwrap()
553    }
554    pub fn try_substitute(&self) -> Result<T, GenericsMismatch>
555    where
556        T: TyVisitable + Clone,
557    {
558        match self.trait_self {
559            None => self.val.clone().try_substitute(self.generics),
560            Some(trait_self) => self
561                .val
562                .clone()
563                .try_substitute_with_self(self.generics, trait_self),
564        }
565    }
566
567    pub fn iter<Item: 'a>(&self) -> impl Iterator<Item = Substituted<'a, Item>>
568    where
569        &'a T: IntoIterator<Item = &'a Item>,
570    {
571        self.val.into_iter().map(move |x| self.rebind(x))
572    }
573}
574
575/// A value of type `T` bound by the generic parameters of item
576/// `item`. Used when dealing with multiple items at a time, to
577/// ensure we don't mix up generics.
578///
579/// To get the value, use `under_binder_of` or `subst_for`.
580#[derive(Debug, Copy, Clone)]
581pub struct ItemBinder<ItemId, T> {
582    pub item_id: ItemId,
583    val: T,
584}
585
586impl<ItemId, T> ItemBinder<ItemId, T>
587where
588    ItemId: Debug + Copy + PartialEq,
589{
590    pub fn new(item_id: ItemId, val: T) -> Self {
591        Self { item_id, val }
592    }
593
594    pub fn as_ref(&self) -> ItemBinder<ItemId, &T> {
595        ItemBinder {
596            item_id: self.item_id,
597            val: &self.val,
598        }
599    }
600
601    pub fn map_bound<U>(self, f: impl FnOnce(T) -> U) -> ItemBinder<ItemId, U> {
602        ItemBinder {
603            item_id: self.item_id,
604            val: f(self.val),
605        }
606    }
607
608    fn assert_item_id(&self, item_id: ItemId) {
609        assert_eq!(
610            self.item_id, item_id,
611            "Trying to use item bound for {:?} as if it belonged to {:?}",
612            self.item_id, item_id
613        );
614    }
615
616    /// Assert that the value is bound for item `item_id`, and returns it. This is used when we
617    /// plan to store the returned value inside that item.
618    pub fn under_binder_of(self, item_id: ItemId) -> T {
619        self.assert_item_id(item_id);
620        self.val
621    }
622
623    /// Given generic args for `item_id`, assert that the value is bound for `item_id` and
624    /// substitute it with the provided generic arguments. Because the arguments are bound in the
625    /// context of another item, so it the resulting substituted value.
626    pub fn substitute<OtherItem: Debug + Copy + PartialEq>(
627        self,
628        args: ItemBinder<OtherItem, &GenericArgs>,
629    ) -> ItemBinder<OtherItem, T>
630    where
631        ItemId: Into<ItemId>,
632        T: TyVisitable,
633    {
634        args.map_bound(|args| self.val.substitute(args))
635    }
636}
637
638/// Dummy item identifier that represents the current item when not ambiguous.
639#[derive(Debug, Copy, Clone, PartialEq, Eq)]
640pub struct CurrentItem;
641
642impl<T> ItemBinder<CurrentItem, T> {
643    pub fn under_current_binder(self) -> T {
644        self.val
645    }
646}
charon_lib/ast/type_level/substitute.rs

charon_lib/ast/type_level/
substitute.rs
