visit_XXX_end

Cargo.toml

@@ -22,6 +22,7 @@ full_moon_derive = { path = "./full-moon-derive", version = "0.2.0" }
 generational-arena = "0.2"
 itertools = "0.8"
 lazy_static = "1.3"
+paste = "0.1"
 regex = "1.1"
 serde = { version = "1.0", features = ["derive"], optional = true }
 

full-moon-derive/src/lib.rs

@@ -185,23 +185,35 @@ pub fn derive_visit(input: TokenStream) -> TokenStream {
 
     let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
 
-    let visit_self = match visit_self_hint(&input.attrs) {
-        Some(VisitSelfHint::Skip) => quote! {},
+    let (visit_self, visit_self_end) = match visit_self_hint(&input.attrs) {
+        Some(VisitSelfHint::Skip) => (quote! {}, quote! {}),
         Some(VisitSelfHint::VisitAs(visit_as)) => {
+            let visit_as_end = syn::Ident::new(&format!("visit_{}_end", visit_as), input_ident.span());
             let visit_as = syn::Ident::new(&format!("visit_{}", visit_as), input_ident.span());
-            quote! {
+
+            (quote! {
                 visitor.#visit_as(self);
-            }
+            }, quote! {
+                visitor.#visit_as_end(self);
+            })
         }
         None => {
             // name of self in snake_case
             let ssself = syn::Ident::new(
                 &format!("visit_{}", snake_case(&input_ident.to_string())),
                 input_ident.span(),
             );
-            quote! {
+
+            let ssself_end = syn::Ident::new(
+                &format!("visit_{}_end", snake_case(&input_ident.to_string())),
+                input_ident.span(),
+            );
+
+            (quote! {
                 visitor.#ssself(self);
-            }
+            }, quote! {
+                visitor.#ssself_end(self);
+            })
         }
     };
 
@@ -216,6 +228,7 @@ pub fn derive_visit(input: TokenStream) -> TokenStream {
 
                 #visit_self
                 #expanded
+                #visit_self_end
             }
         }
 
@@ -229,6 +242,7 @@ pub fn derive_visit(input: TokenStream) -> TokenStream {
 
                 #visit_self
                 #expanded
+                #visit_self_end
             }
         }
     };

src/visitors.rs

@@ -49,10 +49,14 @@ macro_rules! create_visitor {
                 ast.nodes().visit(self);
             }
 
-            $(
-                #[allow(missing_docs)]
-                fn $visit_name(&mut self, _node: &$ast_type<'ast>) { }
-            )+
+            paste::item! {
+                $(
+                    #[allow(missing_docs)]
+                    fn $visit_name(&mut self, _node: &$ast_type<'ast>) { }
+                    #[allow(missing_docs)]
+                    fn [<$visit_name _end>](&mut self, _node: &$ast_type<'ast>) { }
+                )+
+            }
 
             $(
                 #[allow(missing_docs)]
@@ -75,10 +79,14 @@ macro_rules! create_visitor {
                 ast.nodes_mut().visit_mut(self);
             }
 
-            $(
-                #[allow(missing_docs)]
-                fn $visit_name(&mut self, _node: &mut $ast_type<'ast>) { }
-            )+
+            paste::item! {
+                $(
+                    #[allow(missing_docs)]
+                    fn $visit_name(&mut self, _node: &mut $ast_type<'ast>) { }
+                    #[allow(missing_docs)]
+                    fn [<$visit_name _end>](&mut self, _node: &mut $ast_type<'ast>) { }
+                )+
+            }
 
             $(
                 #[allow(missing_docs)]

tests/visitors.rs

@@ -113,3 +113,42 @@ fn test_visit_token() {
         vec!["-- bla bla bla", "-- comment here", "-- and here"]
     );
 }
+
+#[test]
+fn test_end_visit() {
+    #[derive(Default)]
+    struct LogVisitor {
+        instructions: usize,
+        if_start_at: usize,
+        if_end_at: usize,
+        called_at: usize,
+    }
+
+    impl Visitor<'_> for LogVisitor {
+        fn visit_if(&mut self, _: &ast::If) {
+            self.instructions += 1;
+            self.if_start_at = self.instructions
+        }
+
+        fn visit_if_end(&mut self, _: &ast::If) {
+            self.instructions += 1;
+            self.if_end_at = self.instructions;
+        }
+
+        fn visit_call(&mut self, _: &ast::Call) {
+            self.instructions += 1;
+            self.called_at = self.instructions;
+        }
+    }
+
+    let mut visitor = LogVisitor::default();
+    visitor.visit_ast(&parse(r#"
+    if true then
+        call()
+    end
+    "#).unwrap());
+
+    assert_eq!(visitor.if_start_at, 1);
+    assert_eq!(visitor.called_at, 2);
+    assert_eq!(visitor.if_end_at, 3);
+}