Support breaks with values inside while expressions (#10566)

Co-authored-by: Johannes Müller <[email protected]>

spec/compiler/codegen/while_spec.cr

@@ -21,6 +21,61 @@ describe "Codegen: while" do
     run("a = 0; while a < 10; a &+= 1; break if a > 5; end; a").to_i.should eq(6)
   end
 
+  it "break with value" do
+    run(%(
+      struct Nil; def to_i!; 0; end; end
+
+      a = 0
+      b = while a < 10
+        a &+= 1
+        break a &+ 3
+      end
+      b.to_i!
+      )).to_i.should eq(4)
+  end
+
+  it "conditional break with value" do
+    run(%(
+      struct Nil; def to_i!; 0; end; end
+
+      a = 0
+      b = while a < 10
+        a &+= 1
+        break a &+ 3 if a > 5
+      end
+      b.to_i!
+      )).to_i.should eq(9)
+  end
+
+  it "break with value, condition fails" do
+    run(%(
+      a = while 1 == 2
+        break 1
+      end
+      a.nil?
+      )).to_b.should be_true
+  end
+
+  it "endless break with value" do
+    run(%(
+      a = 0
+      while true
+        a &+= 1
+        break a &+ 3
+      end
+      )).to_i.should eq(4)
+  end
+
+  it "endless conditional break with value" do
+    run(%(
+      a = 0
+      while true
+        a &+= 1
+        break a &+ 3 if a > 5
+      end
+      )).to_i.should eq(9)
+  end
+
   it "codegens endless while" do
     codegen "while true; end"
   end

spec/compiler/semantic/while_spec.cr

@@ -6,11 +6,37 @@ describe "Semantic: while" do
   end
 
   it "types while with break without value" do
-    assert_type("while true; break; end") { nil_type }
+    assert_type("while 1; break; end") { nil_type }
   end
 
   it "types while with break with value" do
-    assert_type("while true; break 1; end") { nil_type }
+    assert_type("while 1; break 'a'; end") { nilable char }
+  end
+
+  it "types while with multiple breaks with value" do
+    assert_type(%(
+      while 1
+        break 'a' if 1
+        break "", 123 if 1
+      end
+      )) { nilable union_of(char, tuple_of([string, int32])) }
+  end
+
+  it "types endless while with break without value" do
+    assert_type("while true; break; end") { nil_type }
+  end
+
+  it "types endless while with break with value" do
+    assert_type("while true; break 1; end") { int32 }
+  end
+
+  it "types endless while with multiple breaks with value" do
+    assert_type(%(
+      while true
+        break 'a' if 1
+        break "", 123 if 1
+      end
+      )) { union_of(char, tuple_of([string, int32])) }
   end
 
   it "reports break cannot be used outside a while" do

src/compiler/crystal/codegen/codegen.cr

@@ -796,37 +796,57 @@ module Crystal
       set_ensure_exception_handler(node)
 
       with_cloned_context do
-        while_block, body_block, exit_block = new_blocks "while", "body", "exit"
+        cond = node.cond.single_expression
+        endless_while = cond.true_literal?
 
-        context.while_block = while_block
-        context.while_exit_block = exit_block
-        context.break_phi = nil
-        context.next_phi = nil
+        if endless_while
+          while_block = new_block "while"
 
-        br while_block
+          Phi.open(self, node, @needs_value) do |phi|
+            context.while_block = while_block
+            context.break_phi = phi
+            context.next_phi = nil
 
-        position_at_end while_block
+            br while_block
 
-        request_value do
-          set_current_debug_location node.cond if @debug.line_numbers?
-          codegen_cond_branch node.cond, body_block, exit_block
-        end
+            position_at_end while_block
 
-        position_at_end body_block
+            request_value(false) do
+              accept node.body
+            end
+            br while_block
+          end
+        else
+          while_block, body_block, fail_block = new_blocks "while", "body", "fail"
 
-        request_value(false) do
-          accept node.body
-        end
-        br while_block
+          Phi.open(self, node, @needs_value) do |phi|
+            context.while_block = while_block
+            context.break_phi = phi
+            context.next_phi = nil
 
-        position_at_end exit_block
+            br while_block
 
-        if node.no_returns?
-          unreachable
-        else
-          @last = llvm_nil
+            position_at_end while_block
+
+            request_value do
+              set_current_debug_location node.cond if @debug.line_numbers?
+              codegen_cond_branch node.cond, body_block, fail_block
+            end
+
+            position_at_end body_block
+
+            request_value(false) do
+              accept node.body
+            end
+            br while_block
+
+            position_at_end fail_block
+
+            phi.add llvm_nil, @program.nil, last: true
+          end
         end
       end
+
       false
     end
 
@@ -854,20 +874,28 @@ module Crystal
       set_current_debug_location(node) if @debug.line_numbers?
       node_type = accept_control_expression(node)
 
-      if break_phi = context.break_phi
-        old_last = @last
-        execute_ensures_until(node.target.as(Call))
-        @last = old_last
+      case target = node.target
+      when Call
+        if break_phi = context.break_phi
+          old_last = @last
+          execute_ensures_until(target)
+          @last = old_last
 
-        break_phi.add @last, node_type
-      elsif while_exit_block = context.while_exit_block
-        execute_ensures_until(node.target.as(While))
-        br while_exit_block
-      else
-        node.raise "BUG: unknown exit for break"
+          break_phi.add @last, node_type
+          return false
+        end
+      when While
+        if break_phi = context.break_phi
+          old_last = @last
+          execute_ensures_until(target)
+          @last = old_last
+
+          break_phi.add @last, node_type
+          return false
+        end
       end
 
-      false
+      node.raise "BUG: unknown exit for break"
     end
 
     def visit(node : Next)
@@ -878,15 +906,15 @@ module Crystal
       when Block
         if next_phi = context.next_phi
           old_last = @last
-          execute_ensures_until(target.as(Block))
+          execute_ensures_until(target)
           @last = old_last
 
           next_phi.add @last, node_type
           return false
         end
       when While
         if while_block = context.while_block
-          execute_ensures_until(target.as(While))
+          execute_ensures_until(target)
           br while_block
           return false
         end
@@ -1497,7 +1525,6 @@ module Crystal
 
           context.break_phi = old.return_phi
           context.next_phi = phi
-          context.while_exit_block = nil
           context.closure_parent_context = block_context.closure_parent_context
 
           @needs_value = true

src/compiler/crystal/codegen/context.cr

@@ -11,7 +11,6 @@ class Crystal::CodeGenVisitor
     property break_phi : Phi?
     property next_phi : Phi?
     property while_block : LLVM::BasicBlock?
-    property while_exit_block : LLVM::BasicBlock?
     property! block : Block
     property! block_context : Context
     property closure_vars : Array(MetaVar)?
@@ -44,7 +43,6 @@ class Crystal::CodeGenVisitor
       context.break_phi = @break_phi
       context.next_phi = @next_phi
       context.while_block = @while_block
-      context.while_exit_block = @while_exit_block
       if block = @block
         context.block = block
       end

src/compiler/crystal/semantic/main_visitor.cr

@@ -2110,7 +2110,7 @@ module Crystal
         filter_vars TypeFilters.not(cond_type_filters)
       end
 
-      node.type = @program.nil
+      node.bind_to(@program.nil_var) unless endless_while
 
       false
     end
@@ -2293,6 +2293,7 @@ module Crystal
 
         break_vars = (target_while.break_vars ||= [] of MetaVars)
         break_vars.push @vars.dup
+        target_while.bind_to(node_exp_or_nil_literal(node))
       else
         if @typed_def.try &.captured_block?
           node.raise "can't break from captured block, try using `next`."