Fix for #flat_map

lib/steep/ast/types/any.rb

@@ -28,6 +28,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [1]
         end

lib/steep/ast/types/boolean.rb

@@ -28,6 +28,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [0]
         end

lib/steep/ast/types/bot.rb

@@ -28,6 +28,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [2]
         end

lib/steep/ast/types/class.rb

@@ -30,6 +30,8 @@ def free_variables()
           @fvs = Set.new([self])
         end
 
+        include Helper::NoChild
+
         def level
           [0]
         end

lib/steep/ast/types/helper.rb

@@ -21,6 +21,14 @@ def free_variables()
             @fvs ||= Set.new
           end
         end
+
+        module NoChild
+          def each_child(&block)
+            unless block
+              enum_for :each_child
+            end
+          end
+        end
       end
     end
   end

lib/steep/ast/types/instance.rb

@@ -26,6 +26,8 @@ def free_variables()
           @fvs = Set.new([self])
         end
 
+        include Helper::NoChild
+
         def to_s
           "instance"
         end

lib/steep/ast/types/intersection.rb

@@ -72,6 +72,10 @@ def free_variables()
 
         include Helper::ChildrenLevel
 
+        def each_child(&block)
+          types.each(&block)
+        end
+
         def level
           [0] + level_of_children(types)
         end

lib/steep/ast/types/literal.rb

@@ -31,6 +31,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [0]
         end

lib/steep/ast/types/logic.rb

@@ -13,6 +13,8 @@ def free_variables
             @fvs ||= Set[]
           end
 
+          include Helper::NoChild
+
           def hash
             self.class.hash
           end

lib/steep/ast/types/name.rb

@@ -72,6 +72,10 @@ def free_variables
             end
           end
 
+          def each_child(&block)
+            args.each(&block)
+          end
+
           include Helper::ChildrenLevel
 
           def level
@@ -98,6 +102,8 @@ def to_s
           def with_location(new_location)
             self.class.new(name: name, location: new_location)
           end
+
+          include Helper::NoChild
         end
 
         class Instance < Applying

lib/steep/ast/types/nil.rb

@@ -28,6 +28,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [0]
         end

lib/steep/ast/types/proc.rb

@@ -90,6 +90,15 @@ def back_type
         def block_required?
           block && !block.optional?
         end
+
+        def each_child(&block)
+          if block_given?
+            type.each_child(&block)
+            self.block&.type&.each_child(&block)
+          else
+            enum_for :each_child
+          end
+        end
       end
     end
   end

lib/steep/ast/types/record.rb

@@ -42,6 +42,10 @@ def free_variables()
 
         include Helper::ChildrenLevel
 
+        def each_child(&block)
+          elements.each_value(&block)
+        end
+
         def level
           [0] + level_of_children(elements.values)
         end

lib/steep/ast/types/self.rb

@@ -22,6 +22,8 @@ def to_s
           "self"
         end
 
+        include Helper::NoChild
+
         def subst(s)
           s.self_type or raise "Unexpected substitution: #{inspect}"
         end

lib/steep/ast/types/top.rb

@@ -28,6 +28,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [2]
         end

lib/steep/ast/types/tuple.rb

@@ -40,6 +40,10 @@ def free_variables()
 
         include Helper::ChildrenLevel
 
+        def each_child(&block)
+          types.each(&block)
+        end
+
         def level
           [0] + level_of_children(types)
         end

lib/steep/ast/types/union.rb

@@ -70,6 +70,10 @@ def free_variables
           end
         end
 
+        def each_child(&block)
+          types.each(&block)
+        end
+
         include Helper::ChildrenLevel
 
         def level

lib/steep/ast/types/var.rb

@@ -52,6 +52,8 @@ def free_variables()
           @fvs ||= Set.new([name])
         end
 
+        include Helper::NoChild
+
         def level
           [0]
         end

lib/steep/ast/types/void.rb

@@ -28,6 +28,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [0]
         end

lib/steep/interface/function.rb

@@ -952,6 +952,10 @@ def subst(s)
         )
       end
 
+      def each_child(&block)
+        each_type(&block)
+      end
+
       def each_type(&block)
         if block_given?
           params.each_type(&block)

lib/steep/subtyping/check.rb

@@ -347,7 +347,7 @@ def check_type0(relation)
 
         when relation.super_type.is_a?(AST::Types::Union)
           Any(relation) do |result|
-            relation.super_type.types.each do |super_type|
+            relation.super_type.types.sort_by {|ty| (path = hole_path(ty)) ? -path.size : 1 }.each do |super_type|
               rel = Relation.new(sub_type: relation.sub_type, super_type: super_type)
               result.add(rel) do
                 check_type(rel)
@@ -357,7 +357,7 @@ def check_type0(relation)
 
         when relation.sub_type.is_a?(AST::Types::Intersection)
           Any(relation) do |result|
-            relation.sub_type.types.each do |sub_type|
+            relation.sub_type.types.sort_by {|ty| (path = hole_path(ty)) ? -path.size : 1 }.each do |sub_type|
               rel = Relation.new(sub_type: sub_type, super_type: relation.super_type)
               result.add(rel) do
                 check_type(rel)
@@ -974,6 +974,24 @@ def match_params(name, relation)
       def expand_alias(type, &block)
         factory.expand_alias(type, &block)
       end
+
+      # Returns the shortest type paths for one of the _unknown_ type variables.
+      # Returns nil if there is no path.
+      def hole_path(type, path = [])
+        case type
+        when AST::Types::Var
+          if constraints.unknown?(type.name)
+            [type]
+          else
+            nil
+          end
+        else
+          paths = type.each_child.map do |ty|
+            hole_path(ty, path)&.unshift(ty)
+          end
+          paths.compact.min_by(&:size)
+        end
+      end
     end
   end
 end

test/type_construction_test.rb

@@ -8419,4 +8419,26 @@ def foo(&block)
       end
     end
   end
+
+  def test_flat_map
+    with_checker(<<-RBS) do |checker|
+class FlatMap
+  def flat_map: [A] () { (String) -> (A | Array[A]) } -> Array[A]
+end
+    RBS
+
+      source = parse_ruby(<<-'RUBY')
+# @type var a: FlatMap
+a = _ = nil 
+a.flat_map {|s| [s] }
+      RUBY
+
+      with_standard_construction(checker, source) do |construction, typing|
+        type, _, _ = construction.synthesize(source.node)
+
+        assert_no_error typing
+        assert_equal parse_type("::Array[::String]"), type
+      end
+    end
+  end
 end