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

charon_lib/ast/
visitor.rs
