Fix case-when typing (and && typing...)

lib/steep/type_construction.rb

@@ -1476,7 +1476,7 @@ def synthesize(node, hint: nil)
                                    .for_branch(right)
                                    .synthesize(right)
 
-            type = if left_type.is_a?(AST::Types::Boolean)
+            type = if check_relation(sub_type: left_type, super_type: AST::Types::Boolean.new).success?
                      union_type(left_type, right_type)
                    else
                      union_type(right_type, AST::Builtin.nil_type)
@@ -1584,6 +1584,18 @@ def synthesize(node, hint: nil)
 
               cond_type, constr = constr.synthesize(cond)
               _, cond_vars = interpreter.decompose_value(cond)
+              unless cond_vars.empty?
+                first_var = cond_vars.to_a[0]
+                var_node = cond.updated(
+                  :lvar,
+                  [
+                    ASTUtils::Labeling::LabeledName.new(name: first_var, label: 0)
+                  ]
+                )
+              else
+                first_var = nil
+                var_node = cond
+              end
 
               when_constr = constr
               whens.each do |clause|
@@ -1593,9 +1605,15 @@ def synthesize(node, hint: nil)
                 test_envs = []
 
                 tests.each do |test|
-                  test_node = test.updated(:send, [test, :===, cond.dup])
+                  test_node = test.updated(:send, [test, :===, var_node])
                   test_type, test_constr = test_constr.synthesize(test_node)
                   truthy_env, falsy_env = interpreter.eval(type: test_type, node: test_node, env: test_constr.context.lvar_env)
+                  truthy_env = cond_vars.inject(truthy_env) do |env, var|
+                    env.assign!(var, node: test_node, type: env[first_var])
+                  end
+                  falsy_env = cond_vars.inject(falsy_env) do |env, var|
+                    env.assign!(var, node: test_node, type: env[first_var])
+                  end
                   test_envs << truthy_env
                   test_constr = test_constr.update_lvar_env { falsy_env }
                 end
@@ -1625,17 +1643,14 @@ def synthesize(node, hint: nil)
               types = branch_pairs.map(&:type)
               constrs = branch_pairs.map(&:constr)
 
-              unless els
-                constrs << when_constr
-              end
-
               if when_constr.context.lvar_env[cond_vars.first].is_a?(AST::Types::Bot)
                 # Exhaustive
                 if els
                   typing.add_error Errors::ElseOnExhaustiveCase.new(node: els, type: cond_type)
                 end
               else
                 unless els
+                  constrs << when_constr
                   types << AST::Builtin.nil_type
                 end
               end

test/type_construction_test.rb

@@ -6482,4 +6482,94 @@ def test_typing_record_nilable_attribute
       end
     end
   end
+
+  def test_type_case_case_when_assignment
+    with_checker do |checker|
+      source = parse_ruby(<<EOF)
+# @type var x: String | Integer
+x = (_ = nil)
+
+case x
+when String
+  a = "String"
+when Integer
+  a = "Integer"
+end
+EOF
+
+      with_standard_construction(checker, source) do |construction, typing|
+        _, _, context = construction.synthesize(source.node)
+
+        assert_no_error typing
+        assert_equal parse_type("::String"), context.lvar_env[:a]
+      end
+    end
+  end
+
+  def test_type_case_case_selector
+    with_checker do |checker|
+      source = parse_ruby(<<RUBY)
+x = ["foo", 2, :baz]
+
+a = case y = z = x[0]
+    when String
+      y + ""
+      z + ""
+      "String"
+    when Integer
+      y + 0
+      z + 0
+      "Integer"
+    when Symbol
+      "Array[String]"
+    end
+RUBY
+
+      with_standard_construction(checker, source) do |construction, typing|
+        _, _, context = construction.synthesize(source.node)
+
+        assert_no_error typing
+        assert_equal parse_type("::String"), context.lvar_env[:a]
+      end
+    end
+  end
+
+  def test_type_if_else_when_assignment
+    with_checker do |checker|
+      source = parse_ruby(<<EOF)
+# @type var x: String | Integer
+x = (_ = nil)
+
+if x.is_a?(String)
+  a = "String"
+else
+  a = "Integer"
+end
+EOF
+
+      with_standard_construction(checker, source) do |construction, typing|
+        _, _, context = construction.synthesize(source.node)
+
+        assert_no_error typing
+        assert_equal parse_type("::String"), context.lvar_env[:a]
+      end
+    end
+  end
+
+  def test_bool_and_or
+    with_checker do |checker|
+      source = parse_ruby(<<EOF)
+# @type var x: bool
+
+x = 30.is_a?(Integer) && true
+x = 30.is_a?(String) || false
+EOF
+
+      with_standard_construction(checker, source) do |construction, typing|
+        _, _, context = construction.synthesize(source.node)
+
+        assert_no_error typing
+      end
+    end
+  end
 end