1 files changed, 19 insertions, 3 deletions

diff --git a/source/opt/cfg.cpp b/source/opt/cfg.cpp
index ac0fcc36..a0248d54 100644
--- a/source/opt/cfg.cpp
+++ b/source/opt/cfg.cpp
@@ -74,6 +74,12 @@ void CFG::RemoveNonExistingEdges(uint32_t blk_id) {
 
 void CFG::ComputeStructuredOrder(Function* func, BasicBlock* root,
                                  std::list<BasicBlock*>* order) {
+  ComputeStructuredOrder(func, root, nullptr, order);
+}
+
+void CFG::ComputeStructuredOrder(Function* func, BasicBlock* root,
+                                 BasicBlock* end,
+                                 std::list<BasicBlock*>* order) {
   assert(module_->context()->get_feature_mgr()->HasCapability(
              SpvCapabilityShader) &&
          "This only works on structured control flow");
@@ -81,7 +87,8 @@ void CFG::ComputeStructuredOrder(Function* func, BasicBlock* root,
   // Compute structured successors and do DFS.
   ComputeStructuredSuccessors(func);
   auto ignore_block = [](cbb_ptr) {};
-  auto ignore_edge = [](cbb_ptr, cbb_ptr) {};
+  auto terminal = [end](cbb_ptr bb) { return bb == end; };
+
   auto get_structured_successors = [this](const BasicBlock* b) {
     return &(block2structured_succs_[b]);
   };
@@ -92,7 +99,7 @@ void CFG::ComputeStructuredOrder(Function* func, BasicBlock* root,
     order->push_front(const_cast<BasicBlock*>(b));
   };
   CFA<BasicBlock>::DepthFirstTraversal(root, get_structured_successors,
-                                       ignore_block, post_order, ignore_edge);
+                                       ignore_block, post_order, terminal);
 }
 
 void CFG::ForEachBlockInPostOrder(BasicBlock* bb,
@@ -205,7 +212,7 @@ BasicBlock* CFG::SplitLoopHeader(BasicBlock* bb) {
   // Find the back edge
   BasicBlock* latch_block = nullptr;
   Function::iterator latch_block_iter = header_it;
-  while (++latch_block_iter != fn->end()) {
+  for (; latch_block_iter != fn->end(); ++latch_block_iter) {
     // If blocks are in the proper order, then the only branch that appears
     // after the header is the latch.
     if (std::find(pred.begin(), pred.end(), latch_block_iter->id()) !=
@@ -237,6 +244,15 @@ BasicBlock* CFG::SplitLoopHeader(BasicBlock* bb) {
     context->set_instr_block(inst, new_header);
   });
 
+  // If |bb| was the latch block, the branch back to the header is not in
+  // |new_header|.
+  if (latch_block == bb) {
+    if (new_header->ContinueBlockId() == bb->id()) {
+      new_header->GetLoopMergeInst()->SetInOperand(1, {new_header_id});
+    }
+    latch_block = new_header;
+  }
+
   // Adjust the OpPhi instructions as needed.
   bb->ForEachPhiInst([latch_block, bb, new_header, context](Instruction* phi) {
     std::vector<uint32_t> preheader_phi_ops;