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

charon_lib/ast/
visitor.rs
