Merge pull request #1188 from Shopify/vs/performance_improvements

Various performance optimizations for trackers

lib/tapioca/gem/listeners/foreign_constants.rb

@@ -24,7 +24,7 @@ def on_scope(event)
           # The way we identify these "foreign constants" is by asking the mixin tracker which
           # constants have mixed in the current module that we are handling. We add all the
           # constants that we discover to the pipeline to be processed.
-          Runtime::Trackers::Mixin.constants_with_mixin(mixin).each_value do |location_info|
+          Runtime::Trackers::Mixin.constants_with_mixin(mixin).each do |mixin_type, location_info|
             location_info.each do |constant, location|
               next unless mixed_in_by_gem?(location)
 
@@ -35,10 +35,11 @@ def on_scope(event)
               # base constant. Then, generate RBIs as if the base constant is extending the mixin,
               # which is functionally equivalent to including or prepending to the singleton class.
               if !name && constant.singleton_class?
-                attached_class = attached_class_of(constant)
+                attached_class = Runtime::Trackers::Mixin.resolve_to_attached_class(constant, mixin, mixin_type)
                 next unless attached_class
 
                 constant = attached_class
+                name = @pipeline.name_of(constant)
               end
 
               @pipeline.push_foreign_constant(name, constant) if name

lib/tapioca/gem/listeners/mixins.rb

@@ -73,11 +73,7 @@ def add_mixins(tree, constant, mods, mixin_type)
           ).returns(T::Boolean)
         end
         def mixed_in_by_gem?(constant, mixin, mixin_type)
-          mixin_location =
-            T.cast(
-              Runtime::Trackers::Mixin.constants_with_mixin(mixin).dig(mixin_type, constant),
-              T.nilable(String),
-            )
+          mixin_location = Runtime::Trackers::Mixin.mixin_location(mixin, mixin_type, constant)
 
           return true if mixin_location.nil?
 

lib/tapioca/runtime/trackers/constant_definition.rb

@@ -43,9 +43,10 @@ class ConstantLocation < T::Struct
 
         @creturn_tracepoint = TracePoint.trace(:c_return) do |tp|
           next unless tp.method_id == :new
-          next unless Module === tp.return_value
 
           key = tp.return_value
+          next unless Module === key
+
           loc = build_constant_location(tp, caller_locations)
           (@class_files[key] ||= Set.new) << loc
         end
@@ -58,8 +59,8 @@ def disable!
           end
 
           def build_constant_location(tp, locations)
-            file = resolve_loc(caller_locations)
-            lineno = file == File.realpath(tp.path) ? tp.lineno : 0
+            file = resolve_loc(locations)
+            lineno = File.identical?(file, tp.path) ? tp.lineno : 0
 
             ConstantLocation.new(path: file, lineno: lineno)
           end

lib/tapioca/runtime/trackers/mixin.rb

@@ -31,24 +31,49 @@ def with_disabled_registration(&block)
             with_disabled_tracker(&block)
           end
 
-          sig do
-            params(
-              constant: Module,
-              mixin: Module,
-              mixin_type: Type,
-            ).void
-          end
+          sig { params(constant: Module, mixin: Module, mixin_type: Type).void }
           def register(constant, mixin, mixin_type)
             return unless enabled?
 
             location = Reflection.resolve_loc(caller_locations)
 
-            constants = constants_with_mixin(mixin)
-            constants.fetch(mixin_type).store(constant, location)
+            register_with_location(constant, mixin, mixin_type, location)
+          end
+
+          def resolve_to_attached_class(constant, mixin, mixin_type)
+            attached_class = Reflection.attached_class_of(constant)
+            return unless attached_class
+
+            if mixin_type == Type::Include || mixin_type == Type::Prepend
+              location = mixin_location(mixin, mixin_type, constant)
+              register_with_location(constant, mixin, Type::Extend, T.must(location))
+            end
+
+            attached_class
           end
 
           sig { params(mixin: Module).returns(T::Hash[Type, T::Hash[Module, String]]) }
           def constants_with_mixin(mixin)
+            find_or_initialize_mixin_lookup(mixin)
+          end
+
+          sig { params(mixin: Module, mixin_type: Type, constant: Module).returns(T.nilable(String)) }
+          def mixin_location(mixin, mixin_type, constant)
+            find_or_initialize_mixin_lookup(mixin).dig(mixin_type, constant)
+          end
+
+          private
+
+          sig { params(constant: Module, mixin: Module, mixin_type: Type, location: String).void }
+          def register_with_location(constant, mixin, mixin_type, location)
+            return unless @enabled
+
+            constants = find_or_initialize_mixin_lookup(mixin)
+            constants.fetch(mixin_type).store(constant, location)
+          end
+
+          sig { params(mixin: Module).returns(T::Hash[Type, T::Hash[Module, String]]) }
+          def find_or_initialize_mixin_lookup(mixin)
             @mixins_to_constants[mixin] ||= {
               Type::Prepend => {}.compare_by_identity,
               Type::Include => {}.compare_by_identity,
@@ -70,8 +95,6 @@ def prepend_features(constant)
         Tapioca::Runtime::Trackers::Mixin::Type::Prepend,
       )
 
-      register_extend_on_attached_class(constant) if constant.singleton_class?
-
       super
     end
 
@@ -82,8 +105,6 @@ def append_features(constant)
         Tapioca::Runtime::Trackers::Mixin::Type::Include,
       )
 
-      register_extend_on_attached_class(constant) if constant.singleton_class?
-
       super
     end
 
@@ -95,21 +116,5 @@ def extend_object(obj)
       ) if Module === obj
       super
     end
-
-    private
-
-    # Including or prepending on a singleton class is functionally equivalent to extending the
-    # attached class. Registering the mixin as an extend on the attached class ensures that
-    # this mixin can be found whether searching for an include/prepend on the singleton class
-    # or an extend on the attached class.
-    def register_extend_on_attached_class(constant)
-      attached_class = Tapioca::Runtime::Reflection.attached_class_of(constant)
-
-      Tapioca::Runtime::Trackers::Mixin.register(
-        T.cast(attached_class, Module),
-        self,
-        Tapioca::Runtime::Trackers::Mixin::Type::Extend,
-      ) if attached_class
-    end
   end)
 end

spec/tapioca/cli/gem_spec.rb

@@ -1323,6 +1323,20 @@ module Foo; end
 
           result = @project.tapioca("gem foo")
           assert_empty_stderr(result)
+
+          assert_project_file_equal("sorbet/rbi/gems/[email protected]", <<~RBI)
+            # typed: true
+
+            # DO NOT EDIT MANUALLY
+            # This is an autogenerated file for types exported from the `foo` gem.
+            # Please instead update this file by running `bin/tapioca gem foo`.
+
+            class Bar
+              extend ::Foo
+            end
+
+            module Foo; end
+          RBI
         end
 
         it "must not load engines in the application" do