prototype rb: Extract instance variables and class variables

lib/rbs/prototype/rb.rb

@@ -126,6 +126,7 @@ def process(node, decls:, comments:, context:)
             process child, decls: kls.members, comments: comments, context: new_ctx
           end
           remove_unnecessary_accessibility_methods! kls.members
+          sort_members! kls.members
 
         when :MODULE
           module_name, *module_body = node.children
@@ -195,6 +196,11 @@ def process(node, decls:, comments:, context:)
 
           decls.push member unless decls.include?(member)
 
+          new_ctx = Context.initial.tap { |ctx| ctx.singleton = kind == :singleton }
+          each_node def_body.children do |child|
+            process child, decls: decls, comments: comments, context: new_ctx
+          end
+
         when :ALIAS
           new_name, old_name = node.children.map { |c| literal_to_symbol(c) }
           member = AST::Members::Alias.new(
@@ -377,6 +383,32 @@ def process(node, decls:, comments:, context:)
             comment: comments[node.first_lineno - 1]
           )
 
+        when :IASGN
+          if context.singleton || !context.namespace.empty?
+            member = AST::Members::ClassVariable.new(
+              name: "@#{node.children.first}".to_sym,
+              type: Types::Bases::Any.new(location: nil),
+              location: nil,
+              comment: comments[node.first_lineno - 1]
+            )
+          else
+            member = AST::Members::InstanceVariable.new(
+              name: node.children.first,
+              type: Types::Bases::Any.new(location: nil),
+              location: nil,
+              comment: comments[node.first_lineno - 1]
+            )
+          end
+          decls.push member unless decls.include?(member)
+
+        when :CVASGN
+          member = AST::Members::ClassVariable.new(
+            name: node.children.first,
+            type: Types::Bases::Any.new(location: nil),
+            location: nil,
+            comment: comments[node.first_lineno - 1]
+          )
+          decls.push member unless decls.include?(member)
         else
           process_children(node, decls: decls, comments: comments, context: context)
         end
@@ -743,6 +775,16 @@ def find_def_index_by_name(decls, name)
           ]
         end
       end
+
+      def sort_members!(decls)
+        i = 0
+        orders = {
+          AST::Members::ClassVariable => -2,
+          AST::Members::InstanceVariable => -1,
+        }
+        decls.sort_by! { |decl| [orders.fetch(decl.class, 0), i += 1] }
+      end
     end
   end
 end
+

sig/prototype/rb.rbs

@@ -83,6 +83,8 @@ module RBS
       def is_accessibility?: (decl) -> bool
 
       def find_def_index_by_name: (Array[decl] decls, Symbol name) -> [Integer, AST::Members::MethodDefinition | AST::Members::AttrReader | AST::Members::AttrWriter]?
+
+      def sort_members!: (Array[decl] decls) -> void
     end
   end
 end

test/rbs/rb_prototype_test.rb

@@ -914,6 +914,76 @@ def hello: () -> untyped
     RBS
   end
 
+  def test_instance_variables
+    parser = RB.new
+
+    rb = <<-EOR
+class Hello
+  def message(message)
+    # comment for ivar
+    @message = message
+  end
+end
+    EOR
+
+    parser.parse(rb)
+
+    assert_write parser.decls, <<-EOF
+class Hello
+  # comment for ivar
+  @message: untyped
+
+  def message: (untyped message) -> untyped
+end
+    EOF
+  end
+
+  def test_class_variables
+    parser = RB.new
+
+    rb = <<-EOR
+class Hello
+  def self.foo
+    @foo = 42
+  end
+
+  class << self
+    def bar
+      @bar = 42
+    end
+  end
+
+  @baz = 42
+
+  def message(message)
+    # comment for cvar
+    @@message = message
+  end
+end
+    EOR
+
+    parser.parse(rb)
+
+    assert_write parser.decls, <<-EOF
+class Hello
+  @@foo: untyped
+
+  @@bar: untyped
+
+  @@baz: untyped
+
+  # comment for cvar
+  @@message: untyped
+
+  def self.foo: () -> untyped
+
+  def self.bar: () -> untyped
+
+  def message: (untyped message) -> untyped
+end
+    EOF
+  end
+
   def test_literal_to_type
     parser = RBS::Prototype::RB.new
     [