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

charon_lib/ast/
visitor.rs
