Update OpPhi instructions after splitting block. (#1783)

In the merge return pass, we will split a block, but not update the phi
instructions that reference the block.  Since the branch in the original
block is now part of the block with the new id, the phi nodes must be
updated.

This commit will change this.

I have also considered other places where an id of a basic block could
be referenced, and I don't think any of them need to change.

1) Branch and merge instructions: These jump to the start of the
original block, and so we want them to jump to the block that uses the
original id.  Nothing needs to change.

2) Names and decorations: I don't think it matters with block keeps the
name, and there are no decorations that apply to basic blocks.

Fixes #1736.

3 files changed, 68 insertions, 4 deletions

diff --git a/source/opt/basic_block.cpp b/source/opt/basic_block.cpp
index 5005061b..9d101ba2 100644
--- a/source/opt/basic_block.cpp
+++ b/source/opt/basic_block.cpp
@@ -214,7 +214,30 @@ BasicBlock* BasicBlock::SplitBasicBlock(IRContext* context, uint32_t label_id,
   new_block->insts_.Splice(new_block->end(), &insts_, iter, end());
   new_block->SetParent(GetParent());
 
+  context->AnalyzeDefUse(new_block->GetLabelInst());
+
+  // Update the phi nodes in the successor blocks to reference the new block id.
+  const_cast<const BasicBlock*>(new_block)->ForEachSuccessorLabel(
+      [new_block, this, context](const uint32_t label) {
+        BasicBlock* target_bb = context->get_instr_block(label);
+        target_bb->ForEachPhiInst(
+            [this, new_block, context](Instruction* phi_inst) {
+              bool changed = false;
+              for (uint32_t i = 1; i < phi_inst->NumInOperands(); i += 2) {
+                if (phi_inst->GetSingleWordInOperand(i) == this->id()) {
+                  changed = true;
+                  phi_inst->SetInOperand(i, {new_block->id()});
+                }
+              }
+
+              if (changed) {
+                context->UpdateDefUse(phi_inst);
+              }
+            });
+      });
+
   if (context->AreAnalysesValid(IRContext::kAnalysisInstrToBlockMapping)) {
+    context->set_instr_block(new_block->GetLabelInst(), new_block);
     new_block->ForEachInst([new_block, context](Instruction* inst) {
       context->set_instr_block(inst, new_block);
     });
diff --git a/source/opt/merge_return_pass.cpp b/source/opt/merge_return_pass.cpp
index d84c1aed..82ded202 100644
--- a/source/opt/merge_return_pass.cpp
+++ b/source/opt/merge_return_pass.cpp
@@ -321,7 +321,7 @@ void MergeReturnPass::PredicateBlock(
   // If this block is a header block, predicate the entire structured subgraph.
   // This can act recursively.
 
-  // If |block| is a loop head, then the back edge must jump to the original
+  // If |block| is a loop header, then the back edge must jump to the original
   // code, not the new header.
   if (block->GetLoopMergeInst()) {
     cfg()->SplitLoopHeader(block);
@@ -364,9 +364,7 @@ void MergeReturnPass::PredicateBlock(
   predicated->insert(new_merge);
   new_merge->SetParent(function_);
 
-  // Register the new labels.
-  get_def_use_mgr()->AnalyzeInstDef(old_body->GetLabelInst());
-  context()->set_instr_block(old_body->GetLabelInst(), old_body);
+  // Register the new label.
   get_def_use_mgr()->AnalyzeInstDef(new_merge->GetLabelInst());
   context()->set_instr_block(new_merge->GetLabelInst(), new_merge);
 
diff --git a/test/opt/pass_merge_return_test.cpp b/test/opt/pass_merge_return_test.cpp
index 9da13ebe..52411eb2 100644
--- a/test/opt/pass_merge_return_test.cpp
+++ b/test/opt/pass_merge_return_test.cpp
@@ -443,6 +443,49 @@ OpFunctionEnd
 
   SinglePassRunAndMatch<MergeReturnPass>(before, false);
 }
+
+TEST_F(MergeReturnPassTest, SplitBlockUsedInPhi) {
+  const std::string before =
+      R"(
+; CHECK: OpFunction
+; CHECK-NEXT: OpLabel
+; CHECK: OpSelectionMerge [[merge1:%\w+]] None
+; CHECK: [[merge1]] = OpLabel
+; CHECK: OpBranchConditional %{{\w+}} %{{\w+}} [[old_merge:%\w+]]
+; CHECK: [[old_merge]] = OpLabel
+; CHECK-NEXT: OpSelectionMerge [[merge2:%\w+]]
+; CHECK-NEXT: OpBranchConditional %false [[side_node:%\w+]] [[merge2]]
+; CHECK: [[merge2]] = OpLabel
+; CHECK-NEXT: OpPhi %bool %false [[old_merge]] %true [[side_node]]
+               OpCapability Addresses
+               OpCapability Shader
+               OpCapability Linkage
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %1 "simple_shader"
+       %void = OpTypeVoid
+       %bool = OpTypeBool
+      %false = OpConstantFalse %bool
+       %true = OpConstantTrue %bool
+          %6 = OpTypeFunction %void
+          %1 = OpFunction %void None %6
+          %7 = OpLabel
+               OpSelectionMerge %8 None
+               OpBranchConditional %false %9 %8
+          %9 = OpLabel
+               OpReturn
+          %8 = OpLabel
+               OpSelectionMerge %10 None
+               OpBranchConditional %false %11 %10
+         %11 = OpLabel
+               OpBranch %10
+         %10 = OpLabel
+         %12 = OpPhi %bool %false %8 %true %11
+               OpReturn
+               OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<MergeReturnPass>(before, false);
+}
 #endif
 
 TEST_F(MergeReturnPassTest, StructuredControlFlowBothMergeAndHeader) {