1 //===- DFAJumpThreading.cpp - Threads a switch statement inside a loop ----===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // Transform each threading path to effectively jump thread the DFA. For
10 // example, the CFG below could be transformed as follows, where the cloned
11 // blocks unconditionally branch to the next correct case based on what is
12 // identified in the analysis.
13 //
14 //          sw.bb                        sw.bb
15 //        /   |   \                    /   |   \
16 //   case1  case2  case3          case1  case2  case3
17 //        \   |   /                 |      |      |
18 //       determinator            det.2   det.3  det.1
19 //        br sw.bb                /        |        \
20 //                          sw.bb.2     sw.bb.3     sw.bb.1
21 //                           br case2    br case3    br case1§
22 //
23 // Definitions and Terminology:
24 //
25 // * Threading path:
26 //   a list of basic blocks, the exit state, and the block that determines
27 //   the next state, for which the following notation will be used:
28 //   < path of BBs that form a cycle > [ state, determinator ]
29 //
30 // * Predictable switch:
31 //   The switch variable is always a known constant so that all conditional
32 //   jumps based on switch variable can be converted to unconditional jump.
33 //
34 // * Determinator:
35 //   The basic block that determines the next state of the DFA.
36 //
37 // Representing the optimization in C-like pseudocode: the code pattern on the
38 // left could functionally be transformed to the right pattern if the switch
39 // condition is predictable.
40 //
41 //  X = A                       goto A
42 //  for (...)                   A:
43 //    switch (X)                  ...
44 //      case A                    goto B
45 //        X = B                 B:
46 //      case B                    ...
47 //        X = C                   goto C
48 //
49 // The pass first checks that switch variable X is decided by the control flow
50 // path taken in the loop; for example, in case B, the next value of X is
51 // decided to be C. It then enumerates through all paths in the loop and labels
52 // the basic blocks where the next state is decided.
53 //
54 // Using this information it creates new paths that unconditionally branch to
55 // the next case. This involves cloning code, so it only gets triggered if the
56 // amount of code duplicated is below a threshold.
57 //
58 //===----------------------------------------------------------------------===//
59 
60 #include "llvm/Transforms/Scalar/DFAJumpThreading.h"
61 #include "llvm/ADT/APInt.h"
62 #include "llvm/ADT/DenseMap.h"
63 #include "llvm/ADT/DepthFirstIterator.h"
64 #include "llvm/ADT/SmallSet.h"
65 #include "llvm/ADT/Statistic.h"
66 #include "llvm/Analysis/AssumptionCache.h"
67 #include "llvm/Analysis/CodeMetrics.h"
68 #include "llvm/Analysis/LoopIterator.h"
69 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
70 #include "llvm/Analysis/TargetTransformInfo.h"
71 #include "llvm/IR/CFG.h"
72 #include "llvm/IR/Constants.h"
73 #include "llvm/IR/IntrinsicInst.h"
74 #include "llvm/IR/Verifier.h"
75 #include "llvm/InitializePasses.h"
76 #include "llvm/Pass.h"
77 #include "llvm/Support/CommandLine.h"
78 #include "llvm/Support/Debug.h"
79 #include "llvm/Transforms/Scalar.h"
80 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
81 #include "llvm/Transforms/Utils/Cloning.h"
82 #include "llvm/Transforms/Utils/SSAUpdaterBulk.h"
83 #include "llvm/Transforms/Utils/ValueMapper.h"
84 #include <algorithm>
85 #include <deque>
86 
87 using namespace llvm;
88 
89 #define DEBUG_TYPE "dfa-jump-threading"
90 
91 STATISTIC(NumTransforms, "Number of transformations done");
92 STATISTIC(NumCloned, "Number of blocks cloned");
93 STATISTIC(NumPaths, "Number of individual paths threaded");
94 
95 static cl::opt<bool>
96     ClViewCfgBefore("dfa-jump-view-cfg-before",
97                     cl::desc("View the CFG before DFA Jump Threading"),
98                     cl::Hidden, cl::init(false));
99 
100 static cl::opt<unsigned> MaxPathLength(
101     "dfa-max-path-length",
102     cl::desc("Max number of blocks searched to find a threading path"),
103     cl::Hidden, cl::init(20));
104 
105 static cl::opt<unsigned>
106     CostThreshold("dfa-cost-threshold",
107                   cl::desc("Maximum cost accepted for the transformation"),
108                   cl::Hidden, cl::init(50));
109 
110 namespace {
111 
112 class SelectInstToUnfold {
113   SelectInst *SI;
114   PHINode *SIUse;
115 
116 public:
117   SelectInstToUnfold(SelectInst *SI, PHINode *SIUse) : SI(SI), SIUse(SIUse) {}
118 
119   SelectInst *getInst() { return SI; }
120   PHINode *getUse() { return SIUse; }
121 
122   explicit operator bool() const { return SI && SIUse; }
123 };
124 
125 void unfold(DomTreeUpdater *DTU, SelectInstToUnfold SIToUnfold,
126             std::vector<SelectInstToUnfold> *NewSIsToUnfold,
127             std::vector<BasicBlock *> *NewBBs);
128 
129 class DFAJumpThreading {
130 public:
131   DFAJumpThreading(AssumptionCache *AC, DominatorTree *DT,
132                    TargetTransformInfo *TTI, OptimizationRemarkEmitter *ORE)
133       : AC(AC), DT(DT), TTI(TTI), ORE(ORE) {}
134 
135   bool run(Function &F);
136 
137 private:
138   void
139   unfoldSelectInstrs(DominatorTree *DT,
140                      const SmallVector<SelectInstToUnfold, 4> &SelectInsts) {
141     DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
142     SmallVector<SelectInstToUnfold, 4> Stack;
143     for (SelectInstToUnfold SIToUnfold : SelectInsts)
144       Stack.push_back(SIToUnfold);
145 
146     while (!Stack.empty()) {
147       SelectInstToUnfold SIToUnfold = Stack.back();
148       Stack.pop_back();
149 
150       std::vector<SelectInstToUnfold> NewSIsToUnfold;
151       std::vector<BasicBlock *> NewBBs;
152       unfold(&DTU, SIToUnfold, &NewSIsToUnfold, &NewBBs);
153 
154       // Put newly discovered select instructions into the work list.
155       for (const SelectInstToUnfold &NewSIToUnfold : NewSIsToUnfold)
156         Stack.push_back(NewSIToUnfold);
157     }
158   }
159 
160   AssumptionCache *AC;
161   DominatorTree *DT;
162   TargetTransformInfo *TTI;
163   OptimizationRemarkEmitter *ORE;
164 };
165 
166 class DFAJumpThreadingLegacyPass : public FunctionPass {
167 public:
168   static char ID; // Pass identification
169   DFAJumpThreadingLegacyPass() : FunctionPass(ID) {}
170 
171   void getAnalysisUsage(AnalysisUsage &AU) const override {
172     AU.addRequired<AssumptionCacheTracker>();
173     AU.addRequired<DominatorTreeWrapperPass>();
174     AU.addPreserved<DominatorTreeWrapperPass>();
175     AU.addRequired<TargetTransformInfoWrapperPass>();
176     AU.addRequired<OptimizationRemarkEmitterWrapperPass>();
177   }
178 
179   bool runOnFunction(Function &F) override {
180     if (skipFunction(F))
181       return false;
182 
183     AssumptionCache *AC =
184         &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
185     DominatorTree *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
186     TargetTransformInfo *TTI =
187         &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
188     OptimizationRemarkEmitter *ORE =
189         &getAnalysis<OptimizationRemarkEmitterWrapperPass>().getORE();
190 
191     return DFAJumpThreading(AC, DT, TTI, ORE).run(F);
192   }
193 };
194 } // end anonymous namespace
195 
196 char DFAJumpThreadingLegacyPass::ID = 0;
197 INITIALIZE_PASS_BEGIN(DFAJumpThreadingLegacyPass, "dfa-jump-threading",
198                       "DFA Jump Threading", false, false)
199 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
200 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
201 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
202 INITIALIZE_PASS_DEPENDENCY(OptimizationRemarkEmitterWrapperPass)
203 INITIALIZE_PASS_END(DFAJumpThreadingLegacyPass, "dfa-jump-threading",
204                     "DFA Jump Threading", false, false)
205 
206 // Public interface to the DFA Jump Threading pass
207 FunctionPass *llvm::createDFAJumpThreadingPass() {
208   return new DFAJumpThreadingLegacyPass();
209 }
210 
211 namespace {
212 
213 /// Create a new basic block and sink \p SIToSink into it.
214 void createBasicBlockAndSinkSelectInst(
215     DomTreeUpdater *DTU, SelectInst *SI, PHINode *SIUse, SelectInst *SIToSink,
216     BasicBlock *EndBlock, StringRef NewBBName, BasicBlock **NewBlock,
217     BranchInst **NewBranch, std::vector<SelectInstToUnfold> *NewSIsToUnfold,
218     std::vector<BasicBlock *> *NewBBs) {
219   assert(SIToSink->hasOneUse());
220   assert(NewBlock);
221   assert(NewBranch);
222   *NewBlock = BasicBlock::Create(SI->getContext(), NewBBName,
223                                  EndBlock->getParent(), EndBlock);
224   NewBBs->push_back(*NewBlock);
225   *NewBranch = BranchInst::Create(EndBlock, *NewBlock);
226   SIToSink->moveBefore(*NewBranch);
227   NewSIsToUnfold->push_back(SelectInstToUnfold(SIToSink, SIUse));
228   DTU->applyUpdates({{DominatorTree::Insert, *NewBlock, EndBlock}});
229 }
230 
231 /// Unfold the select instruction held in \p SIToUnfold by replacing it with
232 /// control flow.
233 ///
234 /// Put newly discovered select instructions into \p NewSIsToUnfold. Put newly
235 /// created basic blocks into \p NewBBs.
236 ///
237 /// TODO: merge it with CodeGenPrepare::optimizeSelectInst() if possible.
238 void unfold(DomTreeUpdater *DTU, SelectInstToUnfold SIToUnfold,
239             std::vector<SelectInstToUnfold> *NewSIsToUnfold,
240             std::vector<BasicBlock *> *NewBBs) {
241   SelectInst *SI = SIToUnfold.getInst();
242   PHINode *SIUse = SIToUnfold.getUse();
243   BasicBlock *StartBlock = SI->getParent();
244   BasicBlock *EndBlock = SIUse->getParent();
245   BranchInst *StartBlockTerm =
246       dyn_cast<BranchInst>(StartBlock->getTerminator());
247 
248   assert(StartBlockTerm && StartBlockTerm->isUnconditional());
249   assert(SI->hasOneUse());
250 
251   // These are the new basic blocks for the conditional branch.
252   // At least one will become an actual new basic block.
253   BasicBlock *TrueBlock = nullptr;
254   BasicBlock *FalseBlock = nullptr;
255   BranchInst *TrueBranch = nullptr;
256   BranchInst *FalseBranch = nullptr;
257 
258   // Sink select instructions to be able to unfold them later.
259   if (SelectInst *SIOp = dyn_cast<SelectInst>(SI->getTrueValue())) {
260     createBasicBlockAndSinkSelectInst(DTU, SI, SIUse, SIOp, EndBlock,
261                                       "si.unfold.true", &TrueBlock, &TrueBranch,
262                                       NewSIsToUnfold, NewBBs);
263   }
264   if (SelectInst *SIOp = dyn_cast<SelectInst>(SI->getFalseValue())) {
265     createBasicBlockAndSinkSelectInst(DTU, SI, SIUse, SIOp, EndBlock,
266                                       "si.unfold.false", &FalseBlock,
267                                       &FalseBranch, NewSIsToUnfold, NewBBs);
268   }
269 
270   // If there was nothing to sink, then arbitrarily choose the 'false' side
271   // for a new input value to the PHI.
272   if (!TrueBlock && !FalseBlock) {
273     FalseBlock = BasicBlock::Create(SI->getContext(), "si.unfold.false",
274                                     EndBlock->getParent(), EndBlock);
275     NewBBs->push_back(FalseBlock);
276     BranchInst::Create(EndBlock, FalseBlock);
277     DTU->applyUpdates({{DominatorTree::Insert, FalseBlock, EndBlock}});
278   }
279 
280   // Insert the real conditional branch based on the original condition.
281   // If we did not create a new block for one of the 'true' or 'false' paths
282   // of the condition, it means that side of the branch goes to the end block
283   // directly and the path originates from the start block from the point of
284   // view of the new PHI.
285   BasicBlock *TT = EndBlock;
286   BasicBlock *FT = EndBlock;
287   if (TrueBlock && FalseBlock) {
288     // A diamond.
289     TT = TrueBlock;
290     FT = FalseBlock;
291 
292     // Update the phi node of SI.
293     SIUse->removeIncomingValue(StartBlock, /* DeletePHIIfEmpty = */ false);
294     SIUse->addIncoming(SI->getTrueValue(), TrueBlock);
295     SIUse->addIncoming(SI->getFalseValue(), FalseBlock);
296 
297     // Update any other PHI nodes in EndBlock.
298     for (PHINode &Phi : EndBlock->phis()) {
299       if (&Phi != SIUse) {
300         Phi.addIncoming(Phi.getIncomingValueForBlock(StartBlock), TrueBlock);
301         Phi.addIncoming(Phi.getIncomingValueForBlock(StartBlock), FalseBlock);
302       }
303     }
304   } else {
305     BasicBlock *NewBlock = nullptr;
306     Value *SIOp1 = SI->getTrueValue();
307     Value *SIOp2 = SI->getFalseValue();
308 
309     // A triangle pointing right.
310     if (!TrueBlock) {
311       NewBlock = FalseBlock;
312       FT = FalseBlock;
313     }
314     // A triangle pointing left.
315     else {
316       NewBlock = TrueBlock;
317       TT = TrueBlock;
318       std::swap(SIOp1, SIOp2);
319     }
320 
321     // Update the phi node of SI.
322     for (unsigned Idx = 0; Idx < SIUse->getNumIncomingValues(); ++Idx) {
323       if (SIUse->getIncomingBlock(Idx) == StartBlock)
324         SIUse->setIncomingValue(Idx, SIOp1);
325     }
326     SIUse->addIncoming(SIOp2, NewBlock);
327 
328     // Update any other PHI nodes in EndBlock.
329     for (auto II = EndBlock->begin(); PHINode *Phi = dyn_cast<PHINode>(II);
330          ++II) {
331       if (Phi != SIUse)
332         Phi->addIncoming(Phi->getIncomingValueForBlock(StartBlock), NewBlock);
333     }
334   }
335   StartBlockTerm->eraseFromParent();
336   BranchInst::Create(TT, FT, SI->getCondition(), StartBlock);
337   DTU->applyUpdates({{DominatorTree::Insert, StartBlock, TT},
338                      {DominatorTree::Insert, StartBlock, FT}});
339 
340   // The select is now dead.
341   SI->eraseFromParent();
342 }
343 
344 struct ClonedBlock {
345   BasicBlock *BB;
346   uint64_t State; ///< \p State corresponds to the next value of a switch stmnt.
347 };
348 
349 typedef std::deque<BasicBlock *> PathType;
350 typedef std::vector<PathType> PathsType;
351 typedef SmallPtrSet<const BasicBlock *, 8> VisitedBlocks;
352 typedef std::vector<ClonedBlock> CloneList;
353 
354 // This data structure keeps track of all blocks that have been cloned.  If two
355 // different ThreadingPaths clone the same block for a certain state it should
356 // be reused, and it can be looked up in this map.
357 typedef DenseMap<BasicBlock *, CloneList> DuplicateBlockMap;
358 
359 // This map keeps track of all the new definitions for an instruction. This
360 // information is needed when restoring SSA form after cloning blocks.
361 typedef DenseMap<Instruction *, std::vector<Instruction *>> DefMap;
362 
363 inline raw_ostream &operator<<(raw_ostream &OS, const PathType &Path) {
364   OS << "< ";
365   for (const BasicBlock *BB : Path) {
366     std::string BBName;
367     if (BB->hasName())
368       raw_string_ostream(BBName) << BB->getName();
369     else
370       raw_string_ostream(BBName) << BB;
371     OS << BBName << " ";
372   }
373   OS << ">";
374   return OS;
375 }
376 
377 /// ThreadingPath is a path in the control flow of a loop that can be threaded
378 /// by cloning necessary basic blocks and replacing conditional branches with
379 /// unconditional ones. A threading path includes a list of basic blocks, the
380 /// exit state, and the block that determines the next state.
381 struct ThreadingPath {
382   /// Exit value is DFA's exit state for the given path.
383   uint64_t getExitValue() const { return ExitVal; }
384   void setExitValue(const ConstantInt *V) {
385     ExitVal = V->getZExtValue();
386     IsExitValSet = true;
387   }
388   bool isExitValueSet() const { return IsExitValSet; }
389 
390   /// Determinator is the basic block that determines the next state of the DFA.
391   const BasicBlock *getDeterminatorBB() const { return DBB; }
392   void setDeterminator(const BasicBlock *BB) { DBB = BB; }
393 
394   /// Path is a list of basic blocks.
395   const PathType &getPath() const { return Path; }
396   void setPath(const PathType &NewPath) { Path = NewPath; }
397 
398   void print(raw_ostream &OS) const {
399     OS << Path << " [ " << ExitVal << ", " << DBB->getName() << " ]";
400   }
401 
402 private:
403   PathType Path;
404   uint64_t ExitVal;
405   const BasicBlock *DBB = nullptr;
406   bool IsExitValSet = false;
407 };
408 
409 #ifndef NDEBUG
410 inline raw_ostream &operator<<(raw_ostream &OS, const ThreadingPath &TPath) {
411   TPath.print(OS);
412   return OS;
413 }
414 #endif
415 
416 struct MainSwitch {
417   MainSwitch(SwitchInst *SI, OptimizationRemarkEmitter *ORE) {
418     if (isPredictable(SI)) {
419       Instr = SI;
420     } else {
421       ORE->emit([&]() {
422         return OptimizationRemarkMissed(DEBUG_TYPE, "SwitchNotPredictable", SI)
423                << "Switch instruction is not predictable.";
424       });
425     }
426   }
427 
428   virtual ~MainSwitch() = default;
429 
430   SwitchInst *getInstr() const { return Instr; }
431   const SmallVector<SelectInstToUnfold, 4> getSelectInsts() {
432     return SelectInsts;
433   }
434 
435 private:
436   /// Do a use-def chain traversal. Make sure the value of the switch variable
437   /// is always a known constant. This means that all conditional jumps based on
438   /// switch variable can be converted to unconditional jumps.
439   bool isPredictable(const SwitchInst *SI) {
440     std::deque<Instruction *> Q;
441     SmallSet<Value *, 16> SeenValues;
442     SelectInsts.clear();
443 
444     Value *FirstDef = SI->getOperand(0);
445     auto *Inst = dyn_cast<Instruction>(FirstDef);
446 
447     // If this is a function argument or another non-instruction, then give up.
448     // We are interested in loop local variables.
449     if (!Inst)
450       return false;
451 
452     // Require the first definition to be a PHINode
453     if (!isa<PHINode>(Inst))
454       return false;
455 
456     LLVM_DEBUG(dbgs() << "\tisPredictable() FirstDef: " << *Inst << "\n");
457 
458     Q.push_back(Inst);
459     SeenValues.insert(FirstDef);
460 
461     while (!Q.empty()) {
462       Instruction *Current = Q.front();
463       Q.pop_front();
464 
465       if (auto *Phi = dyn_cast<PHINode>(Current)) {
466         for (Value *Incoming : Phi->incoming_values()) {
467           if (!isPredictableValue(Incoming, SeenValues))
468             return false;
469           addInstToQueue(Incoming, Q, SeenValues);
470         }
471         LLVM_DEBUG(dbgs() << "\tisPredictable() phi: " << *Phi << "\n");
472       } else if (SelectInst *SelI = dyn_cast<SelectInst>(Current)) {
473         if (!isValidSelectInst(SelI))
474           return false;
475         if (!isPredictableValue(SelI->getTrueValue(), SeenValues) ||
476             !isPredictableValue(SelI->getFalseValue(), SeenValues)) {
477           return false;
478         }
479         addInstToQueue(SelI->getTrueValue(), Q, SeenValues);
480         addInstToQueue(SelI->getFalseValue(), Q, SeenValues);
481         LLVM_DEBUG(dbgs() << "\tisPredictable() select: " << *SelI << "\n");
482         if (auto *SelIUse = dyn_cast<PHINode>(SelI->user_back()))
483           SelectInsts.push_back(SelectInstToUnfold(SelI, SelIUse));
484       } else {
485         // If it is neither a phi nor a select, then we give up.
486         return false;
487       }
488     }
489 
490     return true;
491   }
492 
493   bool isPredictableValue(Value *InpVal, SmallSet<Value *, 16> &SeenValues) {
494     if (SeenValues.find(InpVal) != SeenValues.end())
495       return true;
496 
497     if (isa<ConstantInt>(InpVal))
498       return true;
499 
500     // If this is a function argument or another non-instruction, then give up.
501     if (!isa<Instruction>(InpVal))
502       return false;
503 
504     return true;
505   }
506 
507   void addInstToQueue(Value *Val, std::deque<Instruction *> &Q,
508                       SmallSet<Value *, 16> &SeenValues) {
509     if (SeenValues.find(Val) != SeenValues.end())
510       return;
511     if (Instruction *I = dyn_cast<Instruction>(Val))
512       Q.push_back(I);
513     SeenValues.insert(Val);
514   }
515 
516   bool isValidSelectInst(SelectInst *SI) {
517     if (!SI->hasOneUse())
518       return false;
519 
520     Instruction *SIUse = dyn_cast<Instruction>(SI->user_back());
521     // The use of the select inst should be either a phi or another select.
522     if (!SIUse && !(isa<PHINode>(SIUse) || isa<SelectInst>(SIUse)))
523       return false;
524 
525     BasicBlock *SIBB = SI->getParent();
526 
527     // Currently, we can only expand select instructions in basic blocks with
528     // one successor.
529     BranchInst *SITerm = dyn_cast<BranchInst>(SIBB->getTerminator());
530     if (!SITerm || !SITerm->isUnconditional())
531       return false;
532 
533     if (isa<PHINode>(SIUse) &&
534         SIBB->getSingleSuccessor() != cast<Instruction>(SIUse)->getParent())
535       return false;
536 
537     // If select will not be sunk during unfolding, and it is in the same basic
538     // block as another state defining select, then cannot unfold both.
539     for (SelectInstToUnfold SIToUnfold : SelectInsts) {
540       SelectInst *PrevSI = SIToUnfold.getInst();
541       if (PrevSI->getTrueValue() != SI && PrevSI->getFalseValue() != SI &&
542           PrevSI->getParent() == SI->getParent())
543         return false;
544     }
545 
546     return true;
547   }
548 
549   SwitchInst *Instr = nullptr;
550   SmallVector<SelectInstToUnfold, 4> SelectInsts;
551 };
552 
553 struct AllSwitchPaths {
554   AllSwitchPaths(const MainSwitch *MSwitch, OptimizationRemarkEmitter *ORE)
555       : Switch(MSwitch->getInstr()), SwitchBlock(Switch->getParent()),
556         ORE(ORE) {}
557 
558   std::vector<ThreadingPath> &getThreadingPaths() { return TPaths; }
559   unsigned getNumThreadingPaths() { return TPaths.size(); }
560   SwitchInst *getSwitchInst() { return Switch; }
561   BasicBlock *getSwitchBlock() { return SwitchBlock; }
562 
563   void run() {
564     VisitedBlocks Visited;
565     PathsType LoopPaths = paths(SwitchBlock, Visited, /* PathDepth = */ 1);
566     StateDefMap StateDef = getStateDefMap();
567 
568     for (PathType Path : LoopPaths) {
569       ThreadingPath TPath;
570 
571       const BasicBlock *PrevBB = Path.back();
572       for (const BasicBlock *BB : Path) {
573         if (StateDef.count(BB) != 0) {
574           const PHINode *Phi = dyn_cast<PHINode>(StateDef[BB]);
575           assert(Phi && "Expected a state-defining instr to be a phi node.");
576 
577           const Value *V = Phi->getIncomingValueForBlock(PrevBB);
578           if (const ConstantInt *C = dyn_cast<const ConstantInt>(V)) {
579             TPath.setExitValue(C);
580             TPath.setDeterminator(BB);
581             TPath.setPath(Path);
582           }
583         }
584 
585         // Switch block is the determinator, this is the final exit value.
586         if (TPath.isExitValueSet() && BB == Path.front())
587           break;
588 
589         PrevBB = BB;
590       }
591 
592       if (TPath.isExitValueSet())
593         TPaths.push_back(TPath);
594     }
595   }
596 
597 private:
598   // Value: an instruction that defines a switch state;
599   // Key: the parent basic block of that instruction.
600   typedef DenseMap<const BasicBlock *, const PHINode *> StateDefMap;
601 
602   PathsType paths(BasicBlock *BB, VisitedBlocks &Visited,
603                   unsigned PathDepth) const {
604     PathsType Res;
605 
606     // Stop exploring paths after visiting MaxPathLength blocks
607     if (PathDepth > MaxPathLength) {
608       ORE->emit([&]() {
609         return OptimizationRemarkAnalysis(DEBUG_TYPE, "MaxPathLengthReached",
610                                           Switch)
611                << "Exploration stopped after visiting MaxPathLength="
612                << ore::NV("MaxPathLength", MaxPathLength) << " blocks.";
613       });
614       return Res;
615     }
616 
617     Visited.insert(BB);
618 
619     // Some blocks have multiple edges to the same successor, and this set
620     // is used to prevent a duplicate path from being generated
621     SmallSet<BasicBlock *, 4> Successors;
622     for (BasicBlock *Succ : successors(BB)) {
623       if (!Successors.insert(Succ).second)
624         continue;
625 
626       // Found a cycle through the SwitchBlock
627       if (Succ == SwitchBlock) {
628         Res.push_back({BB});
629         continue;
630       }
631 
632       // We have encountered a cycle, do not get caught in it
633       if (Visited.contains(Succ))
634         continue;
635 
636       PathsType SuccPaths = paths(Succ, Visited, PathDepth + 1);
637       for (PathType Path : SuccPaths) {
638         PathType NewPath(Path);
639         NewPath.push_front(BB);
640         Res.push_back(NewPath);
641       }
642     }
643     // This block could now be visited again from a different predecessor. Note
644     // that this will result in exponential runtime. Subpaths could possibly be
645     // cached but it takes a lot of memory to store them.
646     Visited.erase(BB);
647     return Res;
648   }
649 
650   /// Walk the use-def chain and collect all the state-defining instructions.
651   StateDefMap getStateDefMap() const {
652     StateDefMap Res;
653 
654     Value *FirstDef = Switch->getOperand(0);
655 
656     assert(isa<PHINode>(FirstDef) && "After select unfolding, all state "
657                                      "definitions are expected to be phi "
658                                      "nodes.");
659 
660     SmallVector<PHINode *, 8> Stack;
661     Stack.push_back(dyn_cast<PHINode>(FirstDef));
662     SmallSet<Value *, 16> SeenValues;
663 
664     while (!Stack.empty()) {
665       PHINode *CurPhi = Stack.back();
666       Stack.pop_back();
667 
668       Res[CurPhi->getParent()] = CurPhi;
669       SeenValues.insert(CurPhi);
670 
671       for (Value *Incoming : CurPhi->incoming_values()) {
672         if (Incoming == FirstDef || isa<ConstantInt>(Incoming) ||
673             SeenValues.find(Incoming) != SeenValues.end()) {
674           continue;
675         }
676 
677         assert(isa<PHINode>(Incoming) && "After select unfolding, all state "
678                                          "definitions are expected to be phi "
679                                          "nodes.");
680 
681         Stack.push_back(cast<PHINode>(Incoming));
682       }
683     }
684 
685     return Res;
686   }
687 
688   SwitchInst *Switch;
689   BasicBlock *SwitchBlock;
690   OptimizationRemarkEmitter *ORE;
691   std::vector<ThreadingPath> TPaths;
692 };
693 
694 struct TransformDFA {
695   TransformDFA(AllSwitchPaths *SwitchPaths, DominatorTree *DT,
696                AssumptionCache *AC, TargetTransformInfo *TTI,
697                OptimizationRemarkEmitter *ORE,
698                SmallPtrSet<const Value *, 32> EphValues)
699       : SwitchPaths(SwitchPaths), DT(DT), AC(AC), TTI(TTI), ORE(ORE),
700         EphValues(EphValues) {}
701 
702   void run() {
703     if (isLegalAndProfitableToTransform()) {
704       createAllExitPaths();
705       NumTransforms++;
706     }
707   }
708 
709 private:
710   /// This function performs both a legality check and profitability check at
711   /// the same time since it is convenient to do so. It iterates through all
712   /// blocks that will be cloned, and keeps track of the duplication cost. It
713   /// also returns false if it is illegal to clone some required block.
714   bool isLegalAndProfitableToTransform() {
715     CodeMetrics Metrics;
716     SwitchInst *Switch = SwitchPaths->getSwitchInst();
717 
718     // Note that DuplicateBlockMap is not being used as intended here. It is
719     // just being used to ensure (BB, State) pairs are only counted once.
720     DuplicateBlockMap DuplicateMap;
721 
722     for (ThreadingPath &TPath : SwitchPaths->getThreadingPaths()) {
723       PathType PathBBs = TPath.getPath();
724       uint64_t NextState = TPath.getExitValue();
725       const BasicBlock *Determinator = TPath.getDeterminatorBB();
726 
727       // Update Metrics for the Switch block, this is always cloned
728       BasicBlock *BB = SwitchPaths->getSwitchBlock();
729       BasicBlock *VisitedBB = getClonedBB(BB, NextState, DuplicateMap);
730       if (!VisitedBB) {
731         Metrics.analyzeBasicBlock(BB, *TTI, EphValues);
732         DuplicateMap[BB].push_back({BB, NextState});
733       }
734 
735       // If the Switch block is the Determinator, then we can continue since
736       // this is the only block that is cloned and we already counted for it.
737       if (PathBBs.front() == Determinator)
738         continue;
739 
740       // Otherwise update Metrics for all blocks that will be cloned. If any
741       // block is already cloned and would be reused, don't double count it.
742       auto DetIt = std::find(PathBBs.begin(), PathBBs.end(), Determinator);
743       for (auto BBIt = DetIt; BBIt != PathBBs.end(); BBIt++) {
744         BB = *BBIt;
745         VisitedBB = getClonedBB(BB, NextState, DuplicateMap);
746         if (VisitedBB)
747           continue;
748         Metrics.analyzeBasicBlock(BB, *TTI, EphValues);
749         DuplicateMap[BB].push_back({BB, NextState});
750       }
751 
752       if (Metrics.notDuplicatable) {
753         LLVM_DEBUG(dbgs() << "DFA Jump Threading: Not jump threading, contains "
754                           << "non-duplicatable instructions.\n");
755         ORE->emit([&]() {
756           return OptimizationRemarkMissed(DEBUG_TYPE, "NonDuplicatableInst",
757                                           Switch)
758                  << "Contains non-duplicatable instructions.";
759         });
760         return false;
761       }
762 
763       if (Metrics.convergent) {
764         LLVM_DEBUG(dbgs() << "DFA Jump Threading: Not jump threading, contains "
765                           << "convergent instructions.\n");
766         ORE->emit([&]() {
767           return OptimizationRemarkMissed(DEBUG_TYPE, "ConvergentInst", Switch)
768                  << "Contains convergent instructions.";
769         });
770         return false;
771       }
772     }
773 
774     unsigned DuplicationCost = 0;
775 
776     unsigned JumpTableSize = 0;
777     TTI->getEstimatedNumberOfCaseClusters(*Switch, JumpTableSize, nullptr,
778                                           nullptr);
779     if (JumpTableSize == 0) {
780       // Factor in the number of conditional branches reduced from jump
781       // threading. Assume that lowering the switch block is implemented by
782       // using binary search, hence the LogBase2().
783       unsigned CondBranches =
784           APInt(32, Switch->getNumSuccessors()).ceilLogBase2();
785       DuplicationCost = Metrics.NumInsts / CondBranches;
786     } else {
787       // Compared with jump tables, the DFA optimizer removes an indirect branch
788       // on each loop iteration, thus making branch prediction more precise. The
789       // more branch targets there are, the more likely it is for the branch
790       // predictor to make a mistake, and the more benefit there is in the DFA
791       // optimizer. Thus, the more branch targets there are, the lower is the
792       // cost of the DFA opt.
793       DuplicationCost = Metrics.NumInsts / JumpTableSize;
794     }
795 
796     LLVM_DEBUG(dbgs() << "\nDFA Jump Threading: Cost to jump thread block "
797                       << SwitchPaths->getSwitchBlock()->getName()
798                       << " is: " << DuplicationCost << "\n\n");
799 
800     if (DuplicationCost > CostThreshold) {
801       LLVM_DEBUG(dbgs() << "Not jump threading, duplication cost exceeds the "
802                         << "cost threshold.\n");
803       ORE->emit([&]() {
804         return OptimizationRemarkMissed(DEBUG_TYPE, "NotProfitable", Switch)
805                << "Duplication cost exceeds the cost threshold (cost="
806                << ore::NV("Cost", DuplicationCost)
807                << ", threshold=" << ore::NV("Threshold", CostThreshold) << ").";
808       });
809       return false;
810     }
811 
812     ORE->emit([&]() {
813       return OptimizationRemark(DEBUG_TYPE, "JumpThreaded", Switch)
814              << "Switch statement jump-threaded.";
815     });
816 
817     return true;
818   }
819 
820   /// Transform each threading path to effectively jump thread the DFA.
821   void createAllExitPaths() {
822     DomTreeUpdater DTU(*DT, DomTreeUpdater::UpdateStrategy::Eager);
823 
824     // Move the switch block to the end of the path, since it will be duplicated
825     BasicBlock *SwitchBlock = SwitchPaths->getSwitchBlock();
826     for (ThreadingPath &TPath : SwitchPaths->getThreadingPaths()) {
827       LLVM_DEBUG(dbgs() << TPath << "\n");
828       PathType NewPath(TPath.getPath());
829       NewPath.push_back(SwitchBlock);
830       TPath.setPath(NewPath);
831     }
832 
833     // Transform the ThreadingPaths and keep track of the cloned values
834     DuplicateBlockMap DuplicateMap;
835     DefMap NewDefs;
836 
837     SmallSet<BasicBlock *, 16> BlocksToClean;
838     for (BasicBlock *BB : successors(SwitchBlock))
839       BlocksToClean.insert(BB);
840 
841     for (ThreadingPath &TPath : SwitchPaths->getThreadingPaths()) {
842       createExitPath(NewDefs, TPath, DuplicateMap, BlocksToClean, &DTU);
843       NumPaths++;
844     }
845 
846     // After all paths are cloned, now update the last successor of the cloned
847     // path so it skips over the switch statement
848     for (ThreadingPath &TPath : SwitchPaths->getThreadingPaths())
849       updateLastSuccessor(TPath, DuplicateMap, &DTU);
850 
851     // For each instruction that was cloned and used outside, update its uses
852     updateSSA(NewDefs);
853 
854     // Clean PHI Nodes for the newly created blocks
855     for (BasicBlock *BB : BlocksToClean)
856       cleanPhiNodes(BB);
857   }
858 
859   /// For a specific ThreadingPath \p Path, create an exit path starting from
860   /// the determinator block.
861   ///
862   /// To remember the correct destination, we have to duplicate blocks
863   /// corresponding to each state. Also update the terminating instruction of
864   /// the predecessors, and phis in the successor blocks.
865   void createExitPath(DefMap &NewDefs, ThreadingPath &Path,
866                       DuplicateBlockMap &DuplicateMap,
867                       SmallSet<BasicBlock *, 16> &BlocksToClean,
868                       DomTreeUpdater *DTU) {
869     uint64_t NextState = Path.getExitValue();
870     const BasicBlock *Determinator = Path.getDeterminatorBB();
871     PathType PathBBs = Path.getPath();
872 
873     // Don't select the placeholder block in front
874     if (PathBBs.front() == Determinator)
875       PathBBs.pop_front();
876 
877     auto DetIt = std::find(PathBBs.begin(), PathBBs.end(), Determinator);
878     auto Prev = std::prev(DetIt);
879     BasicBlock *PrevBB = *Prev;
880     for (auto BBIt = DetIt; BBIt != PathBBs.end(); BBIt++) {
881       BasicBlock *BB = *BBIt;
882       BlocksToClean.insert(BB);
883 
884       // We already cloned BB for this NextState, now just update the branch
885       // and continue.
886       BasicBlock *NextBB = getClonedBB(BB, NextState, DuplicateMap);
887       if (NextBB) {
888         updatePredecessor(PrevBB, BB, NextBB, DTU);
889         PrevBB = NextBB;
890         continue;
891       }
892 
893       // Clone the BB and update the successor of Prev to jump to the new block
894       BasicBlock *NewBB = cloneBlockAndUpdatePredecessor(
895           BB, PrevBB, NextState, DuplicateMap, NewDefs, DTU);
896       DuplicateMap[BB].push_back({NewBB, NextState});
897       BlocksToClean.insert(NewBB);
898       PrevBB = NewBB;
899     }
900   }
901 
902   /// Restore SSA form after cloning blocks.
903   ///
904   /// Each cloned block creates new defs for a variable, and the uses need to be
905   /// updated to reflect this. The uses may be replaced with a cloned value, or
906   /// some derived phi instruction. Note that all uses of a value defined in the
907   /// same block were already remapped when cloning the block.
908   void updateSSA(DefMap &NewDefs) {
909     SSAUpdaterBulk SSAUpdate;
910     SmallVector<Use *, 16> UsesToRename;
911 
912     for (auto KV : NewDefs) {
913       Instruction *I = KV.first;
914       BasicBlock *BB = I->getParent();
915       std::vector<Instruction *> Cloned = KV.second;
916 
917       // Scan all uses of this instruction to see if it is used outside of its
918       // block, and if so, record them in UsesToRename.
919       for (Use &U : I->uses()) {
920         Instruction *User = cast<Instruction>(U.getUser());
921         if (PHINode *UserPN = dyn_cast<PHINode>(User)) {
922           if (UserPN->getIncomingBlock(U) == BB)
923             continue;
924         } else if (User->getParent() == BB) {
925           continue;
926         }
927 
928         UsesToRename.push_back(&U);
929       }
930 
931       // If there are no uses outside the block, we're done with this
932       // instruction.
933       if (UsesToRename.empty())
934         continue;
935       LLVM_DEBUG(dbgs() << "DFA-JT: Renaming non-local uses of: " << *I
936                         << "\n");
937 
938       // We found a use of I outside of BB.  Rename all uses of I that are
939       // outside its block to be uses of the appropriate PHI node etc.  See
940       // ValuesInBlocks with the values we know.
941       unsigned VarNum = SSAUpdate.AddVariable(I->getName(), I->getType());
942       SSAUpdate.AddAvailableValue(VarNum, BB, I);
943       for (Instruction *New : Cloned)
944         SSAUpdate.AddAvailableValue(VarNum, New->getParent(), New);
945 
946       while (!UsesToRename.empty())
947         SSAUpdate.AddUse(VarNum, UsesToRename.pop_back_val());
948 
949       LLVM_DEBUG(dbgs() << "\n");
950     }
951     // SSAUpdater handles phi placement and renaming uses with the appropriate
952     // value.
953     SSAUpdate.RewriteAllUses(DT);
954   }
955 
956   /// Clones a basic block, and adds it to the CFG.
957   ///
958   /// This function also includes updating phi nodes in the successors of the
959   /// BB, and remapping uses that were defined locally in the cloned BB.
960   BasicBlock *cloneBlockAndUpdatePredecessor(BasicBlock *BB, BasicBlock *PrevBB,
961                                              uint64_t NextState,
962                                              DuplicateBlockMap &DuplicateMap,
963                                              DefMap &NewDefs,
964                                              DomTreeUpdater *DTU) {
965     ValueToValueMapTy VMap;
966     BasicBlock *NewBB = CloneBasicBlock(
967         BB, VMap, ".jt" + std::to_string(NextState), BB->getParent());
968     NewBB->moveAfter(BB);
969     NumCloned++;
970 
971     for (Instruction &I : *NewBB) {
972       // Do not remap operands of PHINode in case a definition in BB is an
973       // incoming value to a phi in the same block. This incoming value will
974       // be renamed later while restoring SSA.
975       if (isa<PHINode>(&I))
976         continue;
977       RemapInstruction(&I, VMap,
978                        RF_IgnoreMissingLocals | RF_NoModuleLevelChanges);
979       if (AssumeInst *II = dyn_cast<AssumeInst>(&I))
980         AC->registerAssumption(II);
981     }
982 
983     updateSuccessorPhis(BB, NewBB, NextState, VMap, DuplicateMap);
984     updatePredecessor(PrevBB, BB, NewBB, DTU);
985     updateDefMap(NewDefs, VMap);
986 
987     // Add all successors to the DominatorTree
988     SmallPtrSet<BasicBlock *, 4> SuccSet;
989     for (auto *SuccBB : successors(NewBB)) {
990       if (SuccSet.insert(SuccBB).second)
991         DTU->applyUpdates({{DominatorTree::Insert, NewBB, SuccBB}});
992     }
993     SuccSet.clear();
994     return NewBB;
995   }
996 
997   /// Update the phi nodes in BB's successors.
998   ///
999   /// This means creating a new incoming value from NewBB with the new
1000   /// instruction wherever there is an incoming value from BB.
1001   void updateSuccessorPhis(BasicBlock *BB, BasicBlock *ClonedBB,
1002                            uint64_t NextState, ValueToValueMapTy &VMap,
1003                            DuplicateBlockMap &DuplicateMap) {
1004     std::vector<BasicBlock *> BlocksToUpdate;
1005 
1006     // If BB is the last block in the path, we can simply update the one case
1007     // successor that will be reached.
1008     if (BB == SwitchPaths->getSwitchBlock()) {
1009       SwitchInst *Switch = SwitchPaths->getSwitchInst();
1010       BasicBlock *NextCase = getNextCaseSuccessor(Switch, NextState);
1011       BlocksToUpdate.push_back(NextCase);
1012       BasicBlock *ClonedSucc = getClonedBB(NextCase, NextState, DuplicateMap);
1013       if (ClonedSucc)
1014         BlocksToUpdate.push_back(ClonedSucc);
1015     }
1016     // Otherwise update phis in all successors.
1017     else {
1018       for (BasicBlock *Succ : successors(BB)) {
1019         BlocksToUpdate.push_back(Succ);
1020 
1021         // Check if a successor has already been cloned for the particular exit
1022         // value. In this case if a successor was already cloned, the phi nodes
1023         // in the cloned block should be updated directly.
1024         BasicBlock *ClonedSucc = getClonedBB(Succ, NextState, DuplicateMap);
1025         if (ClonedSucc)
1026           BlocksToUpdate.push_back(ClonedSucc);
1027       }
1028     }
1029 
1030     // If there is a phi with an incoming value from BB, create a new incoming
1031     // value for the new predecessor ClonedBB. The value will either be the same
1032     // value from BB or a cloned value.
1033     for (BasicBlock *Succ : BlocksToUpdate) {
1034       for (auto II = Succ->begin(); PHINode *Phi = dyn_cast<PHINode>(II);
1035            ++II) {
1036         Value *Incoming = Phi->getIncomingValueForBlock(BB);
1037         if (Incoming) {
1038           if (isa<Constant>(Incoming)) {
1039             Phi->addIncoming(Incoming, ClonedBB);
1040             continue;
1041           }
1042           Value *ClonedVal = VMap[Incoming];
1043           if (ClonedVal)
1044             Phi->addIncoming(ClonedVal, ClonedBB);
1045           else
1046             Phi->addIncoming(Incoming, ClonedBB);
1047         }
1048       }
1049     }
1050   }
1051 
1052   /// Sets the successor of PrevBB to be NewBB instead of OldBB. Note that all
1053   /// other successors are kept as well.
1054   void updatePredecessor(BasicBlock *PrevBB, BasicBlock *OldBB,
1055                          BasicBlock *NewBB, DomTreeUpdater *DTU) {
1056     // When a path is reused, there is a chance that predecessors were already
1057     // updated before. Check if the predecessor needs to be updated first.
1058     if (!isPredecessor(OldBB, PrevBB))
1059       return;
1060 
1061     Instruction *PrevTerm = PrevBB->getTerminator();
1062     for (unsigned Idx = 0; Idx < PrevTerm->getNumSuccessors(); Idx++) {
1063       if (PrevTerm->getSuccessor(Idx) == OldBB) {
1064         OldBB->removePredecessor(PrevBB, /* KeepOneInputPHIs = */ true);
1065         PrevTerm->setSuccessor(Idx, NewBB);
1066       }
1067     }
1068     DTU->applyUpdates({{DominatorTree::Delete, PrevBB, OldBB},
1069                        {DominatorTree::Insert, PrevBB, NewBB}});
1070   }
1071 
1072   /// Add new value mappings to the DefMap to keep track of all new definitions
1073   /// for a particular instruction. These will be used while updating SSA form.
1074   void updateDefMap(DefMap &NewDefs, ValueToValueMapTy &VMap) {
1075     for (auto Entry : VMap) {
1076       Instruction *Inst =
1077           dyn_cast<Instruction>(const_cast<Value *>(Entry.first));
1078       if (!Inst || !Entry.second || isa<BranchInst>(Inst) ||
1079           isa<SwitchInst>(Inst)) {
1080         continue;
1081       }
1082 
1083       Instruction *Cloned = dyn_cast<Instruction>(Entry.second);
1084       if (!Cloned)
1085         continue;
1086 
1087       if (NewDefs.find(Inst) == NewDefs.end())
1088         NewDefs[Inst] = {Cloned};
1089       else
1090         NewDefs[Inst].push_back(Cloned);
1091     }
1092   }
1093 
1094   /// Update the last branch of a particular cloned path to point to the correct
1095   /// case successor.
1096   ///
1097   /// Note that this is an optional step and would have been done in later
1098   /// optimizations, but it makes the CFG significantly easier to work with.
1099   void updateLastSuccessor(ThreadingPath &TPath,
1100                            DuplicateBlockMap &DuplicateMap,
1101                            DomTreeUpdater *DTU) {
1102     uint64_t NextState = TPath.getExitValue();
1103     BasicBlock *BB = TPath.getPath().back();
1104     BasicBlock *LastBlock = getClonedBB(BB, NextState, DuplicateMap);
1105 
1106     // Note multiple paths can end at the same block so check that it is not
1107     // updated yet
1108     if (!isa<SwitchInst>(LastBlock->getTerminator()))
1109       return;
1110     SwitchInst *Switch = cast<SwitchInst>(LastBlock->getTerminator());
1111     BasicBlock *NextCase = getNextCaseSuccessor(Switch, NextState);
1112 
1113     std::vector<DominatorTree::UpdateType> DTUpdates;
1114     SmallPtrSet<BasicBlock *, 4> SuccSet;
1115     for (BasicBlock *Succ : successors(LastBlock)) {
1116       if (Succ != NextCase && SuccSet.insert(Succ).second)
1117         DTUpdates.push_back({DominatorTree::Delete, LastBlock, Succ});
1118     }
1119 
1120     Switch->eraseFromParent();
1121     BranchInst::Create(NextCase, LastBlock);
1122 
1123     DTU->applyUpdates(DTUpdates);
1124   }
1125 
1126   /// After cloning blocks, some of the phi nodes have extra incoming values
1127   /// that are no longer used. This function removes them.
1128   void cleanPhiNodes(BasicBlock *BB) {
1129     // If BB is no longer reachable, remove any remaining phi nodes
1130     if (pred_empty(BB)) {
1131       std::vector<PHINode *> PhiToRemove;
1132       for (auto II = BB->begin(); PHINode *Phi = dyn_cast<PHINode>(II); ++II) {
1133         PhiToRemove.push_back(Phi);
1134       }
1135       for (PHINode *PN : PhiToRemove) {
1136         PN->replaceAllUsesWith(UndefValue::get(PN->getType()));
1137         PN->eraseFromParent();
1138       }
1139       return;
1140     }
1141 
1142     // Remove any incoming values that come from an invalid predecessor
1143     for (auto II = BB->begin(); PHINode *Phi = dyn_cast<PHINode>(II); ++II) {
1144       std::vector<BasicBlock *> BlocksToRemove;
1145       for (BasicBlock *IncomingBB : Phi->blocks()) {
1146         if (!isPredecessor(BB, IncomingBB))
1147           BlocksToRemove.push_back(IncomingBB);
1148       }
1149       for (BasicBlock *BB : BlocksToRemove)
1150         Phi->removeIncomingValue(BB);
1151     }
1152   }
1153 
1154   /// Checks if BB was already cloned for a particular next state value. If it
1155   /// was then it returns this cloned block, and otherwise null.
1156   BasicBlock *getClonedBB(BasicBlock *BB, uint64_t NextState,
1157                           DuplicateBlockMap &DuplicateMap) {
1158     CloneList ClonedBBs = DuplicateMap[BB];
1159 
1160     // Find an entry in the CloneList with this NextState. If it exists then
1161     // return the corresponding BB
1162     auto It = llvm::find_if(ClonedBBs, [NextState](const ClonedBlock &C) {
1163       return C.State == NextState;
1164     });
1165     return It != ClonedBBs.end() ? (*It).BB : nullptr;
1166   }
1167 
1168   /// Helper to get the successor corresponding to a particular case value for
1169   /// a switch statement.
1170   BasicBlock *getNextCaseSuccessor(SwitchInst *Switch, uint64_t NextState) {
1171     BasicBlock *NextCase = nullptr;
1172     for (auto Case : Switch->cases()) {
1173       if (Case.getCaseValue()->getZExtValue() == NextState) {
1174         NextCase = Case.getCaseSuccessor();
1175         break;
1176       }
1177     }
1178     if (!NextCase)
1179       NextCase = Switch->getDefaultDest();
1180     return NextCase;
1181   }
1182 
1183   /// Returns true if IncomingBB is a predecessor of BB.
1184   bool isPredecessor(BasicBlock *BB, BasicBlock *IncomingBB) {
1185     return llvm::find(predecessors(BB), IncomingBB) != pred_end(BB);
1186   }
1187 
1188   AllSwitchPaths *SwitchPaths;
1189   DominatorTree *DT;
1190   AssumptionCache *AC;
1191   TargetTransformInfo *TTI;
1192   OptimizationRemarkEmitter *ORE;
1193   SmallPtrSet<const Value *, 32> EphValues;
1194   std::vector<ThreadingPath> TPaths;
1195 };
1196 
1197 bool DFAJumpThreading::run(Function &F) {
1198   LLVM_DEBUG(dbgs() << "\nDFA Jump threading: " << F.getName() << "\n");
1199 
1200   if (F.hasOptSize()) {
1201     LLVM_DEBUG(dbgs() << "Skipping due to the 'minsize' attribute\n");
1202     return false;
1203   }
1204 
1205   if (ClViewCfgBefore)
1206     F.viewCFG();
1207 
1208   SmallVector<AllSwitchPaths, 2> ThreadableLoops;
1209   bool MadeChanges = false;
1210 
1211   for (BasicBlock &BB : F) {
1212     auto *SI = dyn_cast<SwitchInst>(BB.getTerminator());
1213     if (!SI)
1214       continue;
1215 
1216     LLVM_DEBUG(dbgs() << "\nCheck if SwitchInst in BB " << BB.getName()
1217                       << " is predictable\n");
1218     MainSwitch Switch(SI, ORE);
1219 
1220     if (!Switch.getInstr())
1221       continue;
1222 
1223     LLVM_DEBUG(dbgs() << "\nSwitchInst in BB " << BB.getName() << " is a "
1224                       << "candidate for jump threading\n");
1225     LLVM_DEBUG(SI->dump());
1226 
1227     unfoldSelectInstrs(DT, Switch.getSelectInsts());
1228     if (!Switch.getSelectInsts().empty())
1229       MadeChanges = true;
1230 
1231     AllSwitchPaths SwitchPaths(&Switch, ORE);
1232     SwitchPaths.run();
1233 
1234     if (SwitchPaths.getNumThreadingPaths() > 0) {
1235       ThreadableLoops.push_back(SwitchPaths);
1236 
1237       // For the time being limit this optimization to occurring once in a
1238       // function since it can change the CFG significantly. This is not a
1239       // strict requirement but it can cause buggy behavior if there is an
1240       // overlap of blocks in different opportunities. There is a lot of room to
1241       // experiment with catching more opportunities here.
1242       break;
1243     }
1244   }
1245 
1246   SmallPtrSet<const Value *, 32> EphValues;
1247   if (ThreadableLoops.size() > 0)
1248     CodeMetrics::collectEphemeralValues(&F, AC, EphValues);
1249 
1250   for (AllSwitchPaths SwitchPaths : ThreadableLoops) {
1251     TransformDFA Transform(&SwitchPaths, DT, AC, TTI, ORE, EphValues);
1252     Transform.run();
1253     MadeChanges = true;
1254   }
1255 
1256 #ifdef EXPENSIVE_CHECKS
1257   assert(DT->verify(DominatorTree::VerificationLevel::Full));
1258   verifyFunction(F, &dbgs());
1259 #endif
1260 
1261   return MadeChanges;
1262 }
1263 
1264 } // end anonymous namespace
1265 
1266 /// Integrate with the new Pass Manager
1267 PreservedAnalyses DFAJumpThreadingPass::run(Function &F,
1268                                             FunctionAnalysisManager &AM) {
1269   AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F);
1270   DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(F);
1271   TargetTransformInfo &TTI = AM.getResult<TargetIRAnalysis>(F);
1272   OptimizationRemarkEmitter ORE(&F);
1273 
1274   if (!DFAJumpThreading(&AC, &DT, &TTI, &ORE).run(F))
1275     return PreservedAnalyses::all();
1276 
1277   PreservedAnalyses PA;
1278   PA.preserve<DominatorTreeAnalysis>();
1279   return PA;
1280 }
1281