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

charon_lib/ast/
visitor.rs
