hax_rust_engine/ast/
visitors.rs

1//! Syntax tree traversals to walk a shared or mutable borrow of the syntax tree
2//! of Hax. The visitors are generated using the [`derive_generic_visitor`]
3//! library.
4//!
5//! This module provides visitors of different flavors of visitors, and visitor
6//! wrappers that can enhance the default behavior of a visitor.
7//!
8//! We provide four main visitors.
9//!  - [`AstVisitor`] and [`AstVisitorMut`]: visitor that never early exit.
10//!  - [`AstEarlyExitVisitor`] and [`AstEarlyExitVisitorMut`]: visitor that can early exit.
11//!
12//! Each trait provides methods `visit_expr`, `visit_ty`, etc. enabling easy AST
13//! traversal.
14//!
15//! Importantly, we also provide visitor wrappers that enhance visitors with
16//! common useful behavior. See the module [`wrappers`] for more information.
17
18use super::*;
19use derive_generic_visitor::*;
20
21pub mod wrappers {
22    //! This module provides a visitor wrappers, or transformer of visitors.
23    //! Such wrappers transform the behavior of a visitor.
24    //!
25    //! For example, [`SpanWrapper`] takes care of keeping track of [`Span`]s
26    //! while travesing an AST.
27
28    use std::ops::Deref;
29
30    use super::{infallible::AstVisitable as AstVisitableInfallible, *};
31    use diagnostics::*;
32
33    /// A visitor wrapper that tracks span while visiting the AST. Whenever an
34    /// AST node that carries a span is visited, using this wrapper, the ambient
35    /// span is mutated and accessible via the `HasSpan` trait.
36    pub struct SpanWrapper<'a, V>(pub &'a mut V);
37
38    impl<'a, V: HasSpan> SpanWrapper<'a, V> {
39        /// Performs a spanned action: calls the function `action` on
40        /// `ast_fragment`, with the contextual span information in `self` being
41        /// the span found in `ast_fragment`.
42        fn spanned_action<T: Deref, U>(
43            &mut self,
44            ast_fragment: T,
45            action: impl Fn(&mut Self, T) -> U,
46        ) -> U
47        where
48            T::Target: HasSpan,
49        {
50            let span_before = self.0.span();
51            *self.0.span_mut() = ast_fragment.span();
52            // Perform the provided action on `ast_fragment` with `ast_fragment`'s span as contextual span.
53            let result = action(self, ast_fragment);
54            *self.0.span_mut() = span_before;
55            result
56        }
57    }
58
59    impl<'a, V: AstVisitorMut + HasSpan> AstVisitorMut for SpanWrapper<'a, V> {
60        fn visit_inner<T>(&mut self, x: &mut T)
61        where
62            T: AstVisitableInfallible,
63            T: for<'s> DriveMut<'s, AstVisitableInfallibleWrapper<Self>>,
64        {
65            x.drive_map(self.0)
66        }
67        fn visit_item(&mut self, x: &mut Item) {
68            self.spanned_action(x, Self::visit_inner)
69        }
70        fn visit_expr(&mut self, x: &mut Expr) {
71            self.spanned_action(x, Self::visit_inner)
72        }
73        fn visit_pat(&mut self, x: &mut Pat) {
74            self.spanned_action(x, Self::visit_inner)
75        }
76        fn visit_guard(&mut self, x: &mut Guard) {
77            self.spanned_action(x, Self::visit_inner)
78        }
79        fn visit_arm(&mut self, x: &mut Arm) {
80            self.spanned_action(x, Self::visit_inner)
81        }
82        fn visit_impl_item(&mut self, x: &mut ImplItem) {
83            self.spanned_action(x, Self::visit_inner)
84        }
85        fn visit_trait_item(&mut self, x: &mut TraitItem) {
86            self.spanned_action(x, Self::visit_inner)
87        }
88        fn visit_generic_param(&mut self, x: &mut GenericParam) {
89            self.spanned_action(x, Self::visit_inner)
90        }
91        fn visit_attribute(&mut self, x: &mut Attribute) {
92            self.spanned_action(x, Self::visit_inner)
93        }
94        fn visit_spanned_ty(&mut self, x: &mut SpannedTy) {
95            self.spanned_action(x, Self::visit_inner)
96        }
97    }
98
99    /// A visitor wrapper that automatically collects errors in `ErrorNode`s.
100    /// Coupled with the trait `VisitorWithErrors`, this provides an `error`
101    /// method on a visitor that can be used to throw errors, which will be
102    /// automatically inlined in the AST on the closest error-capable node.
103    pub struct ErrorWrapper<'a, V>(pub &'a mut V);
104
105    /// An opaque error vault. This is the state manipulated by the visitor wrapper [`ErrorWrapper`].
106    /// It is purposefully not-inspectable.
107    #[derive(Default)]
108    pub struct ErrorVault(Vec<Diagnostic>);
109    impl ErrorVault {
110        fn add(&mut self, diagnostic: Diagnostic) {
111            self.0.push(diagnostic);
112        }
113    }
114
115    /// Helper struct that contains error-handling related state information.
116    /// This is used internally by [`setup_error_handling_struct`].
117    pub struct ErrorHandlingState(pub Span, pub ErrorVault);
118    impl Default for ErrorHandlingState {
119        fn default() -> Self {
120            Self(Span::dummy(), Default::default())
121        }
122    }
123
124    #[macro_export]
125    /// Use this macro in an implementation of `AstVisitorMut` to get automatic spans and error handling.
126    macro_rules! setup_error_handling_impl {
127        () => {
128            fn visit<'a, T: $crate::ast::visitors::AstVisitableInfallible>(
129                &'a mut self,
130                x: &mut T,
131            ) {
132                $crate::ast::visitors::wrappers::SpanWrapper(
133                    &mut $crate::ast::visitors::wrappers::ErrorWrapper(self),
134                )
135                .visit(x)
136            }
137        };
138    }
139    pub use setup_error_handling_impl;
140
141    /// Mark a visitor with a specific diagnostic context.
142    pub trait VisitorWithContext {
143        /// Returns the diagnostic context for this visitor.
144        fn context(&self) -> Context;
145    }
146
147    impl<T: HasSpan> HasSpan for ErrorWrapper<'_, T> {
148        fn span(&self) -> Span {
149            self.0.span()
150        }
151
152        fn span_mut(&mut self) -> &mut Span {
153            self.0.span_mut()
154        }
155    }
156
157    /// A visitor that can throw errors. It should be used in combination with
158    /// `ErrorWrapper`, which will take care of bubbling error up to the nearest
159    /// parent capable of representing errors. For instance, if you error out in
160    /// a literal, the error will be represented in the parent expression or the
161    /// parent type, as nodes [`ExprKind::Error`] or [`TyKind ::Error`].
162    pub trait VisitorWithErrors: HasSpan + VisitorWithContext {
163        /// Projects the error vault.
164        fn error_vault(&mut self) -> &mut ErrorVault;
165        /// Send an error.
166        fn error(&mut self, node: impl Into<Fragment>, kind: DiagnosticInfoKind) {
167            let context = self.context();
168            let span = self.span();
169            self.error_vault().add(Diagnostic::new(
170                node,
171                DiagnosticInfo {
172                    context,
173                    span,
174                    kind,
175                },
176            ));
177        }
178    }
179
180    impl<'a, V: VisitorWithErrors> ErrorWrapper<'a, V> {
181        fn error_handled_action<
182            T: FallibleAstNode + Clone + std::fmt::Debug + Into<Fragment>,
183            U,
184        >(
185            &mut self,
186            x: &mut T,
187            action: impl Fn(&mut Self, &mut T) -> U,
188        ) -> U {
189            let diagnostics_snapshot = self.0.error_vault().0.clone();
190            self.0.error_vault().0.clear();
191            let result = action(self, x);
192            let diagnostics: Vec<_> = self.0.error_vault().0.drain(..).collect();
193            if !diagnostics.is_empty() {
194                x.set_error(ErrorNode {
195                    fragment: Box::new(x.clone().into()),
196                    diagnostics,
197                });
198            }
199            self.0.error_vault().0 = diagnostics_snapshot;
200            result
201        }
202    }
203
204    impl<'a, V: AstVisitorMut + VisitorWithErrors> AstVisitorMut for ErrorWrapper<'a, V> {
205        fn visit_inner<T>(&mut self, x: &mut T)
206        where
207            T: AstVisitableInfallible,
208            T: for<'s> DriveMut<'s, AstVisitableInfallibleWrapper<Self>>,
209        {
210            x.drive_map(self.0)
211        }
212        fn visit_item(&mut self, x: &mut Item) {
213            self.error_handled_action(x, Self::visit_inner)
214        }
215        fn visit_pat(&mut self, x: &mut Pat) {
216            self.error_handled_action(x, Self::visit_inner)
217        }
218        fn visit_expr(&mut self, x: &mut Expr) {
219            self.error_handled_action(x, Self::visit_inner)
220        }
221        fn visit_ty(&mut self, x: &mut Ty) {
222            self.error_handled_action(x, Self::visit_inner)
223        }
224    }
225}
226
227#[hax_rust_engine_macros::replace(AstNodes => include(VisitableAstNodes))]
228mod replaced {
229    use super::*;
230    pub mod infallible {
231        use super::*;
232
233        #[visitable_group(
234            visitor(drive_map(
235                /// An mutable visitor that visits the AST for hax.
236                /// 
237                /// ```rust,ignore
238                /// use crate::ast::{diagnostics::*, visitors::*};
239                /// #[setup_error_handling_struct]
240                /// #[derive(Default)]
241                /// struct MyVisitor;
242                ///
243                /// impl VisitorWithContext for MyVisitor {
244                ///     fn context(&self) -> Context {
245                ///         Context::Import
246                ///     }
247                /// }
248                ///
249                /// impl AstVisitorMut for MyVisitor {
250                ///     setup_error_handling_impl!();
251                /// }
252                /// 
253                /// // MyVisitor::visit(my_ast_node)
254                /// ```
255                &mut AstVisitorMut
256            ), infallible),
257            visitor(drive(
258                /// An immutable visitor that visits the AST for hax.
259                &AstVisitor
260            ), infallible),
261            skip(
262                String, bool, char, hax_frontend_exporter::Span,
263            ),
264            drive(
265                for<T: AstVisitable> Box<T>, for<T: AstVisitable> Option<T>, for<T: AstVisitable> Vec<T>,
266                for<A: AstVisitable, B: AstVisitable> (A, B),
267                for<A: AstVisitable, B: AstVisitable, C: AstVisitable> (A, B, C),
268                usize
269            ),
270            override(AstNodes),
271            override_skip(
272                Span, Fragment, GlobalId, Diagnostic,
273            ),
274        )]
275        /// Helper trait to drive visitor.
276        pub trait AstVisitable {}
277    }
278
279    #[allow(missing_docs)]
280    pub mod fallible {
281        use super::*;
282
283        #[visitable_group(
284            visitor(drive(
285                /// An immutable visitor that can exit early.
286                &AstEarlyExitVisitor
287            )),
288            visitor(drive_mut(
289                /// An immutable visitor that can exit early and mutate the AST fragments.
290                &mut AstEarlyExitVisitorMut
291            )),
292            skip(
293                String, bool, char, hax_frontend_exporter::Span,
294            ),
295            drive(
296                for<T: AstVisitable> Box<T>, for<T: AstVisitable> Option<T>, for<T: AstVisitable> Vec<T>,
297                for<A: AstVisitable, B: AstVisitable> (A, B),
298                for<A: AstVisitable, B: AstVisitable, C: AstVisitable> (A, B, C),
299                usize
300            ),
301            override(AstNodes),
302            override_skip(
303                Span, Fragment, GlobalId, Diagnostic,
304            ),
305        )]
306        /// Helper trait to drive visitor.
307        pub trait AstVisitable {}
308    }
309
310    /// This modules provides `dyn` compatible trait for visitors.
311    pub mod dyn_compatible {
312        use super::*;
313
314        macro_rules! derive_erased_ast_visitors {
315            ({$($attrs:tt)*}, $name: ident, $helper: ident, $($ty:ty),*) => {
316                $($attrs)*
317                pub trait $name<'a>: $($helper<'a, $ty> + )* {}
318            };
319        }
320
321        macro_rules! render_path {
322            ($head:ident) => {stringify!($head)};
323            ($head:ident $(::$tail:ident)*) => {
324                concat!(stringify!($head), "::", render_path!($($tail)::*))
325            };
326        }
327
328        macro_rules! make_dyn_compatible {
329            ($($visitable_trait:ident)::*, $($visitor_trait:ident)::*, $helper_name: ident, $name: ident, mut:{$($mut:tt)?}, super:{$($super:ident)::*}, $ret:ty) => {
330                #[doc = concat!("A [dyn-compatible](https://doc.rust-lang.org/reference/items/traits.html#dyn-compatibility) trait similar to [`", render_path!($($visitor_trait)::*),"`].")]
331                #[doc = concat!("This trait provides one `visit` method to visit a given type `T` with a given visitor.")]
332                pub trait $helper_name<'a, T: ?Sized>: $($super)::* {
333                    /// Visit a value with the visitor.
334                    fn visit(&mut self, _: &'a $($mut)? T) -> $ret;
335                }
336
337                impl<'a, T: $($visitable_trait)::*, V: $($visitor_trait)::*> $helper_name<'a, T> for V {
338                    fn visit(&mut self, e: &'a $($mut)? T) -> $ret {
339                        <Self as $($visitor_trait)::*>::visit(self, e)
340                    }
341                }
342                derive_erased_ast_visitors!({
343                    #[doc = concat!("A [dyn-compatible](https://doc.rust-lang.org/reference/items/traits.html#dyn-compatibility) trait similar to [`", render_path!($($visitor_trait)::*),"`].")]
344                    #[doc = concat!("This trait is empty, but it implies a super bound for every type in the AST, so that you can use [`", stringify!($helper_name), "::visit", "`] with the entire AST.")]
345                }, $name, $helper_name, AstNodes);
346
347                impl<'a, V: $($visitor_trait)::*> $name<'a> for V {}
348            };
349        }
350
351        make_dyn_compatible!(
352            infallible::AstVisitable,
353            infallible::AstVisitorMut,
354            AstVisitableMut,
355            AstVisitorMut,
356            mut:{mut},
357            super:{},
358            ()
359        );
360        make_dyn_compatible!(
361            infallible::AstVisitable,
362            infallible::AstVisitor,
363            AstVisitable,
364            AstVisitor,
365            mut:{},
366            super:{},
367            ()
368        );
369
370        make_dyn_compatible!(
371            fallible::AstVisitable,
372            fallible::AstEarlyExitVisitorMut,
373            AstEarlyExitVisitableMut,
374            AstEarlyExitVisitorMut,
375            mut:{mut},
376            super:{Visitor},
377            ControlFlow<<Self as Visitor>::Break>
378        );
379        make_dyn_compatible!(
380            fallible::AstVisitable,
381            fallible::AstEarlyExitVisitor,
382            AstEarlyExitVisitable,
383            AstEarlyExitVisitor,
384            mut:{},
385            super:{Visitor},
386            ControlFlow<<Self as Visitor>::Break>
387        );
388    }
389}
390
391pub use replaced::dyn_compatible;
392use replaced::{fallible, infallible};
393
394pub use fallible::{
395    AstEarlyExitVisitor, AstEarlyExitVisitorMut, AstVisitable as AstVisitableFallible,
396    AstVisitableWrapper,
397};
398pub use hax_rust_engine_macros::setup_error_handling_struct;
399pub use infallible::{
400    AstVisitable as AstVisitableInfallible, AstVisitableInfallibleWrapper, AstVisitor,
401    AstVisitorMut,
402};
403pub use wrappers::{VisitorWithContext, VisitorWithErrors, setup_error_handling_impl};
404
405#[test]
406fn double_literals_in_ast() {
407    use crate::ast::diagnostics::*;
408
409    #[setup_error_handling_struct]
410    #[derive(Default)]
411    struct DoubleU8Literals;
412
413    impl VisitorWithContext for DoubleU8Literals {
414        fn context(&self) -> Context {
415            Context::Import
416        }
417    }
418
419    impl AstVisitorMut for DoubleU8Literals {
420        setup_error_handling_impl!();
421
422        fn visit_literal(&mut self, x: &mut Literal) {
423            let Literal::Int { value, .. } = x else {
424                return;
425            };
426            let Ok(n): Result<u8, _> = str::parse(value) else {
427                return self.error(
428                    x.clone(),
429                    DiagnosticInfoKind::AssertionFailure {
430                        details: "Bad literal".into(),
431                    },
432                );
433            };
434            let n = (n as u16) * 2;
435            if n >= u8::MAX as u16 {
436                return self.error(
437                    x.clone(),
438                    DiagnosticInfoKind::AssertionFailure {
439                        details: "Literal too big".into(),
440                    },
441                );
442            }
443            *value = Symbol::new(&format!("{}", n));
444        }
445    }
446
447    // Syntax helpers
448    let int_kind = IntKind {
449        size: IntSize::S8,
450        signedness: Signedness::Signed,
451    };
452    let mk_lit = |n: isize| Literal::Int {
453        value: Symbol::new(&format!("{}", n)),
454        negative: false,
455        kind: int_kind.clone(),
456    };
457    let meta = Metadata {
458        span: Span::dummy(),
459        attributes: vec![],
460    };
461    let mk_lit_expr = |n| Expr {
462        kind: Box::new(ExprKind::Literal(mk_lit(n))),
463        ty: Ty(Box::new(TyKind::Primitive(PrimitiveTy::Int(
464            int_kind.clone(),
465        )))),
466        meta: meta.clone(),
467    };
468    let mk_array = |exprs| Expr {
469        kind: Box::new(ExprKind::Array(exprs)),
470        ty: Ty(Box::new(TyKind::RawPointer)), // wrong type, but this is not important for this test.
471        meta: meta.clone(),
472    };
473    let mut lit_expr_200 = mk_lit_expr(200);
474
475    // Creates the expression `[50u8, 100u8, 200u8]`: the last one cannot be doubled, and will cause an error.
476    let mut e = mk_array(vec![
477        mk_lit_expr(50),
478        mk_lit_expr(100),
479        lit_expr_200.clone(),
480    ]);
481
482    // Visit the expression.
483    DoubleU8Literals::default().visit(&mut e);
484
485    // Transform `lit_expr_200` into the error `DoubleU8Literal` should produce
486    lit_expr_200.set_error(ErrorNode {
487        fragment: Box::new(lit_expr_200.clone().into()),
488        diagnostics: vec![Diagnostic::new(
489            mk_lit(200),
490            DiagnosticInfo {
491                span: lit_expr_200.span(),
492                context: Context::Import,
493                kind: DiagnosticInfoKind::AssertionFailure {
494                    details: "Literal too big".into(),
495                },
496            },
497        )],
498    });
499
500    // Check that the visitor works as expected
501    assert_eq!(
502        e,
503        mk_array(vec![mk_lit_expr(100), mk_lit_expr(200), lit_expr_200])
504    );
505}
hax_rust_engine/ast/visitors.rs

hax_rust_engine/ast/
visitors.rs
