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::*;
20use hax_lib_macros_types::AttrPayload;
21
22pub mod wrappers {
23    //! This module provides a visitor wrappers, or transformer of visitors.
24    //! Such wrappers transform the behavior of a visitor.
25    //!
26    //! For example, [`SpanWrapper`] takes care of keeping track of [`Span`]s
27    //! while travesing an AST.
28
29    use std::ops::Deref;
30
31    use super::{infallible::AstVisitable as AstVisitableInfallible, *};
32    use diagnostics::*;
33
34    /// A visitor wrapper that tracks span while visiting the AST. Whenever an
35    /// AST node that carries a span is visited, using this wrapper, the ambient
36    /// span is mutated and accessible via the `HasSpan` trait.
37    pub struct SpanWrapper<'a, V>(pub &'a mut V);
38
39    impl<'a, V: HasSpan> SpanWrapper<'a, V> {
40        /// Performs a spanned action: calls the function `action` on
41        /// `ast_fragment`, with the contextual span information in `self` being
42        /// the span found in `ast_fragment`.
43        fn spanned_action<T: Deref, U>(
44            &mut self,
45            ast_fragment: T,
46            action: impl Fn(&mut Self, T) -> U,
47        ) -> U
48        where
49            T::Target: HasSpan,
50        {
51            let span_before = self.0.span();
52            *self.0.span_mut() = ast_fragment.span();
53            // Perform the provided action on `ast_fragment` with `ast_fragment`'s span as contextual span.
54            let result = action(self, ast_fragment);
55            *self.0.span_mut() = span_before;
56            result
57        }
58    }
59
60    impl<'a, V: AstVisitorMut + HasSpan> AstVisitorMut for SpanWrapper<'a, V> {
61        fn visit_inner<T>(&mut self, x: &mut T)
62        where
63            T: AstVisitableInfallible,
64            T: for<'s> DriveMut<'s, AstVisitableInfallibleWrapper<Self>>,
65        {
66            x.drive_map(self.0)
67        }
68        fn visit_item(&mut self, x: &mut Item) {
69            self.spanned_action(x, Self::visit_inner)
70        }
71        fn visit_expr(&mut self, x: &mut Expr) {
72            self.spanned_action(x, Self::visit_inner)
73        }
74        fn visit_pat(&mut self, x: &mut Pat) {
75            self.spanned_action(x, Self::visit_inner)
76        }
77        fn visit_guard(&mut self, x: &mut Guard) {
78            self.spanned_action(x, Self::visit_inner)
79        }
80        fn visit_arm(&mut self, x: &mut Arm) {
81            self.spanned_action(x, Self::visit_inner)
82        }
83        fn visit_impl_item(&mut self, x: &mut ImplItem) {
84            self.spanned_action(x, Self::visit_inner)
85        }
86        fn visit_trait_item(&mut self, x: &mut TraitItem) {
87            self.spanned_action(x, Self::visit_inner)
88        }
89        fn visit_generic_param(&mut self, x: &mut GenericParam) {
90            self.spanned_action(x, Self::visit_inner)
91        }
92        fn visit_attribute(&mut self, x: &mut Attribute) {
93            self.spanned_action(x, Self::visit_inner)
94        }
95        fn visit_spanned_ty(&mut self, x: &mut SpannedTy) {
96            self.spanned_action(x, Self::visit_inner)
97        }
98    }
99
100    /// A visitor wrapper that automatically collects errors in `ErrorNode`s.
101    /// Coupled with the trait `VisitorWithErrors`, this provides an `error`
102    /// method on a visitor that can be used to throw errors, which will be
103    /// automatically inlined in the AST on the closest error-capable node.
104    pub struct ErrorWrapper<'a, V>(pub &'a mut V);
105
106    /// An opaque error vault. This is the state manipulated by the visitor wrapper [`ErrorWrapper`].
107    /// It is purposefully not-inspectable.
108    #[derive(Default)]
109    pub struct ErrorVault(Vec<Diagnostic>);
110    impl ErrorVault {
111        fn add(&mut self, diagnostic: Diagnostic) {
112            self.0.push(diagnostic);
113        }
114    }
115
116    /// Helper struct that contains error-handling related state information.
117    /// This is used internally by [`setup_error_handling_struct`].
118    pub struct ErrorHandlingState(pub Span, pub ErrorVault);
119    impl Default for ErrorHandlingState {
120        fn default() -> Self {
121            Self(Span::dummy(), Default::default())
122        }
123    }
124
125    #[macro_export]
126    /// Use this macro in an implementation of `AstVisitorMut` to get automatic spans and error handling.
127    macro_rules! setup_error_handling_impl {
128        () => {
129            fn visit<T: $crate::ast::visitors::AstVisitableInfallible>(&mut self, x: &mut T) {
130                $crate::ast::visitors::wrappers::SpanWrapper(
131                    &mut $crate::ast::visitors::wrappers::ErrorWrapper(self),
132                )
133                .visit(x)
134            }
135        };
136    }
137    pub use setup_error_handling_impl;
138
139    /// Mark a visitor with a specific diagnostic context.
140    pub trait VisitorWithContext {
141        /// Returns the diagnostic context for this visitor.
142        fn context(&self) -> Context;
143    }
144
145    impl<T: HasSpan> HasSpan for ErrorWrapper<'_, T> {
146        fn span(&self) -> Span {
147            self.0.span()
148        }
149
150        fn span_mut(&mut self) -> &mut Span {
151            self.0.span_mut()
152        }
153    }
154
155    /// A visitor that can throw errors. It should be used in combination with
156    /// `ErrorWrapper`, which will take care of bubbling error up to the nearest
157    /// parent capable of representing errors. For instance, if you error out in
158    /// a literal, the error will be represented in the parent expression or the
159    /// parent type, as nodes [`ExprKind::Error`] or [`TyKind ::Error`].
160    pub trait VisitorWithErrors: HasSpan + VisitorWithContext {
161        /// Projects the error vault.
162        fn error_vault(&mut self) -> &mut ErrorVault;
163        /// Send an error.
164        fn error(&mut self, node: impl Into<Fragment>, kind: DiagnosticInfoKind) {
165            let context = self.context();
166            let span = self.span();
167            self.error_vault().add(Diagnostic::new(
168                node,
169                DiagnosticInfo {
170                    context,
171                    span,
172                    kind,
173                },
174            ));
175        }
176    }
177
178    impl<'a, V: VisitorWithErrors> ErrorWrapper<'a, V> {
179        fn error_handled_action<
180            T: FallibleAstNode + Clone + std::fmt::Debug + Into<Fragment>,
181            U,
182        >(
183            &mut self,
184            x: &mut T,
185            action: impl Fn(&mut Self, &mut T) -> U,
186        ) -> U {
187            let diagnostics_snapshot = self.0.error_vault().0.clone();
188            self.0.error_vault().0.clear();
189            let result = action(self, x);
190            let diagnostics: Vec<_> = self.0.error_vault().0.drain(..).collect();
191            if !diagnostics.is_empty() {
192                x.set_error(ErrorNode {
193                    fragment: Box::new(x.clone().into()),
194                    diagnostics,
195                });
196            }
197            self.0.error_vault().0 = diagnostics_snapshot;
198            result
199        }
200    }
201
202    impl<'a, V: AstVisitorMut + VisitorWithErrors> AstVisitorMut for ErrorWrapper<'a, V> {
203        fn visit_inner<T>(&mut self, x: &mut T)
204        where
205            T: AstVisitableInfallible,
206            T: for<'s> DriveMut<'s, AstVisitableInfallibleWrapper<Self>>,
207        {
208            x.drive_map(self.0)
209        }
210        fn visit_item(&mut self, x: &mut Item) {
211            self.error_handled_action(x, Self::visit_inner)
212        }
213        fn visit_pat(&mut self, x: &mut Pat) {
214            self.error_handled_action(x, Self::visit_inner)
215        }
216        fn visit_expr(&mut self, x: &mut Expr) {
217            self.error_handled_action(x, Self::visit_inner)
218        }
219        fn visit_ty(&mut self, x: &mut Ty) {
220            self.error_handled_action(x, Self::visit_inner)
221        }
222    }
223}
224
225#[hax_rust_engine_macros::replace(AstNodes => include(VisitableAstNodes))]
226mod replaced {
227    use super::*;
228    pub mod infallible {
229        use super::*;
230
231        #[visitable_group(
232            visitor(drive_map(
233                /// An mutable visitor that visits the AST for hax.
234                ///
235                /// ```rust,ignore
236                /// use crate::ast::{diagnostics::*, visitors::*};
237                /// #[setup_error_handling_struct]
238                /// #[derive(Default)]
239                /// struct MyVisitor;
240                ///
241                /// impl VisitorWithContext for MyVisitor {
242                ///     fn context(&self) -> Context {
243                ///         Context::Import
244                ///     }
245                /// }
246                ///
247                /// impl AstVisitorMut for MyVisitor {
248                ///     setup_error_handling_impl!();
249                /// }
250                ///
251                /// // MyVisitor::visit(my_ast_node)
252                /// ```
253                &mut AstVisitorMut
254            ), infallible),
255            visitor(drive(
256                /// An immutable visitor that visits the AST for hax.
257                &AstVisitor
258            ), infallible),
259            skip(
260                String, bool, char, hax_frontend_exporter::Span,
261                for<T> crate::interning::Interned<T>,
262            ),
263            drive(
264                for<T: AstVisitable> Box<T>, for<T: AstVisitable> Option<T>, for<T: AstVisitable> Vec<T>,
265                for<A: AstVisitable, B: AstVisitable> (A, B),
266                for<A: AstVisitable, B: AstVisitable, C: AstVisitable> (A, B, C),
267                usize
268            ),
269            override(AstNodes),
270            override_skip(
271                Span, Fragment, GlobalId, Diagnostic, AttrPayload,
272            ),
273        )]
274        /// Helper trait to drive visitor.
275        pub trait AstVisitable {}
276    }
277
278    #[allow(missing_docs)]
279    pub mod fallible {
280        use super::*;
281
282        #[visitable_group(
283            visitor(drive(
284                /// An immutable visitor that can exit early.
285                &AstEarlyExitVisitor
286            )),
287            visitor(drive_mut(
288                /// An immutable visitor that can exit early and mutate the AST fragments.
289                &mut AstEarlyExitVisitorMut
290            )),
291            skip(
292                String, bool, char, hax_frontend_exporter::Span,
293                for<T> crate::interning::Interned<T>,
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, AttrPayload,
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
