[CodeExtractor] Fix iterator invalidation in findOrCreateBlockForHoisting.

[fhahn]

Summary

By replacing branches to CommonExitBlock, we remove the node from
CommonExitBlock's predecessors, invalidating the iterator. The problem
is exposed when the common exit block has multiple predecessors and
needs to sink lifetime info. The modification in the test case trigger
the issue.

[fhahn]

ping

[davidxl]

lgtm

lib/Transforms/Utils/CodeExtractor.cpp

312

Move ++ into the for clause?

[fhahn]

lib/Transforms/Utils/CodeExtractor.cpp

312

I think we have to increment the iterator at the beginning of the loop, as replaceUsesOfWith may remove the element PI references and then the iterator would become invalid and Transforms/CodeExtractor/live_shrink_hoist.ll crashes.

[fhahn]

Thank you very much for having a look :)

lib/Transforms/Utils/CodeExtractor.cpp

#ifndef NDEBUG
   // If there are any phi nodes, the single pred either exists or has already
   // be created before code extraction.
   assert(!getFirstPHI(CommonExitBlock) && "Phi not expected");
 #endif
 
   BasicBlock *NewExitBlock = CommonExitBlock->splitBasicBlock(
       CommonExitBlock->getFirstNonPHI()->getIterator());
 
-  for (auto *Pred : predecessors(CommonExitBlock)) {
+  for (auto PI = pred_begin(CommonExitBlock), PE = pred_end(CommonExitBlock);
+       PI != PE;) {
+    BasicBlock *Pred = *PI++;
     if (Blocks.count(Pred))
       continue;
     Pred->getTerminator()->replaceUsesOfWith(CommonExitBlock, NewExitBlock);
   }
   // Now add the old exit block to the outline region.
   Blocks.insert(CommonExitBlock);
   return CommonExitBlock;
 }

test/Transforms/CodeExtractor/live_shrink_hoist.ll

-; RUN: opt -S -partial-inliner -max-num-inline-blocks=2 -skip-partial-inlining-cost-analysis  < %s |   FileCheck %s
+; RUN: opt -S -partial-inliner -max-num-inline-blocks=3 -skip-partial-inlining-cost-analysis  < %s |   FileCheck %s
 ; RUN: opt -S -passes=partial-inliner -max-num-inline-blocks=2  -skip-partial-inlining-cost-analysis < %s   | FileCheck %s
 
 %class.A = type { i32 }
 
 @cond = local_unnamed_addr global i32 0, align 4
 
 ; Function Attrs: uwtable
 define void @_Z3foov() local_unnamed_addr  {
 bb:
   %tmp = alloca %class.A, align 4
   %tmp1 = bitcast %class.A* %tmp to i8*
   call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %tmp1) 
   %tmp2 = load i32, i32* @cond, align 4, !tbaa !2
   %tmp3 = icmp eq i32 %tmp2, 0
   br i1 %tmp3, label %bb4, label %bb9
 
 bb4:                                              ; preds = %bb
+  %foo = icmp eq i32 %tmp2, 0
+  br i1 %foo, label %bb5, label %bb9
+
+bb5:                                              ; preds = %bb4
   call void @_ZN1A7memfuncEv(%class.A* nonnull %tmp)
   %tmp5 = getelementptr inbounds %class.A, %class.A* %tmp, i64 0, i32 0
   %tmp6 = load i32, i32* %tmp5, align 4, !tbaa !6
   %tmp7 = icmp sgt i32 %tmp6, 0
   br i1 %tmp7, label %bb9, label %bb8
 
 bb8:                                              ; preds = %bb4
   call void @_ZN1A7memfuncEv(%class.A* nonnull %tmp)