charon_lib/ast/
visitor.rs

1//! Defines two overrideable visitor traits that can be used to conveniently traverse the whole
2//! contents of an item. This is useful when e.g. dealing with types, which show up pretty much
3//! everywhere in the ast.
4//!
5//! The crate defines two traits:
6//! - `AstVisitable` is a trait implemented by all the types that can be visited by this;
7//! - `VisitAst[Mut]` is a (pair of) visitor trait(s) that can be implemented by visitors.
8//!   To define a visitor, implement `VisitAst[Mut]` and override the methods you need. Calling
9//!   `x.drive[_mut](&mut visitor)` will then traverse `x`, calling the visitor methods on all the
10//!   subvalues encountered.
11//!
12//! Underneath it all, this uses `derive_generic_visitor::Drive[Mut]` to do the actual visiting.
13use std::mem;
14use std::path::PathBuf;
15use std::{any::Any, hash::Hash};
16
17use crate::ast::from_rustc;
18use crate::ast::from_rustc::{
19    AttributeKind, DeprecatedSince, Deprecation, Ident, LangItem, OptimizeAttr,
20};
21use crate::ast::*;
22use crate::ids::{Idx, IndexVec};
23use derive_generic_visitor::*;
24
25/// An overrideable visitor trait that can be used to conveniently traverse the whole contents of
26/// an item. This is useful when e.g. dealing with types, which show up pretty much everywhere in
27/// the ast.
28///
29/// This defines three traits:
30/// - `AstVisitable` is a trait implemented by all the types listed below; it has a
31/// `drive[_mut]` method that takes a `VisitAst[Mut]` visitor and calls its methods on all
32/// the relevant subvalues of `self` encountered.
33/// - `VisitAst[Mut]` is a (pair of) visitor trait(s) that can be implemented by visitors. To
34/// define a visitor, implement `VisitAst[Mut]` and override the methods you need.
35///
36/// This trait has a `drive[_mut]` method that knows how to drive a `VisitAst[Mut]` visitor. This
37/// trait is implemented for all the listed types. If listed as `override`, the corresponding
38/// visitor trait has an overrideable method to visit this type. If listed as `drive`, the type
39/// will only be visited by recursing into its contents.
40///
41/// Morally this represents the predicate `for<V: VisitAst[Mut]> Self:
42/// Drive[Mut]<AstVisitableWrapper<V>>`
43#[visitable_group(
44    // Defines the `Visit[Mut]` traits and the `drive[_mut]` method that drives them.
45    visitor(drive(&VisitAst)),
46    visitor(drive_mut(&mut VisitAstMut)),
47    visitor(drive_two(&two ZipAst)),
48    // Types that are skipped by normal visitors but compared for equality by `ZipAst`.
49    skip_but_eq(
50        (), String, PathBuf, bool, char, i128, u8, u32, u64, u128, usize, ustr::Ustr,
51        crate::options::CliOpts,
52        Abi, BuiltinImplData, Byte, DeprecatedSince, DropKind, Error, FileName,
53        GlobalKind, ItemOpacity, LangItem, LifetimeMutability, OptimizeAttr, OverflowMode,
54        ReprOptions, Variance,
55        std::ops::RangeInclusive<IntegerValue>,
56        WithRetag, BuiltinPathElem, BranchId,
57    ),
58    // Types that are completely skipped, even by `ZipAst`.
59    skip(
60        DeclarationGroup, PredicateOrigin, TargetInfo,
61        llbc_ast::BlockId, llbc_ast::StatementId,
62    ),
63    // Types that we unconditionally explore.
64    drive(
65        Assert, AttributeKind, BinderKind, BinOp, BorrowckStatement, BorrowKind, BuiltinAdt, BuiltinAssertKind,
66        Call, CastKind, ClosureInfo, ClosureKind, ConstGenericParam, ConstGenericVarId,
67        Deprecation, Disambiguator, DynPredicate, Field, FieldId, File, FloatTy, FloatValue,
68        FnOperand, FnPtrKind, InlineAttr, IntegerTy, IntTy, UIntTy, ScalarTy,
69        Ident, from_rustc::InlineAttr,
70        llbc_ast::ExprBody, llbc_ast::StatementKind,
71        Loc, Locals, NullOp, Operand, PathElem, PlaceKind,
72        RawAttribute, RefKind, RegionId, RegionParam, IntegerValue, TraitItemName, TraitMethodId, AssocTypeId, AssocConstId, AssocItemId, MaybeAssocItemId,
73        TranslatedCrate, TypeDeclKind, TypeParam, TypePattern, TypeVarId,
74        ullbc_ast::BlockData, ullbc_ast::BlockId, ullbc_ast::ExprBody, ullbc_ast::StatementKind,
75        ullbc_ast::TerminatorKind, SwitchData, SwitchScrutinee,
76        UnOp, UnsizingMetadata, Local, Variant, VariantId, LocalId, Layout, VariantLayout,
77        Discriminator,
78        Size, OffsetExpr, OffsetGuarantee,
79        PtrMetadata,
80        SpanData, SerializedSpan,
81        ItemByVal, VTableField, AssocItemNames,
82        for<Id: AstVisitable> DeclRef<Id>, ItemId,
83        for<T: AstVisitable> Box<T>,
84        for<T: AstVisitable> Option<T>,
85        for<A: AstVisitable, B: AstVisitable> (A, B),
86        for<A: AstVisitable, B: AstVisitable, C: AstVisitable> (A, B, C),
87        for<A: AstVisitable, B: AstVisitable> Result<A, B>,
88        for<A: AstVisitable, B: AstVisitable> OutlivesPred<A, B>,
89        for<T: AstVisitable> Vec<T>,
90        for<T: AstVisitable + HashConsable> HashConsed<T>,
91        for<I: Idx, T: AstVisitable> IndexMap<I, T>,
92        for<I: Idx, T: AstVisitable> IndexVec<I, T>,
93    ),
94    // Types for which we call the corresponding `visit_$ty` method, which by default explores the
95    // type but can be overridden.
96    override(
97        FunDeclId, GlobalDeclId, TypeDeclId, TraitDeclId, TraitImplId, FileId,
98        TypeDeclRef, FunDeclRef, GlobalDeclRef, TraitDeclRef, TraitImplRef, ImplElem,
99        FunDecl, GlobalDecl, TypeDecl, TraitDecl, TraitImpl,
100        ItemMeta, Name, Span, Attribute, AttrInfo,
101        TypeSource, FunSource, GlobalSource, TraitDeclSource, TraitImplSource,
102        TraitAssocTy, TraitAssocConst, TraitMethod, TraitAssocTyImpl,
103        DeBruijnId, Ty, TyKind, Region, TraitRef, TraitRefContents, TraitRefKind,
104        GenericArgs, GenericParams, TraitParam, TraitClauseId, TraitTypeConstraint,
105        for<T: AstVisitable + Idx> DeBruijnVar<T>,
106        for<T: AstVisitable> RegionBinder<T>,
107        for<T: AstVisitable> Binder<T>,
108        llbc_block: llbc_ast::Block, llbc_statement: llbc_ast::Statement,
109        ullbc_statement: ullbc_ast::Statement, ullbc_terminator: ullbc_ast::Terminator,
110        AbortKind, AggregateKind, FnPtr, FunSig,
111        ConstantExpr, ConstantExprKind, InhabitedPredicate, InhabitedPredicateKind, SizeExpr,
112        SizeExprKind, MetadataValue, Place, ProjectionElem, Rvalue, Body,
113    )
114)]
115pub trait AstVisitable: Any {
116    /// The name of the type, used for debug logging.
117    fn name(&self) -> &'static str {
118        std::any::type_name::<Self>()
119    }
120    /// Visit all occurrences of that type inside `self`, in pre-order traversal.
121    fn dyn_visit<T: AstVisitable>(&self, f: impl FnMut(&T))
122    where
123        Self: Sized,
124    {
125        let _ = VisitAst::visit(&mut DynVisitor::new_shared::<T>(f), self);
126    }
127    /// Visit all occurrences of that type inside `self`, in pre-order traversal.
128    fn dyn_visit_mut<T: AstVisitable>(&mut self, f: impl FnMut(&mut T))
129    where
130        Self: Sized,
131    {
132        let _ = VisitAstMut::visit(&mut DynVisitor::new_mut::<T>(f), self);
133    }
134}
135
136/// Manual impl that visits the keys and values
137impl<K: AstVisitable + Hash + Eq, T: AstVisitable> AstVisitable for SeqHashMap<K, T> {
138    fn drive<V: VisitAst>(&self, v: &mut V) -> ControlFlow<V::Break> {
139        for (k, x) in self {
140            v.visit(k)?;
141            v.visit(x)?;
142        }
143        Continue(())
144    }
145    fn drive_mut<V: VisitAstMut>(&mut self, v: &mut V) -> ControlFlow<V::Break> {
146        for (mut k, mut x) in mem::take(self) {
147            v.visit(&mut k)?;
148            v.visit(&mut x)?;
149            self.insert(k, x);
150        }
151        Continue(())
152    }
153    fn drive_two<V: ZipAst>(&self, other: &Self, v: &mut V) -> ControlFlow<V::Break> {
154        if self.len() != other.len() {
155            return Break(Default::default());
156        }
157        for ((key, value), (other_key, other_value)) in self.iter().zip(other) {
158            v.visit(key, other_key)?;
159            v.visit(value, other_value)?;
160        }
161        Continue(())
162    }
163}
164impl<K: BodyVisitable + Hash + Eq, T: BodyVisitable> BodyVisitable for SeqHashMap<K, T> {
165    fn drive_body<V: VisitBody>(&self, v: &mut V) -> ControlFlow<V::Break> {
166        for (k, x) in self {
167            v.visit(k)?;
168            v.visit(x)?;
169        }
170        Continue(())
171    }
172    fn drive_body_mut<V: VisitBodyMut>(&mut self, v: &mut V) -> ControlFlow<V::Break> {
173        for (mut k, mut x) in mem::take(self) {
174            v.visit(&mut k)?;
175            v.visit(&mut x)?;
176            self.insert(k, x);
177        }
178        Continue(())
179    }
180}
181
182/// A smaller visitor group just for function bodies. This explores statements, places and
183/// operands, but does not recurse into types.
184///
185/// This defines three traits:
186/// - `BodyVisitable` is a trait implemented by all the types listed below; it has a
187/// `drive_body[_mut]` method that takes a `VisitBody[Mut]` visitor and calls its methods on all
188/// the relevant subvalues of `self` encountered.
189/// - `VisitBody[Mut]` is a (pair of) visitor trait(s) that can be implemented by visitors. To
190/// define a visitor, implement `VisitBody[Mut]` and override the methods you need.
191///
192/// Morally this represents the predicate `for<V: VisitBody[Mut]> Self:
193/// Drive[Mut]<BodyVisitableWrapper<V>>`
194#[visitable_group(
195    // Defines the `VisitBody[Mut]` traits and the `drive_body[_mut]` method that drives them.
196    visitor(drive_body(&VisitBody)),
197    visitor(drive_body_mut(&mut VisitBodyMut)),
198    // Types that are ignored when encountered.
199    skip(
200        AbortKind, BinOp, BorrowKind, BranchId, BuiltinAssertKind, ConstantExpr, FieldId,
201        TypeDeclRef, FunDeclId, FunDeclRef, FnPtrKind, GenericArgs, GlobalDeclRef, IntegerTy, IntTy, UIntTy,
202        NullOp, RefKind, IntegerValue, Span, Ty, TypeDeclId,  UnOp, VariantId,
203        TraitRef, ScalarTy, Region, RegionId, (), String, PathBuf, bool, u32, usize,
204        DropKind, Error, Variance, WithRetag, BuiltinAdt, BuiltinPathElem,
205        llbc_ast::BlockId, llbc_ast::StatementId,
206    ),
207    // Types that we unconditionally explore.
208    drive(
209        Assert, BorrowckStatement, PlaceKind,
210        llbc_ast::ExprBody, llbc_ast::StatementKind,
211        ullbc_ast::BlockData, ullbc_ast::ExprBody, ullbc_ast::StatementKind,
212        ullbc_ast::TerminatorKind, SwitchData, SwitchScrutinee,
213        Body, Local,
214        for<T: BodyVisitable> Box<T>,
215        for<T: BodyVisitable> Option<T>,
216        for<T: BodyVisitable, E: BodyVisitable> Result<T, E>,
217        for<A: BodyVisitable, B: BodyVisitable> (A, B),
218        for<A: BodyVisitable, B: BodyVisitable, C: BodyVisitable> (A, B, C),
219        for<T: BodyVisitable> Vec<T>,
220        for<I: Idx, T: BodyVisitable> IndexMap<I, T>,
221        for<I: Idx, T: BodyVisitable> IndexVec<I, T>,
222    ),
223    // Types for which we call the corresponding `visit_$ty` method, which by default explores the
224    // type but can be overridden.
225    override(
226        AggregateKind, Call, FnOperand, FnPtr,
227        Operand, Place, ProjectionElem, Rvalue, Locals, LocalId,
228        llbc_block: llbc_ast::Block,
229        llbc_statement: llbc_ast::Statement,
230        ullbc_statement: ullbc_ast::Statement,
231        ullbc_terminator: ullbc_ast::Terminator,
232        ullbc_block_id: ullbc_ast::BlockId,
233    )
234)]
235pub trait BodyVisitable: Any {
236    /// Visit all occurrences of that type inside `self`, in pre-order traversal.
237    fn dyn_visit_in_body<T: BodyVisitable>(&self, f: impl FnMut(&T))
238    where
239        Self: Sized,
240    {
241        let _ = VisitBody::visit(&mut DynVisitor::new_shared::<T>(f), self);
242    }
243
244    /// Visit all occurrences of that type inside `self`, in pre-order traversal.
245    fn dyn_visit_in_body_mut<T: BodyVisitable>(&mut self, f: impl FnMut(&mut T))
246    where
247        Self: Sized,
248    {
249        let _ = VisitBodyMut::visit(&mut DynVisitor::new_mut::<T>(f), self);
250    }
251}
252
253/// Ast and body visitor that uses dynamic dispatch to call the provided function on the visited
254/// values of the right type.
255#[derive(Visitor)]
256pub struct DynVisitor<F> {
257    enter: F,
258}
259impl DynVisitor<()> {
260    pub fn new_shared<T: Any>(mut f: impl FnMut(&T)) -> DynVisitor<impl FnMut(&dyn Any)> {
261        let enter = move |x: &dyn Any| {
262            if let Some(x) = x.downcast_ref::<T>() {
263                f(x);
264            }
265        };
266        DynVisitor { enter }
267    }
268    pub fn new_mut<T: Any>(mut f: impl FnMut(&mut T)) -> DynVisitor<impl FnMut(&mut dyn Any)> {
269        let enter = move |x: &mut dyn Any| {
270            if let Some(x) = x.downcast_mut::<T>() {
271                f(x);
272            }
273        };
274        DynVisitor { enter }
275    }
276}
277impl<F> VisitAst for DynVisitor<F>
278where
279    F: FnMut(&dyn Any),
280{
281    fn visit<T: AstVisitable>(&mut self, x: &T) -> ControlFlow<Self::Break> {
282        (self.enter)(x);
283        x.drive(self)?;
284        Continue(())
285    }
286}
287impl<F> VisitAstMut for DynVisitor<F>
288where
289    F: FnMut(&mut dyn Any),
290{
291    fn visit<T: AstVisitable>(&mut self, x: &mut T) -> ControlFlow<Self::Break> {
292        (self.enter)(x);
293        x.drive_mut(self)?;
294        Continue(())
295    }
296}
297impl<F> VisitBody for DynVisitor<F>
298where
299    F: FnMut(&dyn Any),
300{
301    fn visit<T: BodyVisitable>(&mut self, x: &T) -> ControlFlow<Self::Break> {
302        (self.enter)(x);
303        x.drive_body(self)?;
304        Continue(())
305    }
306}
307impl<F> VisitBodyMut for DynVisitor<F>
308where
309    F: FnMut(&mut dyn Any),
310{
311    fn visit<T: BodyVisitable>(&mut self, x: &mut T) -> ControlFlow<Self::Break> {
312        (self.enter)(x);
313        x.drive_body_mut(self)?;
314        Continue(())
315    }
316}
317
318pub use wrappers::*;
319mod wrappers {
320    //! This module defines a bunch of visitor wrappers, in the model described in the `derive_generic_visitor` crate.
321    //! Each such wrapper is a non-recursive visitor; the only thing it does is that its `.visit()`
322    //! method calls into the appropriate `visit_foo` of the wrapper, then continues visiting with
323    //! the wrapped visitor.
324    //!
325    //! To use such a wrapper, just override the `visit` method of your visitor to call
326    //! `TheWrapper::new(self).visit(x)`. This will integrate the wrapper into the normal behavior of
327    //! your visitor.
328    //!
329    //! Each wrapper interacts with its wrapped visitor via a trait. To be able to use several
330    //! wrappers at once, they must implement the wrapper-specific trait themselves and forward to
331    //! their wrappee visitor. It's a bit annoying as that potentially requires N^2 impls. I don't
332    //! know of a better design.
333    use std::mem;
334
335    use crate::ast::*;
336    use derive_generic_visitor::*;
337
338    /// Struct that we use to be able to use our visitor wrappers with each other to share
339    /// functionality, while still making the wrappers composable. We can implement e.g.
340    /// `VisitorWithItem for DontLeakImplDetails<Wrapper<V>>` while still retaining the capacity to
341    /// implement `impl<V: VisitorWithItem> VisitorWithItem for Wrapper<V>` that forwards to the
342    /// inner visitor.
343    #[repr(transparent)]
344    pub struct DontLeakImplDetails<V>(V);
345
346    impl<V> DontLeakImplDetails<V> {
347        pub fn new(v: &mut V) -> &mut Self {
348            // SAFETY: `repr(transparent)`
349            unsafe { std::mem::transmute(v) }
350        }
351        pub fn inner(&mut self) -> &mut V {
352            // SAFETY: `repr(transparent)`
353            unsafe { std::mem::transmute(self) }
354        }
355    }
356
357    impl<V: Visitor> Visitor for DontLeakImplDetails<V> {
358        type Break = V::Break;
359    }
360    impl<V: VisitAst> VisitAst for DontLeakImplDetails<V> {
361        /// Just forward to the wrapped visitor.
362        fn visit_inner<T>(&mut self, x: &T) -> ControlFlow<Self::Break>
363        where
364            T: AstVisitable,
365        {
366            x.drive(self.inner())
367        }
368    }
369    impl<V: VisitAstMut> VisitAstMut for DontLeakImplDetails<V> {
370        /// Just forward to the wrapped visitor.
371        fn visit_inner<T>(&mut self, x: &mut T) -> ControlFlow<Self::Break>
372        where
373            T: AstVisitable,
374        {
375            x.drive_mut(self.inner())
376        }
377    }
378
379    /// Visitor wrapper that tracks the depth of binders. To use it, make a visitor that implements
380    /// `VisitorWithBinderDepth` and override its `visit` function as follows:
381    /// ```ignore
382    /// impl VisitAst for MyVisitor {
383    ///     fn visit<'a, T: AstVisitable>(&'a mut self, x: &T) -> ControlFlow<Self::Break> {
384    ///         VisitWithBinderDepth::new(self).visit(x)
385    ///     }
386    ///     ...
387    /// }
388    /// ```
389    #[repr(transparent)]
390    pub struct VisitWithBinderDepth<V>(V);
391
392    impl<V: VisitorWithBinderDepth> VisitWithBinderDepth<V> {
393        pub fn new(v: &mut V) -> &mut Self {
394            // SAFETY: `repr(transparent)`
395            unsafe { std::mem::transmute(v) }
396        }
397        pub fn inner(&mut self) -> &mut V {
398            // SAFETY: `repr(transparent)`
399            unsafe { std::mem::transmute(self) }
400        }
401    }
402
403    pub trait VisitorWithBinderDepth {
404        fn binder_depth_mut(&mut self) -> &mut DeBruijnId;
405    }
406
407    impl<V: Visitor> Visitor for VisitWithBinderDepth<V> {
408        type Break = V::Break;
409    }
410    impl<V: VisitAst + VisitorWithBinderDepth> VisitAst for VisitWithBinderDepth<V> {
411        fn visit_inner<T>(&mut self, x: &T) -> ControlFlow<Self::Break>
412        where
413            T: AstVisitable,
414        {
415            x.drive(self.inner())
416        }
417        fn enter_region_binder<T: AstVisitable>(&mut self, _: &RegionBinder<T>) {
418            let binder_depth = self.0.binder_depth_mut();
419            *binder_depth = binder_depth.incr()
420        }
421        fn exit_region_binder<T: AstVisitable>(&mut self, _: &RegionBinder<T>) {
422            let binder_depth = self.0.binder_depth_mut();
423            *binder_depth = binder_depth.decr()
424        }
425        fn enter_binder<T: AstVisitable>(&mut self, _: &Binder<T>) {
426            let binder_depth = self.0.binder_depth_mut();
427            *binder_depth = binder_depth.incr()
428        }
429        fn exit_binder<T: AstVisitable>(&mut self, _: &Binder<T>) {
430            let binder_depth = self.0.binder_depth_mut();
431            *binder_depth = binder_depth.decr()
432        }
433    }
434    impl<V: VisitAstMut + VisitorWithBinderDepth> VisitAstMut for VisitWithBinderDepth<V> {
435        fn visit_inner<T>(&mut self, x: &mut T) -> ControlFlow<Self::Break>
436        where
437            T: AstVisitable,
438        {
439            x.drive_mut(self.inner())
440        }
441        fn enter_region_binder<T: AstVisitable>(&mut self, _: &mut RegionBinder<T>) {
442            let binder_depth = self.0.binder_depth_mut();
443            *binder_depth = binder_depth.incr()
444        }
445        fn exit_region_binder<T: AstVisitable>(&mut self, _: &mut RegionBinder<T>) {
446            let binder_depth = self.0.binder_depth_mut();
447            *binder_depth = binder_depth.decr()
448        }
449        fn enter_binder<T: AstVisitable>(&mut self, _: &mut Binder<T>) {
450            let binder_depth = self.0.binder_depth_mut();
451            *binder_depth = binder_depth.incr()
452        }
453        fn exit_binder<T: AstVisitable>(&mut self, _: &mut Binder<T>) {
454            let binder_depth = self.0.binder_depth_mut();
455            *binder_depth = binder_depth.decr()
456        }
457    }
458
459    /// Visitor wrapper that tracks the variance of the current position. To use it, make a visitor
460    /// that implements [`VisitorWithVariance`] and delegate its `visit` method to this wrapper.
461    pub struct VisitWithVariance<'a, 'ctx, V> {
462        inner: &'a mut V,
463        krate: &'ctx TranslatedCrate,
464    }
465
466    impl<'a, 'ctx, V: VisitorWithVariance> VisitWithVariance<'a, 'ctx, V> {
467        pub fn new(inner: &'a mut V, krate: &'ctx TranslatedCrate) -> Self {
468            Self { inner, krate }
469        }
470
471        fn compose(outer: Variance, inner: Variance) -> Variance {
472            match (outer, inner) {
473                (Variance::Unknown, _) | (_, Variance::Unknown) => Variance::Unknown,
474                (Variance::Bivariant, _) | (_, Variance::Bivariant) => Variance::Bivariant,
475                (Variance::Invariant, _) | (_, Variance::Invariant) => Variance::Invariant,
476                (Variance::Covariant, variance) => variance,
477                (Variance::Contravariant, Variance::Covariant) => Variance::Contravariant,
478                (Variance::Contravariant, Variance::Contravariant) => Variance::Covariant,
479            }
480        }
481    }
482
483    pub trait VisitorWithVariance {
484        fn ambient_variance_mut(&mut self) -> &mut Variance;
485    }
486
487    impl<V: Visitor> Visitor for VisitWithVariance<'_, '_, V> {
488        type Break = V::Break;
489    }
490
491    impl<V: VisitAst + VisitorWithVariance> VisitWithVariance<'_, '_, V> {
492        fn with(
493            &mut self,
494            variance: Variance,
495            f: impl FnOnce(&mut Self) -> ControlFlow<V::Break>,
496        ) -> ControlFlow<V::Break> {
497            let old = *self.inner.ambient_variance_mut();
498            *self.inner.ambient_variance_mut() = Self::compose(old, variance);
499            let result = f(self);
500            *self.inner.ambient_variance_mut() = old;
501            result
502        }
503
504        fn visit_with<T: AstVisitable>(
505            &mut self,
506            variance: Variance,
507            value: &T,
508        ) -> ControlFlow<V::Break> {
509            self.with(variance, |this| this.visit(value))
510        }
511
512        fn visit_inner_with(&mut self, variance: Variance, ty: &Ty) -> ControlFlow<V::Break> {
513            self.with(variance, |this| this.visit_inner(ty))
514        }
515    }
516
517    impl<V: VisitAst + VisitorWithVariance> VisitAst for VisitWithVariance<'_, '_, V> {
518        fn visit_inner<T: AstVisitable>(&mut self, value: &T) -> ControlFlow<Self::Break> {
519            value.drive(self.inner)
520        }
521
522        fn visit_ty(&mut self, ty: &Ty) -> ControlFlow<Self::Break> {
523            match ty.kind() {
524                TyKind::Ref(region, inner, kind) => {
525                    self.visit(region)?;
526                    match kind {
527                        RefKind::Shared => self.visit(inner)?,
528                        RefKind::Mut => self.visit_with(Variance::Invariant, inner)?,
529                    }
530                }
531                TyKind::RawPtr(inner, RefKind::Shared) => self.visit(inner)?,
532                TyKind::RawPtr(inner, RefKind::Mut) => {
533                    self.visit_with(Variance::Invariant, inner)?;
534                }
535                TyKind::FnPtr(..)
536                | TyKind::Adt(..)
537                | TyKind::Array(..)
538                | TyKind::Slice(..)
539                | TyKind::Pattern(..) => {
540                    self.visit_inner(ty)?;
541                }
542                TyKind::TraitType(..)
543                | TyKind::DynTrait(..)
544                | TyKind::FnDef(..)
545                | TyKind::PtrMetadata(..) => {
546                    self.visit_inner_with(Variance::Invariant, ty)?;
547                }
548                TyKind::TypeVar(..) | TyKind::Scalar(..) | TyKind::Never | TyKind::Error(..) => {}
549            }
550            Continue(())
551        }
552
553        fn visit_fun_sig(&mut self, sig: &FunSig) -> ControlFlow<Self::Break> {
554            for input in &sig.inputs {
555                self.visit_with(Variance::Contravariant, input)?;
556            }
557            self.visit(&sig.output)?;
558            Continue(())
559        }
560
561        fn visit_type_decl_ref(&mut self, type_ref: &TypeDeclRef) -> ControlFlow<Self::Break> {
562            if let Some(decl) = self.krate.type_decls.get(type_ref.id) {
563                let params = &decl.generics;
564                for (param, region) in params.regions.iter().zip(&type_ref.generics.regions) {
565                    self.visit_with(param.variance, region)?;
566                }
567                for (param, ty) in params.types.iter().zip(&type_ref.generics.types) {
568                    self.visit_with(param.variance, ty)?;
569                }
570            } else {
571                self.visit_with(Variance::Covariant, &type_ref.generics.regions)?;
572                self.visit_with(Variance::Covariant, &type_ref.generics.types)?;
573            }
574            Continue(())
575        }
576    }
577
578    /// Visitor wrapper that adds item-generic `enter_item` and `exit_item` methods.
579    #[repr(transparent)]
580    pub struct VisitWithItem<V>(V);
581
582    impl<V> VisitWithItem<V> {
583        pub fn new(v: &mut V) -> &mut Self {
584            // SAFETY: `repr(transparent)`
585            unsafe { std::mem::transmute(v) }
586        }
587        pub fn inner(&mut self) -> &mut V {
588            // SAFETY: `repr(transparent)`
589            unsafe { std::mem::transmute(self) }
590        }
591    }
592
593    pub trait VisitorWithItem: VisitAst {
594        fn enter_item(&mut self, _item: ItemRef<'_>) {}
595        fn exit_item(&mut self, _item: ItemRef<'_>) {}
596        fn visit_item(&mut self, item: ItemRef<'_>) -> ControlFlow<Self::Break> {
597            self.enter_item(item);
598            item.drive(self)?;
599            self.exit_item(item);
600            Continue(())
601        }
602    }
603    pub trait VisitorWithItemMut: VisitAstMut {
604        fn enter_item(&mut self, _item: ItemRefMut<'_>) {}
605        fn exit_item(&mut self, _item: ItemRefMut<'_>) {}
606        fn visit_item(&mut self, mut item: ItemRefMut<'_>) -> ControlFlow<Self::Break> {
607            self.enter_item(item.reborrow());
608            item.drive_mut(self)?;
609            self.exit_item(item);
610            Continue(())
611        }
612    }
613
614    impl<V: Visitor> Visitor for VisitWithItem<V> {
615        type Break = V::Break;
616    }
617    impl<V: VisitAst + VisitorWithItem> VisitAst for VisitWithItem<V> {
618        fn visit_inner<T>(&mut self, x: &T) -> ControlFlow<Self::Break>
619        where
620            T: AstVisitable,
621        {
622            x.drive(self.inner())
623        }
624        fn visit_fun_decl(&mut self, x: &FunDecl) -> ControlFlow<Self::Break> {
625            self.0.visit_item(ItemRef::Fun(x))
626        }
627        fn visit_type_decl(&mut self, x: &TypeDecl) -> ControlFlow<Self::Break> {
628            self.0.visit_item(ItemRef::Type(x))
629        }
630        fn visit_global_decl(&mut self, x: &GlobalDecl) -> ControlFlow<Self::Break> {
631            self.0.visit_item(ItemRef::Global(x))
632        }
633        fn visit_trait_decl(&mut self, x: &TraitDecl) -> ControlFlow<Self::Break> {
634            self.0.visit_item(ItemRef::TraitDecl(x))
635        }
636        fn visit_trait_impl(&mut self, x: &TraitImpl) -> ControlFlow<Self::Break> {
637            self.0.visit_item(ItemRef::TraitImpl(x))
638        }
639    }
640    impl<V: VisitAstMut + VisitorWithItemMut> VisitAstMut for VisitWithItem<V> {
641        fn visit_inner<T>(&mut self, x: &mut T) -> ControlFlow<Self::Break>
642        where
643            T: AstVisitable,
644        {
645            x.drive_mut(self.inner())
646        }
647        fn visit_fun_decl(&mut self, x: &mut FunDecl) -> ControlFlow<Self::Break> {
648            self.0.visit_item(ItemRefMut::Fun(x))
649        }
650        fn visit_type_decl(&mut self, x: &mut TypeDecl) -> ControlFlow<Self::Break> {
651            self.0.visit_item(ItemRefMut::Type(x))
652        }
653        fn visit_global_decl(&mut self, x: &mut GlobalDecl) -> ControlFlow<Self::Break> {
654            self.0.visit_item(ItemRefMut::Global(x))
655        }
656        fn visit_trait_decl(&mut self, x: &mut TraitDecl) -> ControlFlow<Self::Break> {
657            self.0.visit_item(ItemRefMut::TraitDecl(x))
658        }
659        fn visit_trait_impl(&mut self, x: &mut TraitImpl) -> ControlFlow<Self::Break> {
660            self.0.visit_item(ItemRefMut::TraitImpl(x))
661        }
662    }
663
664    /// Visitor wrapper that catches references to top-level items.
665    #[repr(transparent)]
666    pub struct VisitWithItemRef<V>(V);
667
668    impl<V> VisitWithItemRef<V> {
669        pub fn new(v: &mut V) -> &mut Self {
670            // SAFETY: `repr(transparent)`
671            unsafe { std::mem::transmute(v) }
672        }
673        pub fn inner(&mut self) -> &mut V {
674            // SAFETY: `repr(transparent)`
675            unsafe { std::mem::transmute(self) }
676        }
677    }
678
679    pub trait VisitorWithItemRef: VisitAst {
680        fn enter_item_ref(&mut self, _item_id: ItemId, _args: &GenericArgs) {}
681        fn exit_item_ref(&mut self, _item_id: ItemId, _args: &GenericArgs) {}
682        fn visit_item_ref(
683            &mut self,
684            item_id: ItemId,
685            args: &GenericArgs,
686        ) -> ControlFlow<Self::Break> {
687            self.enter_item_ref(item_id, args);
688            self.visit_inner(args)?;
689            self.exit_item_ref(item_id, args);
690            Continue(())
691        }
692    }
693    pub trait VisitorWithItemRefMut: VisitAstMut {
694        fn enter_item_ref(&mut self, _item_id: ItemId, _args: &mut GenericArgs) {}
695        fn exit_item_ref(&mut self, _item_id: ItemId, _args: &mut GenericArgs) {}
696        fn visit_item_ref(
697            &mut self,
698            item_id: ItemId,
699            args: &mut GenericArgs,
700        ) -> ControlFlow<Self::Break> {
701            self.enter_item_ref(item_id, args);
702            self.visit_inner(args)?;
703            self.exit_item_ref(item_id, args);
704            Continue(())
705        }
706    }
707
708    impl<V: Visitor> Visitor for VisitWithItemRef<V> {
709        type Break = V::Break;
710    }
711    impl<V: VisitAst + VisitorWithItemRef> VisitAst for VisitWithItemRef<V> {
712        fn visit_inner<T>(&mut self, x: &T) -> ControlFlow<Self::Break>
713        where
714            T: AstVisitable,
715        {
716            x.drive(self.inner())
717        }
718        fn visit_type_decl_ref(&mut self, x: &TypeDeclRef) -> ControlFlow<Self::Break> {
719            self.0.visit_item_ref(ItemId::Type(x.id), &x.generics)
720        }
721        fn visit_fun_decl_ref(&mut self, x: &FunDeclRef) -> ControlFlow<Self::Break> {
722            self.0.visit_item_ref(ItemId::Fun(x.id), &x.generics)
723        }
724        fn visit_global_decl_ref(&mut self, x: &GlobalDeclRef) -> ControlFlow<Self::Break> {
725            self.0.visit_item_ref(ItemId::Global(x.id), &x.generics)
726        }
727        fn visit_trait_decl_ref(&mut self, x: &TraitDeclRef) -> ControlFlow<Self::Break> {
728            self.0.visit_item_ref(ItemId::TraitDecl(x.id), &x.generics)
729        }
730        fn visit_trait_impl_ref(&mut self, x: &TraitImplRef) -> ControlFlow<Self::Break> {
731            self.0.visit_item_ref(ItemId::TraitImpl(x.id), &x.generics)
732        }
733        fn visit_fn_ptr(&mut self, x: &FnPtr) -> ControlFlow<Self::Break> {
734            match x.kind.as_ref() {
735                FnPtrKind::Fun(id) => self.0.visit_item_ref(ItemId::Fun(*id), &x.generics),
736                FnPtrKind::Trait(..) => self.visit_inner(x),
737            }
738        }
739    }
740    impl<V: VisitAstMut + VisitorWithItemRefMut> VisitAstMut for VisitWithItemRef<V> {
741        fn visit_inner<T>(&mut self, x: &mut T) -> ControlFlow<Self::Break>
742        where
743            T: AstVisitable,
744        {
745            x.drive_mut(self.inner())
746        }
747        fn visit_type_decl_ref(&mut self, x: &mut TypeDeclRef) -> ControlFlow<Self::Break> {
748            self.0.visit_item_ref(ItemId::Type(x.id), &mut x.generics)
749        }
750        fn visit_fun_decl_ref(&mut self, x: &mut FunDeclRef) -> ControlFlow<Self::Break> {
751            self.0.visit_item_ref(ItemId::Fun(x.id), &mut x.generics)
752        }
753        fn visit_global_decl_ref(&mut self, x: &mut GlobalDeclRef) -> ControlFlow<Self::Break> {
754            self.0.visit_item_ref(ItemId::Global(x.id), &mut x.generics)
755        }
756        fn visit_trait_decl_ref(&mut self, x: &mut TraitDeclRef) -> ControlFlow<Self::Break> {
757            self.0
758                .visit_item_ref(ItemId::TraitDecl(x.id), &mut x.generics)
759        }
760        fn visit_trait_impl_ref(&mut self, x: &mut TraitImplRef) -> ControlFlow<Self::Break> {
761            self.0
762                .visit_item_ref(ItemId::TraitImpl(x.id), &mut x.generics)
763        }
764        fn visit_fn_ptr(&mut self, x: &mut FnPtr) -> ControlFlow<Self::Break> {
765            match x.kind.as_ref() {
766                FnPtrKind::Fun(id) => self.0.visit_item_ref(ItemId::Fun(*id), &mut x.generics),
767                FnPtrKind::Trait(..) => self.visit_inner(x),
768            }
769        }
770    }
771
772    /// Visitor wrapper that tracks the stack of binders seen so far. See [`VisitWithBinderDepth`] for how to use.
773    #[repr(transparent)]
774    pub struct VisitWithBinderStack<V>(V);
775
776    impl<V: VisitorWithBinderStack> VisitWithBinderStack<V> {
777        // Helper
778        fn wrap(v: &mut V) -> &mut Self {
779            // SAFETY: `repr(transparent)`
780            unsafe { std::mem::transmute(v) }
781        }
782        pub fn new(v: &mut V) -> &mut VisitWithItem<DontLeakImplDetails<Self>> {
783            // Use the `WithItem` wrapper to simplify the implementation of this wrapper. We use
784            // `DontLeakImplDetails` to use the specific `VisitorWithItem` impl we care about
785            // instead of the one that forwards to the `VisitorWithItem` of the containted `V`.
786            VisitWithItem::new(DontLeakImplDetails::new(Self::wrap(v)))
787        }
788        pub fn inner(&mut self) -> &mut V {
789            // SAFETY: `repr(transparent)`
790            unsafe { std::mem::transmute(self) }
791        }
792    }
793
794    pub trait VisitorWithBinderStack {
795        fn binder_stack_mut(&mut self) -> &mut BindingStack<GenericParams>;
796    }
797
798    impl<V: VisitAst + VisitorWithBinderStack> VisitorWithItem
799        for DontLeakImplDetails<VisitWithBinderStack<V>>
800    {
801        fn enter_item(&mut self, item: ItemRef<'_>) {
802            self.0
803                .0
804                .binder_stack_mut()
805                .push(item.generic_params().clone());
806        }
807        fn exit_item(&mut self, _item: ItemRef<'_>) {
808            self.0.0.binder_stack_mut().pop();
809        }
810    }
811    impl<V: VisitAstMut + VisitorWithBinderStack> VisitorWithItemMut
812        for DontLeakImplDetails<VisitWithBinderStack<V>>
813    {
814        fn enter_item(&mut self, item: ItemRefMut<'_>) {
815            self.0
816                .0
817                .binder_stack_mut()
818                .push(item.as_ref().generic_params().clone());
819        }
820        fn exit_item(&mut self, _item: ItemRefMut<'_>) {
821            self.0.0.binder_stack_mut().pop();
822        }
823    }
824
825    impl<V: Visitor> Visitor for VisitWithBinderStack<V> {
826        type Break = V::Break;
827    }
828    impl<V: VisitAst + VisitorWithBinderStack> VisitAst for VisitWithBinderStack<V> {
829        fn visit_inner<T>(&mut self, x: &T) -> ControlFlow<Self::Break>
830        where
831            T: AstVisitable,
832        {
833            x.drive(self.inner())
834        }
835        fn visit_binder<T: AstVisitable>(
836            &mut self,
837            binder: &Binder<T>,
838        ) -> ControlFlow<Self::Break> {
839            self.0.binder_stack_mut().push(binder.params.clone());
840            self.visit_inner(binder)?;
841            self.0.binder_stack_mut().pop();
842            Continue(())
843        }
844        fn visit_region_binder<T: AstVisitable>(
845            &mut self,
846            binder: &RegionBinder<T>,
847        ) -> ControlFlow<Self::Break> {
848            self.0.binder_stack_mut().push(GenericParams {
849                regions: binder.regions.clone(),
850                ..Default::default()
851            });
852            self.visit_inner(binder)?;
853            self.0.binder_stack_mut().pop();
854            Continue(())
855        }
856    }
857    impl<V: VisitAstMut + VisitorWithBinderStack> VisitAstMut for VisitWithBinderStack<V> {
858        fn visit_inner<T>(&mut self, x: &mut T) -> ControlFlow<Self::Break>
859        where
860            T: AstVisitable,
861        {
862            x.drive_mut(self.inner())
863        }
864        fn visit_binder<T: AstVisitable>(
865            &mut self,
866            binder: &mut Binder<T>,
867        ) -> ControlFlow<Self::Break> {
868            self.0.binder_stack_mut().push(binder.params.clone());
869            self.visit_inner(binder)?;
870            self.0.binder_stack_mut().pop();
871            Continue(())
872        }
873        fn visit_region_binder<T: AstVisitable>(
874            &mut self,
875            binder: &mut RegionBinder<T>,
876        ) -> ControlFlow<Self::Break> {
877            self.0.binder_stack_mut().push(GenericParams {
878                regions: binder.regions.clone(),
879                ..Default::default()
880            });
881            self.visit_inner(binder)?;
882            self.0.binder_stack_mut().pop();
883            Continue(())
884        }
885    }
886
887    /// Visitor wrapper that tracks the current span. See [`VisitWithBinderDepth`] for how to use.
888    #[repr(transparent)]
889    pub struct VisitWithSpan<V>(V);
890
891    impl<V: VisitorWithSpan> VisitWithSpan<V> {
892        // Helper
893        fn wrap(v: &mut V) -> &mut Self {
894            // SAFETY: `repr(transparent)`
895            unsafe { std::mem::transmute(v) }
896        }
897        pub fn new(v: &mut V) -> &mut VisitWithItem<DontLeakImplDetails<Self>> {
898            // Use the `WithItem` wrapper to simplify the implementation of this wrapper. We use
899            // `DontLeakImplDetails` to use the specific `VisitorWithItem` impl we care about
900            // instead of the one that forwards to the `VisitorWithItem` of the containted `V`.
901            VisitWithItem::new(DontLeakImplDetails::new(Self::wrap(v)))
902        }
903        pub fn inner(&mut self) -> &mut V {
904            // SAFETY: `repr(transparent)`
905            unsafe { std::mem::transmute(self) }
906        }
907    }
908
909    pub trait VisitorWithSpan {
910        fn current_span(&mut self) -> &mut Span;
911    }
912
913    impl<V: VisitAst + VisitorWithSpan> VisitorWithItem for DontLeakImplDetails<VisitWithSpan<V>> {
914        fn visit_item(&mut self, item: ItemRef<'_>) -> ControlFlow<Self::Break> {
915            let old_span = mem::replace(self.0.0.current_span(), item.item_meta().span);
916            item.drive(self)?;
917            *self.0.0.current_span() = old_span;
918            Continue(())
919        }
920    }
921    impl<V: VisitAstMut + VisitorWithSpan> VisitorWithItemMut
922        for DontLeakImplDetails<VisitWithSpan<V>>
923    {
924        fn visit_item(&mut self, mut item: ItemRefMut<'_>) -> ControlFlow<Self::Break> {
925            let span = item.as_ref().item_meta().span;
926            let old_span = mem::replace(self.0.0.current_span(), span);
927            item.drive_mut(self)?;
928            *self.0.0.current_span() = old_span;
929            Continue(())
930        }
931    }
932
933    impl<V: Visitor> Visitor for VisitWithSpan<V> {
934        type Break = V::Break;
935    }
936    impl<V: VisitAst + VisitorWithSpan> VisitWithSpan<V> {
937        fn visit_inner_track_span<T>(&mut self, x: &T, span: Span) -> ControlFlow<V::Break>
938        where
939            T: AstVisitable,
940            T: for<'s> derive_generic_visitor::Drive<'s, AstVisitableWrapper<Self>>,
941        {
942            let old_span = mem::replace(self.0.current_span(), span);
943            self.visit_inner(x)?;
944            *self.0.current_span() = old_span;
945            Continue(())
946        }
947    }
948    impl<V: VisitAstMut + VisitorWithSpan> VisitWithSpan<V> {
949        fn visit_inner_mut_track_span<T>(&mut self, x: &mut T, span: Span) -> ControlFlow<V::Break>
950        where
951            T: AstVisitable,
952            T: for<'s> derive_generic_visitor::DriveMut<'s, AstVisitableWrapper<Self>>,
953        {
954            let old_span = mem::replace(self.0.current_span(), span);
955            self.visit_inner(x)?;
956            *self.0.current_span() = old_span;
957            Continue(())
958        }
959    }
960    impl<V: VisitAst + VisitorWithSpan> VisitAst for VisitWithSpan<V> {
961        fn visit_inner<T>(&mut self, x: &T) -> ControlFlow<Self::Break>
962        where
963            T: AstVisitable,
964        {
965            x.drive(self.inner())
966        }
967        fn visit_trait_param(&mut self, x: &TraitParam) -> ControlFlow<Self::Break> {
968            match x.span {
969                Some(span) => self.visit_inner_track_span(x, span),
970                None => self.visit_inner(x),
971            }
972        }
973        fn visit_ullbc_statement(&mut self, x: &ullbc_ast::Statement) -> ControlFlow<Self::Break> {
974            self.visit_inner_track_span(x, x.span)
975        }
976        fn visit_ullbc_terminator(
977            &mut self,
978            x: &ullbc_ast::Terminator,
979        ) -> ControlFlow<Self::Break> {
980            self.visit_inner_track_span(x, x.span)
981        }
982        fn visit_llbc_statement(&mut self, x: &llbc_ast::Statement) -> ControlFlow<Self::Break> {
983            self.visit_inner_track_span(x, x.span)
984        }
985        fn visit_llbc_block(&mut self, x: &llbc_ast::Block) -> ControlFlow<Self::Break> {
986            self.visit_inner_track_span(x, x.span)
987        }
988    }
989    impl<V: VisitAstMut + VisitorWithSpan> VisitAstMut for VisitWithSpan<V> {
990        fn visit_inner<T>(&mut self, x: &mut T) -> ControlFlow<Self::Break>
991        where
992            T: AstVisitable,
993        {
994            x.drive_mut(self.inner())
995        }
996        fn visit_trait_param(&mut self, x: &mut TraitParam) -> ControlFlow<Self::Break> {
997            match x.span {
998                Some(span) => self.visit_inner_mut_track_span(x, span),
999                None => self.visit_inner(x),
1000            }
1001        }
1002        fn visit_ullbc_statement(
1003            &mut self,
1004            x: &mut ullbc_ast::Statement,
1005        ) -> ControlFlow<Self::Break> {
1006            self.visit_inner_mut_track_span(x, x.span)
1007        }
1008        fn visit_ullbc_terminator(
1009            &mut self,
1010            x: &mut ullbc_ast::Terminator,
1011        ) -> ControlFlow<Self::Break> {
1012            self.visit_inner_mut_track_span(x, x.span)
1013        }
1014        fn visit_llbc_statement(
1015            &mut self,
1016            x: &mut llbc_ast::Statement,
1017        ) -> ControlFlow<Self::Break> {
1018            self.visit_inner_mut_track_span(x, x.span)
1019        }
1020        fn visit_llbc_block(&mut self, x: &mut llbc_ast::Block) -> ControlFlow<Self::Break> {
1021            self.visit_inner_mut_track_span(x, x.span)
1022        }
1023    }
1024
1025    /// Combo impls to be able to use some wrappers together.
1026    impl<V: VisitorWithSpan> VisitorWithSpan for VisitWithBinderStack<V> {
1027        fn current_span(&mut self) -> &mut Span {
1028            self.0.current_span()
1029        }
1030    }
1031    impl<V: VisitorWithSpan> VisitorWithSpan for VisitWithItem<V> {
1032        fn current_span(&mut self) -> &mut Span {
1033            self.0.current_span()
1034        }
1035    }
1036    impl<V: VisitorWithSpan> VisitorWithSpan for DontLeakImplDetails<V> {
1037        fn current_span(&mut self) -> &mut Span {
1038            self.0.current_span()
1039        }
1040    }
1041    impl<V: VisitorWithBinderDepth> VisitorWithBinderDepth for VisitWithItemRef<V> {
1042        fn binder_depth_mut(&mut self) -> &mut DeBruijnId {
1043            self.0.binder_depth_mut()
1044        }
1045    }
1046}
charon_lib/ast/visitor.rs

charon_lib/ast/
visitor.rs
