Merge pull request #293 from Shopify/uk-method-generation-improvements

Stop generating more methods than necessary

Author: paracycle

lib/tapioca/gem/listeners/methods.rb

@@ -56,10 +56,30 @@ def compile_directly_owned_methods(
         def compile_method(tree, symbol_name, constant, method, visibility = RBI::Public.new)
           return unless method
           return unless method_owned_by_constant?(method, constant)
-          return if @pipeline.symbol_in_payload?(symbol_name) && !@pipeline.method_in_gem?(method)
 
-          signature = lookup_signature_of(method)
-          method = signature.method if signature #: UnboundMethod
+          begin
+            signature = signature_of!(method)
+            method = signature.method if signature #: UnboundMethod
+
+            case @pipeline.method_definition_in_gem(method.name, constant)
+            when Pipeline::MethodUnknown
+              # This means that this is a C-method. Thus, we want to
+              # skip it only if the constant is an ignored one, since
+              # that probably means that we've hit a C-method for a
+              # core type.
+              return if @pipeline.symbol_in_payload?(symbol_name)
+            when Pipeline::MethodNotInGem
+              # Do not process this method, if it is not defined by the current gem
+              return
+            end
+          rescue SignatureBlockError => error
+            @pipeline.error_handler.call(<<~MSG)
+              Unable to compile signature for method: #{method.owner}##{method.name}
+                Exception raised when loading signature: #{error.cause.inspect}
+            MSG
+
+            signature = nil
+          end
 
           method_name = method.name.to_s
           return unless valid_method_name?(method_name)
@@ -192,18 +212,6 @@ def initialize_method_for(constant)
         def ignore?(event)
           event.is_a?(Tapioca::Gem::ForeignScopeNodeAdded)
         end
-
-        #: (UnboundMethod method) -> untyped
-        def lookup_signature_of(method)
-          signature_of!(method)
-        rescue LoadError, StandardError => error
-          @pipeline.error_handler.call(<<~MSG)
-            Unable to compile signature for method: #{method.owner}##{method.name}
-              Exception raised when loading signature: #{error.inspect}
-          MSG
-
-          nil
-        end
       end
     end
   end

lib/tapioca/gem/listeners/source_location.rb

@@ -24,20 +24,24 @@ def on_scope(event)
           # constants that are defined by multiple gems.
           locations = Runtime::Trackers::ConstantDefinition.locations_for(event.constant)
           location = locations.find do |loc|
-            Pathname.new(loc.path).realpath.to_s.include?(@pipeline.gem.full_gem_path)
+            Pathname.new(loc.file).realpath.to_s.include?(@pipeline.gem.full_gem_path)
           end
 
           # The location may still be nil in some situations, like constant aliases (e.g.: MyAlias = OtherConst). These
           # are quite difficult to attribute a correct location, given that the source location points to the original
           # constants and not the alias
-          add_source_location_comment(event.node, location.path, location.lineno) unless location.nil?
+          add_source_location_comment(event.node, location.file, location.line) unless location.nil?
         end
 
         # @override
         #: (MethodNodeAdded event) -> void
         def on_method(event)
-          file, line = event.method.source_location
-          add_source_location_comment(event.node, file, line)
+          definition = @pipeline.method_definition_in_gem(event.method.name, event.constant)
+
+          if Pipeline::MethodInGemWithLocation === definition
+            loc = definition.location
+            add_source_location_comment(event.node, loc.file, loc.line)
+          end
         end
 
         #: (RBI::NodeWithComments node, String? file, Integer? line) -> void

lib/tapioca/gem/pipeline.rb

@@ -107,35 +107,66 @@ def symbol_in_payload?(symbol_name)
         @payload_symbols.include?(symbol_name)
       end
 
-      # this looks something like:
-      # "(eval at /path/to/file.rb:123)"
-      # and we are just interested in the "/path/to/file.rb" part
-      EVAL_SOURCE_FILE_PATTERN = /\(eval at (.+):\d+\)/ #: Regexp
-
       #: ((String | Symbol) name) -> bool
       def constant_in_gem?(name)
-        return true unless Object.respond_to?(:const_source_location)
+        loc = const_source_location(name)
 
-        source_file, _ = Object.const_source_location(name)
-        return true unless source_file
-        # If the source location of the constant is "(eval)", all bets are off.
-        return true if source_file == "(eval)"
+        # If the source location of the constant isn't available or is "(eval)", all bets are off.
+        return true if loc.nil? || loc.file.nil? || loc.file == "(eval)"
 
-        # Ruby 3.3 adds automatic definition of source location for evals if
-        # `file` and `line` arguments are not provided. This results in the source
-        # file being something like `(eval at /path/to/file.rb:123)`. We try to parse
-        # this string to get the actual source file.
-        source_file = source_file.sub(EVAL_SOURCE_FILE_PATTERN, "\\1")
+        gem.contains_path?(loc.file)
+      end
 
-        gem.contains_path?(source_file)
+      class MethodDefinitionLookupResult
+        extend T::Helpers
+        abstract!
       end
 
-      #: (UnboundMethod method) -> bool
-      def method_in_gem?(method)
-        source_location = method.source_location&.first
-        return false if source_location.nil?
+      # The method doesn't seem to exist
+      class MethodUnknown < MethodDefinitionLookupResult; end
+
+      # The method is not defined in the gem
+      class MethodNotInGem < MethodDefinitionLookupResult; end
+
+      # The method probably defined in the gem but doesn't have a source location
+      class MethodInGemWithoutLocation < MethodDefinitionLookupResult; end
+
+      # The method defined in gem and has a source location
+      class MethodInGemWithLocation < MethodDefinitionLookupResult
+        extend T::Sig
+
+        #: Runtime::SourceLocation
+        attr_reader :location
+
+        #: (Runtime::SourceLocation location) -> void
+        def initialize(location)
+          @location = location
+          super()
+        end
+      end
+
+      #: (Symbol method_name, Module owner) -> MethodDefinitionLookupResult
+      def method_definition_in_gem(method_name, owner)
+        definitions = Tapioca::Runtime::Trackers::MethodDefinition.method_definitions_for(method_name, owner)
+
+        # If the source location of the method isn't available, signal that by returning nil.
+        return MethodUnknown.new if definitions.empty?
+
+        # Look up the first entry that matches a file in the gem.
+        found = definitions.find { |loc| @gem.contains_path?(loc.file) }
+
+        unless found
+          # If the source location of the method is "(eval)", err on the side of caution and include the method.
+          found = definitions.find { |loc| loc.file == "(eval)" }
+          # However, we can just return true to signal that the method should be included.
+          # We can't provide a source location for it, but we want it to be included in the gem RBI.
+          return MethodInGemWithoutLocation.new if found
+        end
+
+        # If we searched but couldn't find a source location in the gem, return false to signal that.
+        return MethodNotInGem.new unless found
 
-        @gem.contains_path?(source_location)
+        MethodInGemWithLocation.new(found)
       end
 
       # Helpers

lib/tapioca/runtime/reflection.rb

@@ -1,6 +1,8 @@
 # typed: strict
 # frozen_string_literal: true
 
+require "tapioca/runtime/source_location"
+
 # On Ruby 3.2 or newer, Class defines an attached_object method that returns the
 # attached class of a singleton class without iterating ObjectSpace. On older
 # versions of Ruby, we fall back to iterating ObjectSpace.
@@ -125,15 +127,19 @@ def qualified_name_of(constant)
         end
       end
 
+      SignatureBlockError = Class.new(Tapioca::Error)
+
       #: ((UnboundMethod | Method) method) -> untyped
       def signature_of!(method)
         T::Utils.signature_for_method(method)
+      rescue LoadError, StandardError
+        Kernel.raise SignatureBlockError
       end
 
       #: ((UnboundMethod | Method) method) -> untyped
       def signature_of(method)
         signature_of!(method)
-      rescue LoadError, StandardError
+      rescue SignatureBlockError
         nil
       end
 
@@ -169,23 +175,46 @@ def descendants_of(klass)
         T.unsafe(result)
       end
 
+      #: ((String | Symbol) constant_name) -> SourceLocation?
+      def const_source_location(constant_name)
+        return unless Object.respond_to?(:const_source_location)
+
+        file, line = Object.const_source_location(constant_name)
+
+        SourceLocation.from_loc([file, line]) if file && line
+      end
+
       # Examines the call stack to identify the closest location where a "require" is performed
       # by searching for the label "<top (required)>" or "block in <class:...>" in the
       # case of an ActiveSupport.on_load hook. If none is found, it returns the location
       # labeled "<main>", which is the original call site.
-      #: (Array[Thread::Backtrace::Location]? locations) -> String
+      #: (Array[Thread::Backtrace::Location]? locations) -> SourceLocation?
       def resolve_loc(locations)
-        return "" unless locations
+        return unless locations
 
+        # Find the location of the closest file load, which should give us the location of the file that
+        # triggered the definition.
         resolved_loc = locations.find do |loc|
           label = loc.label
           next unless label
 
           REQUIRED_FROM_LABELS.include?(label) || label.start_with?("block in <class:")
         end
-        return "" unless resolved_loc
+        return unless resolved_loc
+
+        resolved_loc_path = resolved_loc.absolute_path || resolved_loc.path
+
+        # Find the location of the last frame in this file to get the most accurate line number.
+        resolved_loc = locations.find { |loc| loc.absolute_path == resolved_loc_path }
+        return unless resolved_loc
+
+        # If the last operation was a `require`, and we have no more frames,
+        # we are probably dealing with a C-method.
+        return if locations.first&.label == "require"
+
+        file = resolved_loc.absolute_path || resolved_loc.path || ""
 
-        resolved_loc.absolute_path || ""
+        SourceLocation.from_loc([file, resolved_loc.lineno])
       end
 
       #: (Module constant) -> Set[String]

lib/tapioca/runtime/source_location.rb

@@ -0,0 +1,44 @@
+# typed: true
+# frozen_string_literal: true
+
+module Tapioca
+  module Runtime
+    class SourceLocation
+      # this looks something like:
+      # "(eval at /path/to/file.rb:123)"
+      # and we are interested in the "/path/to/file.rb" and "123" parts
+      EVAL_SOURCE_FILE_PATTERN = /^\(eval at (?<file>.+):(?<line>\d+)\)/ #: Regexp
+
+      #: String
+      attr_reader :file
+
+      #: Integer
+      attr_reader :line
+
+      def initialize(file:, line:)
+        # Ruby 3.3 adds automatic definition of source location for evals if
+        # `file` and `line` arguments are not provided. This results in the source
+        # file being something like `(eval at /path/to/file.rb:123)`. We try to parse
+        # this string to get the actual source file.
+        eval_pattern_match = EVAL_SOURCE_FILE_PATTERN.match(file)
+        if eval_pattern_match
+          file = eval_pattern_match[:file]
+          line = eval_pattern_match[:line].to_i
+        end
+
+        @file = file
+        @line = line
+      end
+
+      # force all callers to use the from_loc method
+      private_class_method :new
+
+      class << self
+        #: ([String?, Integer?]? loc) -> SourceLocation?
+        def from_loc(loc)
+          new(file: loc.first, line: loc.last) if loc&.first && loc.last
+        end
+      end
+    end
+  end
+end

lib/tapioca/runtime/trackers.rb

@@ -52,3 +52,4 @@ def register_tracker(tracker)
 require "tapioca/runtime/trackers/constant_definition"
 require "tapioca/runtime/trackers/autoload"
 require "tapioca/runtime/trackers/required_ancestor"
+require "tapioca/runtime/trackers/method_definition"

lib/tapioca/runtime/trackers/constant_definition.rb

@@ -13,12 +13,7 @@ module ConstantDefinition
         extend Reflection
         extend T::Sig
 
-        class ConstantLocation < T::Struct
-          const :lineno, Integer
-          const :path, String
-        end
-
-        @class_files = {}.compare_by_identity
+        @class_files = {}.compare_by_identity #: Hash[Module, Set[SourceLocation]]
 
         # Immediately activated upon load. Observes class/module definition.
         @class_tracepoint = TracePoint.trace(:class) do |tp|
@@ -28,14 +23,17 @@ class ConstantLocation < T::Struct
 
           path = tp.path
           if File.exist?(path)
-            loc = build_constant_location(tp, caller_locations)
+            loc = build_source_location(tp, caller_locations)
           else
             caller_location = T.must(caller_locations)
               .find { |loc| loc.path && File.exist?(loc.path) }
 
             next unless caller_location
 
-            loc = ConstantLocation.new(path: caller_location.absolute_path || "", lineno: caller_location.lineno)
+            loc = SourceLocation.from_loc([
+              caller_location.absolute_path || "",
+              caller_location.lineno,
+            ])
           end
 
           (@class_files[key] ||= Set.new) << loc
@@ -47,31 +45,37 @@ class ConstantLocation < T::Struct
           key = tp.return_value
           next unless Module === key
 
-          loc = build_constant_location(tp, caller_locations)
+          loc = build_source_location(tp, caller_locations)
           (@class_files[key] ||= Set.new) << loc
         end
 
         class << self
+          extend T::Sig
+
           def disable!
             @class_tracepoint.disable
             @creturn_tracepoint.disable
             super
           end
 
-          def build_constant_location(tp, locations)
-            file = resolve_loc(locations)
-            lineno = File.identical?(file, tp.path) ? tp.lineno : 0
+          def build_source_location(tp, locations)
+            loc = resolve_loc(locations)
+            file = loc&.file
+            line = loc&.line
+            lineno = file && File.identical?(file, tp.path) ? tp.lineno : (line || 0)
 
-            ConstantLocation.new(path: file, lineno: lineno)
+            SourceLocation.from_loc([file || "", lineno])
           end
 
           # Returns the files in which this class or module was opened. Doesn't know
           # about situations where the class was opened prior to +require+ing,
           # or where metaprogramming was used via +eval+, etc.
+          #: (Module klass) -> Set[String]
           def files_for(klass)
-            locations_for(klass).map(&:path).to_set
+            locations_for(klass).map(&:file).to_set
           end
 
+          #: (Module klass) -> Set[SourceLocation]
           def locations_for(klass)
             @class_files.fetch(klass, Set.new)
           end