spir-v builder: fix loop-scoped variable not found in continuing block

Make sure variables from the loop block remain in scope for
continuing block. Note that we need to do this because the continuing
block is a sibling of the loop body block in the AST, rather than a
child.

Added test.

Fixed: tint:526
Change-Id: If622995e3aac4cd3c06c2dbd87ffcaa36b0f09c5
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/43680
Commit-Queue: David Neto <dneto@google.com>
Reviewed-by: David Neto <dneto@google.com>

src/writer/spirv/builder.cc

@@ -1766,13 +1766,18 @@ uint32_t Builder::GenerateBinaryExpression(ast::BinaryExpression* expr) {
 
 bool Builder::GenerateBlockStatement(const ast::BlockStatement* stmt) {
   scope_stack_.push_scope();
+  auto result = GenerateBlockStatementWithoutScoping(stmt);
+  scope_stack_.pop_scope();
+  return result;
+}
+
+bool Builder::GenerateBlockStatementWithoutScoping(
+    const ast::BlockStatement* stmt) {
   for (auto* block_stmt : *stmt) {
     if (!GenerateStatement(block_stmt)) {
       return false;
     }
   }
-  scope_stack_.pop_scope();
-
   return true;
 }
 
@@ -2690,7 +2695,12 @@ bool Builder::GenerateLoopStatement(ast::LoopStatement* stmt) {
   if (!GenerateLabel(body_block_id)) {
     return false;
   }
-  if (!GenerateBlockStatement(stmt->body())) {
+
+  // We need variables from the body to be visible in the continuing block, so
+  // manage scope outside of GenerateBlockStatement.
+  scope_stack_.push_scope();
+
+  if (!GenerateBlockStatementWithoutScoping(stmt->body())) {
     return false;
   }
 
@@ -2705,9 +2715,12 @@ bool Builder::GenerateLoopStatement(ast::LoopStatement* stmt) {
   if (!GenerateLabel(continue_block_id)) {
     return false;
   }
-  if (!GenerateBlockStatement(stmt->continuing())) {
+  if (!GenerateBlockStatementWithoutScoping(stmt->continuing())) {
     return false;
   }
+
+  scope_stack_.pop_scope();
+
   if (!push_function_inst(spv::Op::OpBranch, {Operand::Int(loop_header_id)})) {
     return false;
   }

src/writer/spirv/builder.h

@@ -216,10 +216,14 @@ class Builder {
   /// @param assign the statement to generate
   /// @returns true if the statement was successfully generated
   bool GenerateAssignStatement(ast::AssignmentStatement* assign);
-  /// Generates a block statement
+  /// Generates a block statement, wrapped in a push/pop scope
   /// @param stmt the statement to generate
   /// @returns true if the statement was successfully generated
   bool GenerateBlockStatement(const ast::BlockStatement* stmt);
+  /// Generates a block statement
+  /// @param stmt the statement to generate
+  /// @returns true if the statement was successfully generated
+  bool GenerateBlockStatementWithoutScoping(const ast::BlockStatement* stmt);
   /// Generates a break statement
   /// @param stmt the statement to generate
   /// @returns true if the statement was successfully generated

src/writer/spirv/builder_loop_test.cc

@@ -133,6 +133,47 @@ OpBranch %5
 )");
 }
 
+TEST_F(BuilderTest, Loop_WithBodyVariableAccessInContinuing) {
+  // loop {
+  //   var a : i32;
+  //   continuing {
+  //     a = 3;
+  //   }
+  // }
+
+  auto* var = Var("a", ty.i32(), ast::StorageClass::kFunction);
+  auto* var_decl = WrapInStatement(var);
+  auto* body = create<ast::BlockStatement>(ast::StatementList{var_decl});
+  auto* continuing = create<ast::BlockStatement>(
+      ast::StatementList{create<ast::AssignmentStatement>(Expr("a"), Expr(3))});
+
+  auto* loop = create<ast::LoopStatement>(body, continuing);
+  WrapInFunction(loop);
+
+  spirv::Builder& b = Build();
+
+  b.push_function(Function{});
+  EXPECT_TRUE(b.GenerateLoopStatement(loop)) << b.error();
+
+  EXPECT_EQ(DumpInstructions(b.types()), R"(%7 = OpTypeInt 32 1
+%6 = OpTypePointer Function %7
+%8 = OpConstantNull %7
+%9 = OpConstant %7 3
+)");
+  EXPECT_EQ(DumpInstructions(b.functions()[0].instructions()),
+            R"(OpBranch %1
+%1 = OpLabel
+OpLoopMerge %2 %3 None
+OpBranch %4
+%4 = OpLabel
+OpBranch %3
+%3 = OpLabel
+OpStore %5 %9
+OpBranch %1
+%2 = OpLabel
+)");
+}
+
 TEST_F(BuilderTest, Loop_WithContinue) {
   // loop {
   //   continue;