From 7d855574555c5954ec6cf5bd025efcade238f0e9 Mon Sep 17 00:00:00 2001
From: Soutaro Matsumoto <matsumoto@soutaro.com>
Date: Thu, 6 Jul 2023 16:31:40 +0900
Subject: [PATCH] Fix `#partition_union`

---
 lib/steep/ast/types/factory.rb | 13 +++++++++----
 test/type_factory_test.rb      | 27 +++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/lib/steep/ast/types/factory.rb b/lib/steep/ast/types/factory.rb
index bb45a06ee..d308060c9 100644
--- a/lib/steep/ast/types/factory.rb
+++ b/lib/steep/ast/types/factory.rb
@@ -358,16 +358,21 @@ def partition_union(type)
               partition_union(unfold)
             end
           when AST::Types::Union
-            falsy_types, truthy_types = type.types.partition do |type|
-              (type.is_a?(AST::Types::Literal) && type.value == false) ||
-                type.is_a?(AST::Types::Nil)
+            truthy_types = [] #: Array[AST::Types::t]
+            falsy_types = [] #: Array[AST::Types::t]
+
+            type.types.each do |type|
+              truthy, falsy = partition_union(type)
+
+              truthy_types << truthy if truthy
+              falsy_types << falsy if falsy
             end
 
             [
               truthy_types.empty? ? nil : AST::Types::Union.build(types: truthy_types),
               falsy_types.empty? ? nil : AST::Types::Union.build(types: falsy_types)
             ]
-          when AST::Types::Any, AST::Types::Boolean
+          when AST::Types::Any, AST::Types::Boolean, AST::Types::Top
             [type, type]
           when AST::Types::Nil
             [nil, type]
diff --git a/test/type_factory_test.rb b/test/type_factory_test.rb
index 00452191b..8b7f556d0 100644
--- a/test/type_factory_test.rb
+++ b/test/type_factory_test.rb
@@ -395,4 +395,31 @@ class Bar
       end
     end
   end
+
+  def test_partition_union__top
+    with_factory() do |factory|
+      factory.partition_union(factory.type(parse_type("top"))).tap do |truthy, falsy|
+        assert_equal factory.type(parse_type("top")), truthy
+        assert_equal factory.type(parse_type("top")), falsy
+      end
+    end
+  end
+
+  def test_partition_union__boolish
+    with_factory() do |factory|
+      factory.partition_union(factory.type(parse_type("::boolish"))).tap do |truthy, falsy|
+        assert_equal factory.type(parse_type("top")), truthy
+        assert_equal factory.type(parse_type("top")), falsy
+      end
+    end
+  end
+
+  def test_partition_union__bool_union
+    with_factory() do |factory|
+      factory.partition_union(factory.type(parse_type("bool | ::Symbol"))).tap do |truthy, falsy|
+        assert_equal factory.type(parse_type("bool | ::Symbol")), truthy
+        assert_equal factory.type(parse_type("bool")), falsy
+      end
+    end
+  end
 end