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            ),
262            drive(
263                for<T: AstVisitable> Box<T>, for<T: AstVisitable> Option<T>, for<T: AstVisitable> Vec<T>,
264                for<A: AstVisitable, B: AstVisitable> (A, B),
265                for<A: AstVisitable, B: AstVisitable, C: AstVisitable> (A, B, C),
266                usize
267            ),
268            override(AstNodes),
269            override_skip(
270                Span, Fragment, GlobalId, Diagnostic, AttrPayload,
271            ),
272        )]
273        /// Helper trait to drive visitor.
274        pub trait AstVisitable {}
275    }
276
277    #[allow(missing_docs)]
278    pub mod fallible {
279        use super::*;
280
281        #[visitable_group(
282            visitor(drive(
283                /// An immutable visitor that can exit early.
284                &AstEarlyExitVisitor
285            )),
286            visitor(drive_mut(
287                /// An immutable visitor that can exit early and mutate the AST fragments.
288                &mut AstEarlyExitVisitorMut
289            )),
290            skip(
291                String, bool, char, hax_frontend_exporter::Span,
292            ),
293            drive(
294                for<T: AstVisitable> Box<T>, for<T: AstVisitable> Option<T>, for<T: AstVisitable> Vec<T>,
295                for<A: AstVisitable, B: AstVisitable> (A, B),
296                for<A: AstVisitable, B: AstVisitable, C: AstVisitable> (A, B, C),
297                usize
298            ),
299            override(AstNodes),
300            override_skip(
301                Span, Fragment, GlobalId, Diagnostic, AttrPayload,
302            ),
303        )]
304        /// Helper trait to drive visitor.
305        pub trait AstVisitable {}
306    }
307
308    /// This modules provides `dyn` compatible trait for visitors.
309    pub mod dyn_compatible {
310        use super::*;
311
312        macro_rules! derive_erased_ast_visitors {
313            ({$($attrs:tt)*}, $name: ident, $helper: ident, $($ty:ty),*) => {
314                $($attrs)*
315                pub trait $name<'a>: $($helper<'a, $ty> + )* {}
316            };
317        }
318
319        macro_rules! render_path {
320            ($head:ident) => {stringify!($head)};
321            ($head:ident $(::$tail:ident)*) => {
322                concat!(stringify!($head), "::", render_path!($($tail)::*))
323            };
324        }
325
326        macro_rules! make_dyn_compatible {
327            ($($visitable_trait:ident)::*, $($visitor_trait:ident)::*, $helper_name: ident, $name: ident, mut:{$($mut:tt)?}, super:{$($super:ident)::*}, $ret:ty) => {
328                #[doc = concat!("A [dyn-compatible](https://doc.rust-lang.org/reference/items/traits.html#dyn-compatibility) trait similar to [`", render_path!($($visitor_trait)::*),"`].")]
329                #[doc = concat!("This trait provides one `visit` method to visit a given type `T` with a given visitor.")]
330                pub trait $helper_name<'a, T: ?Sized>: $($super)::* {
331                    /// Visit a value with the visitor.
332                    fn visit(&mut self, _: &'a $($mut)? T) -> $ret;
333                }
334
335                impl<'a, T: $($visitable_trait)::*, V: $($visitor_trait)::*> $helper_name<'a, T> for V {
336                    fn visit(&mut self, e: &'a $($mut)? T) -> $ret {
337                        <Self as $($visitor_trait)::*>::visit(self, e)
338                    }
339                }
340                derive_erased_ast_visitors!({
341                    #[doc = concat!("A [dyn-compatible](https://doc.rust-lang.org/reference/items/traits.html#dyn-compatibility) trait similar to [`", render_path!($($visitor_trait)::*),"`].")]
342                    #[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.")]
343                }, $name, $helper_name, AstNodes);
344
345                impl<'a, V: $($visitor_trait)::*> $name<'a> for V {}
346            };
347        }
348
349        make_dyn_compatible!(
350            infallible::AstVisitable,
351            infallible::AstVisitorMut,
352            AstVisitableMut,
353            AstVisitorMut,
354            mut:{mut},
355            super:{},
356            ()
357        );
358        make_dyn_compatible!(
359            infallible::AstVisitable,
360            infallible::AstVisitor,
361            AstVisitable,
362            AstVisitor,
363            mut:{},
364            super:{},
365            ()
366        );
367
368        make_dyn_compatible!(
369            fallible::AstVisitable,
370            fallible::AstEarlyExitVisitorMut,
371            AstEarlyExitVisitableMut,
372            AstEarlyExitVisitorMut,
373            mut:{mut},
374            super:{Visitor},
375            ControlFlow<<Self as Visitor>::Break>
376        );
377        make_dyn_compatible!(
378            fallible::AstVisitable,
379            fallible::AstEarlyExitVisitor,
380            AstEarlyExitVisitable,
381            AstEarlyExitVisitor,
382            mut:{},
383            super:{Visitor},
384            ControlFlow<<Self as Visitor>::Break>
385        );
386    }
387}
388
389pub use replaced::dyn_compatible;
390use replaced::{fallible, infallible};
391
392pub use fallible::{
393    AstEarlyExitVisitor, AstEarlyExitVisitorMut, AstVisitable as AstVisitableFallible,
394    AstVisitableWrapper,
395};
396pub use hax_rust_engine_macros::setup_error_handling_struct;
397pub use infallible::{
398    AstVisitable as AstVisitableInfallible, AstVisitableInfallibleWrapper, AstVisitor,
399    AstVisitorMut,
400};
401pub use wrappers::{VisitorWithContext, VisitorWithErrors, setup_error_handling_impl};
402
403#[test]
404fn double_literals_in_ast() {
405    use crate::ast::diagnostics::*;
406
407    #[setup_error_handling_struct]
408    #[derive(Default)]
409    struct DoubleU8Literals;
410
411    impl VisitorWithContext for DoubleU8Literals {
412        fn context(&self) -> Context {
413            Context::Import
414        }
415    }
416
417    impl AstVisitorMut for DoubleU8Literals {
418        setup_error_handling_impl!();
419
420        fn visit_literal(&mut self, x: &mut Literal) {
421            let Literal::Int { value, .. } = x else {
422                return;
423            };
424            let Ok(n): Result<u8, _> = str::parse(value) else {
425                return self.error(
426                    x.clone(),
427                    DiagnosticInfoKind::AssertionFailure {
428                        details: "Bad literal".into(),
429                    },
430                );
431            };
432            let n = (n as u16) * 2;
433            if n >= u8::MAX as u16 {
434                return self.error(
435                    x.clone(),
436                    DiagnosticInfoKind::AssertionFailure {
437                        details: "Literal too big".into(),
438                    },
439                );
440            }
441            *value = Symbol::new(&format!("{}", n));
442        }
443    }
444
445    // Syntax helpers
446    let int_kind = IntKind {
447        size: IntSize::S8,
448        signedness: Signedness::Signed,
449    };
450    let mk_lit = |n: isize| Literal::Int {
451        value: Symbol::new(&format!("{}", n)),
452        negative: false,
453        kind: int_kind.clone(),
454    };
455    let meta = Metadata {
456        span: Span::dummy(),
457        attributes: vec![],
458    };
459    let mk_lit_expr = |n| Expr {
460        kind: Box::new(ExprKind::Literal(mk_lit(n))),
461        ty: Ty(Box::new(TyKind::Primitive(PrimitiveTy::Int(
462            int_kind.clone(),
463        )))),
464        meta: meta.clone(),
465    };
466    let mk_array = |exprs| Expr {
467        kind: Box::new(ExprKind::Array(exprs)),
468        ty: Ty(Box::new(TyKind::RawPointer)), // wrong type, but this is not important for this test.
469        meta: meta.clone(),
470    };
471    let mut lit_expr_200 = mk_lit_expr(200);
472
473    // Creates the expression `[50u8, 100u8, 200u8]`: the last one cannot be doubled, and will cause an error.
474    let mut e = mk_array(vec![
475        mk_lit_expr(50),
476        mk_lit_expr(100),
477        lit_expr_200.clone(),
478    ]);
479
480    // Visit the expression.
481    DoubleU8Literals::default().visit(&mut e);
482
483    // Transform `lit_expr_200` into the error `DoubleU8Literal` should produce
484    lit_expr_200.set_error(ErrorNode {
485        fragment: Box::new(lit_expr_200.clone().into()),
486        diagnostics: vec![Diagnostic::new(
487            mk_lit(200),
488            DiagnosticInfo {
489                span: lit_expr_200.span(),
490                context: Context::Import,
491                kind: DiagnosticInfoKind::AssertionFailure {
492                    details: "Literal too big".into(),
493                },
494            },
495        )],
496    });
497
498    // Check that the visitor works as expected
499    assert_eq!(
500        e,
501        mk_array(vec![mk_lit_expr(100), mk_lit_expr(200), lit_expr_200])
502    );
503}
hax_rust_engine/ast/visitors.rs

hax_rust_engine/ast/
visitors.rs
