Add support for 'break if' to IR, wgsl-in, and all backends.

src/back/dot/mod.rs

@@ -81,11 +81,15 @@ impl StatementGraph {
                 S::Loop {
                     ref body,
                     ref continuing,
+                    break_if,
                 } => {
                     let body_id = self.add(body);
                     self.flow.push((id, body_id, "body"));
                     let continuing_id = self.add(continuing);
                     self.flow.push((body_id, continuing_id, "continuing"));
+                    if let Some(expr) = break_if {
+                        self.dependencies.push((id, expr, "break if"));
+                    }
                     "Loop"
                 }
                 S::Return { value } => {

src/back/glsl/mod.rs

@@ -1800,15 +1800,24 @@ impl<'a, W: Write> Writer<'a, W> {
             Statement::Loop {
                 ref body,
                 ref continuing,
+                break_if,
             } => {
-                if !continuing.is_empty() {
+                if !continuing.is_empty() || break_if.is_some() {
                     let gate_name = self.namer.call("loop_init");
                     writeln!(self.out, "{}bool {} = true;", level, gate_name)?;
                     writeln!(self.out, "{}while(true) {{", level)?;
                     let l2 = level.next();
+                    let l3 = l2.next();
                     writeln!(self.out, "{}if (!{}) {{", l2, gate_name)?;
                     for sta in continuing {
-                        self.write_stmt(sta, ctx, l2.next())?;
+                        self.write_stmt(sta, ctx, l3)?;
+                    }
+                    if let Some(condition) = break_if {
+                        write!(self.out, "{}if (", l3)?;
+                        self.write_expr(condition, ctx)?;
+                        writeln!(self.out, ") {{")?;
+                        writeln!(self.out, "{}break;", l3.next())?;
+                        writeln!(self.out, "{}}}", l3)?;
                     }
                     writeln!(self.out, "{}}}", l2)?;
                     writeln!(self.out, "{}{} = false;", level.next(), gate_name)?;

src/back/hlsl/writer.rs

@@ -1497,18 +1497,27 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
             Statement::Loop {
                 ref body,
                 ref continuing,
+                break_if,
             } => {
                 let l2 = level.next();
-                if !continuing.is_empty() {
+                if !continuing.is_empty() || break_if.is_some() {
                     let gate_name = self.namer.call("loop_init");
                     writeln!(self.out, "{}bool {} = true;", level, gate_name)?;
                     writeln!(self.out, "{}while(true) {{", level)?;
                     writeln!(self.out, "{}if (!{}) {{", l2, gate_name)?;
+                    let l3 = l2.next();
                     for sta in continuing.iter() {
-                        self.write_stmt(module, sta, func_ctx, l2.next())?;
+                        self.write_stmt(module, sta, func_ctx, l3)?;
                     }
-                    writeln!(self.out, "{}}}", level.next())?;
-                    writeln!(self.out, "{}{} = false;", level.next(), gate_name)?;
+                    if let Some(condition) = break_if {
+                        write!(self.out, "{}if (", l3)?;
+                        self.write_expr(module, condition, func_ctx)?;
+                        writeln!(self.out, ") {{")?;
+                        writeln!(self.out, "{}break;", l3.next())?;
+                        writeln!(self.out, "{}}}", l3)?;
+                    }
+                    writeln!(self.out, "{}}}", l2)?;
+                    writeln!(self.out, "{}{} = false;", l2, gate_name)?;
                 } else {
                     writeln!(self.out, "{}while(true) {{", level)?;
                 }

src/back/msl/writer.rs

@@ -2552,14 +2552,23 @@ impl<W: Write> Writer<W> {
                 crate::Statement::Loop {
                     ref body,
                     ref continuing,
+                    break_if,
                 } => {
-                    if !continuing.is_empty() {
+                    if !continuing.is_empty() || break_if.is_some() {
                         let gate_name = self.namer.call("loop_init");
                         writeln!(self.out, "{}bool {} = true;", level, gate_name)?;
                         writeln!(self.out, "{}while(true) {{", level)?;
                         let lif = level.next();
+                        let lcontinuing = lif.next();
                         writeln!(self.out, "{}if (!{}) {{", lif, gate_name)?;
-                        self.put_block(lif.next(), continuing, context)?;
+                        self.put_block(lcontinuing, continuing, context)?;
+                        if let Some(condition) = break_if {
+                            write!(self.out, "{}if (", lcontinuing)?;
+                            self.put_expression(condition, &context.expression, true)?;
+                            writeln!(self.out, ") {{")?;
+                            writeln!(self.out, "{}break;", lcontinuing.next())?;
+                            writeln!(self.out, "{}}}", lcontinuing)?;
+                        }
                         writeln!(self.out, "{}}}", lif)?;
                         writeln!(self.out, "{}{} = false;", lif, gate_name)?;
                     } else {

src/back/spv/block.rs

@@ -37,6 +37,28 @@ enum ExpressionPointer {
     },
 }
 
+/// The termination statement to be added to the end of the block
+pub enum BlockExit {
+    /// Generates an OpReturn (void return)
+    Return,
+    /// Generates an OpBranch to the specified block
+    Branch {
+        /// The branch target block
+        target: Word,
+    },
+    /// Translates a loop `break if` into an `OpBranchConditional` to the
+    /// merge block if true (the merge block is passed through [`LoopContext::break_id`]
+    /// or else to the loop header (passed through [`preamble_id`])
+    ///
+    /// [`preamble_id`]: Self::BreakIf::preamble_id
+    BreakIf {
+        /// The condition of the `break if`
+        condition: Handle<crate::Expression>,
+        /// The loop header block id
+        preamble_id: Word,
+    },
+}
+
 impl Writer {
     // Flip Y coordinate to adjust for coordinate space difference
     // between SPIR-V and our IR.
@@ -1491,7 +1513,7 @@ impl<'w> BlockContext<'w> {
         &mut self,
         label_id: Word,
         statements: &[crate::Statement],
-        exit_id: Option<Word>,
+        exit: BlockExit,
         loop_context: LoopContext,
     ) -> Result<(), Error> {
         let mut block = Block::new(label_id);
@@ -1508,7 +1530,12 @@ impl<'w> BlockContext<'w> {
                     self.function.consume(block, Instruction::branch(scope_id));
 
                     let merge_id = self.gen_id();
-                    self.write_block(scope_id, block_statements, Some(merge_id), loop_context)?;
+                    self.write_block(
+                        scope_id,
+                        block_statements,
+                        BlockExit::Branch { target: merge_id },
+                        loop_context,
+                    )?;
 
                     block = Block::new(merge_id);
                 }
@@ -1546,10 +1573,20 @@ impl<'w> BlockContext<'w> {
                     );
 
                     if let Some(block_id) = accept_id {
-                        self.write_block(block_id, accept, Some(merge_id), loop_context)?;
+                        self.write_block(
+                            block_id,
+                            accept,
+                            BlockExit::Branch { target: merge_id },
+                            loop_context,
+                        )?;
                     }
                     if let Some(block_id) = reject_id {
-                        self.write_block(block_id, reject, Some(merge_id), loop_context)?;
+                        self.write_block(
+                            block_id,
+                            reject,
+                            BlockExit::Branch { target: merge_id },
+                            loop_context,
+                        )?;
                     }
 
                     block = Block::new(merge_id);
@@ -1611,22 +1648,30 @@ impl<'w> BlockContext<'w> {
                         self.write_block(
                             *label_id,
                             &case.body,
-                            Some(case_finish_id),
+                            BlockExit::Branch {
+                                target: case_finish_id,
+                            },
                             inner_context,
                         )?;
                     }
 
                     // If no default was encountered write a empty block to satisfy the presence of
                     // a block the default label
                     if !reached_default {
-                        self.write_block(default_id, &[], Some(merge_id), inner_context)?;
+                        self.write_block(
+                            default_id,
+                            &[],
+                            BlockExit::Branch { target: merge_id },
+                            inner_context,
+                        )?;
                     }
 
                     block = Block::new(merge_id);
                 }
                 crate::Statement::Loop {
                     ref body,
                     ref continuing,
+                    break_if,
                 } => {
                     let preamble_id = self.gen_id();
                     self.function
@@ -1649,17 +1694,29 @@ impl<'w> BlockContext<'w> {
                     self.write_block(
                         body_id,
                         body,
-                        Some(continuing_id),
+                        BlockExit::Branch {
+                            target: continuing_id,
+                        },
                         LoopContext {
                             continuing_id: Some(continuing_id),
                             break_id: Some(merge_id),
                         },
                     )?;
 
+                    let exit = match break_if {
+                        Some(condition) => BlockExit::BreakIf {
+                            condition,
+                            preamble_id,
+                        },
+                        None => BlockExit::Branch {
+                            target: preamble_id,
+                        },
+                    };
+
                     self.write_block(
                         continuing_id,
                         continuing,
-                        Some(preamble_id),
+                        exit,
                         LoopContext {
                             continuing_id: None,
                             break_id: Some(merge_id),
@@ -1955,19 +2012,30 @@ impl<'w> BlockContext<'w> {
             }
         }
 
-        let termination = match exit_id {
-            Some(id) => Instruction::branch(id),
-            // This can happen if the last branch had all the paths
-            // leading out of the graph (i.e. returning).
-            // Or it may be the end of the self.function.
-            None => match self.ir_function.result {
+        let termination = match exit {
+            // We're generating code for the top-level Block of the function, so we
+            // need to end it with some kind of return instruction.
+            BlockExit::Return => match self.ir_function.result {
                 Some(ref result) if self.function.entry_point_context.is_none() => {
                     let type_id = self.get_type_id(LookupType::Handle(result.ty));
                     let null_id = self.writer.write_constant_null(type_id);
                     Instruction::return_value(null_id)
                 }
                 _ => Instruction::return_void(),
             },
+            BlockExit::Branch { target } => Instruction::branch(target),
+            BlockExit::BreakIf {
+                condition,
+                preamble_id,
+            } => {
+                let condition_id = self.cached[condition];
+
+                Instruction::branch_conditional(
+                    condition_id,
+                    loop_context.break_id.unwrap(),
+                    preamble_id,
+                )
+            }
         };
 
         self.function.consume(block, termination);

src/back/spv/writer.rs

@@ -574,7 +574,12 @@ impl Writer {
         context
             .function
             .consume(prelude, Instruction::branch(main_id));
-        context.write_block(main_id, &ir_function.body, None, LoopContext::default())?;
+        context.write_block(
+            main_id,
+            &ir_function.body,
+            super::block::BlockExit::Return,
+            LoopContext::default(),
+        )?;
 
         // Consume the `BlockContext`, ending its borrows and letting the
         // `Writer` steal back its cached expression table and temp_list.

src/back/wgsl/writer.rs

@@ -908,6 +908,7 @@ impl<W: Write> Writer<W> {
             Statement::Loop {
                 ref body,
                 ref continuing,
+                break_if,
             } => {
                 write!(self.out, "{}", level)?;
                 writeln!(self.out, "loop {{")?;
@@ -917,11 +918,26 @@ impl<W: Write> Writer<W> {
                     self.write_stmt(module, sta, func_ctx, l2)?;
                 }
 
-                if !continuing.is_empty() {
+                // The continuing is optional so we don't need to write it if
+                // it is empty, but the `break if` counts as a continuing statement
+                // so even if `continuing` is empty we must generate it if a
+                // `break if` exists
+                if !continuing.is_empty() || break_if.is_some() {
                     writeln!(self.out, "{}continuing {{", l2)?;
                     for sta in continuing.iter() {
                         self.write_stmt(module, sta, func_ctx, l2.next())?;
                     }
+
+                    // The `break if` is always the last
+                    // statement of the `continuing` block
+                    if let Some(condition) = break_if {
+                        // The trailing space is important
+                        write!(self.out, "{}break if ", l2.next())?;
+                        self.write_expr(module, condition, func_ctx)?;
+                        // Close the `break if` statement
+                        writeln!(self.out, ";")?;
+                    }
+
                     writeln!(self.out, "{}}}", l2)?;
                 }
 

src/front/glsl/parser/functions.rs

@@ -358,6 +358,7 @@ impl<'source> ParsingContext<'source> {
                     Statement::Loop {
                         body: loop_body,
                         continuing: Block::new(),
+                        break_if: None,
                     },
                     meta,
                 );
@@ -411,6 +412,7 @@ impl<'source> ParsingContext<'source> {
                     Statement::Loop {
                         body: loop_body,
                         continuing: Block::new(),
+                        break_if: None,
                     },
                     meta,
                 );
@@ -513,6 +515,7 @@ impl<'source> ParsingContext<'source> {
                     Statement::Loop {
                         body: block,
                         continuing,
+                        break_if: None,
                     },
                     meta,
                 );

src/front/spv/function.rs

@@ -556,7 +556,11 @@ impl<'function> BlockContext<'function> {
                         let continuing = lower_impl(blocks, bodies, continuing);
 
                         block.push(
-                            crate::Statement::Loop { body, continuing },
+                            crate::Statement::Loop {
+                                body,
+                                continuing,
+                                break_if: None,
+                            },
                             crate::Span::default(),
                         )
                     }

src/front/spv/mod.rs

@@ -3565,6 +3565,7 @@ impl<I: Iterator<Item = u32>> Parser<I> {
                 S::Loop {
                     ref mut body,
                     ref mut continuing,
+                    break_if: _,
                 } => {
                     self.patch_statements(body, expressions, fun_parameter_sampling)?;
                     self.patch_statements(continuing, expressions, fun_parameter_sampling)?;