From 0fd4b172f1ed5da5def96157e1ff85ee9b9fc5e0 Mon Sep 17 00:00:00 2001
From: Alexandre Terrasa <alexandre.terrasa@shopify.com>
Date: Mon, 5 Aug 2024 12:53:27 -0400
Subject: [PATCH] Fix parsing of dangling sigs

Signed-off-by: Alexandre Terrasa <alexandre.terrasa@shopify.com>
---
 lib/rbi/parser.rb       |  5 ++++-
 test/rbi/parser_test.rb | 30 ++++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/lib/rbi/parser.rb b/lib/rbi/parser.rb
index ad8aa0c4..e9d31273 100644
--- a/lib/rbi/parser.rb
+++ b/lib/rbi/parser.rb
@@ -185,6 +185,7 @@ def visit_class_node(node)
         current_scope << scope
         @scopes_stack << scope
         visit(node.body)
+        scope.nodes.concat(current_sigs)
         collect_dangling_comments(node)
         @scopes_stack.pop
         @last_node = nil
@@ -269,6 +270,7 @@ def visit_module_node(node)
         current_scope << scope
         @scopes_stack << scope
         visit(node.body)
+        scope.nodes.concat(current_sigs)
         collect_dangling_comments(node)
         @scopes_stack.pop
         @last_node = nil
@@ -278,7 +280,7 @@ def visit_module_node(node)
       def visit_program_node(node)
         @last_node = node
         super
-
+        @tree.nodes.concat(current_sigs)
         collect_orphan_comments
         separate_header_comments
         set_root_tree_loc
@@ -296,6 +298,7 @@ def visit_singleton_class_node(node)
         current_scope << scope
         @scopes_stack << scope
         visit(node.body)
+        scope.nodes.concat(current_sigs)
         collect_dangling_comments(node)
         @scopes_stack.pop
         @last_node = nil
diff --git a/test/rbi/parser_test.rb b/test/rbi/parser_test.rb
index 3298e771..db0b041d 100644
--- a/test/rbi/parser_test.rb
+++ b/test/rbi/parser_test.rb
@@ -153,6 +153,36 @@ def foo; end
       RBI
     end
 
+    def test_parse_dangling_sigs
+      rbi = <<~RBI
+        class Foo
+          sig { void }
+        end
+
+        module Bar
+          class << self
+            sig { void }
+          end
+          sig { void }
+        end
+        sig { void }
+        sig { returns(A) }
+      RBI
+
+      out = Parser.parse_string(rbi)
+      assert_equal(rbi, out.string)
+    end
+
+    def test_parse_sig_standalone
+      rbi = <<~RBI
+        sig { void }
+        sig { returns(A) }
+      RBI
+
+      out = Parser.parse_string(rbi)
+      assert_equal(rbi, out.string)
+    end
+
     def test_parse_methods_with_visibility
       rbi = <<~RBI
         private def m1; end