Change how method definitions are stored and retrieved

Previously we were doing two things wrong:

1. Looking up the method object from the class that the `{singleton_}method_added` call was happening on and indexing method definitions keyed by that was mostly working. However, it was failing in an interesting way when a method was an alias to another method. In that case, the method object would be the same, but the method definition would be different. For example, `alias_method :foo, :bar` would cause the lookup of `foo` method to return the definition of `bar` method. That would cause us to fail generating `foo` method for a gem, if `bar` was defined in another gem.
2. We were storing a single definition for a method, but that proved to be not enough. For example, if a method was defined in a gem and then redefined in application code, we would only store the definition from the application location, and never be able to attribute the method to the gem it was originally defined in. That lead us to not generate that method for the gem. Instead, we now store all definitions of a method in `MethodDefinition` tracker. When looking up a method definition, we now look up all definitions. If there are no definitions, that means the method is probably a C-method, so we return `nil`. If there are definitions, we see if we can find one that matches the gem location. If not, we check for `(eval)` locations, in which case, we include the method but can't return a source location. If all fails, we return `false` to signal that we couldn't find a definition. If we are able to find a definition that matches the gem location, we return that definition.

This new logic is used both in `Tapioca::Gem::Listeners::Methods` and `Tapioca::Gem::Listeners::SourceLocation` listeners. The former uses it to check if the method should be included in the gem RBI, and the latter uses it to add more correct source location information in the comments.

lib/tapioca/gem/listeners/methods.rb

@@ -76,7 +76,7 @@ def compile_method(tree, symbol_name, constant, method, visibility = RBI::Public
             signature = signature_of!(method)
             method = T.let(signature.method, UnboundMethod) if signature
 
-            case @pipeline.method_in_gem?(method)
+            case @pipeline.method_definition_in_gem(method.name, constant)
             when nil
               # This means that this is a C-method. Thus, we want to
               # skip it only if the constant is an ignored one, since
@@ -86,6 +86,8 @@ def compile_method(tree, symbol_name, constant, method, visibility = RBI::Public
             when false
               # Do not process this method, if it is not defined by the current gem
               return
+            else
+              # We want to continue processing this method.
             end
           rescue SignatureBlockError => error
             @pipeline.error_handler.call(<<~MSG)

lib/tapioca/gem/listeners/source_location.rb

@@ -33,8 +33,14 @@ def on_scope(event)
 
         sig { override.params(event: MethodNodeAdded).void }
         def on_method(event)
-          file, line = Tapioca::Runtime::Trackers::MethodDefinition.method_definition_for(event.method)
-          add_source_location_comment(event.node, file, line)
+          definition = @pipeline.method_definition_in_gem(event.method.name, event.constant)
+
+          case definition
+          when NilClass, FalseClass, TrueClass
+            return
+          else
+            add_source_location_comment(event.node, definition.first, definition.last)
+          end
         end
 
         sig { params(node: RBI::NodeWithComments, file: T.nilable(String), line: T.nilable(Integer)).void }

lib/tapioca/gem/pipeline.rb

@@ -126,44 +126,38 @@ 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 = T.let(/\(eval at (.+):\d+\)/, Regexp)
-
       sig { params(name: T.any(String, Symbol)).returns(T::Boolean) }
       def constant_in_gem?(name)
-        return true unless Object.respond_to?(:const_source_location)
-
-        source_file, _ = Object.const_source_location(name)
-        # 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")
+        source_file, _ = const_source_location(name)
 
         # If the source location of the constant isn't available or is "(eval)", all bets are off.
         return true if source_file.nil? || source_file == "(eval)"
 
         gem.contains_path?(source_file)
       end
 
-      sig { params(method: UnboundMethod).returns(T.nilable(T::Boolean)) }
-      def method_in_gem?(method)
-        source_file, _ = Tapioca::Runtime::Trackers::MethodDefinition.method_definition_for(method)
-        # 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")
+      sig { params(method_name: Symbol, owner: Module).returns(T.any([String, Integer], NilClass, T::Boolean)) }
+      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 unless source_file # rubocop:disable Style/ReturnNilInPredicateMethodDefinition
+        return if definitions.empty? # rubocop:disable Style/ReturnNilInPredicateMethodDefinition
+
+        # Look up the first entry that matches a file in the gem.
+        found = definitions.find { |file, _line| @gem.contains_path?(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 { |file, _line| 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 true if found
+        end
 
-        # If the source location of the method is "(eval)", err on the side of caution and include the method.
-        return true if source_file == "(eval)"
+        # If we searched but couldn't find a source location in the gem, return false to signal that.
+        return false unless found
 
-        @gem.contains_path?(source_file)
+        found
       end
 
       # Helpers

lib/tapioca/runtime/reflection.rb

@@ -33,6 +33,10 @@ module Reflection
       UNDEFINED_CONSTANT = T.let(Module.new.freeze, Module)
 
       REQUIRED_FROM_LABELS = T.let(["<top (required)>", "<main>"].freeze, T::Array[String])
+      # 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 = T.let(/\(eval at (.+):\d+\)/, Regexp)
 
       T::Sig::WithoutRuntime.sig { params(constant: BasicObject).returns(T::Boolean) }
       def constant_defined?(constant)
@@ -181,6 +185,21 @@ def descendants_of(klass)
         T.unsafe(result)
       end
 
+      sig { params(constant_name: T.any(String, Symbol)).returns(T.nilable([String, Integer])) }
+      def const_source_location(constant_name)
+        return unless Object.respond_to?(:const_source_location)
+
+        file, line = Object.const_source_location(constant_name)
+
+        # 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.
+        file = file&.sub(EVAL_SOURCE_FILE_PATTERN, "\\1")
+
+        [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)>". If none is found, it returns the location
       # labeled "<main>", which is the original call site.
@@ -201,7 +220,14 @@ def resolve_loc(locations)
         # we are probably dealing with a C-method.
         return if locations.first&.label == "require"
 
-        [resolved_loc.absolute_path || "", resolved_loc.lineno]
+        file = resolved_loc.absolute_path || ""
+        # 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.
+        file = file.sub(EVAL_SOURCE_FILE_PATTERN, "\\1")
+
+        [file, resolved_loc.lineno]
       end
 
       sig { params(constant: Module).returns(T::Set[String]) }

lib/tapioca/runtime/trackers/method_definition.rb

@@ -10,14 +10,14 @@ module MethodDefinition
 
         @method_definitions = T.let(
           {}.compare_by_identity,
-          T::Hash[Module, T::Hash[Symbol, T.nilable([String, Integer])]],
+          T::Hash[Module, T::Hash[Symbol, T::Array[[String, Integer]]]],
         )
 
         class << self
           extend T::Sig
 
-          sig { params(method: T.any(Method, UnboundMethod), locations: T::Array[Thread::Backtrace::Location]).void }
-          def register(method, locations)
+          sig { params(method_name: Symbol, owner: Module, locations: T::Array[Thread::Backtrace::Location]).void }
+          def register(method_name, owner, locations)
             return unless enabled?
             # If Sorbet runtime is redefining a method, it sets this to true.
             # In those cases, we should skip the registration, as the method's original
@@ -27,19 +27,27 @@ def register(method, locations)
             loc = Reflection.resolve_loc(locations)
             return unless loc
 
-            methods_for_owner(method.owner).store(method.name, loc)
+            registrations_for(method_name, owner) << loc
           end
 
-          sig { params(method: T.any(Method, UnboundMethod)).returns(T.nilable([String, Integer])) }
-          def method_definition_for(method)
-            methods_for_owner(method.owner).fetch(method.name, method.source_location)
+          sig { params(method_name: Symbol, owner: Module).returns(T::Array[[String, Integer]]) }
+          def method_definitions_for(method_name, owner)
+            definitions = registrations_for(method_name, owner)
+
+            if definitions.empty?
+              source_loc = owner.instance_method(method_name).source_location
+              definitions = [source_loc] if source_loc
+            end
+
+            definitions
           end
 
           private
 
-          sig { params(owner: Module).returns(T::Hash[Symbol, T.nilable([String, Integer])]) }
-          def methods_for_owner(owner)
-            @method_definitions[owner] ||= {}
+          sig { params(method_name: Symbol, owner: Module).returns(T::Array[[String, Integer]]) }
+          def registrations_for(method_name, owner)
+            owner_lookup = (@method_definitions[owner] ||= {})
+            owner_lookup[method_name] ||= []
           end
         end
       end
@@ -50,18 +58,12 @@ def methods_for_owner(owner)
 class Module
   prepend(Module.new do
     def singleton_method_added(method_name)
-      Tapioca::Runtime::Trackers::MethodDefinition.register(
-        Tapioca::Runtime::Reflection.method_of(self, method_name),
-        caller_locations,
-      )
+      Tapioca::Runtime::Trackers::MethodDefinition.register(method_name, self.singleton_class, caller_locations)
       super
     end
 
     def method_added(method_name)
-      Tapioca::Runtime::Trackers::MethodDefinition.register(
-        instance_method(method_name),
-        caller_locations,
-      )
+      Tapioca::Runtime::Trackers::MethodDefinition.register(method_name, self, caller_locations)
       super
     end
   end)