IR: Add support for break and continue statements (shader-slang#272)

* IR: Add support for break and continue statements

The front-end is already doing the work of connecting this statements to their "parent" statement, so we just needed to build a map from the `Stmt*` to the corresponding `IRBlock*`s to use for break/continue when outputting any loop statement, and then look up in the map for the branch target when outputting a break/continue.

When we get around to adding `switch` statements, the same pattern should work just fine.

I also added support for `do/while` statements in IR codegen, and made sure to exercise those in one of the test cases I added.

There is also an unrelated IR codegen fix for when there is a "bound subscript" on the RHS of an assignment.

* IR: fix handling of do/while and continue

Thanks to @csyonghe for pointing out my mistake in the earlier commit.
I implemented `continue` for `do/while` loops incorrectly, branching to the head of the loop instead of the loop test.
I'll try to blame this mistake on the fact that I never use `do/while` loops because I think they are awful. :)

The fix for that issue wasn't too bad (see `lower-to-ir.cpp`) but it surfaces a much more serious issue: I wasn't actually implementing `continue` correctly *at all* when it comes to generating HLSL/GLSL from the IR (I can't easily make an excuse for that one).
The basic issue at the heart of this is that given an input statement like:

```
for(int ii = 0; ii < N; ii = doSomething(ii))
{ ... }
```

The continue clause (`ii = doSomething(ii)`) could expand into many instructions (across multiple blocks, if we inline), and there is in general no guarantee when we are done that we can package up that code as an expression and spit out a new `for` loop (the same basic argument applies to a `do { ... } while(someComplexExpression())`.

So, if we assume that in general we have to generate a full *statement* for the `continue` clause, what can we emit?

- We could try to "outline" the continue code into its own function, so that we can call it from an expression. That could work, but has high implementation complexity.

- We could introduce additional `bool` variables for control flow, outputting something like:
  ```
  bool useContinueBlock = false;
  for(;;) {
    if(useContinueBlock) { <CONTINUE CODE>; }
    useContinueBlock = true;
    <LOOP TEST>
    <LOOP BODY>
  }
  ```
  This works but user might balk at the extra variable we introduce.

- We could duplicate the code at each continue site. That is, we emit the loop as:
  ```
  for(;;) {
    <LOOP TEST>
    <LOOP BODY>
    <CONTINUE CODE>
  }
  ```
  but then whenever we'd like to emit `continue;` we instead emit `{ <CONTINUE CODE>; continue; }`.
  This doesn't introduce any extra variables, but it causes code duplication (limited, if we don't have too many `continue` sites, and the continue clause is small - which are the common cases).

When I was initially working on the IR codegen I picked that last option just because it is what `fxc` seems to do, but I neglected to actually *implement* the special-case codegen for a `continue` instruction. This change addresses that (see `emit.cpp`).

Finally, once things were fixed the `continue` test case produced the results Yong told me to expect, but it also produced a warning from the downstream HLSL compiler ("hey, your loop doesn't ever actually *loop*!"), so I reworked the test back to one that actually loops (but still tests `continue`).

As a final aside in this essay of a commit message: the current IR representation of control flow uses special-case instructions for various cases of unconditional branch (and two variations on `if`), but these are not strictly necessary, and a future change will hopefully clean it up. The biggest catch in doing that is that it will require the IR->source codegen to carefully track which blocks represent which kinds of branch targets in context (e.g., you can't assume that a `continue` that nees the special handling above will appear as a distinct kind of instruction).

Author: Tim Foley

source/slang/emit.cpp

@@ -97,6 +97,8 @@ struct SharedEmitContext
     Dictionary<IRValue*, UInt> mapIRValueToID;
 
     HashSet<Decl*> irDeclsVisited;
+
+    Dictionary<IRBlock*, IRBlock*> irMapContinueTargetToLoopHead;
 };
 
 struct EmitContext
@@ -5555,10 +5557,10 @@ emitDeclImpl(decl, nullptr);
 
                         emit("for(;;)\n{\n");
 
-                        // TODO: Okay, we *said* we'd do this special
-                        // handling of the `continue` sites, but
-                        // we aren't actually setting anything up here...
-                        //
+                        // Register information so that `continue` sites
+                        // can do the right thing:
+                        ctx->shared->irMapContinueTargetToLoopHead.Add(continueBlock, targetBlock);
+
 
                         emitIRStmtsForBlocks(
                             ctx,
@@ -5579,7 +5581,28 @@ emitDeclImpl(decl, nullptr);
                 return;
 
             case kIROp_continue:
-                emit("continue;\n");
+                // With out current strategy for outputting loops,
+                // just outputting an AST-level `continue` here won't
+                // actually execute the statements in the continue block.
+                //
+                // Instead, we have to manually output those statements
+                // directly here, and *then* do an AST-level `continue`.
+                //
+                // This leads to code duplication when we have multiple
+                // `continue` sites in the original program, but it avoids
+                // introducing additional temporaries for control flow.
+                {
+                    auto continueInst = (IRContinue*) terminator;
+                    auto targetBlock = continueInst->getTargetBlock();
+                    IRBlock* loopHead = nullptr;
+                    ctx->shared->irMapContinueTargetToLoopHead.TryGetValue(targetBlock, loopHead);
+                    SLANG_ASSERT(loopHead);
+                    emitIRStmtsForBlocks(
+                        ctx,
+                        targetBlock,
+                        loopHead);
+                    emit("continue;\n");
+                }
                 return;
 
             case kIROp_loopTest:

source/slang/lower-to-ir.cpp

@@ -279,6 +279,12 @@ struct SharedIRGenContext
     // to reference-count these along the way because
     // they need to get stored into a `union` inside `LoweredValInfo`
     List<RefPtr<ExtendedValueInfo>> extValues;
+
+    // Map from an AST-level statement that can be
+    // used as the target of a `break` or `continue`
+    // to the appropriate basic block to jump to.
+    Dictionary<Stmt*, IRBlock*> breakLabels;
+    Dictionary<Stmt*, IRBlock*> continueLabels;
 };
 
 
@@ -1741,8 +1747,10 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
         auto breakLabel = createBlock();
         auto continueLabel = createBlock();
 
-        // TODO: register `loopHead` as the target for a
-        // `continue` statement.
+        // Register the `break` and `continue` labels so
+        // that we can find them for nested statements.
+        context->shared->breakLabels.Add(stmt, breakLabel);
+        context->shared->continueLabels.Add(stmt, continueLabel);
 
         // Emit the branch that will start out loop,
         // and then insert the block for the head.
@@ -1807,8 +1815,10 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
         // jumps to the head of hte loop.
         auto continueLabel = loopHead;
 
-        // TODO: register appropriate targets for
-        // break/continue statements.
+        // Register the `break` and `continue` labels so
+        // that we can find them for nested statements.
+        context->shared->breakLabels.Add(stmt, breakLabel);
+        context->shared->continueLabels.Add(stmt, continueLabel);
 
         // Emit the branch that will start out loop,
         // and then insert the block for the head.
@@ -1847,6 +1857,66 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
         insertBlock(breakLabel);
     }
 
+    void visitDoWhileStmt(DoWhileStmt* stmt)
+    {
+        // Generating IR for `do {...} while` statement is similar to a
+        // `while` statement, just with the test in a different place
+
+        auto builder = getBuilder();
+
+        // We will create blocks for the various places
+        // we need to jump to inside the control flow,
+        // including the blocks that will be referenced
+        // by `continue` or `break` statements.
+        auto loopHead = createBlock();
+        auto testLabel = createBlock();
+        auto breakLabel = createBlock();
+
+        // A `continue` inside a `do { ... } while ( ... )` loop always
+        // jumps to the loop test.
+        auto continueLabel = testLabel;
+
+        // Register the `break` and `continue` labels so
+        // that we can find them for nested statements.
+        context->shared->breakLabels.Add(stmt, breakLabel);
+        context->shared->continueLabels.Add(stmt, continueLabel);
+
+        // Emit the branch that will start out loop,
+        // and then insert the block for the head.
+
+        auto loopInst = builder->emitLoop(
+            loopHead,
+            breakLabel,
+            continueLabel);
+
+        addLoopDecorations(loopInst, stmt);
+
+        insertBlock(loopHead);
+
+        // Emit the body of the loop
+        lowerStmt(context, stmt->Statement);
+
+        insertBlock(testLabel);
+
+        // Now that we are within the header block, we
+        // want to emit the expression for the loop condition:
+        if (auto condExpr = stmt->Predicate)
+        {
+            auto irCondition = getSimpleVal(context,
+                lowerRValueExpr(context, condExpr));
+
+            // Now we want to `break` if the loop condition is false,
+            // otherwise we will jump back to the head of the loop.
+            builder->emitLoopTest(
+                irCondition,
+                loopHead,
+                breakLabel);
+        }
+
+        // Finally we insert the label that a `break` will jump to
+        insertBlock(breakLabel);
+    }
+
     void visitExpressionStmt(ExpressionStmt* stmt)
     {
         // The statement evaluates an expression
@@ -1917,6 +1987,39 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
     {
         getBuilder()->emitDiscard();
     }
+
+    void visitBreakStmt(BreakStmt* stmt)
+    {
+        // Semantic checking is responsible for finding
+        // the statement taht this `break` breaks out of
+        auto parentStmt = stmt->parentStmt;
+        SLANG_ASSERT(parentStmt);
+
+        // We just need to look up the basic block that
+        // corresponds to the break label for that statement,
+        // and then emit an instruction to jump to it.
+        IRBlock* targetBlock;
+        context->shared->breakLabels.TryGetValue(parentStmt, targetBlock);
+        SLANG_ASSERT(targetBlock);
+        getBuilder()->emitBreak(targetBlock);
+    }
+
+    void visitContinueStmt(ContinueStmt* stmt)
+    {
+        // Semantic checking is responsible for finding
+        // the loop that this `continue` statement continues
+        auto parentStmt = stmt->parentStmt;
+        SLANG_ASSERT(parentStmt);
+
+
+        // We just need to look up the basic block that
+        // corresponds to the continue label for that statement,
+        // and then emit an instruction to jump to it.
+        IRBlock* targetBlock;
+        context->shared->continueLabels.TryGetValue(parentStmt, targetBlock);
+        SLANG_ASSERT(targetBlock);
+        getBuilder()->emitContinue(targetBlock);
+    }
 };
 
 void lowerStmt(
@@ -1947,6 +2050,7 @@ void assign(
         case LoweredValInfo::Flavor::Simple:
         case LoweredValInfo::Flavor::Ptr:
         case LoweredValInfo::Flavor::SwizzledLValue:
+        case LoweredValInfo::Flavor::BoundSubscript:
             {
                 builder->emitStore(
                     left.val,

tests/compute/break-stmt.slang

@@ -0,0 +1,27 @@
+//TEST(compute):COMPARE_COMPUTE:
+//TEST_INPUT:ubuffer(data=[0 1 2 3], stride=4):dxbinding(0),glbinding(0),out
+
+// Test that `break` from a loop works.
+
+int test(int inVal)
+{
+	int ii = 0;
+	for(;;)
+	{
+		if(ii >= inVal)
+			break;
+	    ii++;
+	}
+	return -ii;
+}
+
+RWStructuredBuffer<int> outputBuffer : register(u0);
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+	uint tid = dispatchThreadID.x;
+	int inVal = outputBuffer[tid];
+	int outVal = test(inVal);
+	outputBuffer[tid] = outVal;
+}

tests/compute/break-stmt.slang.expected.txt

@@ -0,0 +1,4 @@
+0
+FFFFFFFF
+FFFFFFFE
+FFFFFFFD

tests/compute/continue-stmt.slang

@@ -0,0 +1,32 @@
+//TEST(compute):COMPARE_COMPUTE:
+//TEST_INPUT:ubuffer(data=[0 1 2 3], stride=4):dxbinding(0),glbinding(0),out
+
+// Test that `break` from a loop works.
+
+int test(int inVal)
+{
+	int ii = 0;
+	do
+	{
+		if(ii < inVal)
+		{
+		    ii++;
+			continue;
+		}
+		break;
+	}
+	while(true);
+
+	return -ii;
+}
+
+RWStructuredBuffer<int> outputBuffer : register(u0);
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+	uint tid = dispatchThreadID.x;
+	int inVal = outputBuffer[tid];
+	int outVal = test(inVal);
+	outputBuffer[tid] = outVal;
+}

tests/compute/continue-stmt.slang.expected.txt

@@ -0,0 +1,4 @@
+0
+FFFFFFFF
+FFFFFFFE
+FFFFFFFD