Compile RBIs for constants not owned by the current gem

lib/tapioca/gem/events.rb

@@ -42,6 +42,23 @@ def initialize(symbol, constant)
       end
     end
 
+    class NotOwnedConstantFound < Event
+      extend T::Sig
+
+      sig { returns(String) }
+      attr_reader :symbol
+
+      sig { returns(Module).checked(:never) }
+      attr_reader :constant
+
+      sig { params(symbol: String, constant: Module).void.checked(:never) }
+      def initialize(symbol, constant)
+        super()
+        @symbol = symbol
+        @constant = constant
+      end
+    end
+
     class NodeAdded < Event
       extend T::Helpers
       extend T::Sig

lib/tapioca/gem/listeners.rb

@@ -13,4 +13,5 @@
 require "tapioca/gem/listeners/sorbet_signatures"
 require "tapioca/gem/listeners/sorbet_type_variables"
 require "tapioca/gem/listeners/subconstants"
+require "tapioca/gem/listeners/not_owned_constants"
 require "tapioca/gem/listeners/yard_doc"

lib/tapioca/gem/listeners/not_owned_constants.rb

@@ -0,0 +1,32 @@
+# typed: strict
+# frozen_string_literal: true
+
+module Tapioca
+  module Gem
+    module Listeners
+      class NotOwnedConstants < Base
+        extend T::Sig
+
+        include Runtime::Reflection
+
+        private
+
+        sig { override.params(event: ScopeNodeAdded).void }
+        def on_scope(event)
+          mixin = event.constant
+
+          # There are cases where we want to compile RBI for constants not owned by the current
+          # gem - e.g. constants defined in one gem that then have a mixin applied to them in
+          # another gem. This is especially common for Ruby standard library classes.
+          #
+          # We can identify these cases by keeping track of which classes have mixins, and then
+          # adding those classes to the pipeline.
+          Runtime::Trackers::Mixin.constants_with_mixin(mixin).each do |constant|
+            name = @pipeline.name_of(constant)
+            @pipeline.push_not_owned_constant(name, constant) if name
+          end
+        end
+      end
+    end
+  end
+end

lib/tapioca/gem/pipeline.rb

@@ -42,6 +42,7 @@ def initialize(gem, include_doc: false)
         @node_listeners << Gem::Listeners::Subconstants.new(self)
         @node_listeners << Gem::Listeners::YardDoc.new(self) if include_doc
         @node_listeners << Gem::Listeners::RemoveEmptyPayloadScopes.new(self)
+        @node_listeners << Gem::Listeners::NotOwnedConstants.new(self)
       end
 
       sig { returns(RBI::Tree) }
@@ -60,6 +61,11 @@ def push_constant(symbol, constant)
         @events << Gem::ConstantFound.new(symbol, constant)
       end
 
+      sig { params(symbol: String, constant: Module).void.checked(:never) }
+      def push_not_owned_constant(symbol, constant)
+        @events << Gem::NotOwnedConstantFound.new(symbol, constant)
+      end
+
       sig { params(symbol: String, constant: Module, node: RBI::Const).void.checked(:never) }
       def push_const(symbol, constant, node)
         @events << Gem::ConstNodeAdded.new(symbol, constant, node)
@@ -124,7 +130,7 @@ def dispatch(event)
         case event
         when Gem::SymbolFound
           on_symbol(event)
-        when Gem::ConstantFound
+        when Gem::ConstantFound, Gem::NotOwnedConstantFound
           on_constant(event)
         when Gem::NodeAdded
           on_node(event)
@@ -142,7 +148,7 @@ def on_symbol(event)
         push_constant(symbol, constant) if constant
       end
 
-      sig { params(event: Gem::ConstantFound).void.checked(:never) }
+      sig { params(event: T.any(Gem::ConstantFound, Gem::NotOwnedConstantFound)).void.checked(:never) }
       def on_constant(event)
         name = event.symbol
 
@@ -156,7 +162,12 @@ def on_constant(event)
         return if T::Enum === constant # T::Enum instances are defined via `compile_enums`
 
         mark_seen(name)
-        compile_constant(name, constant)
+
+        if event.is_a?(Gem::NotOwnedConstantFound)
+          compile_not_owned_constant(name, T.cast(constant, Module))
+        else
+          compile_constant(name, constant)
+        end
       end
 
       sig { params(event: Gem::NodeAdded).void }
@@ -166,6 +177,11 @@ def on_node(event)
 
       # Compile
 
+      sig { params(symbol: String, constant: Module).void }
+      def compile_not_owned_constant(symbol, constant)
+        compile_module(symbol, constant, not_owned_constant: true)
+      end
+
       sig { params(symbol: String, constant: BasicObject).void.checked(:never) }
       def compile_constant(symbol, constant)
         case constant
@@ -228,9 +244,9 @@ def compile_object(name, value)
         @root << node
       end
 
-      sig { params(name: String, constant: Module).void }
-      def compile_module(name, constant)
-        return unless defined_in_gem?(constant, strict: false)
+      sig { params(name: String, constant: Module, not_owned_constant: T::Boolean).void }
+      def compile_module(name, constant, not_owned_constant: false)
+        return unless defined_in_gem?(constant, strict: false) || not_owned_constant
         return if Tapioca::TypeVariableModule === constant
 
         scope =

lib/tapioca/runtime/trackers/mixin.rb

@@ -7,7 +7,8 @@ module Trackers
       module Mixin
         extend T::Sig
 
-        @mixin_map = {}.compare_by_identity
+        @constants_to_mixin_locations = {}.compare_by_identity
+        @mixins_to_constants = {}.compare_by_identity
 
         class Type < T::Enum
           enums do
@@ -20,26 +21,35 @@ class Type < T::Enum
         sig do
           params(
             constant: Module,
-            mod: Module,
+            mixin: Module,
             mixin_type: Type,
             locations: T.nilable(T::Array[Thread::Backtrace::Location])
           ).void
         end
-        def self.register(constant, mod, mixin_type, locations)
+        def self.register(constant, mixin, mixin_type, locations)
           locations ||= []
           locations.map!(&:absolute_path).uniq!
+
           locs = mixin_locations_for(constant)
-          locs.fetch(mixin_type).store(mod, T.cast(locations, T::Array[String]))
+          locs.fetch(mixin_type).store(mixin, T.cast(locations, T::Array[String]))
+
+          constants = constants_with_mixin(mixin)
+          constants << (constant)
         end
 
         sig { params(constant: Module).returns(T::Hash[Type, T::Hash[Module, T::Array[String]]]) }
         def self.mixin_locations_for(constant)
-          @mixin_map[constant] ||= {
+          @constants_to_mixin_locations[constant] ||= {
             Type::Prepend => {}.compare_by_identity,
             Type::Include => {}.compare_by_identity,
             Type::Extend => {}.compare_by_identity,
           }
         end
+
+        sig { params(mixin: Module).returns(T::Array[Module]) }
+        def self.constants_with_mixin(mixin)
+          @mixins_to_constants[mixin] ||= []
+        end
       end
     end
   end

spec/tapioca/cli/gem_spec.rb

@@ -1038,6 +1038,42 @@ def foo(a, b, c, d, e, f, g, h); end
           assert_empty_stderr(result)
           assert_success_status(result)
         end
+
+        it "must generate RBIs for constants defined in a different gem but with mixins in this gem" do
+          foo = mock_gem("foo", "0.0.1") do
+            write("lib/foo.rb", <<~RBI)
+              class Foo; end
+            RBI
+          end
+
+          bar = mock_gem("bar", "0.0.2") do
+            write("lib/bar.rb", <<~RBI)
+              module Bar; end
+
+              Foo.prepend(Bar)
+            RBI
+          end
+
+          @project.require_mock_gem(foo)
+          @project.require_mock_gem(bar)
+          @project.bundle_install
+
+          @project.tapioca("gem bar")
+
+          assert_project_file_equal("sorbet/rbi/gems/bar@0.0.2.rbi", <<~RBI)
+            # typed: true
+
+            # DO NOT EDIT MANUALLY
+            # This is an autogenerated file for types exported from the `bar` gem.
+            # Please instead update this file by running `bin/tapioca gem bar`.
+
+            module Bar; end
+
+            class Foo
+              include ::Bar
+            end
+          RBI
+        end
       end
 
       describe "sync" do

spec/tapioca/gem/pipeline_spec.rb

@@ -471,6 +471,118 @@ def to_foo(base = T.unsafe(nil)); end
       assert_equal(output, compile)
     end
 
+    it "compiles extensions to core types without adding methods" do
+      add_ruby_file("foo.rb", <<~RUBY)
+        class Foo
+          def to_s
+            "Foo"
+          end
+
+          def bar
+            "bar"
+          end
+
+          module Bar; end
+        end
+      RUBY
+
+      add_ruby_file("ext.rb", <<~RUBY)
+        class String
+          include Foo::Bar
+        end
+
+        class Hash
+          extend Foo::Bar
+        end
+
+        class Array
+          prepend Foo::Bar
+        end
+      RUBY
+
+      output = template(<<~RBI)
+        class Array
+          include ::Foo::Bar
+          include ::Enumerable
+          include ::JSON::Ext::Generator::GeneratorMethods::Array
+        end
+
+        class Foo
+          def bar; end
+          def to_s; end
+        end
+
+        module Foo::Bar; end
+
+        class Hash
+          include ::Enumerable
+          include ::JSON::Ext::Generator::GeneratorMethods::Hash
+          extend ::Foo::Bar
+        end
+
+        class String
+          include ::Comparable
+          include ::JSON::Ext::Generator::GeneratorMethods::String
+          include ::Foo::Bar
+          extend ::JSON::Ext::Generator::GeneratorMethods::String::Extend
+        end
+      RBI
+
+      assert_equal(output, compile)
+    end
+
+    it "compiles extensions to core types via #extend, #include, and #prepend methods" do
+      add_ruby_file("foo.rb", <<~RUBY)
+        class Foo
+          def to_s
+            "Foo"
+          end
+
+          def bar
+            "bar"
+          end
+
+          module Bar; end
+        end
+      RUBY
+
+      add_ruby_file("ext.rb", <<~RUBY)
+        String.include(Foo::Bar)
+        Hash.extend(Foo::Bar)
+        Array.prepend(Foo::Bar)
+      RUBY
+
+      output = template(<<~RBI)
+        class Array
+          include ::Foo::Bar
+          include ::Enumerable
+          include ::JSON::Ext::Generator::GeneratorMethods::Array
+        end
+
+        class Foo
+          def bar; end
+          def to_s; end
+        end
+
+        module Foo::Bar; end
+
+        class Hash
+          include ::Enumerable
+          include ::JSON::Ext::Generator::GeneratorMethods::Hash
+          extend ::Foo::Bar
+        end
+
+        class String
+          include ::Comparable
+          include ::JSON::Ext::Generator::GeneratorMethods::String
+          include ::Foo::Bar
+          extend ::JSON::Ext::Generator::GeneratorMethods::String::Extend
+        end
+      RBI
+
+      assert_equal(output, compile)
+    end
+
     it "compiles without annotations" do
       add_ruby_file("bar.rb", <<~RUBY)
         class Bar