[MachineOutliner] Cleanup: move findCandidates out of suffix tree

[paquette]

Summary

Doing some cleanup on the outliner in preparation for some functional changes. This patch moves findCandidates out of the suffix tree. (No test attached since it's just some tidying that doesn't impact the behaviour of the outliner.) Also adds a couple FIXMEs for future cleanup/improvements.

It's much easier to see what's going on when benefit calculations and queries are in the same place. Because the suffix tree ought not to know target-specific details, it makes sense to yank this out of the tree.

[MatzeB]

Looks like NFC to me.

lib/CodeGen/MachineOutliner.cpp

1060–1062

can't you just make this "OUTLINE_FUNCTION_"?

[paquette]

Even better: put findCandidates actually *inside* the MachineOutliner class for consistency.

[paquette]

lib/CodeGen/MachineOutliner.cpp

1060–1062

I think that sounds a little strange, since the function was "outlined"... But it's a character shorter so I'll go ahead and change it.

[MatzeB]

In D35960#823324, @paquette wrote:

Even better: put findCandidates actually *inside* the MachineOutliner class for consistency.

Still LGTM. And if it's just moving code around inside a file you authored, just do it without phabricator review. Post-commit review is fine.

[paquette]

In D35960#823330, @MatzeB wrote:

In D35960#823324, @paquette wrote:

Even better: put findCandidates actually *inside* the MachineOutliner class for consistency.

Still LGTM. And if it's just moving code around inside a file you authored, just do it without phabricator review. Post-commit review is fine.

Awesome, thanks!

[MatzeB]

lib/CodeGen/MachineOutliner.cpp

1060–1062

OUTLINED is fine a made a typo.
My point was to do << "AB" instead of << "A" << "B"

[paquette]

lib/CodeGen/MachineOutliner.cpp

1060–1062

Ah, okay! It's silly that it wasn't that before. I'll fix that.

[MatzeB]

Looks like this was committed in r309334

lib/CodeGen/MachineOutliner.cpp

struct OutlinedFunction {
   /// The number of instructions this function would save.
   unsigned Benefit = 0;
 
   /// \brief Set to true if candidates for this outlined function should be
   /// replaced with tail calls to this OutlinedFunction.
   bool IsTailCall = false;
 
   OutlinedFunction(size_t Name, size_t OccurrenceCount,
-                   const std::vector<unsigned> &Sequence,
-                   unsigned Benefit, bool IsTailCall)
+                   const std::vector<unsigned> &Sequence, unsigned Benefit,
+                   bool IsTailCall)
       : Name(Name), OccurrenceCount(OccurrenceCount), Sequence(Sequence),
-        Benefit(Benefit), IsTailCall(IsTailCall)
-        {}
+        Benefit(Benefit), IsTailCall(IsTailCall) {}
 };
 
 /// Represents an undefined index in the suffix tree.
 const size_t EmptyIdx = -1;
 
 /// A node in a suffix tree which represents a substring or suffix.
 ///
 /// Each node has either no children or at least two children, with the root
struct SuffixTreeNode {
   ///
   /// This has two major purposes in the suffix tree. The first is as a
   /// shortcut in Ukkonen's construction algorithm. One of the things that
   /// Ukkonen's algorithm does to achieve linear-time construction is
   /// keep track of which node the next insert should be at. This makes each
   /// insert O(1), and there are a total of O(N) inserts. The suffix link
   /// helps with inserting children of internal nodes.
   ///
-  /// Say we add a child to an internal node with associated mapping S. The 
+  /// Say we add a child to an internal node with associated mapping S. The
   /// next insertion must be at the node representing S - its first character.
   /// This is given by the way that we iteratively build the tree in Ukkonen's
   /// algorithm. The main idea is to look at the suffixes of each prefix in the
   /// string, starting with the longest suffix of the prefix, and ending with
   /// the shortest. Therefore, if we keep pointers between such nodes, we can
   /// move to the next insertion point in O(1) time. If we don't, then we'd
   /// have to query from the root, which takes O(N) time. This would make the
   /// construction algorithm O(N^2) rather than O(N).

 ///
 /// The suffix tree is implemented using Ukkonen's algorithm for linear-time
 /// suffix tree construction. Ukkonen's algorithm is explained in more detail
 /// in the paper by Esko Ukkonen "On-line construction of suffix trees. The
 /// paper is available at
 ///
 /// https://www.cs.helsinki.fi/u/ukkonen/SuffixT1withFigs.pdf
 class SuffixTree {
-private:
+public:
+  /// Stores each leaf node in the tree.
+  ///
+  /// This is used for finding outlining candidates.
+  std::vector<SuffixTreeNode *> LeafVector;
+
   /// Each element is an integer representing an instruction in the module.
   ArrayRef<unsigned> Str;
 
+private:
   /// Maintains each node in the tree.
   SpecificBumpPtrAllocator<SuffixTreeNode> NodeAllocator;
 
   /// The root of the suffix tree.
   ///
   /// The root represents the empty string. It is maintained by the
   /// \p NodeAllocator like every other node in the tree.
   SuffixTreeNode *Root = nullptr;
 
-  /// Stores each leaf node in the tree.
-  ///
-  /// This is used for finding outlining candidates.
-  std::vector<SuffixTreeNode *> LeafVector;
-
   /// Maintains the end indices of the internal nodes in the tree.
   ///
   /// Each internal node is guaranteed to never have its end index change
   /// during the construction algorithm; however, leaves must be updated at
   /// every step. Therefore, we need to store leaf end indices by reference
   /// to avoid updating O(N) leaves at every step of construction. Thus,
   /// every internal node must be allocated its own end index.
   BumpPtrAllocator InternalEndIdxAllocator;
private:
   /// \param Edge The label on the edge leaving \p Parent to this node.
   ///
   /// \returns A pointer to the allocated leaf node.
   SuffixTreeNode *insertLeaf(SuffixTreeNode &Parent, size_t StartIdx,
                              unsigned Edge) {
 
     assert(StartIdx <= LeafEndIdx && "String can't start after it ends!");
 
-    SuffixTreeNode *N = new (NodeAllocator.Allocate()) SuffixTreeNode(StartIdx, 
-                                                                   &LeafEndIdx,
-                                                                       nullptr,
-                                                                      &Parent);
+    SuffixTreeNode *N = new (NodeAllocator.Allocate())
+        SuffixTreeNode(StartIdx, &LeafEndIdx, nullptr, &Parent);
     Parent.Children[Edge] = N;
 
     return N;
   }
 
   /// Allocate an internal node and add it to the tree.
   ///
   /// \param Parent The parent of this node. Only null when allocating the root.
   /// \param StartIdx The start index of this node's associated string.
   /// \param EndIdx The end index of this node's associated string.
   /// \param Edge The label on the edge leaving \p Parent to this node.
   ///
   /// \returns A pointer to the allocated internal node.
   SuffixTreeNode *insertInternalNode(SuffixTreeNode *Parent, size_t StartIdx,
                                      size_t EndIdx, unsigned Edge) {
 
     assert(StartIdx <= EndIdx && "String can't start after it ends!");
     assert(!(!Parent && StartIdx != EmptyIdx) &&
            "Non-root internal nodes must have parents!");
 
     size_t *E = new (InternalEndIdxAllocator) size_t(EndIdx);
-    SuffixTreeNode *N = new (NodeAllocator.Allocate()) SuffixTreeNode(StartIdx,
-                                                                      E,
-                                                                      Root,
-                                                                      Parent);
+    SuffixTreeNode *N = new (NodeAllocator.Allocate())
+        SuffixTreeNode(StartIdx, E, Root, Parent);
     if (Parent)
       Parent->Children[Edge] = N;
 
     return N;
   }
 
   /// \brief Set the suffix indices of the leaves to the start indices of their
   /// respective suffixes. Also stores each leaf in \p LeafVector at its
   /// respective suffix index.
   ///
   /// \param[in] CurrNode The node currently being visited.
   /// \param CurrIdx The current index of the string being visited.
   void setSuffixIndices(SuffixTreeNode &CurrNode, size_t CurrIdx) {
 
     bool IsLeaf = CurrNode.Children.size() == 0 && !CurrNode.isRoot();
 
     // Store the length of the concatenation of all strings from the root to
     // this node.
     if (!CurrNode.isRoot()) {
       if (CurrNode.ConcatLen == 0)
         CurrNode.ConcatLen = CurrNode.size();
 
       if (CurrNode.Parent)
         CurrNode.ConcatLen += CurrNode.Parent->ConcatLen;
     }
 
     // Traverse the tree depth-first.
     for (auto &ChildPair : CurrNode.Children) {
       assert(ChildPair.second && "Node had a null child!");
-      setSuffixIndices(*ChildPair.second,
-                       CurrIdx + ChildPair.second->size());
+      setSuffixIndices(*ChildPair.second, CurrIdx + ChildPair.second->size());
     }
 
     // Is this node a leaf?
     if (IsLeaf) {
       // If yes, give it a suffix index and bump its parent's occurrence count.
       CurrNode.SuffixIdx = Str.size() - CurrIdx;
       assert(CurrNode.Parent && "CurrNode had no parent!");
       CurrNode.Parent->OccurrenceCount++;
private:
   /// to complete the suffix tree at the current phase.
   ///
   /// \returns The number of suffixes that have not been added at the end of
   /// this step.
   unsigned extend(size_t EndIdx, size_t SuffixesToAdd) {
     SuffixTreeNode *NeedsLink = nullptr;
 
     while (SuffixesToAdd > 0) {
-    
       // Are we waiting to add anything other than just the last character?
       if (Active.Len == 0) {
         // If not, then say the active index is the end index.
         Active.Idx = EndIdx;
       }
 
       assert(Active.Idx <= EndIdx && "Start index can't be after end index!");
 
while (SuffixesToAdd > 0) {
         //
         //   | ABC  ---split--->  | AB
         //   n                    s
         //                     C / \ D
         //                      n   l
 
         // The node s from the diagram
         SuffixTreeNode *SplitNode =
-            insertInternalNode(Active.Node,
-                               NextNode->StartIdx,
-                               NextNode->StartIdx + Active.Len - 1,
-                               FirstChar);
+            insertInternalNode(Active.Node, NextNode->StartIdx,
+                               NextNode->StartIdx + Active.Len - 1, FirstChar);
 
         // Insert the new node representing the new substring into the tree as
         // a child of the split node. This is the node l from the diagram.
         insertLeaf(*SplitNode, EndIdx, LastChar);
 
         // Make the old node a child of the split node and update its start
         // index. This is the node n from the diagram.
         NextNode->StartIdx += Active.Len;
while (SuffixesToAdd > 0) {
         Active.Node = Active.Node->Link;
       }
     }
 
     return SuffixesToAdd;
   }
 
 public:
-
-  /// Find all repeated substrings that satisfy \p BenefitFn.
-  ///
-  /// If a substring appears at least twice, then it must be represented by
-  /// an internal node which appears in at least two suffixes. Each suffix is
-  /// represented by a leaf node. To do this, we visit each internal node in
-  /// the tree, using the leaf children of each internal node. If an internal
-  /// node represents a beneficial substring, then we use each of its leaf
-  /// children to find the locations of its substring.
-  ///
-  /// \param[out] CandidateList Filled with candidates representing each
-  /// beneficial substring.
-  /// \param[out] FunctionList Filled with a list of \p OutlinedFunctions each
-  /// type of candidate.
-  /// \param BenefitFn The function to satisfy.
-  ///
-  /// \returns The length of the longest candidate found.
-  size_t findCandidates(std::vector<Candidate> &CandidateList,
-  std::vector<OutlinedFunction> &FunctionList,
-  const std::function<unsigned(SuffixTreeNode &, size_t, unsigned)>
-  &BenefitFn) {
-
-    CandidateList.clear();
-    FunctionList.clear();
-    size_t FnIdx = 0;
-    size_t MaxLen = 0;
-
-    for (SuffixTreeNode* Leaf : LeafVector) {
-      assert(Leaf && "Leaves in LeafVector cannot be null!");
-      if (!Leaf->IsInTree)
-        continue;
-
-      assert(Leaf->Parent && "All leaves must have parents!");
-      SuffixTreeNode &Parent = *(Leaf->Parent);
-
-      // If it doesn't appear enough, or we already outlined from it, skip it.
-      if (Parent.OccurrenceCount < 2 || Parent.isRoot() || !Parent.IsInTree)
-        continue;
-
-      size_t StringLen = Leaf->ConcatLen - Leaf->size();
-
-      // How many instructions would outlining this string save?
-      unsigned Benefit = BenefitFn(Parent,
-        StringLen, Str[Leaf->SuffixIdx + StringLen - 1]);
-
-      // If it's not beneficial, skip it.
-      if (Benefit < 1)
-        continue;
-
-      if (StringLen > MaxLen)
-        MaxLen = StringLen;
-
-      unsigned OccurrenceCount = 0;
-      for (auto &ChildPair : Parent.Children) {
-        SuffixTreeNode *M = ChildPair.second;
-
-        // Is it a leaf? If so, we have an occurrence of this candidate.
-        if (M && M->IsInTree && M->isLeaf()) {
-          OccurrenceCount++;
-          CandidateList.emplace_back(M->SuffixIdx, StringLen, FnIdx);
-          CandidateList.back().Benefit = Benefit;
-          M->IsInTree = false;
-        }
-      }
-
-      // Save the function for the new candidate sequence.
-      std::vector<unsigned> CandidateSequence;
-      for (unsigned i = Leaf->SuffixIdx; i < Leaf->SuffixIdx + StringLen; i++)
-        CandidateSequence.push_back(Str[i]);
-
-      FunctionList.emplace_back(FnIdx, OccurrenceCount, CandidateSequence,
-                                Benefit, false);
-
-      // Move to the next function.
-      FnIdx++;
-      Parent.IsInTree = false;
-    }
-
-    return MaxLen;
-  }
- 
   /// Construct a suffix tree from a sequence of unsigned integers.
   ///
   /// \param Str The string to construct the suffix tree for.
   SuffixTree(const std::vector<unsigned> &Str) : Str(Str) {
     Root = insertInternalNode(nullptr, EmptyIdx, EmptyIdx, 0);
     Root->IsInTree = true;
     Active.Node = Root;
-    LeafVector = std::vector<SuffixTreeNode*>(Str.size());
+    LeafVector = std::vector<SuffixTreeNode *>(Str.size());
 
     // Keep track of the number of suffixes we have to add of the current
     // prefix.
     size_t SuffixesToAdd = 0;
     Active.Node = Root;
 
     // Construct the suffix tree iteratively on each prefix of the string.
     // PfxEndIdx is the end index of the current prefix.
struct InstructionMapper {
   unsigned mapToLegalUnsigned(MachineBasicBlock::iterator &It) {
 
     // Get the integer for this instruction or give it the current
     // LegalInstrNumber.
     InstrList.push_back(It);
     MachineInstr &MI = *It;
     bool WasInserted;
     DenseMap<MachineInstr *, unsigned, MachineInstrExpressionTrait>::iterator
         ResultIt;
     std::tie(ResultIt, WasInserted) =
         InstructionIntegerMap.insert(std::make_pair(&MI, LegalInstrNumber));
     unsigned MINumber = ResultIt->second;
 
     // There was an insertion.
     if (WasInserted) {
       LegalInstrNumber++;
       IntegerInstructionMap.insert(std::make_pair(MINumber, &MI));
     }
 
     UnsignedVec.push_back(MINumber);
 
     // Make sure we don't overflow or use any integers reserved by the DenseMap.
     if (LegalInstrNumber >= IllegalInstrNumber)
       report_fatal_error("Instruction mapping overflow!");
 
-    assert(LegalInstrNumber != DenseMapInfo<unsigned>::getEmptyKey()
-          && "Tried to assign DenseMap tombstone or empty key to instruction.");
-    assert(LegalInstrNumber != DenseMapInfo<unsigned>::getTombstoneKey()
-          && "Tried to assign DenseMap tombstone or empty key to instruction.");
+    assert(LegalInstrNumber != DenseMapInfo<unsigned>::getEmptyKey() &&
+           "Tried to assign DenseMap tombstone or empty key to instruction.");
+    assert(LegalInstrNumber != DenseMapInfo<unsigned>::getTombstoneKey() &&
+           "Tried to assign DenseMap tombstone or empty key to instruction.");
 
     return MINumber;
   }
 
   /// Maps \p *It to an illegal integer.
   ///
   /// Updates \p InstrList, \p UnsignedVec, and \p IllegalInstrNumber.
   ///
   /// \returns The integer that \p *It was mapped to.
   unsigned mapToIllegalUnsigned(MachineBasicBlock::iterator &It) {
     unsigned MINumber = IllegalInstrNumber;
 
     InstrList.push_back(It);
     UnsignedVec.push_back(IllegalInstrNumber);
     IllegalInstrNumber--;
 
     assert(LegalInstrNumber < IllegalInstrNumber &&
            "Instruction mapping overflow!");
 
-    assert(IllegalInstrNumber !=
-      DenseMapInfo<unsigned>::getEmptyKey() &&
+    assert(IllegalInstrNumber != DenseMapInfo<unsigned>::getEmptyKey() &&
            "IllegalInstrNumber cannot be DenseMap tombstone or empty key!");
 
-    assert(IllegalInstrNumber !=
-      DenseMapInfo<unsigned>::getTombstoneKey() &&
+    assert(IllegalInstrNumber != DenseMapInfo<unsigned>::getTombstoneKey() &&
            "IllegalInstrNumber cannot be DenseMap tombstone or empty key!");
 
     return MINumber;
   }
 
   /// \brief Transforms a \p MachineBasicBlock into a \p vector of \p unsigneds
   /// and appends it to \p UnsignedVec and \p InstrList.
   ///
   /// Two instructions are assigned the same integer if they are identical.
   /// If an instruction is deemed unsafe to outline, then it will be assigned an
   /// unique integer. The resulting mapping is placed into a suffix tree and
   /// queried for candidates.
   ///
   /// \param MBB The \p MachineBasicBlock to be translated into integers.
   /// \param TRI \p TargetRegisterInfo for the module.
   /// \param TII \p TargetInstrInfo for the module.
   void convertToUnsignedVec(MachineBasicBlock &MBB,
                             const TargetRegisterInfo &TRI,
                             const TargetInstrInfo &TII) {
     for (MachineBasicBlock::iterator It = MBB.begin(), Et = MBB.end(); It != Et;
          It++) {
 
       // Keep track of where this instruction is in the module.
-      switch(TII.getOutliningType(*It)) {
-        case TargetInstrInfo::MachineOutlinerInstrType::Illegal:
-          mapToIllegalUnsigned(It);
-          break;
+      switch (TII.getOutliningType(*It)) {
 
-        case TargetInstrInfo::MachineOutlinerInstrType::Legal:
-          mapToLegalUnsigned(It);
-          break;
 
-        case TargetInstrInfo::MachineOutlinerInstrType::Invisible:
-          break;
       }
     }
 
     // After we're done every insertion, uniquely terminate this part of the
     // "string". This makes sure we won't match across basic block or function
     // boundaries since the "end" is encoded uniquely and thus appears in no
     // repeated substring.
     InstrList.push_back(MBB.end());
     UnsignedVec.push_back(IllegalInstrNumber);
     IllegalInstrNumber--;
   }
 
   InstructionMapper() {
     // Make sure that the implementation of DenseMapInfo<unsigned> hasn't
     // changed.
     assert(DenseMapInfo<unsigned>::getEmptyKey() == (unsigned)-1 &&
-                "DenseMapInfo<unsigned>'s empty key isn't -1!");
     assert(DenseMapInfo<unsigned>::getTombstoneKey() == (unsigned)-2 &&
-                "DenseMapInfo<unsigned>'s tombstone key isn't -2!");
   }
 };
 
 /// \brief An interprocedural pass which finds repeated sequences of
 /// instructions and replaces them with calls to functions.
 ///
 /// Each instruction is mapped to an unsigned integer and placed in a string.
 /// The resulting mapping is then placed in a \p SuffixTree. The \p SuffixTree
void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
     ModulePass::getAnalysisUsage(AU);
   }
 
   MachineOutliner() : ModulePass(ID) {
     initializeMachineOutlinerPass(*PassRegistry::getPassRegistry());
   }
 
+  /// Find all repeated substrings that satisfy the outlining cost model.
+  ///
+  /// If a substring appears at least twice, then it must be represented by
+  /// an internal node which appears in at least two suffixes. Each suffix is
+  /// represented by a leaf node. To do this, we visit each internal node in
+  /// the tree, using the leaf children of each internal node. If an internal
+  /// node represents a beneficial substring, then we use each of its leaf
+  /// children to find the locations of its substring.
+  ///
+  /// \param ST A suffix tree to query.
+  /// \param TII TargetInstrInfo for the target.
+  /// \param Mapper Contains outlining mapping information.
+  /// \param[out] CandidateList Filled with candidates representing each
+  /// beneficial substring.
+  /// \param[out] FunctionList Filled with a list of \p OutlinedFunctions each
+  /// type of candidate.
+  ///
+  /// \returns The length of the longest candidate found.
+  size_t findCandidates(SuffixTree &ST, const TargetInstrInfo &TII,
+                        InstructionMapper &Mapper,
+                        std::vector<Candidate> &CandidateList,
+                        std::vector<OutlinedFunction> &FunctionList);
+
   /// \brief Replace the sequences of instructions represented by the
   /// \p Candidates in \p CandidateList with calls to \p MachineFunctions
   /// described in \p FunctionList.
   ///
   /// \param M The module we are outlining from.
   /// \param CandidateList A list of candidates to be outlined.
   /// \param FunctionList A list of functions to be inserted into the module.
   /// \param Mapper Contains the instruction mappings for the module.
struct MachineOutliner : public ModulePass {
   /// \param[out] FunctionList Filled with functions corresponding to each type
   /// of \p Candidate.
   /// \param ST The suffix tree for the module.
   /// \param TII TargetInstrInfo for the module.
   ///
   /// \returns The length of the longest candidate found. 0 if there are none.
   unsigned buildCandidateList(std::vector<Candidate> &CandidateList,
                               std::vector<OutlinedFunction> &FunctionList,
-                              SuffixTree &ST,
-                              InstructionMapper &Mapper,
+                              SuffixTree &ST, InstructionMapper &Mapper,
                               const TargetInstrInfo &TII);
 
   /// \brief Remove any overlapping candidates that weren't handled by the
   /// suffix tree's pruning method.
   ///
   /// Pruning from the suffix tree doesn't necessarily remove all overlaps.
   /// If a short candidate is chosen for outlining, then a longer candidate
   /// which has that short candidate as a suffix is chosen, the tree's pruning
   /// method will not find it. Thus, we need to prune before outlining as well.
   ///
   /// \param[in,out] CandidateList A list of outlining candidates.
   /// \param[in,out] FunctionList A list of functions to be outlined.
   /// \param MaxCandidateLen The length of the longest candidate.
   /// \param TII TargetInstrInfo for the module.
   void pruneOverlaps(std::vector<Candidate> &CandidateList,
                      std::vector<OutlinedFunction> &FunctionList,
-                     unsigned MaxCandidateLen,
-                     const TargetInstrInfo &TII);
+                     unsigned MaxCandidateLen, const TargetInstrInfo &TII);
 
   /// Construct a suffix tree on the instructions in \p M and outline repeated
   /// strings from that tree.
   bool runOnModule(Module &M) override;
 };
 
 } // Anonymous namespace.
 
 char MachineOutliner::ID = 0;
 
 namespace llvm {
 ModulePass *createMachineOutlinerPass() { return new MachineOutliner(); }
+} // namespace llvm
+
+INITIALIZE_PASS(MachineOutliner, DEBUG_TYPE, "Machine Function Outliner", false,
+                false)
+
+size_t
+MachineOutliner::findCandidates(SuffixTree &ST, const TargetInstrInfo &TII,
+                                InstructionMapper &Mapper,
+                                std::vector<Candidate> &CandidateList,
+                                std::vector<OutlinedFunction> &FunctionList) {
+
+  CandidateList.clear();
+  FunctionList.clear();
+  size_t FnIdx = 0;
+  size_t MaxLen = 0;
+
+  // FIXME: Visit internal nodes instead of leaves.
+  for (SuffixTreeNode *Leaf : ST.LeafVector) {
+    assert(Leaf && "Leaves in LeafVector cannot be null!");
+    if (!Leaf->IsInTree)
+      continue;
+
+    assert(Leaf->Parent && "All leaves must have parents!");
+    SuffixTreeNode &Parent = *(Leaf->Parent);
+
+    // If it doesn't appear enough, or we already outlined from it, skip it.
+    if (Parent.OccurrenceCount < 2 || Parent.isRoot() || !Parent.IsInTree)
+      continue;
+
+    // How many instructions would outlining this string save?
+    size_t StringLen = Leaf->ConcatLen - Leaf->size();
+    unsigned EndVal = ST.Str[Leaf->SuffixIdx + StringLen - 1];
+
+    // Determine if this is going to be tail called.
+    // FIXME: The target should decide this. The outlining pass shouldn't care
+    // about things like tail calling. It should be representation-agnostic.
+    MachineInstr *LastInstr = Mapper.IntegerInstructionMap[EndVal];
+    assert(LastInstr && "Last instruction in sequence was unmapped!");
+    bool IsTailCall = LastInstr->isTerminator();
+    unsigned Benefit =
+        TII.getOutliningBenefit(StringLen, Parent.OccurrenceCount, IsTailCall);
+
+    // If it's not beneficial, skip it.
+    if (Benefit < 1)
+      continue;
+
+    if (StringLen > MaxLen)
+      MaxLen = StringLen;
+
+    unsigned OccurrenceCount = 0;
+    for (auto &ChildPair : Parent.Children) {
+      SuffixTreeNode *M = ChildPair.second;
+
+      // Is it a leaf? If so, we have an occurrence of this candidate.
+      if (M && M->IsInTree && M->isLeaf()) {
+        OccurrenceCount++;
+        CandidateList.emplace_back(M->SuffixIdx, StringLen, FnIdx);
+        CandidateList.back().Benefit = Benefit;
+        M->IsInTree = false;
+      }
     }
 
-INITIALIZE_PASS(MachineOutliner, DEBUG_TYPE,
-                "Machine Function Outliner", false, false)
+    // Save the function for the new candidate sequence.
+    std::vector<unsigned> CandidateSequence;
+    for (unsigned i = Leaf->SuffixIdx; i < Leaf->SuffixIdx + StringLen; i++)
+      CandidateSequence.push_back(ST.Str[i]);
+
+    FunctionList.emplace_back(FnIdx, OccurrenceCount, CandidateSequence,
+                              Benefit, false);
+
+    // Move to the next function.
+    FnIdx++;
+    Parent.IsInTree = false;
+  }
+
+  return MaxLen;
+}
 
 void MachineOutliner::pruneOverlaps(std::vector<Candidate> &CandidateList,
                                     std::vector<OutlinedFunction> &FunctionList,
                                     unsigned MaxCandidateLen,
                                     const TargetInstrInfo &TII) {
   // TODO: Experiment with interval trees or other interval-checking structures
   // to lower the time complexity of this function.
   // TODO: Can we do better than the simple greedy choice?
   // Check for overlaps in the range.
   // This is O(MaxCandidateLen * CandidateList.size()).
   for (auto It = CandidateList.begin(), Et = CandidateList.end(); It != Et;
        It++) {
     Candidate &C1 = *It;
     OutlinedFunction &F1 = FunctionList[C1.FunctionIdx];
 
     // If we removed this candidate, skip it.
     if (!C1.InCandidateList)
       continue;
 
     // Is it still worth it to outline C1?
     if (F1.Benefit < 1 || F1.OccurrenceCount < 2) {
       assert(F1.OccurrenceCount > 0 &&
-               "Can't remove OutlinedFunction with no occurrences!");
       F1.OccurrenceCount--;
       C1.InCandidateList = false;
       continue;
     }
 
     // The minimum start index of any candidate that could overlap with this
     // one.
     unsigned FarthestPossibleIdx = 0;
for (auto Sit = It + 1; Sit != Et; Sit++) {
       if (C1.Benefit >= C2.Benefit) {
 
         // C1 is better, so remove C2 and update C2's OutlinedFunction to
         // reflect the removal.
         assert(F2.OccurrenceCount > 0 &&
                "Can't remove OutlinedFunction with no occurrences!");
         F2.OccurrenceCount--;
         F2.Benefit = TII.getOutliningBenefit(F2.Sequence.size(),
-                                             F2.OccurrenceCount,
-                                             F2.IsTailCall
-                                             );
+                                             F2.OccurrenceCount, F2.IsTailCall);
 
         C2.InCandidateList = false;
 
-        DEBUG (
-          dbgs() << "- Removed C2. \n";
-          dbgs() << "--- Num fns left for C2: " << F2.OccurrenceCount << "\n";
-          dbgs() << "--- C2's benefit: " << F2.Benefit << "\n";
-        );
+        DEBUG(dbgs() << "- Removed C2. \n";
+              dbgs() << "--- Num fns left for C2: " << F2.OccurrenceCount
+                     << "\n";
+              dbgs() << "--- C2's benefit: " << F2.Benefit << "\n";);
 
       } else {
         // C2 is better, so remove C1 and update C1's OutlinedFunction to
         // reflect the removal.
         assert(F1.OccurrenceCount > 0 &&
                "Can't remove OutlinedFunction with no occurrences!");
         F1.OccurrenceCount--;
         F1.Benefit = TII.getOutliningBenefit(F1.Sequence.size(),
-                                             F1.OccurrenceCount,
-                                             F1.IsTailCall
-                                             );
+                                             F1.OccurrenceCount, F1.IsTailCall);
         C1.InCandidateList = false;
 
-        DEBUG (
-          dbgs() << "- Removed C1. \n";
-          dbgs() << "--- Num fns left for C1: " << F1.OccurrenceCount << "\n";
-          dbgs() << "--- C1's benefit: " << F1.Benefit << "\n";
-        );
+        DEBUG(dbgs() << "- Removed C1. \n";
+              dbgs() << "--- Num fns left for C1: " << F1.OccurrenceCount
+                     << "\n";
+              dbgs() << "--- C1's benefit: " << F1.Benefit << "\n";);
 
         // C1 is out, so we don't have to compare it against anyone else.
         break;
       }
     }
   }
 }
 
 unsigned
 MachineOutliner::buildCandidateList(std::vector<Candidate> &CandidateList,
                                     std::vector<OutlinedFunction> &FunctionList,
-                                    SuffixTree &ST,
-                                    InstructionMapper &Mapper,
+                                    SuffixTree &ST, InstructionMapper &Mapper,
                                     const TargetInstrInfo &TII) {
 
   std::vector<unsigned> CandidateSequence; // Current outlining candidate.
-  size_t MaxCandidateLen = 0; // Length of the longest candidate.
+  size_t MaxCandidateLen = 0;              // Length of the longest candidate.
 
-  // Function for maximizing query in the suffix tree.
-  // This allows us to define more fine-grained types of things to outline in
-  // the target without putting target-specific info in the suffix tree.
-  auto BenefitFn = [&TII, &Mapper](const SuffixTreeNode &Curr,
-                                   size_t StringLen, unsigned EndVal) {
-
-    // The root represents the empty string.
-    if (Curr.isRoot())
-      return 0u;
-
-    // Is this long enough to outline?
-  // TODO: Let the target decide how "long" a string is in terms of the sizes
-  // of the instructions in the string. For example, if a call instruction
-  // is smaller than a one instruction string, we should outline that string.
-    if (StringLen < 2)
-      return 0u;
-
-    size_t Occurrences = Curr.OccurrenceCount;
-
-    // Anything we want to outline has to appear at least twice.
-    if (Occurrences < 2)
-      return 0u;
-
-    // Check if the last instruction in the sequence is a return.
-    MachineInstr *LastInstr =
-    Mapper.IntegerInstructionMap[EndVal];
-    assert(LastInstr && "Last instruction in sequence was unmapped!");
-
-    // The only way a terminator could be mapped as legal is if it was safe to
-    // tail call.
-    bool IsTailCall = LastInstr->isTerminator();
-    return TII.getOutliningBenefit(StringLen, Occurrences, IsTailCall);
-  };
-
-  MaxCandidateLen = ST.findCandidates(CandidateList, FunctionList, BenefitFn);
+  MaxCandidateLen =
+      findCandidates(ST, TII, Mapper, CandidateList, FunctionList);
 
   for (auto &OF : FunctionList)
-    OF.IsTailCall = Mapper.
-                    IntegerInstructionMap[OF.Sequence.back()]->isTerminator();
+    OF.IsTailCall =
+        Mapper.IntegerInstructionMap[OF.Sequence.back()]->isTerminator();
 
   // Sort the candidates in decending order. This will simplify the outlining
   // process when we have to remove the candidates from the mapping by
   // allowing us to cut them out without keeping track of an offset.
   std::stable_sort(CandidateList.begin(), CandidateList.end());
 
   return MaxCandidateLen;
 }
 
 MachineFunction *
 MachineOutliner::createOutlinedFunction(Module &M, const OutlinedFunction &OF,
                                         InstructionMapper &Mapper) {
 
   // Create the function name. This should be unique. For now, just hash the
   // module name and include it in the function name plus the number of this
   // function.
   std::ostringstream NameStream;
-  NameStream << "OUTLINED_FUNCTION" << "_" << OF.Name;
+  NameStream << "OUTLINED_FUNCTION"
+             << "_" << OF.Name;
 
   // Create the function using an IR-level function.
   LLVMContext &C = M.getContext();
   Function *F = dyn_cast<Function>(
       M.getOrInsertFunction(NameStream.str(), Type::getVoidTy(C)));
   assert(F && "Function was null!");
 
   // NOTE: If this is linkonceodr, then we can take advantage of linker deduping
   // which gives us better results when we outline from linkonceodr functions.
for (const Candidate &C : CandidateList) {
     MBB->erase(StartIt, EndIt);
 
     OutlinedSomething = true;
 
     // Statistics.
     NumOutlined++;
   }
 
-  DEBUG (
-    dbgs() << "OutlinedSomething = " << OutlinedSomething << "\n";
-  );
+  DEBUG(dbgs() << "OutlinedSomething = " << OutlinedSomething << "\n";);
 
   return OutlinedSomething;
 }
 
 bool MachineOutliner::runOnModule(Module &M) {
 
   // Is there anything in the module at all?
   if (M.empty())
     return false;
 
   MachineModuleInfo &MMI = getAnalysis<MachineModuleInfo>();
-  const TargetSubtargetInfo &STI = MMI.getOrCreateMachineFunction(*M.begin())
-                                      .getSubtarget();
+  const TargetSubtargetInfo &STI =
+      MMI.getOrCreateMachineFunction(*M.begin()).getSubtarget();
   const TargetRegisterInfo *TRI = STI.getRegisterInfo();
   const TargetInstrInfo *TII = STI.getInstrInfo();
 
   InstructionMapper Mapper;
 
   // Build instruction mappings for each function in the module.
   for (Function &F : M) {
     MachineFunction &MF = MMI.getOrCreateMachineFunction(F);