Don't use invalidated iterators in FlattenCFGPass

[kuhar]

Summary

FlattenCFG may erase unnecessary blocks, which also invalidates iterators to those erased blocks.
Before this patch, iterativelyFlattenCFG could try to increment a BB iterator after that BB has been removed and crash.

This patch makes FlattenCFGPass use WeakVH to skip over erased blocks.

[asbirlea]

lgtm.

llvm/trunk/lib/Transforms/Scalar/FlattenCFGPass.cpp

 //===- FlattenCFGPass.cpp - CFG Flatten Pass ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements flattening of CFG.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/AliasAnalysis.h"
-#include "llvm/Transforms/Utils/Local.h"
 #include "llvm/IR/CFG.h"
+#include "llvm/IR/ValueHandle.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Utils/Local.h"
+
 using namespace llvm;
 
 #define DEBUG_TYPE "flattencfg"
 
 namespace {
 struct FlattenCFGPass : public FunctionPass {
   static char ID; // Pass identification, replacement for typeid
 public:

 // Public interface to the FlattenCFG pass
 FunctionPass *llvm::createFlattenCFGPass() { return new FlattenCFGPass(); }
 
 /// iterativelyFlattenCFG - Call FlattenCFG on all the blocks in the function,
 /// iterating until no more changes are made.
 static bool iterativelyFlattenCFG(Function &F, AliasAnalysis *AA) {
   bool Changed = false;
   bool LocalChange = true;
+
+  // Use block handles instead of iterating over function blocks directly
+  // to avoid using iterators invalidated by erasing blocks.
+  std::vector<WeakVH> Blocks;
+  Blocks.reserve(F.size());
+  for (auto &BB : F)
+    Blocks.push_back(&BB);
+
   while (LocalChange) {
     LocalChange = false;
 
-    // Loop over all of the basic blocks and remove them if they are unneeded...
-    //
-    for (Function::iterator BBIt = F.begin(); BBIt != F.end();) {
-      if (FlattenCFG(&*BBIt++, AA)) {
+    // Loop over all of the basic blocks and try to flatten them.
+    for (WeakVH &BlockHandle : Blocks) {
+      // Skip blocks erased by FlattenCFG.
+      if (auto *BB = cast_or_null<BasicBlock>(BlockHandle))
+        if (FlattenCFG(BB, AA))
           LocalChange = true;
-      }
-    }
     Changed |= LocalChange;
   }
   return Changed;
 }
 
 bool FlattenCFGPass::runOnFunction(Function &F) {
   AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
   bool EverChanged = false;
   // iterativelyFlattenCFG can make some blocks dead.
   while (iterativelyFlattenCFG(F, AA)) {
     removeUnreachableBlocks(F);
     EverChanged = true;
   }
   return EverChanged;
 }

llvm/trunk/test/Transforms/Util/flattencfg.ll

b0:                                ; preds = %entry
   br i1 %cond1, label %exit, label %b1
 
 b1:                                       ; preds = %entry, %b0
   br label %exit
 
 exit:                               ; preds = %entry, %b0, %b1
   ret void
 }
+
+; CHECK-LABEL: @test_not_crash2
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    %0 = fcmp ult float %a
+; CHECK-NEXT:    %1 = fcmp ult float %b
+; CHECK-NEXT:    [[COND:%[a-z0-9]+]] = or i1 %0, %1
+; CHECK-NEXT:    br i1 [[COND]], label %bb4, label %bb3
+; CHECK:       bb3:
+; CHECK-NEXT:    br label %bb4
+; CHECK:       bb4:
+; CHECK-NEXT:    ret void
+define void @test_not_crash2(float %a, float %b) #0 {
+entry:
+  %0 = fcmp ult float %a, 1.000000e+00
+  br i1 %0, label %bb0, label %bb1
+
+bb3:                                               ; preds = %bb0
+  br label %bb4
+
+bb4:                                               ; preds = %bb0, %bb3
+  ret void
+
+bb1:                                               ; preds = %entry
+  br label %bb0
+
+bb0:                                               ; preds = %bb1, %entry
+  %1 = fcmp ult float %b, 1.000000e+00
+  br i1 %1, label %bb4, label %bb3
+}
+