tint/uniformity: Fix handling of continuing block

Variables declared inside a loop block were not visible to the
continuing block. Special-case the continuing block by processing it
inside the loop-body block statement, instead of afterwards.

Change-Id: I05bc906bd98b24295dc91116b9ef8d8ef02c3af3
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/114860
Reviewed-by: Antonio Maiorano <amaiorano@google.com>
Commit-Queue: James Price <jrprice@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>

src/tint/resolver/uniformity.cc

@@ -462,6 +462,22 @@ class UniformityGraph {
                         }
                     }
 
+                    auto* parent = sem_.Get(b)->Parent();
+                    auto* loop = parent ? parent->As<sem::LoopStatement>() : nullptr;
+                    if (loop) {
+                        // We've reached the end of a loop body. If there is a continuing block,
+                        // process it before ending the block so that any variables declared in the
+                        // loop body are visible to the continuing block.
+                        if (auto* continuing =
+                                loop->Declaration()->As<ast::LoopStatement>()->continuing) {
+                            auto& loop_body_behavior = sem_.Get(b)->Behaviors();
+                            if (loop_body_behavior.Contains(sem::Behavior::kNext) ||
+                                loop_body_behavior.Contains(sem::Behavior::kContinue)) {
+                                cf = ProcessStatement(cf, continuing);
+                            }
+                        }
+                    }
+
                     if (sem_.Get<sem::FunctionBlockStatement>(b)) {
                         // We've reached the end of the function body.
                         // Add edges from pointer parameter outputs to their current value.
@@ -860,13 +876,11 @@ class UniformityGraph {
                     current_function_->variables.Set(v, in_node);
                 }
 
+                // Note: The continuing block is processed as a special case at the end of
+                // processing the loop body BlockStatement. This is so that variable declarations
+                // inside the loop body are visible to the continuing statement.
                 auto* cf1 = ProcessStatement(cfx, l->body);
-                if (l->continuing) {
+                cfx->AddEdge(cf1);
-                    auto* cf2 = ProcessStatement(cf1, l->continuing);
-                    cfx->AddEdge(cf2);
-                } else {
-                    cfx->AddEdge(cf1);
-                }
                 cfx->AddEdge(cf);
 
                 // Add edges from variable loop input nodes to their values at the end of the loop.

src/tint/resolver/uniformity_test.cc

@@ -1545,6 +1545,71 @@ fn foo() {
     RunTest(src, true);
 }
 
+TEST_F(UniformityAnalysisTest, Loop_NonUniformValueDeclaredInBody_UnreachableContinuing) {
+    std::string src = R"(
+@group(0) @binding(0) var<storage, read_write> non_uniform : i32;
+
+fn foo() {
+  var condition = true;
+  loop {
+    var v = non_uniform;
+    if (condition) {
+      break;
+    } else {
+      break;
+    }
+
+    continuing {
+      if (v == 0) {
+        workgroupBarrier();
+      }
+    }
+  }
+}
+)";
+
+    RunTest(src, true);
+}
+
+TEST_F(UniformityAnalysisTest, Loop_NonUniformValueDeclaredInBody_MaybeReachesContinuing) {
+    std::string src = R"(
+@group(0) @binding(0) var<storage, read_write> non_uniform : i32;
+
+fn foo() {
+  var condition = true;
+  loop {
+    var v = non_uniform;
+    if (condition) {
+      continue;
+    } else {
+      break;
+    }
+
+    continuing {
+      if (v == 0) {
+        workgroupBarrier();
+      }
+    }
+  }
+}
+)";
+
+    RunTest(src, false);
+    EXPECT_EQ(error_,
+              R"(test:16:9 warning: 'workgroupBarrier' must only be called from uniform control flow
+        workgroupBarrier();
+        ^^^^^^^^^^^^^^^^
+
+test:15:7 note: control flow depends on non-uniform value
+      if (v == 0) {
+      ^^
+
+test:7:13 note: reading from read_write storage buffer 'non_uniform' may result in a non-uniform value
+    var v = non_uniform;
+            ^^^^^^^^^^^
+)");
+}
+
 TEST_F(UniformityAnalysisTest, Loop_NonUniformBreakInBody_Reconverge) {
     // Loops reconverge at exit, so test that we can call workgroupBarrier() after a loop that
     // contains a non-uniform conditional break.