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

charon_lib/ast/
visitor.rs
