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

charon_lib/ast/
visitor.rs
