Don't change decorations and names in merge return. (#1777)

When creating a new phi for a value in the function, merge return will
rewrite all uses of an id that are no longer dominated by its
definition.  Uses that are not in a basic block, like OpName or
decorations, are not dominated, but they should not be replaced.

Fixes #1736.

2 files changed, 109 insertions, 9 deletions

diff --git a/source/opt/merge_return_pass.cpp b/source/opt/merge_return_pass.cpp
index 6d955527..d84c1aed 100644
--- a/source/opt/merge_return_pass.cpp
+++ b/source/opt/merge_return_pass.cpp
@@ -198,7 +198,11 @@ void MergeReturnPass::CreatePhiNodesForInst(BasicBlock* merge_block,
     std::vector<Instruction*> users_to_update;
     context()->get_def_use_mgr()->ForEachUser(
         &inst, [&users_to_update, &dom_tree, inst_bb, this](Instruction* user) {
-          if (!dom_tree->Dominates(inst_bb, context()->get_instr_block(user))) {
+          BasicBlock* user_bb = context()->get_instr_block(user);
+          // If |user_bb| is nullptr, then |user| is not in the function.  It is
+          // something like an OpName or decoration, which should not be
+          // replaced with the result of the OpPhi.
+          if (user_bb && !dom_tree->Dominates(inst_bb, user_bb)) {
             users_to_update.push_back(user);
           }
         });
diff --git a/test/opt/pass_merge_return_test.cpp b/test/opt/pass_merge_return_test.cpp
index f597d6c9..9da13ebe 100644
--- a/test/opt/pass_merge_return_test.cpp
+++ b/test/opt/pass_merge_return_test.cpp
@@ -268,11 +268,11 @@ TEST_F(MergeReturnPassTest, StructuredControlFlowWithUnreachableMerge) {
 ; CHECK: OpSelectionMerge [[merge_lab:%\w+]]
 ; CHECK: OpBranchConditional [[cond:%\w+]] [[if_lab:%\w+]] [[then_lab:%\w+]]
 ; CHECK: [[if_lab]] = OpLabel
-; CHECK-Next: OpStore [[var]] [[true]]
-; CHECK-Next: OpBranch
+; CHECK-NEXT: OpStore [[var]] [[true]]
+; CHECK-NEXT: OpBranch
 ; CHECK: [[then_lab]] = OpLabel
-; CHECK-Next: OpStore [[var]] [[true]]
-; CHECK-Next: OpBranch [[merge_lab]]
+; CHECK-NEXT: OpStore [[var]] [[true]]
+; CHECK-NEXT: OpBranch [[merge_lab]]
 ; CHECK: OpReturn
 OpCapability Addresses
 OpCapability Shader
@@ -310,11 +310,10 @@ TEST_F(MergeReturnPassTest, StructuredControlFlowAddPhi) {
 ; CHECK: OpBranchConditional [[cond:%\w+]] [[if_lab:%\w+]] [[then_lab:%\w+]]
 ; CHECK: [[if_lab]] = OpLabel
 ; CHECK-NEXT: [[add:%\w+]] = OpIAdd [[type:%\w+]]
-; CHECK-Next: OpStore [[var]] [[true]]
-; CHECK-Next: OpBranch
+; CHECK-NEXT: OpBranch
 ; CHECK: [[then_lab]] = OpLabel
-; CHECK-Next: OpStore [[var]] [[true]]
-; CHECK-Next: OpBranch [[merge_lab]]
+; CHECK-NEXT: OpStore [[var]] [[true]]
+; CHECK-NEXT: OpBranch [[merge_lab]]
 ; CHECK: [[merge_lab]] = OpLabel
 ; CHECK-NEXT: [[phi:%\w+]] = OpPhi [[type]] [[add]] [[if_lab]] [[undef:%\w+]] [[then_lab]]
 ; CHECK: OpIAdd [[type]] [[phi]] [[phi]]
@@ -347,6 +346,103 @@ OpFunctionEnd
 
   SinglePassRunAndMatch<MergeReturnPass>(before, false);
 }
+
+TEST_F(MergeReturnPassTest, StructuredControlDecoration) {
+  const std::string before =
+      R"(
+; CHECK: OpDecorate [[dec_id:%\w+]] RelaxedPrecision
+; CHECK: [[false:%\w+]] = OpConstantFalse
+; CHECK: [[true:%\w+]] = OpConstantTrue
+; CHECK: OpFunction
+; CHECK: [[var:%\w+]] = OpVariable [[:%\w+]] Function [[false]]
+; CHECK: OpSelectionMerge [[merge_lab:%\w+]]
+; CHECK: OpBranchConditional [[cond:%\w+]] [[if_lab:%\w+]] [[then_lab:%\w+]]
+; CHECK: [[if_lab]] = OpLabel
+; CHECK-NEXT: [[dec_id]] = OpIAdd [[type:%\w+]]
+; CHECK-NEXT: OpBranch
+; CHECK: [[then_lab]] = OpLabel
+; CHECK-NEXT: OpStore [[var]] [[true]]
+; CHECK-NEXT: OpBranch [[merge_lab]]
+; CHECK: [[merge_lab]] = OpLabel
+; CHECK: OpReturn
+OpCapability Addresses
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %6 "simple_shader"
+OpDecorate %11 RelaxedPrecision
+%2 = OpTypeVoid
+%3 = OpTypeBool
+%int = OpTypeInt 32 0
+%int_0 = OpConstant %int 0
+%4 = OpConstantFalse %3
+%1 = OpTypeFunction %2
+%6 = OpFunction %2 None %1
+%7 = OpLabel
+OpSelectionMerge %10 None
+OpBranchConditional %4 %8 %9
+%8 = OpLabel
+%11 = OpIAdd %int %int_0 %int_0
+OpBranch %10
+%9 = OpLabel
+OpReturn
+%10 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<MergeReturnPass>(before, false);
+}
+
+TEST_F(MergeReturnPassTest, StructuredControlDecoration2) {
+  const std::string before =
+      R"(
+; CHECK: OpDecorate [[dec_id:%\w+]] RelaxedPrecision
+; CHECK: [[false:%\w+]] = OpConstantFalse
+; CHECK: [[true:%\w+]] = OpConstantTrue
+; CHECK: OpFunction
+; CHECK: [[var:%\w+]] = OpVariable [[:%\w+]] Function [[false]]
+; CHECK: OpSelectionMerge [[merge_lab:%\w+]]
+; CHECK: OpBranchConditional [[cond:%\w+]] [[if_lab:%\w+]] [[then_lab:%\w+]]
+; CHECK: [[if_lab]] = OpLabel
+; CHECK-NEXT: [[dec_id]] = OpIAdd [[type:%\w+]]
+; CHECK-NEXT: OpBranch
+; CHECK: [[then_lab]] = OpLabel
+; CHECK-NEXT: OpStore [[var]] [[true]]
+; CHECK-NEXT: OpBranch [[merge_lab]]
+; CHECK: [[merge_lab]] = OpLabel
+; CHECK-NEXT: [[phi:%\w+]] = OpPhi [[type]] [[dec_id]] [[if_lab]] [[undef:%\w+]] [[then_lab]]
+; CHECK: OpIAdd [[type]] [[phi]] [[phi]]
+; CHECK: OpReturn
+OpCapability Addresses
+OpCapability Shader
+OpCapability Linkage
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %6 "simple_shader"
+OpDecorate %11 RelaxedPrecision
+%2 = OpTypeVoid
+%3 = OpTypeBool
+%int = OpTypeInt 32 0
+%int_0 = OpConstant %int 0
+%4 = OpConstantFalse %3
+%1 = OpTypeFunction %2
+%6 = OpFunction %2 None %1
+%7 = OpLabel
+OpSelectionMerge %10 None
+OpBranchConditional %4 %8 %9
+%8 = OpLabel
+%11 = OpIAdd %int %int_0 %int_0
+OpBranch %10
+%9 = OpLabel
+OpReturn
+%10 = OpLabel
+%12 = OpIAdd %int %11 %11
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<MergeReturnPass>(before, false);
+}
 #endif
 
 TEST_F(MergeReturnPassTest, StructuredControlFlowBothMergeAndHeader) {