From 66dcf35b6e900fa29f0223377c9d4f5f3ce7b613 Mon Sep 17 00:00:00 2001
From: Soutaro Matsumoto <matsumoto@soutaro.com>
Date: Thu, 13 Apr 2023 16:41:32 +0900
Subject: [PATCH] Set up break contexts correctly for untyped blocks

---
 lib/steep/type_construction.rb | 97 ++++++++++++++++++++--------------
 test/type_construction_test.rb | 24 +++++++++
 2 files changed, 82 insertions(+), 39 deletions(-)

diff --git a/lib/steep/type_construction.rb b/lib/steep/type_construction.rb
index b73eca2d5..286ea75e9 100644
--- a/lib/steep/type_construction.rb
+++ b/lib/steep/type_construction.rb
@@ -3183,46 +3183,65 @@ def type_send(node, send_node:, block_params:, block_body:, unwrap: false, tapp:
       receiver_type = checker.factory.deep_expand_alias(recv_type)
       private = receiver.nil? || receiver.type == :self
 
-      type, constr = case receiver_type
-                     when nil
-                       raise
-
-                     when AST::Types::Any
-                       constr = constr.synthesize_children(node, skips: [receiver])
-                       constr.add_call(
-                         TypeInference::MethodCall::Untyped.new(
-                           node: node,
-                           context: context.call_context,
-                           method_name: method_name
-                         )
-                       )
+      type, constr =
+        case receiver_type
+        when nil
+          raise
 
-                     else
-                       if interface = calculate_interface(receiver_type, private: private)
-                         constr.type_send_interface(
-                           node,
-                           interface: interface,
-                           receiver: receiver,
-                           receiver_type: receiver_type,
-                           method_name: method_name,
-                           arguments: arguments,
-                           block_params: block_params,
-                           block_body: block_body,
-                           tapp: tapp
-                         )
-                       else
-                         constr = constr.synthesize_children(node, skips: [receiver])
-                         constr.add_call(
-                           TypeInference::MethodCall::NoMethodError.new(
-                             node: node,
-                             context: context.call_context,
-                             method_name: method_name,
-                             receiver_type: receiver_type,
-                             error: Diagnostic::Ruby::NoMethod.new(node: node, method: method_name, type: receiver_type)
-                           )
-                         )
-                       end
-                     end
+        when AST::Types::Any
+          case node.type
+          when :block, :numblock
+            # @type var node: Parser::AST::Node & Parser::AST::_BlockNode
+            block_annotations = source.annotations(block: node, factory: checker.factory, context: nesting)
+            block_params or raise
+
+            constr = constr.synthesize_children(node.children[0])
+
+            constr.type_block_without_hint(
+              node: node,
+              block_params: TypeInference::BlockParams.from_node(block_params, annotations: block_annotations),
+              block_annotations: block_annotations,
+              block_body: block_body
+            ) do |error|
+              constr.typing.errors << error
+            end
+          else
+            constr = constr.synthesize_children(node, skips: [receiver])
+          end
+
+          constr.add_call(
+            TypeInference::MethodCall::Untyped.new(
+              node: node,
+              context: context.call_context,
+              method_name: method_name
+            )
+          )
+        else
+          if interface = calculate_interface(receiver_type, private: private)
+            constr.type_send_interface(
+              node,
+              interface: interface,
+              receiver: receiver,
+              receiver_type: receiver_type,
+              method_name: method_name,
+              arguments: arguments,
+              block_params: block_params,
+              block_body: block_body,
+              tapp: tapp
+            )
+          else
+            constr = constr.synthesize_children(node, skips: [receiver])
+            constr.add_call(
+              TypeInference::MethodCall::NoMethodError.new(
+                node: node,
+                context: context.call_context,
+                method_name: method_name,
+                receiver_type: receiver_type,
+                error: Diagnostic::Ruby::NoMethod.new(node: node, method: method_name, type: receiver_type)
+              )
+            )
+          end
+        end
 
       Pair.new(type: type, constr: constr)
     end
diff --git a/test/type_construction_test.rb b/test/type_construction_test.rb
index 341f1075a..9402c5e43 100644
--- a/test/type_construction_test.rb
+++ b/test/type_construction_test.rb
@@ -10372,4 +10372,28 @@ def next
       end
     end
   end
+
+  def test_untyped_call_block
+    with_checker(<<~RBS) do |checker|
+      RBS
+      source = parse_ruby(<<~RUBY)
+        (_ = [1, 2, 3]).each {
+          case
+          when 1
+            next
+          when 2
+            retry
+          when 3
+            break
+          end
+        }
+      RUBY
+
+      with_standard_construction(checker, source) do |construction, typing|
+        type, _, context = construction.synthesize(source.node)
+
+        assert_no_error(typing)
+      end
+    end
+  end
 end