[tint][ir][val] Restrict types in root allowed by kAllowOverrides

Restricts the instruction types that are allowed in the root block
when kAllowOverrides is set. This doesn't do the walk up to make sure
that the specific values/calls are allowed. This supersedes a number
of previous patches that were playing whack-a-mole.

Fixes: 382453612
Fixes: 382422295
Change-Id: Iccd2145afb29bf61b2dc7ccbfd97b08648e39208
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/218354
Commit-Queue: Ryan Harrison <[email protected]>
Auto-Submit: Ryan Harrison <[email protected]>
Reviewed-by: James Price <[email protected]>
Commit-Queue: James Price <[email protected]>

src/tint/lang/core/ir/validator.cc

@@ -1884,20 +1884,25 @@ void Validator::CheckRootBlock(const Block* blk) {
             [&](const core::ir::Construct* c) {
                 if (capabilities_.Contains(Capability::kAllowModuleScopeLets)) {
                     CheckInstruction(c);
+                    CheckOnlyUsedInRootBlock(inst);
                 } else {
                     AddError(inst) << "root block: invalid instruction: " << inst->TypeInfo().name;
                 }
             },
-            [&](const core::ir::Discard* d) {
-                AddError(d) << "root block: invalid instruction: " << d->TypeInfo().name;
-            },
             [&](Default) {
-                // Note, this validation is looser then it could be. There are only certain
-                // expressions and builtins which can be used in an override. We can tighten this up
-                // later to a constrained set of instructions and builtins if necessary.
-                if (capabilities_.Contains(Capability::kAllowOverrides)) {
-                    // If overrides are allowed we can have regular instructions in the root block,
-                    // with the caveat that those instructions can _only_ be used in the root block.
+                // Note, this validation around kAllowOverrides is looser than it could be. There
+                // are only certain expressions and builtins which can be used in an override, but
+                // the current checks are only doing type level checking. To tighten this up will
+                // require walking up the tree to make sure that operands are const/override and
+                // builtins are allowed.
+                if (capabilities_.Contains(Capability::kAllowOverrides) &&
+                    inst->IsAnyOf<core::ir::Unary, core::ir::Binary, core::ir::CoreBuiltinCall,
+                                  core::ir::Convert, core::ir::Swizzle, core::ir::Access,
+                                  core::ir::Bitcast>()) {
+                    CheckInstruction(inst);
+                    // If overrides are allowed we can have certain regular instructions in the root
+                    // block, with the caveat that those instructions can _only_ be used in the root
+                    // block.
                     CheckOnlyUsedInRootBlock(inst);
                 } else {
                     AddError(inst) << "root block: invalid instruction: " << inst->TypeInfo().name;

src/tint/lang/core/ir/validator_test.cc

@@ -3849,9 +3849,8 @@ note: # Disassembly
 )");
 }
 
-TEST_F(IR_ValidatorTest, Discard_OutsideFunction) {
-    auto* d = b.Discard();
-    mod.root_block->Append(d);
+TEST_F(IR_ValidatorTest, Discard_RootBlock) {
+    mod.root_block->Append(b.Discard());
 
     auto res = ir::Validate(mod);
     ASSERT_NE(res, Success);
@@ -3872,6 +3871,55 @@ note: # Disassembly
 )");
 }
 
+TEST_F(IR_ValidatorTest, Terminator_RootBlock) {
+    auto f = b.Function("f", ty.void_());
+    b.Append(f->Block(), [&] { b.Unreachable(); });
+
+    mod.root_block->Append(b.Return(f));
+    mod.root_block->Append(b.Unreachable());
+    mod.root_block->Append(b.TerminateInvocation());
+    auto res = ir::Validate(mod);
+    ASSERT_NE(res, Success);
+    EXPECT_EQ(res.Failure().reason.Str(),
+              R"(:2:3 error: return: root block: invalid instruction: tint::core::ir::Return
+  ret
+  ^^^
+
+:1:1 note: in block
+$B1: {  # root
+^^^
+
+:3:3 error: unreachable: root block: invalid instruction: tint::core::ir::Unreachable
+  unreachable
+  ^^^^^^^^^^^
+
+:1:1 note: in block
+$B1: {  # root
+^^^
+
+:4:3 error: terminate_invocation: root block: invalid instruction: tint::core::ir::TerminateInvocation
+  terminate_invocation
+  ^^^^^^^^^^^^^^^^^^^^
+
+:1:1 note: in block
+$B1: {  # root
+^^^
+
+note: # Disassembly
+$B1: {  # root
+  ret
+  unreachable
+  terminate_invocation
+}
+
+%f = func():void {
+  $B2: {
+    unreachable
+  }
+}
+)");
+}
+
 TEST_F(IR_ValidatorTest, Terminator_HasResult) {
     auto* ret_func = b.Function("ret_func", ty.void_());
     b.Append(ret_func->Block(), [&] {
@@ -4883,6 +4931,34 @@ note: # Disassembly
 )");
 }
 
+TEST_F(IR_ValidatorTest, If_RootBlock) {
+    auto* if_ = b.If(true);
+    if_->True()->Append(b.Unreachable());
+    mod.root_block->Append(if_);
+
+    auto res = ir::Validate(mod, core::ir::Capabilities{core::ir::Capability::kAllowOverrides});
+    ASSERT_NE(res, Success);
+    EXPECT_EQ(res.Failure().reason.Str(),
+              R"(:2:3 error: if: root block: invalid instruction: tint::core::ir::If
+  if true [t: $B2] {  # if_1
+  ^^^^^^^^^^^^^^^^
+
+:1:1 note: in block
+$B1: {  # root
+^^^
+
+note: # Disassembly
+$B1: {  # root
+  if true [t: $B2] {  # if_1
+    $B2: {  # true
+      unreachable
+    }
+  }
+}
+
+)");
+}
+
 TEST_F(IR_ValidatorTest, If_EmptyFalse) {
     auto* f = b.Function("my_func", ty.void_());
 
@@ -5040,6 +5116,34 @@ note: # Disassembly
 )");
 }
 
+TEST_F(IR_ValidatorTest, Loop_RootBlock) {
+    auto* l = b.Loop();
+    l->Body()->Append(b.ExitLoop(l));
+    mod.root_block->Append(l);
+
+    auto res = ir::Validate(mod, core::ir::Capabilities{core::ir::Capability::kAllowOverrides});
+    ASSERT_NE(res, Success);
+    EXPECT_EQ(res.Failure().reason.Str(),
+              R"(:2:3 error: loop: root block: invalid instruction: tint::core::ir::Loop
+  loop [b: $B2] {  # loop_1
+  ^^^^^^^^^^^^^
+
+:1:1 note: in block
+$B1: {  # root
+^^^
+
+note: # Disassembly
+$B1: {  # root
+  loop [b: $B2] {  # loop_1
+    $B2: {  # body
+      exit_loop  # loop_1
+    }
+  }
+}
+
+)");
+}
+
 TEST_F(IR_ValidatorTest, Loop_OnlyBody) {
     auto* f = b.Function("my_func", ty.void_());
 
@@ -5080,6 +5184,35 @@ note: # Disassembly
 )");
 }
 
+TEST_F(IR_ValidatorTest, Switch_RootBlock) {
+    auto* switch_ = b.Switch(1_i);
+    auto* def = b.DefaultCase(switch_);
+    def->Append(b.ExitSwitch(switch_));
+    mod.root_block->Append(switch_);
+
+    auto res = ir::Validate(mod, core::ir::Capabilities{core::ir::Capability::kAllowOverrides});
+    ASSERT_NE(res, Success);
+    EXPECT_EQ(res.Failure().reason.Str(),
+              R"(:2:3 error: switch: root block: invalid instruction: tint::core::ir::Switch
+  switch 1i [c: (default, $B2)] {  # switch_1
+  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+:1:1 note: in block
+$B1: {  # root
+^^^
+
+note: # Disassembly
+$B1: {  # root
+  switch 1i [c: (default, $B2)] {  # switch_1
+    $B2: {  # case
+      exit_switch  # switch_1
+    }
+  }
+}
+
+)");
+}
+
 TEST_F(IR_ValidatorTest, Type_VectorElements) {
     auto* f = b.Function("my_func", ty.void_());