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

charon_lib/ast/
visitor.rs
