tint/uniformity: Fix issues with for-loops

Variables declared in for-loop initializers were not being tracked
properly across iterations as a check was wrongly determining them to
be declared inside the loop body.

Also fixes an issue where variables declared in for-loop initializers
were still considered to be in scope after loop exit.

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

src/tint/resolver/uniformity.cc

@@ -591,20 +591,11 @@ class UniformityGraph {
                 auto& info = current_function_->LoopSwitchInfoFor(parent);
 
                 // Propagate assignments to the loop input nodes.
-                for (auto* var : current_function_->local_var_decls) {
-                    // Skip variables that were declared inside this loop.
-                    if (auto* lv = var->As<sem::LocalVariable>();
-                        lv &&
-                        lv->Statement()->FindFirstParent([&](auto* s) { return s == parent; })) {
-                        continue;
-                    }
-
-                    // Add an edge from the variable's loop input node to its value at this point.
-                    auto in_node = info.var_in_nodes.Find(var);
-                    TINT_ASSERT(Resolver, in_node != nullptr);
-                    auto* out_node = current_function_->variables.Get(var);
-                    if (out_node != *in_node) {
-                        (*in_node)->AddEdge(out_node);
+                for (auto v : info.var_in_nodes) {
+                    auto* in_node = v.value;
+                    auto* out_node = current_function_->variables.Get(v.key);
+                    if (out_node != in_node) {
+                        in_node->AddEdge(out_node);
                     }
                 }
                 return cf;
@@ -677,6 +668,13 @@ class UniformityGraph {
                     current_function_->variables.Set(v.key, v.value);
                 }
 
+                if (f->initializer) {
+                    // Remove variables declared in the for-loop initializer from the current scope.
+                    if (auto* decl = f->initializer->As<ast::VariableDeclStatement>()) {
+                        current_function_->local_var_decls.Remove(sem_.Get(decl->variable));
+                    }
+                }
+
                 current_function_->RemoveLoopSwitchInfoFor(sem_loop);
 
                 if (sem_loop->Behaviors() == sem::Behaviors{sem::Behavior::kNext}) {

src/tint/resolver/uniformity_test.cc

@@ -1939,6 +1939,43 @@ test:12:9 note: reading from read_write storage buffer 'non_uniform' may result
 )");
 }
 
+TEST_F(UniformityAnalysisTest,
+       ForLoop_InitializerVarBecomesNonUniformBeforeConditionalContinue_BarrierAtStart) {
+    // Use a variable declared in a for-loop initializer for a conditional barrier in a loop, assign
+    // a non-uniform value to that variable later in that loop and then execute a continue.
+    // Tests that variables declared in the for-loop initializer are properly tracked.
+    std::string src = R"(
+@group(0) @binding(0) var<storage, read_write> non_uniform : i32;
+
+fn foo() {
+  for (var i = 0; i < 10; i++) {
+    if (i < 5) {
+      workgroupBarrier();
+    }
+    if (true) {
+      i = non_uniform;
+      continue;
+    }
+  }
+}
+)";
+
+    RunTest(src, false);
+    EXPECT_EQ(error_,
+              R"(test:7:7 warning: 'workgroupBarrier' must only be called from uniform control flow
+      workgroupBarrier();
+      ^^^^^^^^^^^^^^^^
+
+test:5:3 note: control flow depends on non-uniform value
+  for (var i = 0; i < 10; i++) {
+  ^^^
+
+test:10:11 note: reading from read_write storage buffer 'non_uniform' may result in a non-uniform value
+      i = non_uniform;
+          ^^^^^^^^^^^
+)");
+}
+
 TEST_F(UniformityAnalysisTest, ForLoop_NonUniformCondition_Reconverge) {
     // Loops reconverge at exit, so test that we can call workgroupBarrier() after a loop that has a
     // non-uniform condition.
@@ -1955,6 +1992,41 @@ fn foo() {
     RunTest(src, true);
 }
 
+TEST_F(UniformityAnalysisTest, ForLoop_VarDeclaredInBody) {
+    // Make sure that we can declare a variable inside the loop body without causing issues for
+    // tracking local variables across iterations.
+    std::string src = R"(
+@group(0) @binding(0) var<storage, read_write> n : i32;
+
+fn foo() {
+  var outer : i32;
+  for (var i = 0; i < n; i = i + 1) {
+    var inner : i32;
+  }
+}
+)";
+
+    RunTest(src, true);
+}
+
+TEST_F(UniformityAnalysisTest, ForLoop_InitializerScope) {
+    // Make sure that variables declared in a for-loop initializer are properly removed from the
+    // local variable list, otherwise a parent control-flow statement will try to add edges to nodes
+    // that no longer exist.
+    std::string src = R"(
+@group(0) @binding(0) var<storage, read_write> n : i32;
+
+fn foo() {
+  if (n == 5) {
+    for (var i = 0; i < n; i = i + 1) {
+    }
+  }
+}
+)";
+
+    RunTest(src, true);
+}
+
 TEST_F(UniformityAnalysisTest, While_CallInside_UniformCondition) {
     std::string src = R"(
 @group(0) @binding(0) var<storage, read> n : i32;