ScheduleTree: introduce virtual "clone" method

The original implementation of schedule trees with "element" classes was
relying on a manual macro-based dynamic dispatch to copy schedule tree
nodes.  Arguably, this is not desirable in the C++ code where the same
behavior can be achieved through language mechanisms (virtual functions
or CRTP).  Since copy constructors cannot be overloaded, introduce an
abstract "clone" method to ScheduleTree and override it in all
subclasses to produce copies of the given node wrapped in a unique
pointer to the base class.  Note that children of the given node are not
copied to remain consistent with the existing uses.  Replace macro-based
dispatch with a call to the virtual function.

Using language featuers for dispatch is less error-prone and converts
run-time errors (e.g., unhandled type in macro-based if/else sequence)
into compile-time errors (e.g., attempting to instantiate an abstract
class because the abstract function was not implemented).  It also
localizes the changes to be made to the ScheduleTree API when
introducing or removing a new node type.

Author: ftynse

tc/core/polyhedral/schedule_tree.cc

@@ -122,31 +122,6 @@ vector<ScheduleTree*> ancestorsInSubTree(
   }
   return res;
 }
-
-static std::unique_ptr<ScheduleTree> makeElem(const ScheduleTree& st) {
-#define ELEM_MAKE_CASE_CSTR(CLASS)                      \
-  else if (st.type_ == CLASS::NodeType) {               \
-    return CLASS::make(static_cast<const CLASS*>(&st)); \
-  }
-
-  if (st.type_ == detail::ScheduleTreeType::None) {
-    LOG(FATAL) << "Hit Error node!";
-  }
-  ELEM_MAKE_CASE_CSTR(ScheduleTreeBand)
-  ELEM_MAKE_CASE_CSTR(ScheduleTreeContext)
-  ELEM_MAKE_CASE_CSTR(ScheduleTreeDomain)
-  ELEM_MAKE_CASE_CSTR(ScheduleTreeExtension)
-  ELEM_MAKE_CASE_CSTR(ScheduleTreeFilter)
-  ELEM_MAKE_CASE_CSTR(ScheduleTreeMapping)
-  ELEM_MAKE_CASE_CSTR(ScheduleTreeSequence)
-  ELEM_MAKE_CASE_CSTR(ScheduleTreeSet)
-  ELEM_MAKE_CASE_CSTR(ScheduleTreeThreadSpecificMarker)
-
-#undef ELEM_MAKE_CASE_CSTR
-
-  LOG(FATAL) << "NYI: ScheduleTree from type: " << static_cast<int>(st.type_);
-  return nullptr;
-}
 } // namespace
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -163,7 +138,7 @@ ScheduleTree::ScheduleTree(const ScheduleTree& st)
 }
 
 ScheduleTreeUPtr ScheduleTree::makeScheduleTree(const ScheduleTree& tree) {
-  return makeElem(tree);
+  return tree.clone();
 }
 
 ScheduleTree* ScheduleTree::child(const vector<size_t>& positions) {

tc/core/polyhedral/schedule_tree.h

@@ -465,6 +465,10 @@ struct ScheduleTree {
   // Note that this function does _not_ output the child trees.
   virtual std::ostream& write(std::ostream& os) const = 0;
 
+  // Clone the current node.
+  // Note that this function does _not_ clone the child trees.
+  virtual ScheduleTreeUPtr clone() const = 0;
+
   //
   // Data members
   //

tc/core/polyhedral/schedule_tree_elem.h

@@ -58,6 +58,9 @@ struct ScheduleTreeContext : public ScheduleTree {
   }
 
   virtual std::ostream& write(std::ostream& os) const override;
+  virtual ScheduleTreeUPtr clone() const override {
+    return make(this);
+  }
 
  public:
   isl::set context_;
@@ -91,6 +94,9 @@ struct ScheduleTreeDomain : public ScheduleTree {
   }
 
   virtual std::ostream& write(std::ostream& os) const override;
+  virtual ScheduleTreeUPtr clone() const override {
+    return make(this);
+  }
 
  public:
   isl::union_set domain_;
@@ -124,6 +130,9 @@ struct ScheduleTreeExtension : public ScheduleTree {
   }
 
   virtual std::ostream& write(std::ostream& os) const override;
+  virtual ScheduleTreeUPtr clone() const override {
+    return make(this);
+  }
 
  public:
   isl::union_map extension_;
@@ -157,6 +166,9 @@ struct ScheduleTreeFilter : public ScheduleTree {
       std::vector<ScheduleTreeUPtr>&& children = {});
 
   virtual std::ostream& write(std::ostream& os) const override;
+  virtual ScheduleTreeUPtr clone() const override {
+    return make(this);
+  }
 
  public:
   isl::union_set filter_;
@@ -195,6 +207,9 @@ struct ScheduleTreeMapping : public ScheduleTree {
       std::vector<ScheduleTreeUPtr>&& children = {});
 
   virtual std::ostream& write(std::ostream& os) const override;
+  virtual ScheduleTreeUPtr clone() const override {
+    return make(this);
+  }
 
  public:
   // Mapping from identifiers to affine functions on domain elements.
@@ -229,6 +244,9 @@ struct ScheduleTreeSequence : public ScheduleTree {
       std::vector<ScheduleTreeUPtr>&& children = {});
 
   virtual std::ostream& write(std::ostream& os) const override;
+  virtual ScheduleTreeUPtr clone() const override {
+    return make(this);
+  }
 };
 
 struct ScheduleTreeSet : public ScheduleTree {
@@ -256,6 +274,9 @@ struct ScheduleTreeSet : public ScheduleTree {
       std::vector<ScheduleTreeUPtr>&& children = {});
 
   virtual std::ostream& write(std::ostream& os) const override;
+  virtual ScheduleTreeUPtr clone() const override {
+    return make(this);
+  }
 };
 
 struct ScheduleTreeBand : public ScheduleTree {
@@ -280,6 +301,9 @@ struct ScheduleTreeBand : public ScheduleTree {
   }
 
   virtual std::ostream& write(std::ostream& os) const override;
+  virtual ScheduleTreeUPtr clone() const override {
+    return make(this);
+  }
 
   // Make a schedule node band from partial schedule.
   // Replace "mupa" by its greatest integer part to ensure that the
@@ -353,6 +377,9 @@ struct ScheduleTreeThreadSpecificMarker : public ScheduleTree {
       std::vector<ScheduleTreeUPtr>&& children = {});
 
   virtual std::ostream& write(std::ostream& os) const override;
+  virtual ScheduleTreeUPtr clone() const override {
+    return make(this);
+  }
 };
 
 bool elemEquals(