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, ConstGenericVar, ConstGenericVarId,
46        Disambiguator, DynPredicate, Field, FieldId, FieldProjKind, FloatTy, FloatValue,
47        FnOperand, FunId, FunIdOrTraitMethodRef, FunSig, ImplElem, IntegerTy, IntTy, UIntTy, Literal, LiteralTy,
48        llbc_ast::Block, llbc_ast::ExprBody, llbc_ast::RawStatement, llbc_ast::Switch,
49        Locals, Name, NullOp, Opaque, Operand, PathElem, PlaceKind, ProjectionElem, RawConstantExpr,
50        RefKind, RegionId, RegionVar, ScalarValue, TraitItemName,
51        TranslatedCrate, TypeDeclKind, TypeId, TypeVar, TypeVarId, llbc_ast::StatementId,
52        ullbc_ast::BlockData, ullbc_ast::BlockId, ullbc_ast::ExprBody, ullbc_ast::RawStatement,
53        ullbc_ast::RawTerminator, ullbc_ast::SwitchTargets, ullbc_ast::Terminator,
54        UnOp, UnsizingMetadata, Local, Variant, VariantId, LocalId, CopyNonOverlapping, Layout, VariantLayout, PtrMetadata, VTable,
55        for<T: AstVisitable> Box<T>,
56        for<T: AstVisitable> Option<T>,
57        for<A: AstVisitable, B: AstVisitable> (A, B),
58        for<A: AstVisitable, B: AstVisitable, C: AstVisitable> (A, B, C),
59        for<A: AstVisitable, B: AstVisitable> Result<A, B>,
60        for<A: AstVisitable, B: AstVisitable> OutlivesPred<A, B>,
61        for<T: AstVisitable> Vec<T>,
62        for<I: Idx, T: AstVisitable> Vector<I, T>,
63    ),
64    // Types for which we call the corresponding `visit_$ty` method, which by default explores the
65    // type but can be overridden.
66    override(
67        DeBruijnId, Ty, TyKind, Region, ConstGeneric, TraitRef, TraitRefKind,
68        TypeDeclRef, FunDeclRef, MaybeBuiltinFunDeclRef, TraitMethodRef, GlobalDeclRef, TraitDeclRef, TraitImplRef,
69        GenericArgs, GenericParams, TraitClause, TraitClauseId, TraitTypeConstraint, Place, Rvalue,
70        for<T: AstVisitable + Idx> DeBruijnVar<T>,
71        for<T: AstVisitable> RegionBinder<T>,
72        for<T: AstVisitable> Binder<T>,
73        llbc_statement: llbc_ast::Statement, ullbc_statement: ullbc_ast::Statement,
74        AggregateKind, FnPtr, ItemKind, ItemMeta, Span, ConstantExpr,
75        FunDeclId, GlobalDeclId, TypeDeclId, TraitDeclId, TraitImplId,
76        FunDecl, GlobalDecl, TypeDecl, TraitDecl, TraitImpl,
77    )
78)]
79pub trait AstVisitable: Any {
80    /// The name of the type, used for debug logging.
81    fn name(&self) -> &'static str {
82        std::any::type_name::<Self>()
83    }
84    /// Visit all occurrences of that type inside `self`, in pre-order traversal.
85    fn dyn_visit<T: AstVisitable>(&self, f: impl FnMut(&T)) {
86        let _ = self.drive(&mut DynVisitor::new_shared::<T>(f));
87    }
88    /// Visit all occurrences of that type inside `self`, in pre-order traversal.
89    fn dyn_visit_mut<T: AstVisitable>(&mut self, f: impl FnMut(&mut T)) {
90        let _ = self.drive_mut(&mut DynVisitor::new_mut::<T>(f));
91    }
92}
93
94/// Manual impl that only visits the values
95impl<K: Any, T: AstVisitable> AstVisitable for HashMap<K, T> {
96    fn drive<V: VisitAst>(&self, v: &mut V) -> ControlFlow<V::Break> {
97        for x in self.values() {
98            v.visit(x)?;
99        }
100        Continue(())
101    }
102    fn drive_mut<V: VisitAstMut>(&mut self, v: &mut V) -> ControlFlow<V::Break> {
103        for x in self.values_mut() {
104            v.visit(x)?;
105        }
106        Continue(())
107    }
108}
109
110/// Manual impl that only visits the values
111impl<K: Any, T: AstVisitable> AstVisitable for IndexMap<K, T> {
112    fn drive<V: VisitAst>(&self, v: &mut V) -> ControlFlow<V::Break> {
113        for x in self.values() {
114            v.visit(x)?;
115        }
116        Continue(())
117    }
118    fn drive_mut<V: VisitAstMut>(&mut self, v: &mut V) -> ControlFlow<V::Break> {
119        for x in self.values_mut() {
120            v.visit(x)?;
121        }
122        Continue(())
123    }
124}
125
126/// A smaller visitor group just for function bodies. This explores statements, places and
127/// operands, but does not recurse into types.
128///
129/// This defines three traits:
130/// - `BodyVisitable` is a trait implemented by all the types listed below; it has a
131/// `drive_body[_mut]` method that takes a `VisitBody[Mut]` visitor and calls its methods on all
132/// the relevant subvalues of `self` encountered.
133/// - `VisitBody[Mut]` is a (pair of) visitor trait(s) that can be implemented by visitors. To
134/// define a visitor, implement `VisitBody[Mut]` and override the methods you need.
135///
136/// Morally this represents the predicate `for<V: VisitBody[Mut]> Self:
137/// Drive[Mut]<BodyVisitableWrapper<V>>`
138#[visitable_group(
139    // Defines the `VisitBody[Mut]` traits and the `drive_body[_mut]` method that drives them.
140    visitor(drive_body(&VisitBody)),
141    visitor(drive_body_mut(&mut VisitBodyMut)),
142    // Types that are ignored when encountered.
143    skip(
144        AbortKind, BinOp, BorrowKind, ConstantExpr, ConstGeneric, FieldId, FieldProjKind,
145        TypeDeclRef, FunDeclId, FunIdOrTraitMethodRef, GenericArgs, GlobalDeclRef, IntegerTy, IntTy, UIntTy,
146        NullOp, RefKind, ScalarValue, Span, Ty, TypeDeclId, TypeId, UnOp, VariantId, LocalId,
147        TraitRef,
148    ),
149    // Types that we unconditionally explore.
150    drive(
151        Assert, PlaceKind,
152        llbc_ast::ExprBody, llbc_ast::RawStatement, llbc_ast::Switch,
153        ullbc_ast::BlockData, ullbc_ast::ExprBody, ullbc_ast::RawStatement,
154        ullbc_ast::RawTerminator, ullbc_ast::SwitchTargets, CopyNonOverlapping,
155        llbc_ast::StatementId,
156        Body, Opaque, Locals, Local,
157        for<T: BodyVisitable> Box<T>,
158        for<T: BodyVisitable> Option<T>,
159        for<T: BodyVisitable, E: BodyVisitable> Result<T, E>,
160        for<A: BodyVisitable, B: BodyVisitable> (A, B),
161        for<A: BodyVisitable, B: BodyVisitable, C: BodyVisitable> (A, B, C),
162        for<T: BodyVisitable> Vec<T>,
163        for<I: Idx, T: BodyVisitable> Vector<I, T>,
164    ),
165    // Types for which we call the corresponding `visit_$ty` method, which by default explores the
166    // type but can be overridden.
167    override(
168        AggregateKind, Call, FnOperand, FnPtr,
169        Operand, Place, ProjectionElem, Rvalue,
170        llbc_block: llbc_ast::Block,
171        llbc_statement: llbc_ast::Statement,
172        ullbc_statement: ullbc_ast::Statement,
173        ullbc_terminator: ullbc_ast::Terminator,
174        ullbc_block_id: ullbc_ast::BlockId,
175    )
176)]
177pub trait BodyVisitable: Any {
178    /// Visit all occurrences of that type inside `self`, in pre-order traversal.
179    fn dyn_visit_in_body<T: BodyVisitable>(&self, f: impl FnMut(&T)) {
180        let _ = self.drive_body(&mut DynVisitor::new_shared::<T>(f));
181    }
182
183    /// Visit all occurrences of that type inside `self`, in pre-order traversal.
184    fn dyn_visit_in_body_mut<T: BodyVisitable>(&mut self, f: impl FnMut(&mut T)) {
185        let _ = self.drive_body_mut(&mut DynVisitor::new_mut::<T>(f));
186    }
187}
188
189/// Ast and body visitor that uses dynamic dispatch to call the provided function on the visited
190/// values of the right type.
191#[derive(Visitor)]
192pub struct DynVisitor<F> {
193    enter: F,
194}
195impl DynVisitor<()> {
196    pub fn new_shared<T: Any>(mut f: impl FnMut(&T)) -> DynVisitor<impl FnMut(&dyn Any)> {
197        let enter = move |x: &dyn Any| {
198            if let Some(x) = x.downcast_ref::<T>() {
199                f(x);
200            }
201        };
202        DynVisitor { enter }
203    }
204    pub fn new_mut<T: Any>(mut f: impl FnMut(&mut T)) -> DynVisitor<impl FnMut(&mut dyn Any)> {
205        let enter = move |x: &mut dyn Any| {
206            if let Some(x) = x.downcast_mut::<T>() {
207                f(x);
208            }
209        };
210        DynVisitor { enter }
211    }
212}
213impl<F> VisitAst for DynVisitor<F>
214where
215    F: FnMut(&dyn Any),
216{
217    fn visit<'a, T: AstVisitable>(&'a mut self, x: &T) -> ControlFlow<Self::Break> {
218        (self.enter)(x);
219        x.drive(self)?;
220        Continue(())
221    }
222}
223impl<F> VisitAstMut for DynVisitor<F>
224where
225    F: FnMut(&mut dyn Any),
226{
227    fn visit<'a, T: AstVisitable>(&'a mut self, x: &mut T) -> ControlFlow<Self::Break> {
228        (self.enter)(x);
229        x.drive_mut(self)?;
230        Continue(())
231    }
232}
233impl<F> VisitBody for DynVisitor<F>
234where
235    F: FnMut(&dyn Any),
236{
237    fn visit<'a, T: BodyVisitable>(&'a mut self, x: &T) -> ControlFlow<Self::Break> {
238        (self.enter)(x);
239        x.drive_body(self)?;
240        Continue(())
241    }
242}
243impl<F> VisitBodyMut for DynVisitor<F>
244where
245    F: FnMut(&mut dyn Any),
246{
247    fn visit<'a, T: BodyVisitable>(&'a mut self, x: &mut T) -> ControlFlow<Self::Break> {
248        (self.enter)(x);
249        x.drive_body_mut(self)?;
250        Continue(())
251    }
252}
charon_lib/ast/visitor.rs

charon_lib/ast/
visitor.rs
