Prevent eliminating case constructs in block merging (#4976)

Fixes #4918

* Prevent block merging from producing an invalid case construct by
  merging a switch target/default with another construct's merge or
  continue block
* This is to satisfy the structural dominance requirement between the
  switch header and the case constructs

2 files changed, 106 insertions, 0 deletions

diff --git a/source/opt/block_merge_util.cpp b/source/opt/block_merge_util.cpp
index 8ae8020a5..83c702ca3 100644
--- a/source/opt/block_merge_util.cpp
+++ b/source/opt/block_merge_util.cpp
@@ -125,6 +125,26 @@ bool CanMergeWithSuccessor(IRContext* context, BasicBlock* block) {
       return false;
     }
   }
+
+  if (succ_is_merge || IsContinue(context, lab_id)) {
+    auto* struct_cfg = context->GetStructuredCFGAnalysis();
+    auto switch_block_id = struct_cfg->ContainingSwitch(block->id());
+    if (switch_block_id) {
+      auto switch_merge_id = struct_cfg->SwitchMergeBlock(switch_block_id);
+      const auto* switch_inst =
+          &*block->GetParent()->FindBlock(switch_block_id)->tail();
+      for (uint32_t i = 1; i < switch_inst->NumInOperands(); i += 2) {
+        auto target_id = switch_inst->GetSingleWordInOperand(i);
+        if (target_id == block->id() && target_id != switch_merge_id) {
+          // Case constructs must be structurally dominated by the OpSwitch.
+          // Since the successor is the merge/continue for another construct,
+          // merging the blocks would break that requirement.
+          return false;
+        }
+      }
+    }
+  }
+
   return true;
 }
 
diff --git a/test/opt/block_merge_test.cpp b/test/opt/block_merge_test.cpp
index 9698fed23..6129bb273 100644
--- a/test/opt/block_merge_test.cpp
+++ b/test/opt/block_merge_test.cpp
@@ -1201,6 +1201,92 @@ OpFunctionEnd
   SinglePassRunAndMatch<BlockMergePass>(text, true);
 }
 
+TEST_F(BlockMergeTest, DontLoseCaseConstruct) {
+  const std::string text = R"(
+OpCapability Shader
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %func "func"
+OpExecutionMode %func LocalSize 1 1 1
+OpName %entry "entry";
+OpName %loop "loop"
+OpName %loop_merge "loop_merge"
+OpName %loop_cont "loop_cont"
+OpName %switch "switch"
+OpName %switch_merge "switch_merge"
+%void = OpTypeVoid
+%bool = OpTypeBool
+%int = OpTypeInt 32 0
+%void_fn = OpTypeFunction %void
+%bool_undef = OpUndef %bool
+%int_undef = OpUndef %int
+%func = OpFunction %void None %void_fn
+%entry = OpLabel
+OpBranch %loop
+%loop = OpLabel
+OpLoopMerge %loop_merge %loop_cont None
+OpBranch %switch
+%switch = OpLabel
+OpSelectionMerge %switch_merge None
+OpSwitch %int_undef %switch_merge 0 %break 1 %break
+%break = OpLabel
+OpBranch %loop_merge
+%switch_merge = OpLabel
+OpBranch %loop_cont
+%loop_cont = OpLabel
+OpBranch %loop
+%loop_merge = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  auto result = SinglePassRunAndDisassemble<opt::BlockMergePass>(
+      text, /* skip_nop = */ true, /* do_validation = */ true);
+  EXPECT_EQ(opt::Pass::Status::SuccessWithoutChange, std::get<1>(result));
+}
+
+TEST_F(BlockMergeTest, DontLoseDefaultCaseConstruct) {
+  const std::string text = R"(
+OpCapability Shader
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %func "func"
+OpExecutionMode %func LocalSize 1 1 1
+OpName %entry "entry";
+OpName %loop "loop"
+OpName %loop_merge "loop_merge"
+OpName %loop_cont "loop_cont"
+OpName %switch "switch"
+OpName %switch_merge "switch_merge"
+%void = OpTypeVoid
+%bool = OpTypeBool
+%int = OpTypeInt 32 0
+%void_fn = OpTypeFunction %void
+%bool_undef = OpUndef %bool
+%int_undef = OpUndef %int
+%func = OpFunction %void None %void_fn
+%entry = OpLabel
+OpBranch %loop
+%loop = OpLabel
+OpLoopMerge %loop_merge %loop_cont None
+OpBranch %switch
+%switch = OpLabel
+OpSelectionMerge %switch_merge None
+OpSwitch %int_undef %cont 0 %switch_merge 1 %switch_merge
+%cont = OpLabel
+OpBranch %loop_cont
+%switch_merge = OpLabel
+OpBranch %loop_merge
+%loop_cont = OpLabel
+OpBranch %loop
+%loop_merge = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  auto result = SinglePassRunAndDisassemble<opt::BlockMergePass>(
+      text, /* skip_nop = */ true, /* do_validation = */ true);
+  EXPECT_EQ(opt::Pass::Status::SuccessWithoutChange, std::get<1>(result));
+}
+
 // TODO(greg-lunarg): Add tests to verify handling of these cases:
 //
 //    More complex control flow