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

charon_lib/ast/
visitor.rs
