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        AbortKind, Assert, BinOp, BorrowKind, BuiltinAssertKind, BuiltinFunId, BuiltinIndexOp, BuiltinTy,
47        Call, CastKind, ClosureInfo, ClosureKind, ConstGenericParam, ConstGenericVarId,
48        Disambiguator, DynPredicate, Field, FieldId, FieldProjKind, 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,
53        TranslatedCrate, TypeDeclKind, TypeId, TypeParam, 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,
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, TraitMethodRef, 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        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 tracks the stack of binders seen so far. See [`VisitWithBinderDepth`] for how to use.
495    #[repr(transparent)]
496    pub struct VisitWithBinderStack<V>(V);
497
498    impl<V: VisitorWithBinderStack> VisitWithBinderStack<V> {
499        // Helper
500        fn wrap(v: &mut V) -> &mut Self {
501            // SAFETY: `repr(transparent)`
502            unsafe { std::mem::transmute(v) }
503        }
504        pub fn new(v: &mut V) -> &mut VisitWithItem<DontLeakImplDetails<Self>> {
505            // Use the `WithItem` wrapper to simplify the implementation of this wrapper. We use
506            // `DontLeakImplDetails` to use the specific `VisitorWithItem` impl we care about
507            // instead of the one that forwards to the `VisitorWithItem` of the containted `V`.
508            VisitWithItem::new(DontLeakImplDetails::new(Self::wrap(v)))
509        }
510        pub fn inner(&mut self) -> &mut V {
511            // SAFETY: `repr(transparent)`
512            unsafe { std::mem::transmute(self) }
513        }
514    }
515
516    pub trait VisitorWithBinderStack {
517        fn binder_stack_mut(&mut self) -> &mut BindingStack<GenericParams>;
518    }
519
520    impl<V: VisitAst + VisitorWithBinderStack> VisitorWithItem
521        for DontLeakImplDetails<VisitWithBinderStack<V>>
522    {
523        fn enter_item(&mut self, item: ItemRef<'_>) {
524            self.0
525                .0
526                .binder_stack_mut()
527                .push(item.generic_params().clone());
528        }
529        fn exit_item(&mut self, _item: ItemRef<'_>) {
530            self.0.0.binder_stack_mut().pop();
531        }
532    }
533    impl<V: VisitAstMut + VisitorWithBinderStack> VisitorWithItemMut
534        for DontLeakImplDetails<VisitWithBinderStack<V>>
535    {
536        fn enter_item(&mut self, item: ItemRefMut<'_>) {
537            self.0
538                .0
539                .binder_stack_mut()
540                .push(item.as_ref().generic_params().clone());
541        }
542        fn exit_item(&mut self, _item: ItemRefMut<'_>) {
543            self.0.0.binder_stack_mut().pop();
544        }
545    }
546
547    impl<V: Visitor> Visitor for VisitWithBinderStack<V> {
548        type Break = V::Break;
549    }
550    impl<V: VisitAst + VisitorWithBinderStack> VisitAst for VisitWithBinderStack<V> {
551        fn visit_inner<T>(&mut self, x: &T) -> ControlFlow<Self::Break>
552        where
553            T: AstVisitable,
554        {
555            x.drive(self.inner())
556        }
557        fn visit_binder<T: AstVisitable>(
558            &mut self,
559            binder: &Binder<T>,
560        ) -> ControlFlow<Self::Break> {
561            self.0.binder_stack_mut().push(binder.params.clone());
562            self.visit_inner(binder)?;
563            self.0.binder_stack_mut().pop();
564            Continue(())
565        }
566        fn visit_region_binder<T: AstVisitable>(
567            &mut self,
568            binder: &RegionBinder<T>,
569        ) -> ControlFlow<Self::Break> {
570            self.0.binder_stack_mut().push(GenericParams {
571                regions: binder.regions.clone(),
572                ..Default::default()
573            });
574            self.visit_inner(binder)?;
575            self.0.binder_stack_mut().pop();
576            Continue(())
577        }
578    }
579    impl<V: VisitAstMut + VisitorWithBinderStack> VisitAstMut for VisitWithBinderStack<V> {
580        fn visit_inner<T>(&mut self, x: &mut T) -> ControlFlow<Self::Break>
581        where
582            T: AstVisitable,
583        {
584            x.drive_mut(self.inner())
585        }
586        fn visit_binder<T: AstVisitable>(
587            &mut self,
588            binder: &mut Binder<T>,
589        ) -> ControlFlow<Self::Break> {
590            self.0.binder_stack_mut().push(binder.params.clone());
591            self.visit_inner(binder)?;
592            self.0.binder_stack_mut().pop();
593            Continue(())
594        }
595        fn visit_region_binder<T: AstVisitable>(
596            &mut self,
597            binder: &mut RegionBinder<T>,
598        ) -> ControlFlow<Self::Break> {
599            self.0.binder_stack_mut().push(GenericParams {
600                regions: binder.regions.clone(),
601                ..Default::default()
602            });
603            self.visit_inner(binder)?;
604            self.0.binder_stack_mut().pop();
605            Continue(())
606        }
607    }
608
609    /// Visitor wrapper that tracks the current span. See [`VisitWithBinderDepth`] for how to use.
610    #[repr(transparent)]
611    pub struct VisitWithSpan<V>(V);
612
613    impl<V: VisitorWithSpan> VisitWithSpan<V> {
614        // Helper
615        fn wrap(v: &mut V) -> &mut Self {
616            // SAFETY: `repr(transparent)`
617            unsafe { std::mem::transmute(v) }
618        }
619        pub fn new(v: &mut V) -> &mut VisitWithItem<DontLeakImplDetails<Self>> {
620            // Use the `WithItem` wrapper to simplify the implementation of this wrapper. We use
621            // `DontLeakImplDetails` to use the specific `VisitorWithItem` impl we care about
622            // instead of the one that forwards to the `VisitorWithItem` of the containted `V`.
623            VisitWithItem::new(DontLeakImplDetails::new(Self::wrap(v)))
624        }
625        pub fn inner(&mut self) -> &mut V {
626            // SAFETY: `repr(transparent)`
627            unsafe { std::mem::transmute(self) }
628        }
629    }
630
631    pub trait VisitorWithSpan {
632        fn current_span(&mut self) -> &mut Span;
633    }
634
635    impl<V: VisitAst + VisitorWithSpan> VisitorWithItem for DontLeakImplDetails<VisitWithSpan<V>> {
636        fn visit_item(&mut self, item: ItemRef<'_>) -> ControlFlow<Self::Break> {
637            let old_span = mem::replace(self.0.0.current_span(), item.item_meta().span);
638            item.drive(self)?;
639            *self.0.0.current_span() = old_span;
640            Continue(())
641        }
642    }
643    impl<V: VisitAstMut + VisitorWithSpan> VisitorWithItemMut
644        for DontLeakImplDetails<VisitWithSpan<V>>
645    {
646        fn visit_item(&mut self, mut item: ItemRefMut<'_>) -> ControlFlow<Self::Break> {
647            let span = item.as_ref().item_meta().span;
648            let old_span = mem::replace(self.0.0.current_span(), span);
649            item.drive_mut(self)?;
650            *self.0.0.current_span() = old_span;
651            Continue(())
652        }
653    }
654
655    impl<V: Visitor> Visitor for VisitWithSpan<V> {
656        type Break = V::Break;
657    }
658    impl<V: VisitAst + VisitorWithSpan> VisitWithSpan<V> {
659        fn visit_inner_track_span<T>(&mut self, x: &T, span: Span) -> ControlFlow<V::Break>
660        where
661            T: AstVisitable,
662            T: for<'s> derive_generic_visitor::Drive<'s, AstVisitableWrapper<Self>>,
663        {
664            let old_span = mem::replace(self.0.current_span(), span);
665            self.visit_inner(x)?;
666            *self.0.current_span() = old_span;
667            Continue(())
668        }
669    }
670    impl<V: VisitAstMut + VisitorWithSpan> VisitWithSpan<V> {
671        fn visit_inner_mut_track_span<T>(&mut self, x: &mut T, span: Span) -> ControlFlow<V::Break>
672        where
673            T: AstVisitable,
674            T: for<'s> derive_generic_visitor::DriveMut<'s, AstVisitableWrapper<Self>>,
675        {
676            let old_span = mem::replace(self.0.current_span(), span);
677            self.visit_inner(x)?;
678            *self.0.current_span() = old_span;
679            Continue(())
680        }
681    }
682    impl<V: VisitAst + VisitorWithSpan> VisitAst for VisitWithSpan<V> {
683        fn visit_inner<T>(&mut self, x: &T) -> ControlFlow<Self::Break>
684        where
685            T: AstVisitable,
686        {
687            x.drive(self.inner())
688        }
689        fn visit_trait_param(&mut self, x: &TraitParam) -> ControlFlow<Self::Break> {
690            match x.span {
691                Some(span) => self.visit_inner_track_span(x, span),
692                None => self.visit_inner(x),
693            }
694        }
695        fn visit_ullbc_statement(&mut self, x: &ullbc_ast::Statement) -> ControlFlow<Self::Break> {
696            self.visit_inner_track_span(x, x.span)
697        }
698        fn visit_ullbc_terminator(
699            &mut self,
700            x: &ullbc_ast::Terminator,
701        ) -> ControlFlow<Self::Break> {
702            self.visit_inner_track_span(x, x.span)
703        }
704        fn visit_llbc_statement(&mut self, x: &llbc_ast::Statement) -> ControlFlow<Self::Break> {
705            self.visit_inner_track_span(x, x.span)
706        }
707        fn visit_llbc_block(&mut self, x: &llbc_ast::Block) -> ControlFlow<Self::Break> {
708            self.visit_inner_track_span(x, x.span)
709        }
710    }
711    impl<V: VisitAstMut + VisitorWithSpan> VisitAstMut for VisitWithSpan<V> {
712        fn visit_inner<T>(&mut self, x: &mut T) -> ControlFlow<Self::Break>
713        where
714            T: AstVisitable,
715        {
716            x.drive_mut(self.inner())
717        }
718        fn visit_trait_param(&mut self, x: &mut TraitParam) -> ControlFlow<Self::Break> {
719            match x.span {
720                Some(span) => self.visit_inner_mut_track_span(x, span),
721                None => self.visit_inner(x),
722            }
723        }
724        fn visit_ullbc_statement(
725            &mut self,
726            x: &mut ullbc_ast::Statement,
727        ) -> ControlFlow<Self::Break> {
728            self.visit_inner_mut_track_span(x, x.span)
729        }
730        fn visit_ullbc_terminator(
731            &mut self,
732            x: &mut ullbc_ast::Terminator,
733        ) -> ControlFlow<Self::Break> {
734            self.visit_inner_mut_track_span(x, x.span)
735        }
736        fn visit_llbc_statement(
737            &mut self,
738            x: &mut llbc_ast::Statement,
739        ) -> ControlFlow<Self::Break> {
740            self.visit_inner_mut_track_span(x, x.span)
741        }
742        fn visit_llbc_block(&mut self, x: &mut llbc_ast::Block) -> ControlFlow<Self::Break> {
743            self.visit_inner_mut_track_span(x, x.span)
744        }
745    }
746
747    /// Combo impl to be able to use `VisitWithSpan` and `VisitWithBinderStack` together.
748    impl<V: VisitorWithSpan> VisitorWithSpan for VisitWithBinderStack<V> {
749        fn current_span(&mut self) -> &mut Span {
750            self.0.current_span()
751        }
752    }
753    impl<V: VisitorWithSpan> VisitorWithSpan for VisitWithItem<V> {
754        fn current_span(&mut self) -> &mut Span {
755            self.0.current_span()
756        }
757    }
758    impl<V: VisitorWithSpan> VisitorWithSpan for DontLeakImplDetails<V> {
759        fn current_span(&mut self) -> &mut Span {
760            self.0.current_span()
761        }
762    }
763}
charon_lib/ast/visitor.rs

charon_lib/ast/
visitor.rs
