From 9e946eb373a824f271c88111c7c77da992cc3017 Mon Sep 17 00:00:00 2001
From: Soutaro Matsumoto <matsumoto@soutaro.com>
Date: Tue, 7 Jan 2025 14:52:11 +0900
Subject: [PATCH] Fix annotation in `when` clause

---
 lib/steep/type_construction.rb |  2 +-
 sig/test/type_check_test.rbs   |  2 ++
 test/type_check_test.rb        | 28 ++++++++++++++++++++++++++++
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/lib/steep/type_construction.rb b/lib/steep/type_construction.rb
index 71a40dbf..ee98400e 100644
--- a/lib/steep/type_construction.rb
+++ b/lib/steep/type_construction.rb
@@ -2026,8 +2026,8 @@ def synthesize(node, hint: nil, condition: false)
                 branch_result =
                   if body
                     when_clause_constr
-                      .for_branch(body)
                       .update_type_env {|env| env.join(*body_envs) }
+                      .for_branch(body)
                       .tap {|constr| typing.cursor_context.set_node_context(body, constr.context) }
                       .synthesize(body, hint: hint)
                   else
diff --git a/sig/test/type_check_test.rbs b/sig/test/type_check_test.rbs
index fe285808..57ef3d8a 100644
--- a/sig/test/type_check_test.rbs
+++ b/sig/test/type_check_test.rbs
@@ -181,4 +181,6 @@ class TypeCheckTest < Minitest::Test
   def test_undeclared_method2: () -> untyped
 
   def test_when__assertion: () -> untyped
+
+  def test_when__type_annotation: () -> untyped
 end
diff --git a/test/type_check_test.rb b/test/type_check_test.rb
index 56e7033b..761b4181 100644
--- a/test/type_check_test.rb
+++ b/test/type_check_test.rb
@@ -2869,4 +2869,32 @@ def test_when__assertion
       YAML
     )
   end
+
+  def test_when__type_annotation
+    run_type_check_test(
+      signatures: {
+        "a.rbs" => <<~RBS
+        RBS
+      },
+      code: {
+        "a.rb" => <<~RUBY
+          a = [1, ""].sample
+
+          case
+          when 1.even?
+            # @type var a: String
+            a + ""
+          when 2.even?
+            # @type var a: Integer
+            a + 1
+          end
+        RUBY
+      },
+      expectations: <<~YAML
+        ---
+        - file: a.rb
+          diagnostics: []
+      YAML
+    )
+  end
 end