Visit TokenStream in generated syntax tree traversals

Author: dtolnay

codegen/src/fold.rs

@@ -207,7 +207,7 @@ fn node(traits: &mut TokenStream, impls: &mut TokenStream, s: &Node, defs: &Defi
         };
     }
 
-    let traits_body = if s.ident == "Span" {
+    let traits_body = if s.ident == "Span" || s.ident == "TokenStream" {
         quote!(i)
     } else {
         quote!(#fold_fn(self, i))
@@ -219,14 +219,16 @@ fn node(traits: &mut TokenStream, impls: &mut TokenStream, s: &Node, defs: &Defi
         }
     });
 
-    impls.extend(quote! {
-        pub fn #fold_fn<F>(f: &mut F, node: #ty) -> #ty
-        where
-            F: Fold + ?Sized,
-        {
-            #fold_impl
-        }
-    });
+    if s.ident != "TokenStream" {
+        impls.extend(quote! {
+            pub fn #fold_fn<F>(f: &mut F, node: #ty) -> #ty
+            where
+                F: Fold + ?Sized,
+            {
+                #fold_impl
+            }
+        });
+    }
 
     if s.ident == "Attribute" {
         let features = cfg::features(&s.features, DocCfg::Ordinary);

codegen/src/gen.rs

@@ -3,7 +3,7 @@ use inflections::Inflect;
 use proc_macro2::{Ident, Span, TokenStream};
 use syn_codegen::{Data, Definitions, Features, Node};
 
-pub const TERMINAL_TYPES: [&str; 2] = ["Ident", "Span"];
+pub const TERMINAL_TYPES: [&str; 3] = ["Ident", "Span", "TokenStream"];
 
 pub fn under_name(name: &str) -> Ident {
     Ident::new(&name.to_snake_case(), Span::call_site())

codegen/src/visit.rs

@@ -183,7 +183,7 @@ fn node(traits: &mut TokenStream, impls: &mut TokenStream, s: &Node, defs: &Defi
         Some(quote!('ast))
     };
 
-    let traits_body = if s.ident == "Span" {
+    let traits_body = if s.ident == "Span" || s.ident == "TokenStream" {
         None
     } else {
         Some(quote! {
@@ -197,14 +197,16 @@ fn node(traits: &mut TokenStream, impls: &mut TokenStream, s: &Node, defs: &Defi
         }
     });
 
-    impls.extend(quote! {
-        pub fn #visit_fn<'ast, V>(v: &mut V, node: &#ast_lifetime #ty)
-        where
-            V: Visit<'ast> + ?Sized,
-        {
-            #visit_impl
-        }
-    });
+    if s.ident != "TokenStream" {
+        impls.extend(quote! {
+            pub fn #visit_fn<'ast, V>(v: &mut V, node: &#ast_lifetime #ty)
+            where
+                V: Visit<'ast> + ?Sized,
+            {
+                #visit_impl
+            }
+        });
+    }
 }
 
 pub fn generate(defs: &Definitions) -> Result<()> {

codegen/src/visit_mut.rs

@@ -186,7 +186,7 @@ fn node(traits: &mut TokenStream, impls: &mut TokenStream, s: &Node, defs: &Defi
         }
     }
 
-    let traits_body = if s.ident == "Span" {
+    let traits_body = if s.ident == "Span" || s.ident == "TokenStream" {
         None
     } else {
         Some(quote! {
@@ -200,14 +200,16 @@ fn node(traits: &mut TokenStream, impls: &mut TokenStream, s: &Node, defs: &Defi
         }
     });
 
-    impls.extend(quote! {
-        pub fn #visit_mut_fn<V>(v: &mut V, node: &mut #ty)
-        where
-            V: VisitMut + ?Sized,
-        {
-            #visit_mut_impl
-        }
-    });
+    if s.ident != "TokenStream" {
+        impls.extend(quote! {
+            pub fn #visit_mut_fn<V>(v: &mut V, node: &mut #ty)
+            where
+                V: VisitMut + ?Sized,
+            {
+                #visit_mut_impl
+            }
+        });
+    }
 
     if s.ident == "Attribute" {
         let features = cfg::features(&s.features, DocCfg::Ordinary);