[Dominators] Introduce DomTreeNodeTraits to allow customization. (NFC)

This patch introduces DomTreeNodeTraits for customization. Clients can implement
DomTreeNodeTraitsCustom to provide custom ParentPtr, getEntryNode and getParent.
There's also a default specialization if DomTreeNodeTraitsCustom is not implemented,
that assume a Function-like NodeT. This is what is used for the existing DominatorTree
and MachineDominatorTree.

The main motivation for this patch is using DominatorTreeBase across all
regions of a VPlan, see D140513.

Reviewed By: kuhar

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

llvm/include/llvm/Support/GenericDomTree.h

@@ -14,7 +14,8 @@
 ///
 /// Unlike ADT/* graph algorithms, generic dominator tree has more requirements
 /// on the graph's NodeRef. The NodeRef should be a pointer and,
-/// NodeRef->getParent() must return the parent node that is also a pointer.
+/// either NodeRef->getParent() must return the parent node that is also a
+/// pointer or DomTreeNodeTraits needs to be specialized.
 ///
 /// FIXME: Maybe GenericDomTree needs a TreeTraits, instead of GraphTraits.
 ///
@@ -220,6 +221,20 @@ template <typename DomTreeT>
 bool Verify(const DomTreeT &DT, typename DomTreeT::VerificationLevel VL);
 }  // namespace DomTreeBuilder
 
+/// Default DomTreeNode traits for NodeT. The default implementation assume a
+/// Function-like NodeT. Can be specialized to support different node types.
+template <typename NodeT> struct DomTreeNodeTraits {
+  using NodeType = NodeT;
+  using NodePtr = NodeT *;
+  using ParentPtr = decltype(std::declval<NodePtr>()->getParent());
+  static_assert(std::is_pointer<ParentPtr>::value,
+                "Currently NodeT's parent must be a pointer type");
+  using ParentType = std::remove_pointer_t<ParentPtr>;
+
+  static NodeT *getEntryNode(ParentPtr Parent) { return &Parent->front(); }
+  static ParentPtr getParent(NodePtr BB) { return BB->getParent(); }
+};
+
 /// Core dominator tree base class.
 ///
 /// This class is a generic template over graph nodes. It is instantiated for
@@ -229,9 +244,10 @@ class DominatorTreeBase {
  public:
   static_assert(std::is_pointer<typename GraphTraits<NodeT *>::NodeRef>::value,
                 "Currently DominatorTreeBase supports only pointer nodes");
-  using NodeType = NodeT;
-  using NodePtr = NodeT *;
-  using ParentPtr = decltype(std::declval<NodeT *>()->getParent());
+  using NodeTrait = DomTreeNodeTraits<NodeT>;
+  using NodeType = typename NodeTrait::NodeType;
+  using NodePtr = typename NodeTrait::NodePtr;
+  using ParentPtr = typename NodeTrait::ParentPtr;
   static_assert(std::is_pointer<ParentPtr>::value,
                 "Currently NodeT's parent must be a pointer type");
   using ParentType = std::remove_pointer_t<ParentPtr>;
@@ -467,13 +483,14 @@ class DominatorTreeBase {
   /// must have tree nodes.
   NodeT *findNearestCommonDominator(NodeT *A, NodeT *B) const {
     assert(A && B && "Pointers are not valid");
-    assert(A->getParent() == B->getParent() &&
+    assert(NodeTrait::getParent(A) == NodeTrait::getParent(B) &&
            "Two blocks are not in same function");
 
     // If either A or B is a entry block then it is nearest common dominator
     // (for forward-dominators).
     if (!isPostDominator()) {
-      NodeT &Entry = A->getParent()->front();
+      NodeT &Entry =
+          *DomTreeNodeTraits<NodeT>::getEntryNode(NodeTrait::getParent(A));
       if (A == &Entry || B == &Entry)
         return &Entry;
     }
@@ -584,8 +601,8 @@ class DominatorTreeBase {
   void insertEdge(NodeT *From, NodeT *To) {
     assert(From);
     assert(To);
-    assert(From->getParent() == Parent);
-    assert(To->getParent() == Parent);
+    assert(NodeTrait::getParent(From) == Parent);
+    assert(NodeTrait::getParent(To) == Parent);
     DomTreeBuilder::InsertEdge(*this, From, To);
   }
 
@@ -602,8 +619,8 @@ class DominatorTreeBase {
   void deleteEdge(NodeT *From, NodeT *To) {
     assert(From);
     assert(To);
-    assert(From->getParent() == Parent);
-    assert(To->getParent() == Parent);
+    assert(NodeTrait::getParent(From) == Parent);
+    assert(NodeTrait::getParent(To) == Parent);
     DomTreeBuilder::DeleteEdge(*this, From, To);
   }
 

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -2111,12 +2111,6 @@ class VPRegionBlock : public VPBlockBase {
     EntryBlock->setParent(this);
   }
 
-  // FIXME: DominatorTreeBase is doing 'A->getParent()->front()'. 'front' is a
-  // specific interface of llvm::Function, instead of using
-  // GraphTraints::getEntryNode. We should add a new template parameter to
-  // DominatorTreeBase representing the Graph type.
-  VPBlockBase &front() const { return *Entry; }
-
   const VPBlockBase *getExiting() const { return Exiting; }
   VPBlockBase *getExiting() { return Exiting; }
 

llvm/lib/Transforms/Vectorize/VPlanCFG.h

@@ -29,10 +29,12 @@ namespace llvm {
 // successors/predecessors but not to the blocks inside the region.
 
 template <> struct GraphTraits<VPBlockBase *> {
+  using ParentPtrTy = VPRegionBlock *;
   using NodeRef = VPBlockBase *;
   using ChildIteratorType = SmallVectorImpl<VPBlockBase *>::iterator;
 
   static NodeRef getEntryNode(NodeRef N) { return N; }
+  static NodeRef getEntryNode(VPRegionBlock *R) { return R->getEntry(); }
 
   static inline ChildIteratorType child_begin(NodeRef N) {
     return N->getSuccessors().begin();

llvm/lib/Transforms/Vectorize/VPlanDominatorTree.h

@@ -22,6 +22,16 @@
 
 namespace llvm {
 
+template <> struct DomTreeNodeTraits<VPBlockBase> {
+  using NodeType = VPBlockBase;
+  using NodePtr = VPBlockBase *;
+  using ParentPtr = VPRegionBlock *;
+
+  static NodePtr getEntryNode(ParentPtr Parent) { return Parent->getEntry(); }
+  static ParentPtr getParent(VPBlockBase *B) { return B->getParent(); }
+};
+
+///
 /// Template specialization of the standard LLVM dominator tree utility for
 /// VPBlockBases.
 using VPDominatorTree = DomTreeBase<VPBlockBase>;