[tint][ir] Validate NextIteration instructions

Change-Id: I3d2e08dc93c744263bdf6b786ec542680183a29d
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/187685
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: James Price <jrprice@google.com>
diff --git a/src/tint/lang/core/ir/validator.cc b/src/tint/lang/core/ir/validator.cc
index f862f95..891fc09 100644
--- a/src/tint/lang/core/ir/validator.cc
+++ b/src/tint/lang/core/ir/validator.cc
@@ -327,13 +327,17 @@
     /// @param b the terminator to validate
     void CheckTerminator(const Terminator* b);
 
+    /// Validates the continue instruction
+    /// @param c the continue to validate
+    void CheckContinue(const Continue* c);
+
     /// Validates the given exit
     /// @param e the exit to validate
     void CheckExit(const Exit* e);
 
-    /// Validates the continue instruction
-    /// @param c the continue to validate
-    void CheckContinue(const Continue* c);
+    /// Validates the next iteration instruction
+    /// @param n the next iteration to validate
+    void CheckNextIteration(const NextIteration* n);
 
     /// Validates the given exit if
     /// @param e the exit if to validate
@@ -1153,14 +1157,14 @@
     // DemoteToHelper) so we can't add validation.
 
     tint::Switch(
-        b,                                                 //
-        [&](const ir::BreakIf*) {},                        //
-        [&](const ir::Continue* c) { CheckContinue(c); },  //
-        [&](const ir::Exit* e) { CheckExit(e); },          //
-        [&](const ir::NextIteration*) {},                  //
-        [&](const ir::Return* ret) { CheckReturn(ret); },  //
-        [&](const ir::TerminateInvocation*) {},            //
-        [&](const ir::Unreachable*) {},                    //
+        b,                                                           //
+        [&](const ir::BreakIf*) {},                                  //
+        [&](const ir::Continue* c) { CheckContinue(c); },            //
+        [&](const ir::Exit* e) { CheckExit(e); },                    //
+        [&](const ir::NextIteration* n) { CheckNextIteration(n); },  //
+        [&](const ir::Return* ret) { CheckReturn(ret); },            //
+        [&](const ir::TerminateInvocation*) {},                      //
+        [&](const ir::Unreachable*) {},                              //
         [&](Default) { AddError(b) << "missing validation"; });
 }
 
@@ -1219,6 +1223,21 @@
         [&](Default) { AddError(e) << "missing validation"; });
 }
 
+void Validator::CheckNextIteration(const NextIteration* n) {
+    auto* loop = n->Loop();
+    if (loop == nullptr) {
+        AddError(n) << "has no associated loop";
+        return;
+    }
+    if (!TransitivelyHolds(loop->Initializer(), n) && !TransitivelyHolds(loop->Continuing(), n)) {
+        if (control_stack_.Any(Eq<const ControlInstruction*>(loop))) {
+            AddError(n) << "must only be called from loop initializer or continuing";
+        } else {
+            AddError(n) << "called outside of associated loop";
+        }
+    }
+}
+
 void Validator::CheckExitIf(const ExitIf* e) {
     if (control_stack_.Back() != e->If()) {
         AddError(e) << "if target jumps over other control instructions";
diff --git a/src/tint/lang/core/ir/validator_test.cc b/src/tint/lang/core/ir/validator_test.cc
index 8e158af..b14f6c3 100644
--- a/src/tint/lang/core/ir/validator_test.cc
+++ b/src/tint/lang/core/ir/validator_test.cc
@@ -2722,6 +2722,98 @@
 )");
 }
 
+TEST_F(IR_ValidatorTest, NextIterationOutsideOfLoop) {
+    auto* f = b.Function("my_func", ty.void_());
+    b.Append(f->Block(), [&] {
+        auto* loop = b.Loop();
+        b.Append(loop->Body(), [&] { b.ExitLoop(loop); });
+        b.NextIteration(loop);
+    });
+
+    auto res = ir::Validate(mod);
+    ASSERT_NE(res, Success);
+    EXPECT_EQ(res.Failure().reason.Str(),
+              R"(:8:5 error: next_iteration: called outside of associated loop
+    next_iteration  # -> $B2
+    ^^^^^^^^^^^^^^
+
+:2:3 note: in block
+  $B1: {
+  ^^^
+
+note: # Disassembly
+%my_func = func():void {
+  $B1: {
+    loop [b: $B2] {  # loop_1
+      $B2: {  # body
+        exit_loop  # loop_1
+      }
+    }
+    next_iteration  # -> $B2
+  }
+}
+)");
+}
+
+TEST_F(IR_ValidatorTest, NextIterationInLoopInit) {
+    auto* f = b.Function("my_func", ty.void_());
+    b.Append(f->Block(), [&] {
+        auto* loop = b.Loop();
+        b.Append(loop->Initializer(), [&] { b.NextIteration(loop); });
+        b.Append(loop->Body(), [&] { b.ExitLoop(loop); });
+        b.Return(f);
+    });
+
+    auto res = ir::Validate(mod);
+    ASSERT_EQ(res, Success);
+}
+
+TEST_F(IR_ValidatorTest, NextIterationInLoopBody) {
+    auto* f = b.Function("my_func", ty.void_());
+    b.Append(f->Block(), [&] {
+        auto* loop = b.Loop();
+        b.Append(loop->Body(), [&] { b.NextIteration(loop); });
+        b.Return(f);
+    });
+
+    auto res = ir::Validate(mod);
+    ASSERT_NE(res, Success);
+    EXPECT_EQ(res.Failure().reason.Str(),
+              R"(:5:9 error: next_iteration: must only be called from loop initializer or continuing
+        next_iteration  # -> $B2
+        ^^^^^^^^^^^^^^
+
+:4:7 note: in block
+      $B2: {  # body
+      ^^^
+
+note: # Disassembly
+%my_func = func():void {
+  $B1: {
+    loop [b: $B2] {  # loop_1
+      $B2: {  # body
+        next_iteration  # -> $B2
+      }
+    }
+    ret
+  }
+}
+)");
+}
+
+TEST_F(IR_ValidatorTest, NextIterationInLoopContinuing) {
+    auto* f = b.Function("my_func", ty.void_());
+    b.Append(f->Block(), [&] {
+        auto* loop = b.Loop();
+        b.Append(loop->Body(), [&] { b.ExitLoop(loop); });
+        b.Append(loop->Continuing(), [&] { b.NextIteration(loop); });
+        b.Return(f);
+    });
+
+    auto res = ir::Validate(mod);
+    ASSERT_EQ(res, Success);
+}
+
 TEST_F(IR_ValidatorTest, ContinuingUseValueBeforeContinue) {
     auto* f = b.Function("my_func", ty.void_());
     auto* value = b.Let("value", 1_i);
diff --git a/src/tint/lang/wgsl/writer/ir_to_program/ir_to_program_test.cc b/src/tint/lang/wgsl/writer/ir_to_program/ir_to_program_test.cc
index 5e601cc..7b9b3cd 100644
--- a/src/tint/lang/wgsl/writer/ir_to_program/ir_to_program_test.cc
+++ b/src/tint/lang/wgsl/writer/ir_to_program/ir_to_program_test.cc
@@ -2093,7 +2093,7 @@
                 b.Append(if2->True(), [&] { b.Return(fn, 1_i); });
                 b.Append(if2->False(), [&] { b.Return(fn, 2_i); });
 
-                b.NextIteration(loop);
+                b.Continue(loop);
             });
         });