Support extend self by prototype rb

lib/rbs/prototype/rb.rb

@@ -1,9 +1,9 @@
 module RBS
   module Prototype
     class RB
-      Context = Struct.new(:module_function, :singleton, keyword_init: true) do
-        def self.initial
-          self.new(module_function: false, singleton: false)
+      Context = Struct.new(:module_function, :singleton, :namespace, keyword_init: true) do
+        def self.initial(namespace: Namespace.root)
+          self.new(module_function: false, singleton: false, namespace: namespace)
         end
 
         def method_kind
@@ -87,8 +87,9 @@ def process(node, decls:, comments:, context:)
 
           decls.push kls
 
+          new_ctx = Context.initial(namespace: context.namespace + kls.name.to_namespace)
           each_node class_body do |child|
-            process child, decls: kls.members, comments: comments, context: Context.initial
+            process child, decls: kls.members, comments: comments, context: new_ctx
           end
           remove_unnecessary_accessibility_methods! kls.members
 
@@ -107,8 +108,9 @@ def process(node, decls:, comments:, context:)
 
           decls.push mod
 
+          new_ctx = Context.initial(namespace: context.namespace + mod.name.to_namespace)
           each_node module_body do |child|
-            process child, decls: mod.members, comments: comments, context: Context.initial
+            process child, decls: mod.members, comments: comments, context: new_ctx
           end
           remove_unnecessary_accessibility_methods! mod.members
 
@@ -119,7 +121,7 @@ def process(node, decls:, comments:, context:)
             RBS.logger.warn "`class <<` syntax with not-self may be compiled to incorrect code: #{this}"
           end
 
-          ctx = Context.initial.tap { |ctx| ctx.singleton = true}
+          ctx = Context.initial.tap { |ctx| ctx.singleton = true }
           process_children(body, decls: decls, comments: comments, context: ctx)
 
         when :DEFN, :DEFS
@@ -183,7 +185,7 @@ def process(node, decls:, comments:, context:)
             end
           when :extend
             args.each do |arg|
-              if (name = const_to_name(arg))
+              if (name = const_to_name(arg, context: context))
                 decls << AST::Members::Extend.new(
                   name: name,
                   args: [],
@@ -320,7 +322,7 @@ def process_children(node, decls:, comments:, context:)
         end
       end
 
-      def const_to_name(node)
+      def const_to_name(node, context: nil)
         case node&.type
         when :CONST
           TypeName.new(name: node.children[0], namespace: Namespace.empty)
@@ -334,6 +336,8 @@ def const_to_name(node)
           TypeName.new(name: node.children[1], namespace: namespace)
         when :COLON3
           TypeName.new(name: node.children[0], namespace: Namespace.root)
+        when :SELF
+          context&.then { |c| c.namespace.to_type_name }
         end
       end
 

test/rbs/rb_prototype_test.rb

@@ -251,6 +251,14 @@ class Hello
   attr_accessor :y, :z
   attr_writer foo, :a, 'b'
 end
+
+module Mod
+  extend self
+
+  module Mod2
+    extend self
+  end
+end
     EOR
 
     parser.parse(rb)
@@ -271,6 +279,14 @@ class Hello
 
   attr_writer b: untyped
 end
+
+module Mod
+  extend ::Mod
+
+  module Mod2
+    extend ::Mod::Mod2
+  end
+end
     EOF
   end