Support auto extending module (via AncestorBuilder)

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/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,
@@ -334,7 +334,32 @@ def one_interface_ancestors(type_name)
           end
       end
 
+      def auto_extended_modules(type_name, resolver)
+        auto_extended_modules = [] #: Array[Definition::Ancestor::Instance]
+
+        mod = env.class_decls[type_name] or raise "Unknown name for include: #{type_name}"
+        mod.decls.each do |mod_decl|
+          mod_decl.decl.annotations.each do |annotation|
+            if annotation.string.start_with? "autoextend:"
+              auto_extended_mod_name = TypeName(annotation.string.split(":", 2)[1])
+              auto_extended_mod_name = resolver.resolve(auto_extended_mod_name, context: mod_decl.context) || auto_extended_mod_name
+
+              auto_extend_source = AST::Members::Extend.new(
+                name: auto_extended_mod_name,
+                args: [],
+                annotations: [],
+                location: nil,
+                comment: nil,
+              )
+              auto_extended_modules << Definition::Ancestor::Instance.new(name: auto_extended_mod_name, args: [], source: auto_extend_source)
+            end
+          end
+        end
+        auto_extended_modules
+      end
+
       def mixin_ancestors0(decl, type_name, align_params:, included_modules:, included_interfaces:, extended_modules:, prepended_modules:, extended_interfaces:)
+        resolver = Resolver::TypeNameResolver.new(env)
         decl.each_mixin do |member|
           case member
           when AST::Members::Include
@@ -348,6 +373,12 @@ def mixin_ancestors0(decl, type_name, align_params:, included_modules:, included
 
               module_name = env.normalize_module_name(module_name)
               included_modules << Definition::Ancestor::Instance.new(name: module_name, args: module_args, source: member)
+
+              if extended_modules
+                auto_extended_modules(module_name, resolver).each do |mod|
+                  extended_modules << mod unless extended_modules.include?(mod)
+                end
+              end
             when member.name.interface? && included_interfaces
               NoMixinFoundError.check!(member.name, env: env, member: member)
 
@@ -375,7 +406,8 @@ def mixin_ancestors0(decl, type_name, align_params:, included_modules:, included
               NoMixinFoundError.check!(member.name, env: env, member: member)
 
               module_name = env.normalize_module_name(module_name)
-              extended_modules << Definition::Ancestor::Instance.new(name: module_name, args: module_args, source: member)
+              mod = Definition::Ancestor::Instance.new(name: module_name, args: module_args, source: member)
+              extended_modules << mod unless extended_modules.include?(mod)
             when member.name.interface? && extended_interfaces
               NoMixinFoundError.check!(member.name, env: env, member: member)
 

sig/ancestor_builder.rbs

@@ -138,6 +138,8 @@ module RBS
 
       def validate_super_class!: (TypeName, Environment::ClassEntry) -> void
 
+      def auto_extended_modules: (TypeName, Resolver::TypeNameResolver) -> Array[Definition::Ancestor::Instance]
+
       def mixin_ancestors: (Environment::ClassEntry | Environment::ModuleEntry,
                             TypeName,
                             included_modules: Array[Definition::Ancestor::Instance]?,

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