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

charon_lib/ast/
visitor.rs
