Inline instance variable annotation in module declaration

lib/rbs/inline_parser.rb

@@ -119,7 +119,27 @@ def visit_class_node(node)
 
         class_decl = AST::Ruby::Declarations::ClassDecl.new(buffer, class_name, node, super_class)
         insert_declaration(class_decl)
-        push_module_nesting(class_decl) do
+        visit_class_or_module_body(class_decl, node)
+      end
+
+      def visit_module_node(node)
+        return if skip_node?(node)
+
+        unless module_name = constant_as_type_name(node.constant_path)
+          diagnostics << Diagnostic::NonConstantModuleName.new(
+            rbs_location(node.constant_path.location),
+            "Module name must be a constant"
+          )
+          return
+        end
+
+        module_decl = AST::Ruby::Declarations::ModuleDecl.new(buffer, module_name, node)
+        insert_declaration(module_decl)
+        visit_class_or_module_body(module_decl, node)
+      end
+
+      def visit_class_or_module_body(decl, node)
+        push_module_nesting(decl) do
           visit_child_nodes(node)
 
           node.child_nodes.each do |child_node|
@@ -137,7 +157,7 @@ def visit_class_node(node)
           block.each_paragraph([]) do |paragraph|
             case paragraph
             when AST::Ruby::Annotations::InstanceVariableAnnotation
-              class_decl.members << AST::Ruby::Members::InstanceVariableMember.new(buffer, paragraph)
+              decl.members << AST::Ruby::Members::InstanceVariableMember.new(buffer, paragraph)
             when Location
               # Skip
             when AST::Ruby::CommentBlock::AnnotationSyntaxError
@@ -150,29 +170,7 @@ def visit_class_node(node)
           report_unused_annotation(*unused_annotations)
         end
 
-        class_decl.members.sort_by! { _1.location.start_line }
-      end
-
-      def visit_module_node(node)
-        return if skip_node?(node)
-
-        unless module_name = constant_as_type_name(node.constant_path)
-          diagnostics << Diagnostic::NonConstantModuleName.new(
-            rbs_location(node.constant_path.location),
-            "Module name must be a constant"
-          )
-          return
-        end
-
-        module_decl = AST::Ruby::Declarations::ModuleDecl.new(buffer, module_name, node)
-        insert_declaration(module_decl)
-        push_module_nesting(module_decl) do
-          visit_child_nodes(node)
-        end
-
-        comments.each_enclosed_block(node) do |block|
-          report_unused_block(block)
-        end
+        decl.members.sort_by! { _1.location.start_line }
       end
 
       def visit_def_node(node)

sig/inline_parser.rbs

@@ -110,6 +110,8 @@ module RBS
 
       def report_unused_block: (AST::Ruby::CommentBlock) -> void
 
+      def visit_class_or_module_body: (module_context, Prism::ClassNode | Prism::ModuleNode) -> void
+
       private
 
       def parse_mixin_call: (Prism::CallNode) -> void

test/rbs/inline_parser_test.rb

@@ -1502,6 +1502,28 @@ def initialize(name, age)
     end
   end
 
+  def test_parse__instance_variable_in_module
+    result = parse(<<~RUBY)
+      module Foo
+        # @rbs @bar: String
+      end
+    RUBY
+
+    assert_empty result.diagnostics
+
+    result.declarations[0].tap do |decl|
+      assert_instance_of RBS::AST::Ruby::Declarations::ModuleDecl, decl
+
+      assert_equal 1, decl.members.size
+
+      decl.members[0].tap do |member|
+        assert_instance_of RBS::AST::Ruby::Members::InstanceVariableMember, member
+        assert_equal :@bar, member.name
+        assert_equal "String", member.type.to_s
+      end
+    end
+  end
+
   def test_error__instance_variable_ignored
     result = parse(<<~RUBY)
       # @rbs @global_decl: String