Add generated Visit trait for AST node traversal

Enable walking the AST by generating a Visit trait with per-node visitor
methods. It uses double dispatch to route each node type to its
corresponding visitor method. This avoids consumers needing to manually
match on Node variants and allows overriding specific visits while
inheriting default behavior for others.

Author: alexcrocha

rust/ruby-rbs/build.rs

@@ -27,6 +27,14 @@ struct Node {
     fields: Option<Vec<NodeField>>,
 }
 
+impl Node {
+    fn variant_name(&self) -> &str {
+        self.rust_name
+            .strip_suffix("Node")
+            .unwrap_or(&self.rust_name)
+    }
+}
+
 fn main() -> Result<(), Box<dyn Error>> {
     let config_path = Path::new(env!("CARGO_MANIFEST_DIR"))
         .join("../../config.yml")
@@ -62,11 +70,12 @@ fn main() -> Result<(), Box<dyn Error>> {
 enum CIdentifier {
     Type,     // foo_bar_t
     Constant, // FOO_BAR
+    Method,   // visit_foo_bar
 }
 
 fn convert_name(name: &str, identifier: CIdentifier) -> String {
     let type_name = name.replace("::", "_");
-    let lowercase = matches!(identifier, CIdentifier::Type);
+    let lowercase = matches!(identifier, CIdentifier::Type | CIdentifier::Method);
     let mut out = String::new();
     let mut prev_is_lower = false;
 
@@ -94,12 +103,67 @@ fn convert_name(name: &str, identifier: CIdentifier) -> String {
         }
     }
 
-    if lowercase {
+    if matches!(identifier, CIdentifier::Type) {
         out.push_str("_t");
     }
     out
 }
 
+fn write_visit_trait(file: &mut File, config: &Config) -> Result<(), Box<dyn std::error::Error>> {
+    writeln!(file, "/// A trait for visiting the AST.")?;
+    writeln!(file, "pub trait Visit {{")?;
+    writeln!(file, "   /// Visits a node.")?;
+    writeln!(file, "   fn visit(&mut self, node: &Node) {{")?;
+    writeln!(file, "       match node {{")?;
+
+    for node in &config.nodes {
+        writeln!(file, "           Node::{}(it) => {{", node.variant_name())?;
+        writeln!(
+            file,
+            "               self.visit_{}_node(it);",
+            convert_name(&node.variant_name(), CIdentifier::Method)
+        )?;
+        writeln!(file, "           }}")?;
+    }
+
+    writeln!(file, "       }}")?;
+    writeln!(file, "   }}")?;
+
+    for node in &config.nodes {
+        writeln!(file)?;
+        writeln!(file, "    /// Visits a `{}` node.", node.name)?;
+        writeln!(
+            file,
+            "    fn visit_{}_node(&mut self, node: &{}Node) {{",
+            convert_name(&node.variant_name(), CIdentifier::Method),
+            node.variant_name()
+        )?;
+        writeln!(
+            file,
+            "        visit_{}_node(self, node);",
+            convert_name(&node.variant_name(), CIdentifier::Method)
+        )?;
+        writeln!(file, "    }}")?;
+    }
+    writeln!(file, "}}")?;
+    writeln!(file)?;
+
+    for node in &config.nodes {
+        writeln!(file, "#[allow(unused_variables)]")?; // TODO: Remove this once all nodes that need visitor are implemented
+        writeln!(
+            file,
+            "pub fn visit_{}_node<V: Visit + ?Sized>(visitor: &mut V, node: &{}Node) {{",
+            convert_name(&node.variant_name(), CIdentifier::Method),
+            node.variant_name()
+        )?;
+        // TODO: Implement the default implementation of public visits
+        writeln!(file, "}}")?;
+        writeln!(file)?;
+    }
+
+    Ok(())
+}
+
 fn generate(config: &Config) -> Result<(), Box<dyn Error>> {
     let out_dir = env::var("OUT_DIR")?;
     let dest_path = Path::new(&out_dir).join("bindings.rs");
@@ -291,18 +355,13 @@ fn generate(config: &Config) -> Result<(), Box<dyn Error>> {
     )?;
     writeln!(file, "        match unsafe {{ (*node).type_ }} {{")?;
     for node in &config.nodes {
-        let variant_name = node
-            .rust_name
-            .strip_suffix("Node")
-            .unwrap_or(&node.rust_name);
-
         let enum_name = convert_name(&node.name, CIdentifier::Constant);
 
         writeln!(
             file,
             "            rbs_node_type::{} => Self::{}({} {{ parser, pointer: node.cast::<{}>() }}),",
             enum_name,
-            variant_name,
+            node.variant_name(),
             node.rust_name,
             convert_name(&node.name, CIdentifier::Type)
         )?;
@@ -314,6 +373,9 @@ fn generate(config: &Config) -> Result<(), Box<dyn Error>> {
     writeln!(file, "        }}")?;
     writeln!(file, "    }}")?;
     writeln!(file, "}}")?;
+    writeln!(file)?;
+
+    write_visit_trait(&mut file, config)?;
 
     Ok(())
 }