From 00a4b65b28e540d77b5f016d67a820d95ffecd58 Mon Sep 17 00:00:00 2001
From: Marco Costa <marco.costa@datadoghq.com>
Date: Tue, 14 Jan 2025 12:48:41 -0800
Subject: [PATCH] Fix self union type checking

---
 lib/steep/subtyping/check.rb | 11 +++++------
 test/type_check_test.rb      | 27 +++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/lib/steep/subtyping/check.rb b/lib/steep/subtyping/check.rb
index fd8a56d2..c6c67f1c 100644
--- a/lib/steep/subtyping/check.rb
+++ b/lib/steep/subtyping/check.rb
@@ -285,11 +285,6 @@ def check_type0(relation)
             )
           end
 
-        when relation.sub_type.is_a?(AST::Types::Self) && !self_type.is_a?(AST::Types::Self)
-          Expand(relation) do
-            check_type(Relation.new(sub_type: self_type, super_type: relation.super_type))
-          end
-
         when relation.sub_type.is_a?(AST::Types::Instance) && !instance_type.is_a?(AST::Types::Instance)
           Expand(relation) do
             check_type(Relation.new(sub_type: instance_type, super_type: relation.super_type))
@@ -419,6 +414,11 @@ def check_type0(relation)
             end
           end
 
+        when relation.sub_type.is_a?(AST::Types::Self) && !self_type.is_a?(AST::Types::Self)
+          Expand(relation) do
+            check_type(Relation.new(sub_type: self_type, super_type: relation.super_type))
+          end
+
         when relation.super_type.is_a?(AST::Types::Name::Interface)
           Expand(relation) do
             check_interface(
@@ -608,7 +608,6 @@ def check_type0(relation)
           Expand(relation) do
             check_type(Relation.new(sub_type: relation.sub_type.back_type, super_type: relation.super_type))
           end
-
         else
           Failure(relation, Result::Failure::UnknownPairError.new(relation: relation))
         end
diff --git a/test/type_check_test.rb b/test/type_check_test.rb
index 61df331f..3eee168d 100644
--- a/test/type_check_test.rb
+++ b/test/type_check_test.rb
@@ -2927,4 +2927,31 @@ def test_tuple_type_assertion
       YAML
     )
   end
+
+  def test_self_type_union_assertion
+    run_type_check_test(
+      signatures: {
+        "a.rbs" => <<~RBS
+          class Foo
+            def bar: (bool) -> self?
+          end
+        RBS
+      },
+      code: {
+        "a.rb" => <<~RUBY
+          class Foo
+            def bar(var)
+              return nil if var
+              self
+            end
+          end
+        RUBY
+      },
+      expectations: <<~YAML
+        ---
+        - file: a.rb
+          diagnostics: []
+      YAML
+    )
+  end
 end