[CodeExtractor] Fix multiple bugs under certain shape of extracted region

Summary:
If the extracted region has multiple exported data flows toward the same BB which is not included in the region, correct resotre instructions and PHI nodes won't be generated inside the exitStub. The solution is simply put the restore instructions right after the definition of output values instead of putting in exitStub.
Unittest for this bug is included.

Author: myhsu

Reviewers: chandlerc, davide, lattner, silvas, davidxl, wmi, kuhar

Subscribers: dberlin, kuhar, mgorny, llvm-commits

Differential Revision: https://reviews.llvm.org/D37902

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

2 files changed, 70 insertions, 0 deletions

diff --git a/unittests/Transforms/Utils/CMakeLists.txt b/unittests/Transforms/Utils/CMakeLists.txt
index 8c09bae5aa3..e2bb0af0f77 100644
--- a/unittests/Transforms/Utils/CMakeLists.txt
+++ b/unittests/Transforms/Utils/CMakeLists.txt
@@ -9,6 +9,7 @@ set(LLVM_LINK_COMPONENTS
 add_llvm_unittest(UtilsTests
   ASanStackFrameLayoutTest.cpp
   Cloning.cpp
+  CodeExtractor.cpp
   FunctionComparator.cpp
   IntegerDivision.cpp
   Local.cpp
diff --git a/unittests/Transforms/Utils/CodeExtractor.cpp b/unittests/Transforms/Utils/CodeExtractor.cpp
new file mode 100644
index 00000000000..c229be6d695
--- /dev/null
+++ b/unittests/Transforms/Utils/CodeExtractor.cpp
@@ -0,0 +1,69 @@
+//===- CodeExtractor.cpp - Unit tests for CodeExtractor -------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/CodeExtractor.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Verifier.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+TEST(CodeExtractor, ExitStub) {
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M(parseAssemblyString(R"invalid(
+    define i32 @foo(i32 %x, i32 %y, i32 %z) {
+    header:
+      %0 = icmp ugt i32 %x, %y
+      br i1 %0, label %body1, label %body2
+
+    body1:
+      %1 = add i32 %z, 2
+      br label %notExtracted
+
+    body2:
+      %2 = mul i32 %z, 7
+      br label %notExtracted
+
+    notExtracted:
+      %3 = phi i32 [ %1, %body1 ], [ %2, %body2 ]
+      %4 = add i32 %3, %x
+      ret i32 %4
+    }
+  )invalid",
+                                                Err, Ctx));
+
+  Function *Func = M->getFunction("foo");
+  SmallVector<BasicBlock *, 3> Candidates;
+  for (auto &BB : *Func) {
+    if (BB.getName() == "body1")
+      Candidates.push_back(&BB);
+    if (BB.getName() == "body2")
+      Candidates.push_back(&BB);
+  }
+  // CodeExtractor requires the first basic block
+  // to dominate all the other ones.
+  Candidates.insert(Candidates.begin(), &Func->getEntryBlock());
+
+  DominatorTree DT(*Func);
+  CodeExtractor CE(Candidates, &DT);
+  EXPECT_TRUE(CE.isEligible());
+
+  Function *Outlined = CE.extractCodeRegion();
+  EXPECT_TRUE(Outlined);
+  EXPECT_FALSE(verifyFunction(*Outlined));
+}
+} // end anonymous namespace