Proposal: Support auto extending module

docs/syntax.md

@@ -780,3 +780,17 @@ _annotation_ ::= `%a{` _annotation-text_ `}`  # Annotation using {}
 
 _annotation-text_ ::= /[^\x00]*/              # Any characters except NUL (and parenthesis)
 ```
+
+#### Auto extending modules
+
+Module having "autoextend:..." annotation is considered as an auto extending module.
+When such auto extending modules are included, the including class will be extended by annotated modules.
+
+```
+%a{autoextend:Mod::ClassMethods}
+module Mod
+  module ClassMethods
+    def foo: () -> void
+  end
+end
+```

lib/rbs/definition_builder.rb

@@ -239,6 +239,27 @@ def build_singleton0(type_name)
             definition.class_variables.merge!(defn.class_variables)
           end
 
+          resolver = Resolver::TypeNameResolver.new(env)
+          singleton_ancestors = ancestor_builder.singleton_ancestors(type_name)
+          ancestor_builder.instance_ancestors(type_name).ancestors.each do |ancestor|
+            next unless RBS::Definition::Ancestor::Instance === ancestor && RBS::AST::Members::Include === ancestor.source
+
+            included_module = env.class_decls[ancestor.name] or raise "Unknown name for instance_ancestors: #{ancestor.name}"
+            included_module.decls.each do |decl|
+              decl.decl.annotations.each do |annotation|
+                if annotation.string.start_with? "autoextend:"
+                  mod_name = TypeName(annotation.string.split(":", 2)[1])
+                  mod_name = resolver.resolve(mod_name, context: decl.context) || mod_name
+
+                  if singleton_ancestors.ancestors.none? { |ancestor| RBS::Definition::Ancestor::Instance === ancestor && RBS::AST::Members::Extend === ancestor.source && ancestor.name == mod_name }
+                    subst = tapp_subst(mod_name, [])
+                    define_instance(definition, mod_name, subst)
+                  end
+                end
+              end
+            end
+          end
+
           one_ancestors.each_extended_module do |mod|
             mod.args.each do |arg|
               validate_type_presence(arg)

lib/rbs/definition_builder/ancestor_builder.rb

@@ -111,7 +111,7 @@ def self.singleton(type_name:, super_class:)
             params: nil,
             super_class: super_class,
             self_types: nil,
-            included_modules: nil,
+            included_modules: [],
             included_interfaces: nil,
             prepended_modules: nil,
             extended_modules: [],
@@ -306,7 +306,7 @@ def one_singleton_ancestors(type_name)
 
         mixin_ancestors(entry,
                         type_name,
-                        included_modules: nil,
+                        included_modules: ancestors.included_modules,
                         included_interfaces: nil,
                         prepended_modules: nil,
                         extended_modules: ancestors.extended_modules,

test/rbs/ancestor_builder_test.rb

@@ -64,7 +64,8 @@ class Hello[X] < Array[Integer]
 
           assert_equal Ancestor::Singleton.new(name: type_name("::Array")),
                        a.super_class
-          assert_nil a.included_modules
+          assert_equal [Ancestor::Instance.new(name: type_name("::Bar"), args: [parse_type("X", variables: [:X])], source: nil)],
+                       a.included_modules
           assert_nil a.included_interfaces
           assert_nil a.prepended_modules
           assert_equal [Ancestor::Instance.new(name: type_name("::Foo"), args: [parse_type("::String")], source: nil)],
@@ -130,7 +131,8 @@ module Hello[X] : _I1[Array[X]]
           assert_equal Ancestor::Instance.new(name: type_name("::Module"), args: [], source: nil),
                        a.super_class
           assert_nil a.self_types
-          assert_nil a.included_modules
+          assert_equal [Ancestor::Instance.new(name: type_name("::M2"), args: [parse_type("X", variables: [:X])], source: nil)],
+                       a.included_modules
           assert_nil a.prepended_modules
           assert_equal [Ancestor::Instance.new(name: type_name("::M1"), args: [parse_type("::String")], source: nil)],
                        a.extended_modules

test/rbs/ancestor_graph_test.rb

@@ -83,7 +83,7 @@ module M
           graph.each_ancestor(InstanceNode("::M")).to_set
         )
         assert_equal(
-          Set[InstanceNode("::B"), InstanceNode("::C")],
+          Set[InstanceNode("::B"), InstanceNode("::C"), SingletonNode("::B"), SingletonNode("::C")],
           graph.each_descendant(InstanceNode("::M")).to_set
         )
 

test/rbs/definition_builder_test.rb

@@ -281,6 +281,37 @@ def get: () -> X
     end
   end
 
+  def test_build_singleton_including_autoextending_module
+    SignatureManager.new do |manager|
+      manager.files[Pathname("foo.rbs")] = <<EOF
+%a{autoextend:Mod::ClassMethods}
+module Mod
+  module ClassMethods
+    def count: -> Integer
+  end
+end
+
+class Class
+  include Mod
+end
+EOF
+      manager.build do |env|
+        builder = DefinitionBuilder.new(env: env)
+
+        builder.build_singleton(type_name("::Class")).yield_self do |definition|
+          assert_instance_of Definition, definition
+
+          assert_equal [:__id__, :count, :initialize, :new, :puts, :respond_to_missing?, :to_i], definition.methods.keys.sort
+          assert_method_definition definition.methods[:__id__], ["() -> ::Integer"]
+          assert_method_definition definition.methods[:initialize], ["() -> void"]
+          assert_method_definition definition.methods[:puts], ["(*untyped) -> nil"]
+          assert_method_definition definition.methods[:respond_to_missing?], ["(::Symbol, bool) -> bool"]
+        end
+      end
+    end
+  end
+
+
   def test_build_instance_module_include_module
     SignatureManager.new do |manager|
       manager.files[Pathname("foo.rbs")] = <<EOF