From 4cd88474849598711ff79e365811c5dc14cff4e1 Mon Sep 17 00:00:00 2001
From: Emily Samp <emily.samp@shopify.com>
Date: Fri, 9 Aug 2024 17:00:46 -0500
Subject: [PATCH] WIP: use a `MethodDefinition` tracker to properly attribute
 methods definitions that are triggered by one gem but defined in another

---
 lib/tapioca/gem/listeners/methods.rb          |  2 +-
 lib/tapioca/gem/pipeline.rb                   | 18 +++++++-
 lib/tapioca/runtime/trackers.rb               |  1 +
 .../runtime/trackers/method_definition.rb     | 43 +++++++++++++++++++
 4 files changed, 61 insertions(+), 3 deletions(-)
 create mode 100644 lib/tapioca/runtime/trackers/method_definition.rb

diff --git a/lib/tapioca/gem/listeners/methods.rb b/lib/tapioca/gem/listeners/methods.rb
index 48ac37ba1..5ad05224f 100644
--- a/lib/tapioca/gem/listeners/methods.rb
+++ b/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_in_gem?(constant, method)
             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
diff --git a/lib/tapioca/gem/pipeline.rb b/lib/tapioca/gem/pipeline.rb
index 138a1a200..f5ba66b03 100644
--- a/lib/tapioca/gem/pipeline.rb
+++ b/lib/tapioca/gem/pipeline.rb
@@ -148,8 +148,8 @@ def constant_in_gem?(name)
         gem.contains_path?(source_file)
       end
 
-      sig { params(method: UnboundMethod).returns(T.nilable(T::Boolean)) }
-      def method_in_gem?(method)
+      sig { params(constant: Module, method: UnboundMethod).returns(T.nilable(T::Boolean)) }
+      def method_in_gem?(constant, method)
         source_file, _ = method.source_location
         # Ruby 3.3 adds automatic definition of source location for evals if
         # `file` and `line` arguments are not provided. This results in the source
@@ -163,6 +163,20 @@ def method_in_gem?(method)
         # 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)"
 
+        return true if @gem.contains_path?(source_file)
+
+        method_definitions = Tapioca::Runtime::Trackers::MethodDefinition.method_definitions_for(constant)
+
+        source_file = method_definitions[method.name]
+
+        # If the source location of the method isn't available, signal that by returning nil.
+        return unless source_file
+
+        # 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)"
+
+        source_file = source_file&.sub(EVAL_SOURCE_FILE_PATTERN, "\\1")
+
         @gem.contains_path?(source_file)
       end
 
diff --git a/lib/tapioca/runtime/trackers.rb b/lib/tapioca/runtime/trackers.rb
index d57f087e3..f1f2aa4e0 100644
--- a/lib/tapioca/runtime/trackers.rb
+++ b/lib/tapioca/runtime/trackers.rb
@@ -56,3 +56,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"
diff --git a/lib/tapioca/runtime/trackers/method_definition.rb b/lib/tapioca/runtime/trackers/method_definition.rb
new file mode 100644
index 000000000..fda6cec8d
--- /dev/null
+++ b/lib/tapioca/runtime/trackers/method_definition.rb
@@ -0,0 +1,43 @@
+# typed: true
+# frozen_string_literal: true
+
+module Tapioca
+  module Runtime
+    module Trackers
+      module MethodDefinition
+        extend Tracker
+        extend T::Sig
+
+        @method_definitions = {}.compare_by_identity
+
+        class << self
+          extend T::Sig
+
+          sig { params(constant: Module, method_name: Symbol).void }
+          def register(constant, method_name)
+            return unless enabled?
+
+            @method_definitions[constant] ||= {}
+            @method_definitions[constant][method_name] = Reflection.resolve_loc(caller_locations)
+          end
+
+          sig { params(constant: Module).returns(T::Hash[Symbol, T::Array[T::Hash[Symbol, String]]]) }
+          def method_definitions_for(constant)
+            @method_definitions[constant] || {}
+          end
+        end
+      end
+    end
+  end
+end
+
+class Module
+  prepend(Module.new do
+    def method_added(method_name)
+      Tapioca::Runtime::Trackers::MethodDefinition.register(self, method_name)
+      super
+    end
+  end)
+end
+
+# TODO: are there other methods I have to override here?