From 403f13913b9e23f4eb0ce907fa4d5badc2be6f31 Mon Sep 17 00:00:00 2001
From: Alex Rocha <alex.rocha@shopify.com>
Date: Fri, 28 Nov 2025 10:04:17 -0800
Subject: [PATCH] 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.
---
 rust/ruby-rbs/build.rs | 78 +++++++++++++++++++++++++++++++++++++-----
 1 file changed, 70 insertions(+), 8 deletions(-)

diff --git a/rust/ruby-rbs/build.rs b/rust/ruby-rbs/build.rs
index ce05d2e62..f30a97c3a 100644
--- a/rust/ruby-rbs/build.rs
+++ b/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 traversing the AST using a visitor")?;
+    writeln!(file, "pub trait Visit {{")?;
+    writeln!(
+        file,
+        "   /// Visit any node of the AST. Generally used to continue traversal"
+    )?;
+    writeln!(file, "   fn visit(&mut self, node: &Node) {{")?;
+    writeln!(file, "       match node {{")?;
+
+    for node in &config.nodes {
+        let node_variant_name = node.variant_name();
+        let method_name = convert_name(node_variant_name, CIdentifier::Method);
+
+        writeln!(file, "           Node::{}(it) => {{", node_variant_name)?;
+        writeln!(file, "               self.visit_{}_node(it);", method_name,)?;
+        writeln!(file, "           }}")?;
+    }
+
+    writeln!(file, "       }}")?;
+    writeln!(file, "   }}")?;
+
+    for node in &config.nodes {
+        let node_variant_name = node.variant_name();
+        let method_name = convert_name(node_variant_name, CIdentifier::Method);
+
+        writeln!(file)?;
+        writeln!(
+            file,
+            "    fn visit_{}_node(&mut self, node: &{}Node) {{",
+            method_name, node_variant_name
+        )?;
+        writeln!(file, "        visit_{}_node(self, node);", method_name)?;
+        writeln!(file, "    }}")?;
+    }
+    writeln!(file, "}}")?;
+    writeln!(file)?;
+
+    for node in &config.nodes {
+        let node_variant_name = node.variant_name();
+        let method_name = convert_name(node_variant_name, CIdentifier::Method);
+
+        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) {{",
+            method_name, node_variant_name
+        )?;
+        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(())
 }