[CodeExtractor] Store outputs at the first valid insertion point

When CodeExtractor outlines values which are used by the original
function, it must store those values in some in-out parameter. This
store instruction must not be inserted in between a PHI and an EH pad
instruction, as that results in invalid IR.

This fixes the following verifier failure seen while outlining within
ObjC methods with live exit values:

  The unwind destination does not have an exception handling instruction!
    %call35 = invoke i8* bitcast (i8* (i8*, i8*, ...)* @objc_msgSend to i8* (i8*, i8*)*)(i8* %exn.adjusted, i8* %1)
            to label %invoke.cont34 unwind label %lpad33, !dbg !4183
  The unwind destination does not have an exception handling instruction!
    invoke void @objc_exception_throw(i8* %call35) #12
            to label %invoke.cont36 unwind label %lpad33, !dbg !4184
  LandingPadInst not the first non-PHI instruction in the block.
    %3 = landingpad { i8*, i32 }
            catch i8* null, !dbg !1411

rdar://46540815

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@348562 91177308-0d34-0410-b5e6-96231b3b80d8

Author: vedantk

lib/Transforms/Utils/CodeExtractor.cpp

@@ -992,17 +992,15 @@ emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
       continue;
 
     // Find proper insertion point.
-    Instruction *InsertPt;
+    BasicBlock::iterator InsertPt;
     // In case OutI is an invoke, we insert the store at the beginning in the
     // 'normal destination' BB. Otherwise we insert the store right after OutI.
     if (auto *InvokeI = dyn_cast<InvokeInst>(OutI))
-      InsertPt = InvokeI->getNormalDest()->getFirstNonPHI();
+      InsertPt = InvokeI->getNormalDest()->getFirstInsertionPt();
+    else if (auto *Phi = dyn_cast<PHINode>(OutI))
+      InsertPt = Phi->getParent()->getFirstInsertionPt();
     else
-      InsertPt = OutI->getNextNode();
-
-    // Let's assume that there is no other guy interleave non-PHI in PHIs.
-    if (isa<PHINode>(InsertPt))
-      InsertPt = InsertPt->getParent()->getFirstNonPHI();
+      InsertPt = std::next(OutI->getIterator());
 
     assert(OAI != newFunction->arg_end() &&
            "Number of output arguments should match "
@@ -1012,13 +1010,13 @@ emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
       Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
       Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), FirstOut + i);
       GetElementPtrInst *GEP = GetElementPtrInst::Create(
-          StructArgTy, &*OAI, Idx, "gep_" + outputs[i]->getName(), InsertPt);
-      new StoreInst(outputs[i], GEP, InsertPt);
+          StructArgTy, &*OAI, Idx, "gep_" + outputs[i]->getName(), &*InsertPt);
+      new StoreInst(outputs[i], GEP, &*InsertPt);
       // Since there should be only one struct argument aggregating
       // all the output values, we shouldn't increment OAI, which always
       // points to the struct argument, in this case.
     } else {
-      new StoreInst(outputs[i], &*OAI, InsertPt);
+      new StoreInst(outputs[i], &*OAI, &*InsertPt);
       ++OAI;
     }
   }
@@ -1379,8 +1377,10 @@ Function *CodeExtractor::extractCodeRegion() {
   if (doesNotReturn)
     newFunction->setDoesNotReturn();
 
-  LLVM_DEBUG(if (verifyFunction(*newFunction))
-             report_fatal_error("verification of newFunction failed!"));
+  LLVM_DEBUG(if (verifyFunction(*newFunction, &errs())) {
+    newFunction->dump();
+    report_fatal_error("verification of newFunction failed!");
+  });
   LLVM_DEBUG(if (verifyFunction(*oldFunction))
              report_fatal_error("verification of oldFunction failed!"));
   return newFunction;

unittests/Transforms/Utils/CodeExtractorTest.cpp

@@ -127,4 +127,72 @@ TEST(CodeExtractor, ExitPHIOnePredFromRegion) {
   EXPECT_FALSE(verifyFunction(*Outlined));
   EXPECT_FALSE(verifyFunction(*Func));
 }
+
+TEST(CodeExtractor, StoreOutputInvokeResultAfterEHPad) {
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M(parseAssemblyString(R"invalid(
+    declare i8 @hoge()
+
+    define i32 @foo() personality i8* null {
+      entry:
+        %call = invoke i8 @hoge()
+                to label %invoke.cont unwind label %lpad
+
+      invoke.cont:                                      ; preds = %entry
+        unreachable
+
+      lpad:                                             ; preds = %entry
+        %0 = landingpad { i8*, i32 }
+                catch i8* null
+        br i1 undef, label %catch, label %finally.catchall
+
+      catch:                                            ; preds = %lpad
+        %call2 = invoke i8 @hoge()
+                to label %invoke.cont2 unwind label %lpad2
+
+      invoke.cont2:                                    ; preds = %catch
+        %call3 = invoke i8 @hoge()
+                to label %invoke.cont3 unwind label %lpad2
+
+      invoke.cont3:                                    ; preds = %invoke.cont2
+        unreachable
+
+      lpad2:                                           ; preds = %invoke.cont2, %catch
+        %ex.1 = phi i8* [ undef, %invoke.cont2 ], [ null, %catch ]
+        %1 = landingpad { i8*, i32 }
+                catch i8* null
+        br label %finally.catchall
+
+      finally.catchall:                                 ; preds = %lpad33, %lpad
+        %ex.2 = phi i8* [ %ex.1, %lpad2 ], [ null, %lpad ]
+        unreachable
+    }
+  )invalid", Err, Ctx));
+
+	if (!M) {
+    Err.print("unit", errs());
+    exit(1);
+  }
+
+  Function *Func = M->getFunction("foo");
+  EXPECT_FALSE(verifyFunction(*Func, &errs()));
+
+  SmallVector<BasicBlock *, 2> ExtractedBlocks{
+    getBlockByName(Func, "catch"),
+    getBlockByName(Func, "invoke.cont2"),
+    getBlockByName(Func, "invoke.cont3"),
+    getBlockByName(Func, "lpad2")
+  };
+
+  DominatorTree DT(*Func);
+  CodeExtractor CE(ExtractedBlocks, &DT);
+  EXPECT_TRUE(CE.isEligible());
+
+  Function *Outlined = CE.extractCodeRegion();
+  EXPECT_TRUE(Outlined);
+  EXPECT_FALSE(verifyFunction(*Outlined, &errs()));
+  EXPECT_FALSE(verifyFunction(*Func, &errs()));
+}
+
 } // end anonymous namespace