Fix structured exit validation

Fixes #3139

* If the header of the construct is also a merge block, jump to the
associated header instead of the immediate dominator
* new test

source/val/construct.cpp

@@ -171,7 +171,18 @@ bool Construct::IsStructuredExit(ValidationState_t& _, BasicBlock* dest) const {
 
     bool seen_switch = false;
     auto header = entry_block();
+    //  The first block in the traversal is either:
+    //  i.  The header block that declares |header| as its merge block.
+    //  ii. The immediate dominator of |header|.
     auto block = header->immediate_dominator();
+    for (auto& use : header->label()->uses()) {
+      if ((use.first->opcode() == SpvOpLoopMerge ||
+           use.first->opcode() == SpvOpSelectionMerge) &&
+          use.second == 1) {
+        block = use.first->block();
+        break;
+      }
+    }
     while (block) {
       auto terminator = block->terminator();
       auto index = terminator - &_.ordered_instructions()[0];
@@ -197,7 +208,9 @@ bool Construct::IsStructuredExit(ValidationState_t& _, BasicBlock* dest) const {
           }
         }
 
-        if (terminator->opcode() == SpvOpSwitch) seen_switch = true;
+        if (terminator->opcode() == SpvOpSwitch) {
+          seen_switch = true;
+        }
 
         // Hit an enclosing loop and didn't break or continue.
         if (merge_inst->opcode() == SpvOpLoopMerge) return false;

source/val/validate_cfg.cpp

@@ -728,10 +728,10 @@ spv_result_t StructuredControlFlowChecks(
     }
 
     Construct::ConstructBlockSet construct_blocks = construct.blocks(function);
+    std::string construct_name, header_name, exit_name;
+    std::tie(construct_name, header_name, exit_name) =
+        ConstructNames(construct.type());
     for (auto block : construct_blocks) {
-      std::string construct_name, header_name, exit_name;
-      std::tie(construct_name, header_name, exit_name) =
-          ConstructNames(construct.type());
       // Check that all exits from the construct are via structured exits.
       for (auto succ : *block->successors()) {
         if (block->reachable() && !construct_blocks.count(succ) &&

test/val/val_cfg_test.cpp

@@ -4187,6 +4187,57 @@ OpFunctionEnd
           "1[%loop], but its merge block 2[%continue] is not"));
 }
 
+TEST_F(ValidateCFG, ExitFromConstructWhoseHeaderIsAMerge) {
+  const std::string text = R"(
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+%void = OpTypeVoid
+%2 = OpTypeFunction %void
+%int = OpTypeInt 32 1
+%4 = OpUndef %int
+%bool = OpTypeBool
+%6 = OpUndef %bool
+%7 = OpFunction %void None %2
+%8 = OpLabel
+OpSelectionMerge %9 None
+OpSwitch %4 %10 0 %11
+%10 = OpLabel
+OpBranch %9
+%11 = OpLabel
+OpBranch %12
+%12 = OpLabel
+OpLoopMerge %13 %14 None
+OpBranch %15
+%15 = OpLabel
+OpSelectionMerge %16 None
+OpSwitch %4 %17 1 %18 2 %19
+%17 = OpLabel
+OpBranch %16
+%18 = OpLabel
+OpBranch %14
+%19 = OpLabel
+OpBranch %16
+%16 = OpLabel
+OpBranch %14
+%14 = OpLabel
+OpBranchConditional %6 %12 %13
+%13 = OpLabel
+OpSelectionMerge %20 None
+OpBranchConditional %6 %21 %20
+%21 = OpLabel
+OpBranch %9
+%20 = OpLabel
+OpBranch %10
+%9 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  CompileSuccessfully(text);
+  EXPECT_EQ(SPV_SUCCESS, ValidateInstructions());
+}
+
 }  // namespace
 }  // namespace val
 }  // namespace spvtools