From 8142fa2879fe99f49270fb5f3cb0d59c4281fc7d Mon Sep 17 00:00:00 2001
From: Soutaro Matsumoto <matsumoto@soutaro.com>
Date: Tue, 8 Mar 2022 20:28:25 +0900
Subject: [PATCH 1/2] Add #each_child helper method

---
 lib/steep/ast/types/any.rb          | 2 ++
 lib/steep/ast/types/boolean.rb      | 2 ++
 lib/steep/ast/types/bot.rb          | 2 ++
 lib/steep/ast/types/class.rb        | 2 ++
 lib/steep/ast/types/helper.rb       | 8 ++++++++
 lib/steep/ast/types/instance.rb     | 2 ++
 lib/steep/ast/types/intersection.rb | 4 ++++
 lib/steep/ast/types/literal.rb      | 2 ++
 lib/steep/ast/types/logic.rb        | 2 ++
 lib/steep/ast/types/name.rb         | 6 ++++++
 lib/steep/ast/types/nil.rb          | 2 ++
 lib/steep/ast/types/proc.rb         | 9 +++++++++
 lib/steep/ast/types/record.rb       | 4 ++++
 lib/steep/ast/types/self.rb         | 2 ++
 lib/steep/ast/types/top.rb          | 2 ++
 lib/steep/ast/types/tuple.rb        | 4 ++++
 lib/steep/ast/types/union.rb        | 4 ++++
 lib/steep/ast/types/var.rb          | 2 ++
 lib/steep/ast/types/void.rb         | 2 ++
 lib/steep/interface/function.rb     | 4 ++++
 20 files changed, 67 insertions(+)

diff --git a/lib/steep/ast/types/any.rb b/lib/steep/ast/types/any.rb
index 56fa6659b..f9123b4e4 100644
--- a/lib/steep/ast/types/any.rb
+++ b/lib/steep/ast/types/any.rb
@@ -28,6 +28,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [1]
         end
diff --git a/lib/steep/ast/types/boolean.rb b/lib/steep/ast/types/boolean.rb
index 9da21689c..70e218103 100644
--- a/lib/steep/ast/types/boolean.rb
+++ b/lib/steep/ast/types/boolean.rb
@@ -28,6 +28,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [0]
         end
diff --git a/lib/steep/ast/types/bot.rb b/lib/steep/ast/types/bot.rb
index 4e4e9787a..f79ab8ecc 100644
--- a/lib/steep/ast/types/bot.rb
+++ b/lib/steep/ast/types/bot.rb
@@ -28,6 +28,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [2]
         end
diff --git a/lib/steep/ast/types/class.rb b/lib/steep/ast/types/class.rb
index 6014baf52..881c801b4 100644
--- a/lib/steep/ast/types/class.rb
+++ b/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
diff --git a/lib/steep/ast/types/helper.rb b/lib/steep/ast/types/helper.rb
index 87195ccf8..8965d897e 100644
--- a/lib/steep/ast/types/helper.rb
+++ b/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
diff --git a/lib/steep/ast/types/instance.rb b/lib/steep/ast/types/instance.rb
index 20329d3e3..920f36af8 100644
--- a/lib/steep/ast/types/instance.rb
+++ b/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
diff --git a/lib/steep/ast/types/intersection.rb b/lib/steep/ast/types/intersection.rb
index ee28b426e..9705801b6 100644
--- a/lib/steep/ast/types/intersection.rb
+++ b/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
diff --git a/lib/steep/ast/types/literal.rb b/lib/steep/ast/types/literal.rb
index e0ea85cf5..99483f302 100644
--- a/lib/steep/ast/types/literal.rb
+++ b/lib/steep/ast/types/literal.rb
@@ -31,6 +31,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [0]
         end
diff --git a/lib/steep/ast/types/logic.rb b/lib/steep/ast/types/logic.rb
index 4b76ab972..e4bd95b2b 100644
--- a/lib/steep/ast/types/logic.rb
+++ b/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
diff --git a/lib/steep/ast/types/name.rb b/lib/steep/ast/types/name.rb
index 364ea9f47..9b1809d9a 100644
--- a/lib/steep/ast/types/name.rb
+++ b/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
diff --git a/lib/steep/ast/types/nil.rb b/lib/steep/ast/types/nil.rb
index 1540a6b8b..bc2840667 100644
--- a/lib/steep/ast/types/nil.rb
+++ b/lib/steep/ast/types/nil.rb
@@ -28,6 +28,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [0]
         end
diff --git a/lib/steep/ast/types/proc.rb b/lib/steep/ast/types/proc.rb
index 1812dae31..738af4453 100644
--- a/lib/steep/ast/types/proc.rb
+++ b/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
diff --git a/lib/steep/ast/types/record.rb b/lib/steep/ast/types/record.rb
index 5bafcc06a..4702e1f31 100644
--- a/lib/steep/ast/types/record.rb
+++ b/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
diff --git a/lib/steep/ast/types/self.rb b/lib/steep/ast/types/self.rb
index 018f5a611..7a4304c9f 100644
--- a/lib/steep/ast/types/self.rb
+++ b/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
diff --git a/lib/steep/ast/types/top.rb b/lib/steep/ast/types/top.rb
index fdd5b9753..46bec38b9 100644
--- a/lib/steep/ast/types/top.rb
+++ b/lib/steep/ast/types/top.rb
@@ -28,6 +28,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [2]
         end
diff --git a/lib/steep/ast/types/tuple.rb b/lib/steep/ast/types/tuple.rb
index cd47af125..d6afb624c 100644
--- a/lib/steep/ast/types/tuple.rb
+++ b/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
diff --git a/lib/steep/ast/types/union.rb b/lib/steep/ast/types/union.rb
index 9bc5b4fcc..75f72dab3 100644
--- a/lib/steep/ast/types/union.rb
+++ b/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
diff --git a/lib/steep/ast/types/var.rb b/lib/steep/ast/types/var.rb
index adff1e61d..e33eb3a5c 100644
--- a/lib/steep/ast/types/var.rb
+++ b/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
diff --git a/lib/steep/ast/types/void.rb b/lib/steep/ast/types/void.rb
index d4ca4b863..887aa6d3e 100644
--- a/lib/steep/ast/types/void.rb
+++ b/lib/steep/ast/types/void.rb
@@ -28,6 +28,8 @@ def to_s
 
         include Helper::NoFreeVariables
 
+        include Helper::NoChild
+
         def level
           [0]
         end
diff --git a/lib/steep/interface/function.rb b/lib/steep/interface/function.rb
index a628c338b..e33fdfa33 100644
--- a/lib/steep/interface/function.rb
+++ b/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)

From e0383cdc1612f80b5bd9039a84beb88239ca0cdc Mon Sep 17 00:00:00 2001
From: Soutaro Matsumoto <matsumoto@soutaro.com>
Date: Tue, 8 Mar 2022 20:59:30 +0900
Subject: [PATCH 2/2] Choose the best type on `Any` result

---
 lib/steep/subtyping/check.rb   | 22 ++++++++++++++++++++--
 test/type_construction_test.rb | 22 ++++++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/lib/steep/subtyping/check.rb b/lib/steep/subtyping/check.rb
index 8af1d63c6..1418a1cc9 100644
--- a/lib/steep/subtyping/check.rb
+++ b/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
diff --git a/test/type_construction_test.rb b/test/type_construction_test.rb
index 57c9c4e57..a8feea06d 100644
--- a/test/type_construction_test.rb
+++ b/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