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

[myhsu]

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

[davide]

I'm a little nervous about updating the dominator manually, still. I wonder if we can use the new API that Jakub implemented (cc: @kuhar / @dberlin )

[• dberlin]

Please don't update it manually.
Pretty much every manual update LLVM does has been proven incorrect over time.
Please just express the added/removed edges if you have to.
If you are cutting out a region, this should not be difficult.

[kuhar]

I have to agree with Danny, I'm not sure this version is correct. It should be fairly easy to use the batch updater here.
If we discover that doing incremental update makes the code too slow, that it's also great, as I still haven't come up a good benchmark to tell me what needs more optimization.

Nit: please run clang-format on your changes.

lib/Transforms/Utils/CodeExtractor.cpp

473

Does this handle the self-loop case?

480

In general, we prefer to update DFS numbers lazily and it's better to just invalidate them.

[kuhar]

lib/Transforms/Utils/CodeExtractor.cpp

771

You can use auto here, as the type is obvious looking at the right hand side.

[myhsu]

It turn out that the DomTree updating issue has actually been fixed in D32308
The unittest also pass after I applying that change
Sorry I didn't work incrementally on master branch at first place

[myhsu]

@kuhar @dberlin
Could you help reviewing my update?
Thanks : )

[chandlerc]

Can you revert the unrelated formatting changes? Otherwise reviewing this is a bit hard.

[myhsu]

@chandlerc
Temporarily revert to version before applying clang-format

[kuhar]

In D37902#878726, @myhsu wrote:

@chandlerc
Temporarily revert to version before applying clang-format

I think the point was to only format your changes, not the surrounding code. You can use the clang-format-diff.py script to to that.

lib/Transforms/Utils/CodeExtractor.cpp

791

Why is it the case that incrementing OAI works here and gets to OutputArgEnd? Same with the GEP::Create, why is passing the OAI pointer there valid if it doesn't get incremented and you use outputs[i] the same time?

This is not clear to me, could you consider adding a comment?
I think that the comment from the previous version of this snippet got lost here.

unittests/Transforms/Utils/CodeExtractor.cpp

57

nit: s/require/requires
Maybe it'd be better to phrase is like CodeExtractor requires the first element to dominate all the other ones?. Isn't an element really a BasicBlock here?

68

Why is FALSE expected here?

[myhsu]

@kuhar
Fix things mentioned in the inline comments

[myhsu]

lib/Transforms/Utils/CodeExtractor.cpp

791

The original code didn't check the OutputArgEnd either, however I add an assertion anyway.

OAI iterator shouldn't be increased in each iteration if AggregateArgs is true. Since under the configuration of using aggregate argument, there would be only one struct argument that accommodate all of the output values in struct fields, and OAI should always point to that struct argument.

The code comment in the original code seems to be not related to the aggregate argument condition, but I add few comment regarding this concern anyway

unittests/Transforms/Utils/CodeExtractor.cpp

68

verifyFunction would return false if no error

[myhsu]

Fix typo

[kuhar]

lib/Transforms/Utils/CodeExtractor.cpp

745–746

nit: s/reference/reference.
Comments should be full sentences and end with ..

770

nit: s/value/value.

782

nit: maybe s/amount/number?

792

Maybe s/increase/increment or even s/increase/advance?

793

nit: s/point/points
s/case/case.

unittests/Transforms/Utils/CodeExtractor.cpp

59

nit: s/ones/ones.

68

Maybe that's obvious for people more familiar with this part of the codebase, but I'd appreciate a comment explaining that here :)

[myhsu]

@kuhar
Fix typos

[kuhar]

Other than the two minor nits, the patch looks good to me.

lib/Transforms/Utils/CodeExtractor.cpp

774

nit: s/point/point.

777

nit: s/PHIs/PHIs.
You can also move this piece of comment to the line above, as you won't hit the 80-column limit (I think :P).

[myhsu]

Fix comment

[myhsu]

In D37902#887576, @kuhar wrote:

Other than the two minor nits, the patch looks good to me.

Thanks @kuhar :)

[myhsu]

@chandlerc @kuhar
I'm not pretty familiar the procedure yet, could someone help me commit the patch or merge it?
Thanks

[kuhar]

In D37902#889231, @myhsu wrote:

@chandlerc @kuhar
I'm not pretty familiar the procedure yet, could someone help me commit the patch or merge it?
Thanks

If you have commit access, you can go ahead and submit the patch. Otherwise, I can submit it for you.
Please take a look here: Obtaining Commit Access.

[myhsu]

In D37902#889297, @kuhar wrote:

In D37902#889231, @myhsu wrote:

@chandlerc @kuhar
I'm not pretty familiar the procedure yet, could someone help me commit the patch or merge it?
Thanks

If you have commit access, you can go ahead and submit the patch. Otherwise, I can submit it for you.
Please take a look here: Obtaining Commit Access.

Since this is my first time committing a patch, I guess I don't have enough record for being granted the commit access
Please submit this changed for me, thanks!

[kuhar]

Can you please rebase the patch so that is applies cleanly to the current ToT?

[myhsu]

Rebase to master

[kuhar]

With your patch, FindPhiPredForUseInBlock becomes unused and it causes the following warning to be emitted:

llvm/lib/Transforms/Utils/CodeExtractor.cpp:657:20: warning: unused function 'FindPhiPredForUseInBlock' [-Wunused-function]
static BasicBlock* FindPhiPredForUseInBlock(Value* Used, BasicBlock* BB) {

And that would make many buildbots very grumpy :)
I think that the function should be now deleted, as it has no users.

[myhsu]

Remove 'FindPhiPredForUsedInBlock' function

lib/Transforms/Utils/CodeExtractor.cpp

 using namespace llvm;
 
 #define DEBUG_TYPE "code-extractor"
 
 // Provide a command-line option to aggregate function arguments into a struct
 // for functions produced by the code extractor. This is useful when converting
 // extracted functions to pthread-based code, as only one argument (void*) can
 // be passed in to pthread_create().
-static cl::opt<bool>
-AggregateArgsOpt("aggregate-extracted-args", cl::Hidden,
-                 cl::desc("Aggregate arguments to code-extracted functions"));
+static cl::opt<bool> AggregateArgsOpt(
+    "aggregate-extracted-args", cl::Hidden,
 
 /// \brief Test whether a block is valid for extraction.
 bool CodeExtractor::isBlockValidForExtraction(const BasicBlock &BB) {
   // Landing pads must be in the function where they were inserted for cleanup.
   if (BB.isEHPad())
     return false;
   // taking the address of a basic block moved to another function is illegal
   if (BB.hasAddressTaken())
if (!CodeExtractor::isBlockValidForExtraction(*BB)) {
       return Result;
     }
   }
 
 #ifndef NDEBUG
   for (SetVector<BasicBlock *>::iterator I = std::next(Result.begin()),
                                          E = Result.end();
        I != E; ++I)
-    for (pred_iterator PI = pred_begin(*I), PE = pred_end(*I);
-         PI != PE; ++PI)
+    for (pred_iterator PI = pred_begin(*I), PE = pred_end(*I); PI != PE; ++PI)
       assert(Result.count(*PI) &&
              "No blocks in this region may have entries from outside the region"
              " except for the first block!");
 #endif
 
   return Result;
 }
 
if (Blocks.count(I->getParent()))
       return true;
   return false;
 }
 
 /// definedInCaller - Return true if the specified value is defined in the
 /// function being code extracted, but not in the region being extracted.
 /// These values must be passed in as live-ins to the function.
 static bool definedInCaller(const SetVector<BasicBlock *> &Blocks, Value *V) {
-  if (isa<Argument>(V)) return true;
+  if (isa<Argument>(V))
+    return true;
   if (Instruction *I = dyn_cast<Instruction>(V))
     if (!Blocks.count(I->getParent()))
       return true;
   return false;
 }
 
 static BasicBlock *getCommonExitBlock(const SetVector<BasicBlock *> &Blocks) {
   BasicBlock *CommonExitBlock = nullptr;

 /// region, we need to split the entry block of the region so that the PHI node
 /// is easier to deal with.
 void CodeExtractor::severSplitPHINodes(BasicBlock *&Header) {
   unsigned NumPredsFromRegion = 0;
   unsigned NumPredsOutsideRegion = 0;
 
   if (Header != &Header->getParent()->getEntryBlock()) {
     PHINode *PN = dyn_cast<PHINode>(Header->begin());
-    if (!PN) return;  // No PHI nodes.
+    if (!PN)
+      return; // No PHI nodes.
 
     // If the header node contains any PHI nodes, check to see if there is more
     // than one entry from outside the region.  If so, we need to sever the
     // header block into two.
     for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
       if (Blocks.count(PN->getIncomingBlock(i)))
         ++NumPredsFromRegion;
       else
         ++NumPredsOutsideRegion;
 
     // If there is one (or fewer) predecessor from outside the region, we don't
     // need to do anything special.
-    if (NumPredsOutsideRegion <= 1) return;
+    if (NumPredsOutsideRegion <= 1)
+      return;
   }
 
   // Otherwise, we need to split the header block into two pieces: one
   // containing PHI nodes merging values from outside of the region, and a
   // second that contains all of the code for the block and merges back any
   // incoming values from inside of the region.
   BasicBlock *NewBB = llvm::SplitBlock(Header, Header->getFirstNonPHI(), DT);
 
   // We only want to code extract the second block now, and it becomes the new
   // header of the region.
   BasicBlock *OldPred = Header;
   Blocks.remove(OldPred);
   Blocks.insert(NewBB);
   Header = NewBB;
 
   // Okay, now we need to adjust the PHI nodes and any branches from within the
   // region to go to the new header block instead of the old header block.
   if (NumPredsFromRegion) {
     PHINode *PN = cast<PHINode>(OldPred->begin());
     // Loop over all of the predecessors of OldPred that are in the region,
     // changing them to branch to NewBB instead.
     for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
       if (Blocks.count(PN->getIncomingBlock(i))) {
         TerminatorInst *TI = PN->getIncomingBlock(i)->getTerminator();
         TI->replaceUsesOfWith(OldPred, NewBB);
       }
 
-    // Okay, everything within the region is now branching to the right block, we
-    // just have to update the PHI nodes now, inserting PHI nodes into NewBB.
+    // Okay, everything within the region is now branching to the right block,
+    // we just have to update the PHI nodes now, inserting PHI nodes into NewBB.
     BasicBlock::iterator AfterPHIs;
     for (AfterPHIs = OldPred->begin(); isa<PHINode>(AfterPHIs); ++AfterPHIs) {
       PHINode *PN = cast<PHINode>(AfterPHIs);
       // Create a new PHI node in the new region, which has an incoming value
       // from OldPred of PN.
       PHINode *NewPN = PHINode::Create(PN->getType(), 1 + NumPredsFromRegion,
                                        PN->getName() + ".ce", &NewBB->front());
       PN->replaceAllUsesWith(NewPN);

 /// constructFunction - make a function based on inputs and outputs, as follows:
 /// f(in0, ..., inN, out0, ..., outN)
 ///
 Function *CodeExtractor::constructFunction(const ValueSet &inputs,
                                            const ValueSet &outputs,
                                            BasicBlock *header,
                                            BasicBlock *newRootNode,
                                            BasicBlock *newHeader,
-                                           Function *oldFunction,
-                                           Module *M) {
+                                           Function *oldFunction, Module *M) {
   DEBUG(dbgs() << "inputs: " << inputs.size() << "\n");
   DEBUG(dbgs() << "outputs: " << outputs.size() << "\n");
 
   // This function returns unsigned, outputs will go back by reference.
   switch (NumExitBlocks) {
   case 0:
-  case 1: RetTy = Type::getVoidTy(header->getContext()); break;
-  case 2: RetTy = Type::getInt1Ty(header->getContext()); break;
-  default: RetTy = Type::getInt16Ty(header->getContext()); break;
+  case 1:
+    RetTy = Type::getVoidTy(header->getContext());
+    break;
+  case 2:
+    RetTy = Type::getInt1Ty(header->getContext());
+    break;
+  default:
+    RetTy = Type::getInt16Ty(header->getContext());
+    break;
   }
 
-  std::vector<Type*> paramTy;
+  std::vector<Type *> paramTy;
 
   // Add the types of the input values to the function's argument list
   for (Value *value : inputs) {
     DEBUG(dbgs() << "value used in func: " << *value << "\n");
     paramTy.push_back(value->getType());
   }
 
   // Add the types of the output values to the function's argument list.
Function *CodeExtractor::constructFunction(const ValueSet &inputs,
   });
 
   StructType *StructTy;
   if (AggregateArgs && (inputs.size() + outputs.size() > 0)) {
     StructTy = StructType::get(M->getContext(), paramTy);
     paramTy.clear();
     paramTy.push_back(PointerType::getUnqual(StructTy));
   }
-  FunctionType *funcType =
-                  FunctionType::get(RetTy, paramTy, false);
+  FunctionType *funcType = FunctionType::get(RetTy, paramTy, false);
 
   // Create the new function
-  Function *newFunction = Function::Create(funcType,
-                                           GlobalValue::InternalLinkage,
-                                           oldFunction->getName() + "_" +
-                                           header->getName(), M);
+  Function *newFunction =
+      Function::Create(funcType, GlobalValue::InternalLinkage,
+                       oldFunction->getName() + "_" + header->getName(), M);
   // If the old function is no-throw, so is the new one.
   if (oldFunction->doesNotThrow())
     newFunction->setDoesNotThrow();
 
   // Inherit the uwtable attribute if we need to.
   if (oldFunction->hasUWTable())
     newFunction->setHasUWTable();
 
if (AggregateArgs) {
       Idx[1] = ConstantInt::get(Type::getInt32Ty(header->getContext()), i);
       TerminatorInst *TI = newFunction->begin()->getTerminator();
       GetElementPtrInst *GEP = GetElementPtrInst::Create(
           StructTy, &*AI, Idx, "gep_" + inputs[i]->getName(), TI);
       RewriteVal = new LoadInst(GEP, "loadgep_" + inputs[i]->getName(), TI);
     } else
       RewriteVal = &*AI++;
 
-    std::vector<User*> Users(inputs[i]->user_begin(), inputs[i]->user_end());
+    std::vector<User *> Users(inputs[i]->user_begin(), inputs[i]->user_end());
     for (User *use : Users)
       if (Instruction *inst = dyn_cast<Instruction>(use))
         if (Blocks.count(inst->getParent()))
           inst->replaceUsesOfWith(inputs[i], RewriteVal);
   }
 
   // Set names for input and output arguments.
   if (!AggregateArgs) {
     AI = newFunction->arg_begin();
     for (unsigned i = 0, e = inputs.size(); i != e; ++i, ++AI)
       AI->setName(inputs[i]->getName());
     for (unsigned i = 0, e = outputs.size(); i != e; ++i, ++AI)
-      AI->setName(outputs[i]->getName()+".out");
+      AI->setName(outputs[i]->getName() + ".out");
   }
 
   // Rewrite branches to basic blocks outside of the loop to new dummy blocks
   // within the new function. This must be done before we lose track of which
   // blocks were originally in the code region.
-  std::vector<User*> Users(header->user_begin(), header->user_end());
+  std::vector<User *> Users(header->user_begin(), header->user_end());
   for (unsigned i = 0, e = Users.size(); i != e; ++i)
     // The BasicBlock which contains the branch is not in the region
     // modify the branch target to a new block
     if (TerminatorInst *TI = dyn_cast<TerminatorInst>(Users[i]))
       if (!Blocks.count(TI->getParent()) &&
           TI->getParent()->getParent() == oldFunction)
         TI->replaceUsesOfWith(header, newHeader);
 
   return newFunction;
 }
 
 /// FindPhiPredForUseInBlock - Given a value and a basic block, find a PHI
 /// that uses the value within the basic block, and return the predecessor
 /// block associated with that use, or return 0 if none is found.
-static BasicBlock* FindPhiPredForUseInBlock(Value* Used, BasicBlock* BB) {
+static BasicBlock *FindPhiPredForUseInBlock(Value *Used, BasicBlock *BB) {
   for (Use &U : Used->uses()) {
-     PHINode *P = dyn_cast<PHINode>(U.getUser());
-     if (P && P->getParent() == BB)
-       return P->getIncomingBlock(U);
   }
 
   return nullptr;
 }
 
 /// emitCallAndSwitchStatement - This method sets up the caller side by adding
 /// the call instruction, splitting any PHI nodes in the header block as
 /// necessary.
-void CodeExtractor::
-emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
-                           ValueSet &inputs, ValueSet &outputs) {
+void CodeExtractor::emitCallAndSwitchStatement(Function *newFunction,
+                                               BasicBlock *codeReplacer,
+                                               ValueSet &inputs,
+                                               ValueSet &outputs) {
   // Emit a call to the new function, passing in: *pointer to struct (if
   // aggregating parameters), or plan inputs and allocated memory for outputs
-  std::vector<Value*> params, StructValues, ReloadOutputs, Reloads;
+  std::vector<Value *> params, StructValues, ReloadOutputs, Reloads;
 
   Module *M = newFunction->getParent();
   LLVMContext &Context = M->getContext();
   const DataLayout &DL = M->getDataLayout();
 
   // Add inputs as params, or to be filled into the struct
   for (Value *input : inputs)
     if (AggregateArgs)
       StructValues.push_back(input);
     else
       params.push_back(input);
 
   // Create allocas for the outputs
   for (Value *output : outputs) {
     if (AggregateArgs) {
       StructValues.push_back(output);
     } else {
       AllocaInst *alloca =
-        new AllocaInst(output->getType(), DL.getAllocaAddrSpace(),
-                       nullptr, output->getName() + ".loc",
+          new AllocaInst(output->getType(), DL.getAllocaAddrSpace(), nullptr,
+                         output->getName() + ".loc",
                          &codeReplacer->getParent()->front().front());
       ReloadOutputs.push_back(alloca);
       params.push_back(alloca);
     }
   }
 
   StructType *StructArgTy = nullptr;
   AllocaInst *Struct = nullptr;
   if (AggregateArgs && (inputs.size() + outputs.size() > 0)) {
-    std::vector<Type*> ArgTypes;
-    for (ValueSet::iterator v = StructValues.begin(),
-           ve = StructValues.end(); v != ve; ++v)
+    std::vector<Type *> ArgTypes;
+    for (ValueSet::iterator v = StructValues.begin(), ve = StructValues.end();
+         v != ve; ++v)
       ArgTypes.push_back((*v)->getType());
 
     // Allocate a struct at the beginning of this function
     StructArgTy = StructType::get(newFunction->getContext(), ArgTypes);
     Struct = new AllocaInst(StructArgTy, DL.getAllocaAddrSpace(), nullptr,
                             "structArg",
                             &codeReplacer->getParent()->front().front());
     params.push_back(Struct);
CallInst *call = CallInst::Create(newFunction, params,
                                     NumExitBlocks > 1 ? "targetBlock" : "");
   codeReplacer->getInstList().push_back(call);
 
   Function::arg_iterator OutputArgBegin = newFunction->arg_begin();
   unsigned FirstOut = inputs.size();
   if (!AggregateArgs)
     std::advance(OutputArgBegin, inputs.size());
 
   // Reload the outputs passed in by reference
+  Function::arg_iterator OAI = OutputArgBegin;
   for (unsigned i = 0, e = outputs.size(); i != e; ++i) {
     Value *Output = nullptr;
     if (AggregateArgs) {
       Value *Idx[2];
       Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
       Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), FirstOut + i);
       GetElementPtrInst *GEP = GetElementPtrInst::Create(
           StructArgTy, Struct, Idx, "gep_reload_" + outputs[i]->getName());
       codeReplacer->getInstList().push_back(GEP);
       Output = GEP;
     } else {
       Output = ReloadOutputs[i];
     }
-    LoadInst *load = new LoadInst(Output, outputs[i]->getName()+".reload");
+    LoadInst *load = new LoadInst(Output, outputs[i]->getName() + ".reload");
     Reloads.push_back(load);
     codeReplacer->getInstList().push_back(load);
-    std::vector<User*> Users(outputs[i]->user_begin(), outputs[i]->user_end());
+    std::vector<User *> Users(outputs[i]->user_begin(), outputs[i]->user_end());
     for (unsigned u = 0, e = Users.size(); u != e; ++u) {
       Instruction *inst = cast<Instruction>(Users[u]);
       if (!Blocks.count(inst->getParent()))
         inst->replaceUsesOfWith(outputs[i], load);
     }
+
+    // Store to argument right after the definition of output value
+    auto *OutI = dyn_cast<Instruction>(outputs[i]);
+    if (!OutI)
+      continue;
+    // Find proper insertion point
+    Instruction *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();
+
+    if (AggregateArgs) {
+      Value *Idx[2];
+      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);
+    } else {
+      new StoreInst(outputs[i], &*OAI, InsertPt);
+      ++OAI;
+    }
   }
 
   // Now we can emit a switch statement using the call as a value.
   SwitchInst *TheSwitch =
       SwitchInst::Create(Constant::getNullValue(Type::getInt16Ty(Context)),
                          codeReplacer, 0, codeReplacer);
 
   // Since there may be multiple exits from the original region, make the new
   // function return an unsigned, switch on that number.  This loop iterates
   // over all of the blocks in the extracted region, updating any terminator
   // instructions in the to-be-extracted region that branch to blocks that are
   // not in the region to be extracted.
-  std::map<BasicBlock*, BasicBlock*> ExitBlockMap;
+  std::map<BasicBlock *, BasicBlock *> ExitBlockMap;
 
   unsigned switchVal = 0;
   for (BasicBlock *Block : Blocks) {
     TerminatorInst *TI = Block->getTerminator();
     for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
       if (!Blocks.count(TI->getSuccessor(i))) {
         BasicBlock *OldTarget = TI->getSuccessor(i);
         // add a new basic block which returns the appropriate value
         BasicBlock *&NewTarget = ExitBlockMap[OldTarget];
         if (!NewTarget) {
           // If we don't already have an exit stub for this non-extracted
           // destination, create one now!
-          NewTarget = BasicBlock::Create(Context,
-                                         OldTarget->getName() + ".exitStub",
-                                         newFunction);
+          NewTarget = BasicBlock::Create(
+              Context, OldTarget->getName() + ".exitStub", newFunction);
           unsigned SuccNum = switchVal++;
 
           Value *brVal = nullptr;
           switch (NumExitBlocks) {
           case 0:
-          case 1: break;  // No value needed.
-          case 2:         // Conditional branch, return a bool
+          case 1:
+            break; // No value needed.
+          case 2:  // Conditional branch, return a bool
             brVal = ConstantInt::get(Type::getInt1Ty(Context), !SuccNum);
             break;
           default:
             brVal = ConstantInt::get(Type::getInt16Ty(Context), SuccNum);
             break;
           }
 
-          ReturnInst *NTRet = ReturnInst::Create(Context, brVal, NewTarget);
+          ReturnInst::Create(Context, brVal, NewTarget);
 
           // Update the switch instruction.
-          TheSwitch->addCase(ConstantInt::get(Type::getInt16Ty(Context),
-                                              SuccNum),
-                             OldTarget);
-
-          // Restore values just before we exit
-          Function::arg_iterator OAI = OutputArgBegin;
-          for (unsigned out = 0, e = outputs.size(); out != e; ++out) {
-            // For an invoke, the normal destination is the only one that is
-            // dominated by the result of the invocation
-            BasicBlock *DefBlock = cast<Instruction>(outputs[out])->getParent();
-
-            bool DominatesDef = true;
-
-            BasicBlock *NormalDest = nullptr;
-            if (auto *Invoke = dyn_cast<InvokeInst>(outputs[out]))
-              NormalDest = Invoke->getNormalDest();
-
-            if (NormalDest) {
-              DefBlock = NormalDest;
-
-              // Make sure we are looking at the original successor block, not
-              // at a newly inserted exit block, which won't be in the dominator
-              // info.
-              for (const auto &I : ExitBlockMap)
-                if (DefBlock == I.second) {
-                  DefBlock = I.first;
-                  break;
-                }
-
-              // In the extract block case, if the block we are extracting ends
-              // with an invoke instruction, make sure that we don't emit a
-              // store of the invoke value for the unwind block.
-              if (!DT && DefBlock != OldTarget)
-                DominatesDef = false;
-            }
-
-            if (DT) {
-              DominatesDef = DT->dominates(DefBlock, OldTarget);
-              
-              // If the output value is used by a phi in the target block,
-              // then we need to test for dominance of the phi's predecessor
-              // instead.  Unfortunately, this a little complicated since we
-              // have already rewritten uses of the value to uses of the reload.
-              BasicBlock* pred = FindPhiPredForUseInBlock(Reloads[out], 
-                                                          OldTarget);
-              if (pred && DT && DT->dominates(DefBlock, pred))
-                DominatesDef = true;
-            }
-
-            if (DominatesDef) {
-              if (AggregateArgs) {
-                Value *Idx[2];
-                Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
-                Idx[1] = ConstantInt::get(Type::getInt32Ty(Context),
-                                          FirstOut+out);
-                GetElementPtrInst *GEP = GetElementPtrInst::Create(
-                    StructArgTy, &*OAI, Idx, "gep_" + outputs[out]->getName(),
-                    NTRet);
-                new StoreInst(outputs[out], GEP, NTRet);
-              } else {
-                new StoreInst(outputs[out], &*OAI, NTRet);
-              }
-            }
-            // Advance output iterator even if we don't emit a store
-            if (!AggregateArgs) ++OAI;
-          }
+          TheSwitch->addCase(
+              ConstantInt::get(Type::getInt16Ty(Context), SuccNum), OldTarget);
         }
 
         // rewrite the original branch instruction with this new target
         TI->setSuccessor(i, NewTarget);
       }
   }
 
   // Now that we've done the deed, simplify the switch instruction.
   Type *OldFnRetTy = TheSwitch->getParent()->getParent()->getReturnType();
   switch (NumExitBlocks) {
   case 0:
     // There are no successors (the block containing the switch itself), which
     // means that previously this was the last part of the function, and hence
     // this should be rewritten as a `ret'
 
     // Check if the function should return a value
     if (OldFnRetTy->isVoidTy()) {
-      ReturnInst::Create(Context, nullptr, TheSwitch);  // Return void
+      ReturnInst::Create(Context, nullptr, TheSwitch); // Return void
     } else if (OldFnRetTy == TheSwitch->getCondition()->getType()) {
       // return what we have
       ReturnInst::Create(Context, TheSwitch->getCondition(), TheSwitch);
     } else {
       // Otherwise we must have code extracted an unwind or something, just
       // return whatever we want.
-      ReturnInst::Create(Context, 
-                         Constant::getNullValue(OldFnRetTy), TheSwitch);
+      ReturnInst::Create(Context, Constant::getNullValue(OldFnRetTy),
+                         TheSwitch);
     }
 
     TheSwitch->eraseFromParent();
     break;
   case 1:
     // Only a single destination, change the switch into an unconditional
     // branch.
     BranchInst::Create(TheSwitch->getSuccessor(1), TheSwitch);
     TheSwitch->eraseFromParent();
     break;
   case 2:
     BranchInst::Create(TheSwitch->getSuccessor(1), TheSwitch->getSuccessor(2),
                        call, TheSwitch);
     TheSwitch->eraseFromParent();
     break;
   default:
     // Otherwise, make the default destination of the switch instruction be one
     // of the other successors.
     TheSwitch->setCondition(call);
     TheSwitch->setDefaultDest(TheSwitch->getSuccessor(NumExitBlocks));
     // Remove redundant case
-    TheSwitch->removeCase(SwitchInst::CaseIt(TheSwitch, NumExitBlocks-1));
+    TheSwitch->removeCase(SwitchInst::CaseIt(TheSwitch, NumExitBlocks - 1));
     break;
   }
 }
 
 void CodeExtractor::moveCodeToFunction(Function *newFunction) {
   Function *oldFunc = (*Blocks.begin())->getParent();
   Function::BasicBlockListType &oldBlocks = oldFunc->getBasicBlockList();
   Function::BasicBlockListType &newBlocks = newFunction->getBasicBlockList();
Function *CodeExtractor::extractCodeRegion() {
 
   // If we have any return instructions in the region, split those blocks so
   // that the return is not in the region.
   splitReturnBlocks();
 
   Function *oldFunction = header->getParent();
 
   // This takes place of the original loop
-  BasicBlock *codeReplacer = BasicBlock::Create(header->getContext(), 
-                                                "codeRepl", oldFunction,
-                                                header);
+  BasicBlock *codeReplacer =
+      BasicBlock::Create(header->getContext(), "codeRepl", oldFunction, header);
 
   // The new function needs a root node because other nodes can branch to the
   // head of the region, but the entry node of a function cannot have preds.
-  BasicBlock *newFuncRoot = BasicBlock::Create(header->getContext(), 
-                                               "newFuncRoot");
+  BasicBlock *newFuncRoot =
+      BasicBlock::Create(header->getContext(), "newFuncRoot");
   newFuncRoot->getInstList().push_back(BranchInst::Create(header));
 
   findAllocas(SinkingCands, HoistingCands, CommonExit);
   assert(HoistingCands.empty() || CommonExit);
 
   // Find inputs to, outputs from the code region.
   findInputsOutputs(inputs, outputs, SinkingCands);
 
for (succ_iterator SI = succ_begin(Block), SE = succ_end(Block); SI != SE;
         }
         ExitBlocks.insert(*SI);
       }
     }
   }
   NumExitBlocks = ExitBlocks.size();
 
   // Construct new function based on inputs/outputs & add allocas for all defs.
-  Function *newFunction = constructFunction(inputs, outputs, header,
-                                            newFuncRoot,
-                                            codeReplacer, oldFunction,
-                                            oldFunction->getParent());
+  Function *newFunction =
+      constructFunction(inputs, outputs, header, newFuncRoot, codeReplacer,
+                        oldFunction, oldFunction->getParent());
 
   // Update the entry count of the function.
   if (BFI) {
     Optional<uint64_t> EntryCount =
         BFI->getProfileCountFromFreq(EntryFreq.getFrequency());
     if (EntryCount.hasValue())
       newFunction->setEntryCount(EntryCount.getValue());
     BFI->setBlockFreq(codeReplacer, EntryFreq.getFrequency());
for (BasicBlock::iterator I = header->begin(); isa<PHINode>(I); ++I) {
     for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
       if (!Blocks.count(PN->getIncomingBlock(i)))
         PN->setIncomingBlock(i, newFuncRoot);
   }
 
   // Look at all successors of the codeReplacer block.  If any of these blocks
   // had PHI nodes in them, we need to update the "from" block to be the code
   // replacer, not the original block in the extracted region.
-  std::vector<BasicBlock*> Succs(succ_begin(codeReplacer),
+  std::vector<BasicBlock *> Succs(succ_begin(codeReplacer),
                                   succ_end(codeReplacer));
   for (unsigned i = 0, e = Succs.size(); i != e; ++i)
     for (BasicBlock::iterator I = Succs[i]->begin(); isa<PHINode>(I); ++I) {
       PHINode *PN = cast<PHINode>(I);
-      std::set<BasicBlock*> ProcessedPreds;
+      std::set<BasicBlock *> ProcessedPreds;
       for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
         if (Blocks.count(PN->getIncomingBlock(i))) {
           if (ProcessedPreds.insert(PN->getIncomingBlock(i)).second)
             PN->setIncomingBlock(i, codeReplacer);
           else {
             // There were multiple entries in the PHI for this block, now there
             // is only one, so remove the duplicated entries.
             PN->removeIncomingValue(i, false);
-            --i; --e;
+            --i;
+            --e;
           }
         }
     }
 
-  DEBUG(if (verifyFunction(*newFunction)) 
+  DEBUG(if (verifyFunction(*newFunction))
             report_fatal_error("verifyFunction failed!"));
   return newFunction;
 }

unittests/Transforms/Utils/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Analysis
+  AsmParser
   Core
   Support
   TransformUtils
   )
 
 add_llvm_unittest(UtilsTests
   ASanStackFrameLayoutTest.cpp
   Cloning.cpp
+  CodeExtractor.cpp
   FunctionComparator.cpp
   IntegerDivision.cpp
   Local.cpp
   OrderedInstructions.cpp
   ValueMapperTest.cpp
   )

unittests/Transforms/Utils/CodeExtractor.cpp

+//===- 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 require first element
+  // should dominate others
+  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