From ed2f375a955657a76726feac09af6d0ba64af37a Mon Sep 17 00:00:00 2001
From: Alex Rocha <alex.rocha@shopify.com>
Date: Thu, 20 Nov 2025 15:48:35 -0800
Subject: [PATCH] Add test demonstrating AST traversal and type checking

Adds test_parse_integer() which parses an integer literal type alias
and traverses the AST (TypeAlias -> LiteralType -> Integer) using
pattern matching to verify node types and extract values.

This validates that the generated node wrappers enable idiomatic Rust
AST traversal with proper type safety.

Also adds Debug derives and fixes memory management by returning
SignatureNode instead of raw pointer, with Drop impl to free parser.
---
 rust/ruby-rbs/build.rs   |  2 ++
 rust/ruby-rbs/src/lib.rs | 38 +++++++++++++++++++++++++++++++++++---
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/rust/ruby-rbs/build.rs b/rust/ruby-rbs/build.rs
index d105f247f..a8bdfccd2 100644
--- a/rust/ruby-rbs/build.rs
+++ b/rust/ruby-rbs/build.rs
@@ -89,6 +89,7 @@ fn generate(config: &Config) -> Result<(), Box<dyn Error>> {
     // TODO: Go through all of the nodes and generate the structs to back them up
     for node in &config.nodes {
         writeln!(file, "#[allow(dead_code)]")?; // TODO: Remove this once all nodes that need parser are implemented
+        writeln!(file, "#[derive(Debug)]")?;
         writeln!(file, "pub struct {} {{", node.rust_name)?;
         writeln!(file, "    parser: *mut rbs_parser_t,")?;
         writeln!(
@@ -188,6 +189,7 @@ fn generate(config: &Config) -> Result<(), Box<dyn Error>> {
     }
 
     // Generate the Node enum to wrap all of the structs
+    writeln!(file, "#[derive(Debug)]")?;
     writeln!(file, "pub enum Node {{")?;
     for node in &config.nodes {
         let variant_name = node
diff --git a/rust/ruby-rbs/src/lib.rs b/rust/ruby-rbs/src/lib.rs
index 20069ebf9..dae6787c2 100644
--- a/rust/ruby-rbs/src/lib.rs
+++ b/rust/ruby-rbs/src/lib.rs
@@ -13,7 +13,7 @@ static INIT: Once = Once::new();
 /// let signature = parse(rbs_code.as_bytes());
 /// assert!(signature.is_ok(), "Failed to parse RBS signature");
 /// ```
-pub fn parse(rbs_code: &[u8]) -> Result<*mut rbs_signature_t, String> {
+pub fn parse(rbs_code: &[u8]) -> Result<SignatureNode, String> {
     unsafe {
         INIT.call_once(|| {
             rbs_constant_pool_init(RBS_GLOBAL_CONSTANT_POOL, 26);
@@ -30,16 +30,27 @@ pub fn parse(rbs_code: &[u8]) -> Result<*mut rbs_signature_t, String> {
         let mut signature: *mut rbs_signature_t = std::ptr::null_mut();
         let result = rbs_parse_signature(parser, &mut signature);
 
-        rbs_parser_free(parser);
+        let signature_node = SignatureNode {
+            parser,
+            pointer: signature,
+        };
 
         if result {
-            Ok(signature)
+            Ok(signature_node)
         } else {
             Err(String::from("Failed to parse RBS signature"))
         }
     }
 }
 
+impl Drop for SignatureNode {
+    fn drop(&mut self) {
+        unsafe {
+            rbs_parser_free(self.parser);
+        }
+    }
+}
+
 pub struct NodeListIter {
     parser: *mut rbs_parser_t,
     current: *mut rbs_node_list_node_t,
@@ -140,6 +151,7 @@ impl RBSLocationList {
     }
 }
 
+#[derive(Debug)]
 pub struct RBSString {
     pointer: *const rbs_string_t,
 }
@@ -223,4 +235,24 @@ mod tests {
         let signature2 = parse(rbs_code2.as_bytes());
         assert!(signature2.is_ok(), "Failed to parse RBS signature");
     }
+
+    #[test]
+    fn test_parse_integer() {
+        let rbs_code = r#"type foo = 1"#;
+        let signature = parse(rbs_code.as_bytes());
+        assert!(signature.is_ok(), "Failed to parse RBS signature");
+
+        let signature_node = signature.unwrap();
+        if let Node::TypeAlias(node) = signature_node.declarations().iter().next().unwrap()
+            && let Node::LiteralType(literal) = node.type_()
+            && let Node::Integer(integer) = literal.literal()
+        {
+            assert_eq!(
+                "1".to_string(),
+                String::from_utf8(integer.string_representation().as_bytes().to_vec()).unwrap()
+            );
+        } else {
+            panic!("No literal type node found");
+        }
+    }
 }