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

charon_lib/ast/
visitor.rs
