Fix signature lookup through prepended methods

lib/tapioca/runtime/reflection.rb

@@ -125,7 +125,36 @@ def qualified_name_of(constant)
 
       #: ((UnboundMethod | Method) method) -> untyped
       def signature_of!(method)
-        T::Utils.signature_for_method(method)
+        # We use `T::Utils.signature_for_method` to find a method's signature. To support
+        # `prepend` in cases like this, we need to call it a second time:
+        #   module Wrapper
+        #     def foo = super
+        #   end
+        #   class Example
+        #     sig { void }
+        #     def foo; end
+        #     prepend Wrapper
+        #   end
+        # Looking up the original `Example#foo` evaluates its `sig`. Because `Wrapper`
+        # was prepended, Ruby now resolves `Example#foo` to `Wrapper#foo`, so Sorbet
+        # stores the signature for `Wrapper#foo`. The first pass causes the store; the
+        # second pass finds `Example#foo`'s signature on `Wrapper#foo`.
+        needs_second_pass = false #: bool
+        2.times do
+          current_method = method #: (UnboundMethod | Method)?
+
+          while current_method
+            needs_second_pass ||= T::Private::Methods.has_sig_block_for_method(current_method)
+            signature = T::Utils.signature_for_method(current_method)
+            return signature if signature
+
+            current_method = current_method.super_method
+          end
+
+          return nil unless needs_second_pass
+        end
+
+        nil
       rescue LoadError, StandardError
         Kernel.raise SignatureBlockError
       end

sorbet/rbi/shims/sorbet.rbi

@@ -25,6 +25,9 @@ module T::Private
   module Methods
     ARG_NOT_PROVIDED = T.let(T.unsafe(nil), Object)
 
+    sig { params(method: T.any(::Method, ::UnboundMethod)).returns(T::Boolean) }
+    def self.has_sig_block_for_method(method); end
+
     class Declaration
       def on_failure; end
       def on_failure=(on_failure); end

spec/tapioca/dsl/compiler_spec.rb

@@ -149,6 +149,60 @@ def some_attribute; end
           assert_equal(expected, rbi_for(:Post))
         end
 
+        it "compiles the nearest signatures through prepended modules" do
+          add_ruby_file("post.rb", <<~RUBY)
+            module FirstPostWrapper
+              def publish(...)
+                super
+              end
+            end
+
+            module SecondPostWrapper
+              extend T::Sig
+
+              sig { params(slug: Symbol).returns(String) }
+              def find(slug)
+                super
+              end
+
+              def publish(*args, **kwargs, &block)
+                super
+              end
+            end
+
+            class Post
+              extend T::Sig
+
+              sig { params(title: String).returns(Integer) }
+              def find(title)
+                title.length
+              end
+
+              sig { params(title: String).returns(Integer) }
+              def publish(title)
+                title.length
+              end
+
+              prepend FirstPostWrapper
+              prepend SecondPostWrapper
+            end
+          RUBY
+
+          expected = <<~RBI
+            # typed: strong
+
+            class Post
+              sig { params(slug: ::Symbol).returns(::String) }
+              def find(slug); end
+
+              sig { params(title: ::String).returns(::Integer) }
+              def publish(title); end
+            end
+          RBI
+
+          assert_equal(expected, rbi_for(:Post))
+        end
+
         it "compiles a class that overrides caller_locations" do
           add_ruby_file("post.rb", <<~RUBY)
             class Post

spec/tapioca/dsl/compilers/active_job_spec.rb

@@ -163,6 +163,43 @@ def perform_now(user_id); end
               RBI
               assert_equal(expected, rbi_for(:NotifyJob))
             end
+
+            it "generates RBS comment signatures through prepended modules" do
+              add_ruby_file("job.rb", <<~RUBY)
+                # typed: strict
+
+                module NotifyJobInstrumentation
+                  def perform(*args, **kwargs, &block)
+                    super
+                  end
+                end
+
+                class NotifyJob < ActiveJob::Base
+                  #: (Integer) -> void
+                  def perform(user_id)
+                    # ...
+                  end
+
+                  prepend NotifyJobInstrumentation
+                end
+              RUBY
+
+              expected = template(<<~RBI)
+                # typed: strong
+
+                class NotifyJob
+                  class << self
+                    sig { params(user_id: ::Integer, block: T.nilable(T.proc.params(job: NotifyJob).void)).returns(T.any(NotifyJob, FalseClass)) }
+                    def perform_later(user_id, &block); end
+
+                    sig { params(user_id: ::Integer).void }
+                    def perform_now(user_id); end
+                  end
+                end
+              RBI
+
+              assert_equal(expected, rbi_for(:NotifyJob))
+            end
           end
         end
       end

spec/tapioca/runtime/reflection_spec.rb

@@ -80,6 +80,46 @@ def bad_method
       def unknown_method
         ' ¯\_(ツ)_/¯ '
       end
+
+      #: (String) -> String
+      def wrapped_method(value)
+        value
+      end
+    end
+
+    module UnsignedSignatureWrapper
+      def unknown_method(...)
+        super
+      end
+
+      def wrapped_method(...)
+        super
+      end
+    end
+
+    SignatureFoo.prepend(UnsignedSignatureWrapper)
+
+    module ParentSignatureWrapper
+      def inherited_wrapped_method(...)
+        super
+      end
+    end
+
+    class ParentSignatureFoo
+      #: -> String
+      def inherited_wrapped_method = "wrapped"
+
+      prepend ParentSignatureWrapper
+    end
+
+    module ChildSignatureWrapper
+      def inherited_wrapped_method(...)
+        super
+      end
+    end
+
+    class ChildSignatureFoo < ParentSignatureFoo
+      prepend ChildSignatureWrapper
     end
 
     class ReflectionSpec < Minitest::Spec
@@ -154,7 +194,33 @@ class ReflectionSpec < Minitest::Spec
 
           it "returns nil when a signature is not defined" do
             method = SignatureFoo.instance_method(:unknown_method)
-            assert_nil(Runtime::Reflection.signature_of(method))
+            calls = []
+
+            signature = T::Utils.stub(:signature_for_method, ->(current_method) do
+              calls << current_method
+              nil
+            end) do
+              Runtime::Reflection.signature_of(method)
+            end
+
+            assert_nil(signature)
+            assert_equal([method, method.super_method], calls)
+          end
+
+          it "returns a signature from a super method when a prepended method has none" do
+            method = SignatureFoo.instance_method(:wrapped_method)
+            signature = Runtime::Reflection.signature_of(method)
+
+            refute_nil(signature)
+            assert_equal("::String", signature.return_type.to_s)
+          end
+
+          it "returns an inherited signature hidden by prepended methods on both classes" do
+            signature = Runtime::Reflection.signature_of(ChildSignatureFoo.instance_method(:inherited_wrapped_method))
+
+            refute_nil(signature)
+            assert_equal(ParentSignatureFoo, signature.method.owner)
+            assert_equal("::String", signature.return_type.to_s)
           end
 
           it "returns nil when a signature block raises an exception" do