Support finding signatures through prepended methods

lib/tapioca/dsl/compiler.rb

@@ -139,10 +139,16 @@ def parameters_types_from_signature(method_def, signature)
         params
       end
 
-      #: (RBI::Scope scope, (Method | UnboundMethod) method_def, ?class_method: bool) -> void
-      def create_method_from_def(scope, method_def, class_method: false)
-        parameters = compile_method_parameters_to_rbi(method_def)
-        return_type = compile_method_return_type_to_rbi(method_def)
+      #: (
+      #|   RBI::Scope scope,
+      #|   (Method | UnboundMethod) method_def,
+      #|   ?class_method: bool,
+      #|   ?lookup_from: Module[top]
+      #| ) -> void
+      def create_method_from_def(scope, method_def, class_method: false, lookup_from: constant)
+        signature = signature_of(method_def, lookup_from: lookup_from)
+        parameters = compile_method_parameters_to_rbi(method_def, signature: signature)
+        return_type = compile_method_return_type_to_rbi(method_def, signature: signature)
 
         scope.create_method(
           method_def.name.to_s,
@@ -152,9 +158,11 @@ def create_method_from_def(scope, method_def, class_method: false)
         )
       end
 
-      #: ((Method | UnboundMethod) method_def) -> Array[RBI::TypedParam]
-      def compile_method_parameters_to_rbi(method_def)
-        signature = signature_of(method_def)
+      #: ((Method | UnboundMethod) method_def, ?signature: untyped) -> Array[RBI::TypedParam]
+      def compile_method_parameters_to_rbi(
+        method_def,
+        signature: signature_of(method_def, lookup_from: constant)
+      )
         method_def = signature.nil? ? method_def : signature.method
         method_types = parameters_types_from_signature(method_def, signature)
 
@@ -188,9 +196,11 @@ def compile_method_parameters_to_rbi(method_def)
         end
       end
 
-      #: ((Method | UnboundMethod) method_def) -> String
-      def compile_method_return_type_to_rbi(method_def)
-        signature = signature_of(method_def)
+      #: ((Method | UnboundMethod) method_def, ?signature: untyped) -> String
+      def compile_method_return_type_to_rbi(
+        method_def,
+        signature: signature_of(method_def, lookup_from: constant)
+      )
         return_type = signature.nil? ? "T.untyped" : name_of_type(signature.return_type)
         sanitize_signature_types(return_type)
       end

lib/tapioca/dsl/compilers/action_controller_helpers.rb

@@ -94,14 +94,14 @@ def decorate
               # Generate a method definition in the helper module for each
               # helper method defined via the `helper_method` call in the controller.
               helpers_module.instance_methods(false).each do |method_name|
-                method = if proxied_helper_methods.include?(method_name)
-                  helper_method_proxy_target(method_name)
+                method, lookup_from = if proxied_helper_methods.include?(method_name)
+                  [helper_method_proxy_target(method_name), constant]
                 else
-                  helpers_module.instance_method(method_name)
+                  [helpers_module.instance_method(method_name), helpers_module]
                 end
 
                 if method
-                  create_method_from_def(helper_methods, method)
+                  create_method_from_def(helper_methods, method, lookup_from: lookup_from)
                 else
                   create_unknown_proxy_method(helper_methods, method_name)
                 end

lib/tapioca/dsl/compilers/graphql_mutation.rb

@@ -43,7 +43,8 @@ def decorate
           return unless constant.method_defined?(:resolve)
 
           method_def = constant.instance_method(:resolve)
-          return if signature_of(method_def) # Skip if the mutation already has an inline sig
+          # Skip if the mutation already has an inline signature.
+          return if signature_of(method_def, lookup_from: constant)
 
           arguments = constant.all_argument_definitions
           return if arguments.empty?

lib/tapioca/dsl/helpers/active_model_type_helper.rb

@@ -65,7 +65,7 @@ def lookup_arg_type_of_method(obj, method)
 
           #: (untyped obj, Symbol method) -> untyped
           def lookup_signature_of_method(obj, method)
-            Runtime::Reflection.signature_of(obj.method(method))
+            Runtime::Reflection.signature_of(obj.method(method), lookup_from: obj)
           rescue NameError
             nil
           end

lib/tapioca/dsl/helpers/graphql_type_helper.rb

@@ -77,7 +77,7 @@ def type_for(type, ignore_nilable_wrapper: false, prepare_method: nil)
             end
           when GraphQL::Schema::Scalar.singleton_class
             method = Runtime::Reflection.method_of(unwrapped_type, :coerce_input)
-            signature = Runtime::Reflection.signature_of(method)
+            signature = Runtime::Reflection.signature_of(method, lookup_from: unwrapped_type)
             return_type = signature&.return_type
 
             # Wrap as non-nilable for required arguments. `coerce_input` supports both
@@ -93,7 +93,7 @@ def type_for(type, ignore_nilable_wrapper: false, prepare_method: nil)
 
           prepared = false
           if prepare_method
-            prepare_signature = Runtime::Reflection.signature_of(prepare_method)
+            prepare_signature = Runtime::Reflection.signature_of(prepare_method, lookup_from: prepare_method.receiver)
             prepare_return_type = prepare_signature&.return_type
             if valid_return_type?(prepare_return_type)
               parsed_type = prepare_return_type&.to_s
@@ -122,7 +122,7 @@ def type_for_constant(constant)
           if constant.method_defined?(:prepare)
             prepare_method = constant.instance_method(:prepare)
 
-            prepare_signature = Runtime::Reflection.signature_of(prepare_method)
+            prepare_signature = Runtime::Reflection.signature_of(prepare_method, lookup_from: constant)
 
             return prepare_signature.return_type&.to_s if valid_return_type?(prepare_signature&.return_type)
           end

lib/tapioca/gem/listeners/methods.rb

@@ -68,7 +68,7 @@ def compile_method(tree, symbol_name, constant, method, visibility = RBI::Public
           return unless method_owned_by_constant?(method, constant)
 
           begin
-            signature = signature_of!(method)
+            signature = signature_of!(method, lookup_from: constant)
             signature ||= inferred_attr_writer_signature(method, constant)
             method = signature.method if signature #: UnboundMethod
 
@@ -198,7 +198,7 @@ def inferred_attr_writer_signature(method, constant)
           reader_method = attr_reader_for_writer(method, constant)
           return unless reader_method
 
-          reader_signature = signature_of(reader_method)
+          reader_signature = signature_of(reader_method, lookup_from: constant)
           return unless reader_signature
 
           build_attr_writer_signature(method, reader_method, reader_signature)
@@ -211,7 +211,7 @@ def attr_reader_for_writer(method, constant)
           return unless method.parameters == [[:req]]
 
           reader_method = T.let(constant.instance_method(method_name.delete_suffix("=").to_sym), UnboundMethod)
-          reader_method = original_method(reader_method)
+          reader_method = original_method(reader_method, constant)
           return unless same_source_location?(method, reader_method)
           return unless method_owned_by_constant?(reader_method, constant)
 
@@ -242,9 +242,9 @@ def build_attr_writer_signature(writer_method, reader_method, reader_signature)
           )
         end
 
-        #: (UnboundMethod method) -> UnboundMethod
-        def original_method(method)
-          T.let(signature_of(method)&.method || method, UnboundMethod)
+        #: (UnboundMethod method, Module[top] lookup_from) -> UnboundMethod
+        def original_method(method, lookup_from)
+          T.let(signature_of(method, lookup_from: lookup_from)&.method || method, UnboundMethod)
         end
 
         #: (UnboundMethod method, UnboundMethod other_method) -> bool

lib/tapioca/runtime/reflection.rb

@@ -19,7 +19,9 @@ module Reflection
       PUBLIC_INSTANCE_METHODS_METHOD = Module.instance_method(:public_instance_methods) #: UnboundMethod
       PROTECTED_INSTANCE_METHODS_METHOD = Module.instance_method(:protected_instance_methods) #: UnboundMethod
       PRIVATE_INSTANCE_METHODS_METHOD = Module.instance_method(:private_instance_methods) #: UnboundMethod
+      INSTANCE_METHOD_METHOD = Module.instance_method(:instance_method) #: UnboundMethod
       METHOD_METHOD = Kernel.instance_method(:method) #: UnboundMethod
+      METHOD_RECEIVER_METHOD = Method.instance_method(:receiver) #: UnboundMethod
       UNDEFINED_CONSTANT = Module.new.freeze #: Module[top]
 
       REQUIRED_FROM_LABELS = ["<top (required)>", "<main>", "<compiled>"].freeze #: Array[String]
@@ -123,20 +125,148 @@ def qualified_name_of(constant)
 
       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
+      #: ((UnboundMethod | Method) method, lookup_from: untyped) -> untyped
+      def signature_of!(method, lookup_from:)
+        # We use `T::Utils.signature_for_method` on a method and the portion of its
+        # `super_method` chain introduced by `prepend`. This finds signatures hidden
+        # by prepended methods without crossing ordinary method implementations. We
+        # check the chain a second time to support `prepend` in cases like this:
+        #   module Wrapper
+        #     def foo = super
+        #   end
+        #   class Example
+        #     sig { void }
+        #     def foo; end
+        #     prepend Wrapper
+        #   end
+        # Asking Sorbet for the signature on 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`.
+        lookup_scope = lookup_scope_for(method, lookup_from)
+        resolved_method = method_from_lookup(method, lookup_from, lookup_scope)
+        unless are_equal?(method.owner, resolved_method.owner)
+          Kernel.raise ArgumentError, "`method` must be looked up from `lookup_from`"
+        end
+        candidates = signature_candidates(method, lookup_scope)
+
+        begin
+          needs_second_pass = false #: bool
+          candidates.each do |current_method|
+            needs_second_pass ||= T::Private::Methods.has_sig_block_for_method(current_method)
+            current_signature = T::Utils.signature_for_method(current_method)
+            return current_signature if current_signature
+          end
+          return nil unless needs_second_pass
+
+          # Evaluating a `sig` lazily may register it under the prepended wrapper,
+          # so resolve the method and inspect the chain again.
+          refreshed_method = method_from_lookup(method, lookup_from, lookup_scope)
+          signature_candidates(refreshed_method, lookup_scope).each do |current_method|
+            current_signature = T::Utils.signature_for_method(current_method)
+            return current_signature if current_signature
+          end
+
+          nil
+        rescue LoadError, StandardError
+          Kernel.raise SignatureBlockError
+        end
       end
 
-      #: ((UnboundMethod | Method) method) -> untyped
-      def signature_of(method)
-        signature_of!(method)
+      #: ((UnboundMethod | Method) method, lookup_from: untyped) -> untyped
+      def signature_of(method, lookup_from:)
+        signature_of!(method, lookup_from: lookup_from)
       rescue SignatureBlockError
         nil
       end
 
+      #: ((UnboundMethod | Method) method, untyped lookup_from) -> Module[top]
+      def lookup_scope_for(method, lookup_from)
+        if Method === method
+          receiver = METHOD_RECEIVER_METHOD.bind_call(method)
+          unless are_equal?(receiver, lookup_from)
+            Kernel.raise ArgumentError, "`lookup_from` must be the method receiver"
+          end
+
+          singleton_class_of(lookup_from)
+        elsif Module === lookup_from
+          lookup_from
+        else
+          Kernel.raise ArgumentError, "`lookup_from` must be a module for an unbound method"
+        end
+      end
+      private :lookup_scope_for
+
+      #: (
+      #|   (UnboundMethod | Method) method,
+      #|   untyped lookup_from,
+      #|   Module[top] lookup_scope
+      #| ) -> (UnboundMethod | Method)
+      def method_from_lookup(method, lookup_from, lookup_scope)
+        if Method === method
+          METHOD_METHOD.bind_call(lookup_from, method.name)
+        else
+          INSTANCE_METHOD_METHOD.bind_call(lookup_scope, method.name)
+        end
+      end
+      private :method_from_lookup
+
+      #: ((UnboundMethod | Method) method, Module[top] lookup_scope) -> Array[UnboundMethod | Method]
+      def signature_candidates(method, lookup_scope)
+        return [method] if are_equal?(method.owner, lookup_scope)
+
+        ancestors = ancestors_of(lookup_scope)
+        prepended_positions = prepended_ancestor_positions(ancestors)
+        candidates = [] #: Array[UnboundMethod | Method]
+        current_method = method #: (UnboundMethod | Method)?
+        first_ancestor_index = 0
+
+        while current_method
+          ancestor_index = (first_ancestor_index...ancestors.length).find do |index|
+            are_equal?(ancestors.fetch(index), current_method.owner)
+          end
+          unless ancestor_index
+            Kernel.raise ArgumentError, "method does not belong to the `lookup_from` ancestor chain"
+          end
+
+          candidates << current_method
+          break unless prepended_positions.fetch(ancestor_index)
+
+          first_ancestor_index = ancestor_index + 1
+          current_method = current_method.super_method
+        end
+
+        candidates
+      end
+      private :signature_candidates
+
+      #: (Array[Module[top]] ancestors) -> Array[bool]
+      def prepended_ancestor_positions(ancestors)
+        prepended = Array.new(ancestors.length, false) #: Array[bool]
+
+        # Mark prepend positions in the flattened ancestor chain. We track positions
+        # instead of module identities because Ruby can include and prepend the same
+        # module at different points in one ancestor chain.
+        ancestors.each_with_index do |ancestor, ancestor_index|
+          prepended_ancestors = ancestors_of(ancestor).take_while do |candidate|
+            !are_equal?(candidate, ancestor)
+          end
+          first_prepend_index = ancestor_index - prepended_ancestors.length
+          next if first_prepend_index.negative?
+          next unless prepended_ancestors.each_with_index.all? do |candidate, index|
+            are_equal?(candidate, ancestors.fetch(first_prepend_index + index))
+          end
+
+          first_prepend_index.upto(ancestor_index - 1) do |index|
+            prepended[index] = true
+          end
+        end
+
+        prepended
+      end
+      private :prepended_ancestor_positions
+
       #: (T::Types::Base type) -> String
       def name_of_type(type)
         type.to_s

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