Generate RBI methods inherited from Struct and other anonymous superclasses

lib/tapioca/gem/listeners/methods.rb

@@ -46,9 +46,12 @@ def compile_directly_owned_methods(
                 method = mod.instance_method(name)
                 method_visibility = visibility
 
-                if method.owner != mod
-                  # Use the visibility of the method `mod` itself defines, ignoring any modules prepended to it
-                  method_visibility = visibility_defined_by_constant(name, mod) || visibility
+                unless are_equal?(method.owner, mod)
+                  # Use the visibility of the method `mod` itself defines, ignoring any modules prepended to it.
+                  # Methods from anonymous superclasses aren't defined by `mod`, so check those superclasses next.
+                  method_visibility = visibility_defined_by_constant(name, mod) ||
+                    visibility_defined_by_anonymous_superclasses(name, mod) ||
+                    visibility
                 end
 
                 vis = case method_visibility
@@ -81,8 +84,12 @@ def compile_method(tree, symbol_name, constant, method, visibility = RBI::Public
             signature = signature_defined_by_constant(method, constant)
             signature ||= inferred_attr_writer_signature(method, constant)
             method = signature.method if signature #: UnboundMethod
+            # Methods from anonymous superclasses aren't tracked for the constant, so use the method's source location
+            # instead of `constant.instance_method`, which can return a prepended module's method or a `sig` wrapper.
+            fallback_source_method = method if @pipeline.anonymous_superclass_owning(method, constant)
+            definition = @pipeline.method_definition_in_gem(method.name, constant, fallback_source_method: fallback_source_method)
 
-            case @pipeline.method_definition_in_gem(method.name, constant)
+            case definition
             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
@@ -186,13 +193,17 @@ def compile_method(tree, symbol_name, constant, method, visibility = RBI::Public
         # It walks up the ancestor tree via the `super_method` method; if any of the super
         # methods are owned by the constant, it means that the constant declares the method,
         # and that super method is returned.
+        #
+        # Methods owned by the constant's anonymous superclasses also count, since those
+        # superclasses are not part of the RBI and their methods would otherwise be lost.
         #: (UnboundMethod method, Module[top] constant) -> UnboundMethod?
         def method_defined_by_constant(method, constant)
           # Widen the type of `method` to be nilable
           method = method #: UnboundMethod?
 
           while method
-            return method if method.owner == constant
+            return method if are_equal?(method.owner, constant)
+            return method if @pipeline.anonymous_superclass_owning(method, constant)
 
             method = method.super_method
           end
@@ -210,14 +221,14 @@ def method_defined_by_constant(method, constant)
         #: (UnboundMethod method, Module[top] constant) -> untyped
         def signature_defined_by_constant(method, constant)
           signature = signature_of!(method)
-          return signature if signature && signature.method.owner == method.owner
+          return signature if signature && are_equal?(signature.method.owner, method.owner)
 
           # Widen the type of `prepended_method` to be nilable
           prepended_method = constant.instance_method(method.name) #: UnboundMethod?
 
-          while prepended_method && prepended_method.owner != method.owner
+          while prepended_method && !are_equal?(prepended_method.owner, method.owner)
             signature = signature_of(prepended_method)
-            return signature if signature && signature.method.owner == method.owner
+            return signature if signature && are_equal?(signature.method.owner, method.owner)
 
             prepended_method = prepended_method.super_method
           end
@@ -246,6 +257,18 @@ def visibility_defined_by_constant(name, constant)
           end
         end
 
+        # Return the visibility of the method the constant's anonymous superclasses define, ignoring any modules
+        # prepended to them.
+        #: (Symbol name, Module[top] constant) -> Symbol?
+        def visibility_defined_by_anonymous_superclasses(name, constant)
+          @pipeline.anonymous_superclasses_of(constant).each do |superclass|
+            visibility = visibility_defined_by_constant(name, superclass)
+            return visibility if visibility
+          end
+
+          nil
+        end
+
         #: (UnboundMethod method, Module[top] constant) -> untyped
         def inferred_attr_writer_signature(method, constant)
           reader_method = attr_reader_for_writer(method, constant)

lib/tapioca/gem/listeners/mixins.rb

@@ -60,12 +60,34 @@ def add_mixins(tree, constant, mods, mixin_type)
         #: (Module[top] constant, Module[top] mixin, Runtime::Trackers::Mixin::Type mixin_type) -> bool
         def mixed_in_by_gem?(constant, mixin, mixin_type)
           mixin_location = Runtime::Trackers::Mixin.mixin_location(mixin, mixin_type, constant)
+          mixin_location ||= anonymous_superclass_mixin_location(constant, mixin, mixin_type)
 
           return true if mixin_location.nil?
 
           @pipeline.gem.contains_path?(mixin_location)
         end
 
+        # Mixins added inside an anonymous superclass are tracked under that superclass, not the constant, so look
+        # for them there. Modules prepended into an anonymous superclass come after the constant in its ancestors
+        # and are listed as includes, so includes are looked up as either an include or a prepend.
+        #: (Module[top] constant, Module[top] mixin, Runtime::Trackers::Mixin::Type mixin_type) -> String?
+        def anonymous_superclass_mixin_location(constant, mixin, mixin_type)
+          mixin_types = if mixin_type == Runtime::Trackers::Mixin::Type::Extend
+            [mixin_type]
+          else
+            [Runtime::Trackers::Mixin::Type::Include, Runtime::Trackers::Mixin::Type::Prepend]
+          end
+
+          @pipeline.anonymous_superclasses_of(constant).each do |superclass|
+            mixin_types.each do |type|
+              location = Runtime::Trackers::Mixin.mixin_location(mixin, type, superclass)
+              return location if location
+            end
+          end
+
+          nil
+        end
+
         #: (String mixin_name) -> bool
         def filtered_mixin?(mixin_name)
           # filter T:: namespace mixins that aren't T::Props
@@ -75,7 +97,14 @@ def filtered_mixin?(mixin_name)
 
         #: (Module[top] constant) -> Array[Module[top]]
         def interesting_ancestors_of(constant)
-          inherited_ancestors = Set.new.compare_by_identity.merge(inherited_ancestors_of(constant))
+          # Drop the ancestors the superclass in the RBI already brings in. `class Foo < Struct.new(:a)` is written as
+          # `class Foo < ::Struct`, so only `Struct`'s ancestors are dropped (it's the superclass of the last anonymous
+          # superclass), and mixins added inside `Struct.new` are kept. Also drop the anonymous superclasses themselves,
+          # or a named struct like `Struct::Foo` shows up as `include ::Struct::Foo`.
+          anonymous_superclasses = @pipeline.anonymous_superclasses_of(constant)
+          inherited_ancestors = Set.new.compare_by_identity.merge(
+            inherited_ancestors_of(anonymous_superclasses.last || constant),
+          ).merge(anonymous_superclasses)
 
           # TODO: There is actually a bug here where this will drop modules that
           # may be included twice. For example:

lib/tapioca/gem/listeners/source_location.rb

@@ -34,7 +34,12 @@ def on_scope(event)
         # @override
         #: (MethodNodeAdded event) -> void
         def on_method(event)
-          definition = @pipeline.method_definition_in_gem(event.method.name, event.constant)
+          anonymous_superclass = @pipeline.anonymous_superclass_owning(event.method, event.constant)
+
+          # Methods from anonymous superclasses are tracked under the anonymous superclass, not the constant,
+          # so use it if it owns the method. Otherwise, fall back to the constant, which is always there.
+          tracked_owner = anonymous_superclass || event.constant
+          definition = @pipeline.method_definition_in_gem(event.method.name, tracked_owner)
 
           if Pipeline::MethodInGemWithLocation === definition
             loc = definition.location

lib/tapioca/gem/pipeline.rb

@@ -26,6 +26,7 @@ def initialize(
         @gem = gem
         @seen = Set.new #: Set[String]
         @alias_namespace = Set.new #: Set[String]
+        @anonymous_superclasses = {}.compare_by_identity #: Hash[Module[top], Array[Module[top]]]
         @error_handler = error_handler
 
         @events = [] #: Array[Gem::Event]
@@ -150,9 +151,13 @@ def initialize(location)
         end
       end
 
-      #: (Symbol method_name, Module[top] owner) -> MethodDefinitionLookupResult
-      def method_definition_in_gem(method_name, owner)
-        definitions = Tapioca::Runtime::Trackers::MethodDefinition.method_definitions_for(method_name, owner)
+      #: (Symbol method_name, Module[top] owner, ?fallback_source_method: UnboundMethod?) -> MethodDefinitionLookupResult
+      def method_definition_in_gem(method_name, owner, fallback_source_method: nil)
+        definitions = Tapioca::Runtime::Trackers::MethodDefinition.method_definitions_for(
+          method_name,
+          owner,
+          fallback_source_method: fallback_source_method,
+        )
 
         # If the source location of the method isn't available, signal that by returning nil.
         return MethodUnknown.new if definitions.empty?
@@ -189,8 +194,90 @@ def name_of(constant)
         name
       end
 
+      # Return the superclasses that `compile_superclass` skips because they can't be written to the RBI, like the
+      # unnamed ones created by `class Foo < Struct.new(:bar)`, `Data.define` or `Class.new`, and named structs like
+      # `Struct.new("Bar", :baz)`, which are written as `Struct`. Abstract and generic superclasses are not included
+      # (see `fold_into_constant?`). The methods and mixins of these superclasses are folded into the constant.
+      #
+      # Only the gem that first defines the constant picks its superclass, so other gems that reopen the constant get
+      # no superclasses. Otherwise, they would repeat the superclass methods in their own RBI.
+      #: (Module[top] constant) -> Array[Module[top]]
+      def anonymous_superclasses_of(constant)
+        @anonymous_superclasses[constant] ||= begin
+          superclasses = [] #: Array[Module[top]]
+
+          if constant.singleton_class?
+            # Singleton classes have no names, so find the anonymous superclasses through the attached class.
+            # Since `Foo.singleton_class.superclass == Foo.superclass.singleton_class`, return the anonymous
+            # superclasses' singleton classes, which define class methods like `new`.
+            singleton_class = constant #: as Class[top]
+            attached_class = attached_class_of(singleton_class)
+
+            if attached_class
+              superclasses = anonymous_superclasses_of(attached_class).map { |klass| singleton_class_of(klass) }
+            end
+          elsif Class === constant
+            superclass = superclass_of(constant) #: Class[top]?
+
+            while superclass && fold_into_constant?(superclass)
+              superclasses << superclass
+              superclass = superclass_of(superclass)
+            end
+
+            # Drop the superclasses if another gem first defined the constant and this gem only reopens it. Most classes
+            # have none, so check that first to skip looking up where the constant was defined.
+            superclasses = [] unless superclasses.empty? || superclass_picked_by_gem?(constant)
+          end
+
+          superclasses
+        end
+      end
+
+      # Return the anonymous superclass of the constant that owns the method, or nil if none does
+      #: (UnboundMethod method, Module[top] constant) -> Module[top]?
+      def anonymous_superclass_owning(method, constant)
+        owner = method.owner
+
+        anonymous_superclasses_of(constant).find { |superclass| are_equal?(superclass, owner) }
+      end
+
       private
 
+      # Whether the gem first defines the constant. Ruby only sets the superclass the first time the class is defined
+      # (`class Foo < Bar`), so that gem picks it, and reopening the class elsewhere can't change it.
+      #: (Module[top] constant) -> bool
+      def superclass_picked_by_gem?(constant)
+        name = name_of(constant)
+        # Without a name, there's no way to tell where the constant was defined, so assume it's this gem
+        return true if name.nil?
+
+        constant_in_gem?(name)
+      end
+
+      # Whether the superclass is missing from the RBI, so its methods and mixins should be folded into the constant
+      #: (Class[top] superclass) -> bool
+      def fold_into_constant?(superclass)
+        # Most superclasses are named and written to the RBI, so check the name first
+        name = name_of(superclass)
+        return false unless name.nil? || named_struct?(superclass, name)
+
+        # Abstract methods and type variables can't be written on the constant, so stop at these superclasses
+        return false if abstract_type_of(superclass)
+        return false if Runtime::GenericTypeRegistry.lookup_type_variables(superclass)
+
+        true
+      end
+
+      # Whether the superclass is a named struct like `Struct::Bar`, which `name_of` writes as `Struct`
+      #: (Class[top] superclass, String name) -> bool
+      def named_struct?(superclass, name)
+        return false unless name == "Struct"
+        return false if are_equal?(superclass, ::Struct)
+
+        # Use `ancestors` instead of `<`, which classes can override
+        ancestors_of(superclass).any? { |ancestor| are_equal?(ancestor, ::Struct) }
+      end
+
       #: (Gemfile::GemSpec gem) -> Set[String]
       def load_bootstrap_symbols(gem)
         engine_symbols = Static::SymbolLoader.engine_symbols(gem)

lib/tapioca/runtime/trackers/method_definition.rb

@@ -24,12 +24,15 @@ def register(method_name, owner, locations)
             registrations_for(method_name, owner) << loc
           end
 
-          #: (Symbol method_name, Module[top] owner) -> Array[SourceLocation]
-          def method_definitions_for(method_name, owner)
+          # Return where `method_name` was defined on `owner`. If no definitions were recorded for `owner` (for example,
+          # the method was defined on an anonymous superclass), fall back to the source location of `fallback_source_method`,
+          # or of the method `owner` resolves to.
+          #: (Symbol method_name, Module[top] owner, ?fallback_source_method: UnboundMethod?) -> Array[SourceLocation]
+          def method_definitions_for(method_name, owner, fallback_source_method: nil)
             definitions = registrations_for(method_name, owner)
 
             if definitions.empty?
-              source_loc = owner.instance_method(method_name).source_location
+              source_loc = (fallback_source_method || owner.instance_method(method_name)).source_location
               definitions = [SourceLocation.from_loc(source_loc)].compact
             end