[Dominators] Change Roots type to SmallVector

[kuhar]

Summary

We can use the template parameter IsPostDom to pick an appropriate SmallVector size to store DomTree roots for dominators and postdominators. Before, the code would always allocate memory with std::vector.

[grosser]

LGTM

[kuhar]

Update the diff to apply cleanly.

include/llvm/Analysis/DominanceFrontier.h

 public:
   using DomSetType = std::set<BlockT *>;                // Dom set for a bb
   using DomSetMapType = std::map<BlockT *, DomSetType>; // Dom set map
 
 protected:
   using BlockTraits = GraphTraits<BlockT *>;
 
   DomSetMapType Frontiers;
-  std::vector<BlockT *> Roots;
+  // Postdominators can have multiple roots.
+  SmallVector<BlockT *, IsPostDom ? 4 : 1> Roots;
   static constexpr bool IsPostDominators = IsPostDom;
 
 public:
   DominanceFrontierBase() = default;
 
   /// getRoots - Return the root blocks of the current CFG.  This may include
   /// multiple blocks if we are computing post dominators.  For forward
   /// dominators, this will always be a single block (the entry node).
-  inline const std::vector<BlockT *> &getRoots() const {
-    return Roots;
-  }
+  const SmallVectorImpl<BlockT *> &getRoots() const { return Roots; }
 
   BlockT *getRoot() const {
     assert(Roots.size() == 1 && "Should always have entry node!");
     return Roots[0];
   }
 
   /// isPostDominator - Returns true if analysis based of postdoms
   bool isPostDominator() const {
private:
   using BlockTraits = GraphTraits<BlockT *>;
 
 public:
   using DomTreeT = DomTreeBase<BlockT>;
   using DomTreeNodeT = DomTreeNodeBase<BlockT>;
   using DomSetType = typename DominanceFrontierBase<BlockT, false>::DomSetType;
 
   void analyze(DomTreeT &DT) {
-    this->Roots = DT.getRoots();
-    assert(this->Roots.size() == 1 &&
+    assert(DT.getRoots().size() == 1 &&
            "Only one entry block for forward domfronts!");
+    this->Roots = {DT.getRoot()};
     calculate(DT, DT[this->Roots[0]]);
   }
 
   const DomSetType &calculate(const DomTreeT &DT, const DomTreeNodeT *Node);
 };
 
 class DominanceFrontier : public ForwardDominanceFrontierBase<BasicBlock> {
 public:

include/llvm/CodeGen/MachineDominanceFrontier.h

  MachineDominanceFrontier &operator=(const MachineDominanceFrontier &) = delete;
 
  static char ID;
 
  MachineDominanceFrontier();
 
  DominanceFrontierBase<MachineBasicBlock, false> &getBase() { return Base; }
 
- inline const std::vector<MachineBasicBlock *> &getRoots() const {
+  const SmallVectorImpl<MachineBasicBlock *> &getRoots() const {
    return Base.getRoots();
   }
 
   MachineBasicBlock *getRoot() const {
     return Base.getRoot();
   }
 
   bool isPostDominator() const {

include/llvm/CodeGen/MachineDominators.h

public:
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   /// getRoots -  Return the root blocks of the current CFG.  This may include
   /// multiple blocks if we are computing post dominators.  For forward
   /// dominators, this will always be a single block (the entry node).
   ///
-  inline const std::vector<MachineBasicBlock*> &getRoots() const {
+  inline const SmallVectorImpl<MachineBasicBlock*> &getRoots() const {
     applySplitCriticalEdges();
     return DT->getRoots();
   }
 
   inline MachineBasicBlock *getRoot() const {
     applySplitCriticalEdges();
     return DT->getRoot();
   }

include/llvm/CodeGen/MachinePostDominators.h

public:
   static char ID;
 
   MachinePostDominatorTree();
 
   ~MachinePostDominatorTree() override;
 
   FunctionPass *createMachinePostDominatorTreePass();
 
-  const std::vector<MachineBasicBlock *> &getRoots() const {
+  const SmallVectorImpl<MachineBasicBlock *> &getRoots() const {
     return DT->getRoots();
   }
 
   MachineDomTreeNode *getRootNode() const {
     return DT->getRootNode();
   }
 
   MachineDomTreeNode *operator[](MachineBasicBlock *BB) const {

include/llvm/Support/GenericDomTree.h

public:
   using NodePtr = NodeT *;
   using ParentPtr = decltype(std::declval<NodeT *>()->getParent());
   static_assert(std::is_pointer<ParentPtr>::value,
                 "Currently NodeT's parent must be a pointer type");
   using ParentType = typename std::remove_pointer<ParentPtr>::type;
   static constexpr bool IsPostDominator = IsPostDom;
 
  protected:
-  std::vector<NodeT *> Roots;
+  // Dominators always have a single root, postdominators can have more.
+  SmallVector<NodeT *, IsPostDom ? 4 : 1> Roots;
 
   using DomTreeNodeMapType =
      DenseMap<NodeT *, std::unique_ptr<DomTreeNodeBase<NodeT>>>;
   DomTreeNodeMapType DomTreeNodes;
   DomTreeNodeBase<NodeT> *RootNode;
   ParentPtr Parent = nullptr;
 
   mutable bool DFSInfoValid = false;
public:
 
   DominatorTreeBase(const DominatorTreeBase &) = delete;
   DominatorTreeBase &operator=(const DominatorTreeBase &) = delete;
 
   /// getRoots - Return the root blocks of the current CFG.  This may include
   /// multiple blocks if we are computing post dominators.  For forward
   /// dominators, this will always be a single block (the entry node).
   ///
-  const std::vector<NodeT *> &getRoots() const { return Roots; }
+  const SmallVectorImpl<NodeT *> &getRoots() const { return Roots; }
 
   /// isPostDominator - Returns true if analysis based of postdoms
   ///
   bool isPostDominator() const { return IsPostDominator; }
 
   /// compare - Return false if the other dominator tree base matches this
   /// dominator tree base. Otherwise return true.
   bool compare(const DominatorTreeBase &Other) const {