TreeNode refactor

apache · Jul 8, 2019 · 4f1381b · 4f1381b
1 parent 3d88c70
commit 4f1381b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc
@@ -571,7 +571,9 @@ struct TagCompare : ConditionNode {
   ~TagCompare() {}
 };
 
-using TreeNodePtr = relay::TreeNodePtr;
+using TreeNodePtr = typename relay::TreeNode<ConditionNodePtr>::pointer;
+using TreeLeafNode = relay::TreeLeafNode<ConditionNodePtr>;
+using TreeLeafFatalNode = relay::TreeLeafFatalNode<ConditionNodePtr>;
 using TreeBranchNode = relay::TreeBranchNode<ConditionNodePtr>;
 
 void CompileTreeNode(TreeNodePtr tree, VMCompiler* compiler) {

diff --git a/src/relay/pass/pass_util.h b/src/relay/pass/pass_util.h
@@ -109,12 +109,16 @@ inline bool IsAtomic(const Expr& e) {
   return e.as<VarNode>() || e.as<OpNode>() || e.as<ConstructorNode>() || e.as<GlobalVarNode>();
 }
 
+template<typename ConditionNodePtr>
 struct TreeNode {
+  typedef std::shared_ptr<TreeNode<ConditionNodePtr>> pointer;
   virtual ~TreeNode() {}
 };
-using TreeNodePtr = std::shared_ptr<TreeNode>;
 
-struct TreeLeafNode : TreeNode {
+template<typename ConditionNodePtr>
+struct TreeLeafNode : TreeNode<ConditionNodePtr> {
+  using TreeNodePtr = typename TreeNode<ConditionNodePtr>::pointer;
+
   Expr body;
 
   explicit TreeLeafNode(Expr body): body(body) {}
@@ -126,7 +130,10 @@ struct TreeLeafNode : TreeNode {
   ~TreeLeafNode() {}
 };
 
-struct TreeLeafFatalNode : TreeNode {
+template<typename ConditionNodePtr>
+struct TreeLeafFatalNode : TreeNode<ConditionNodePtr> {
+  using TreeNodePtr = typename TreeNode<ConditionNodePtr>::pointer;
+
   TreeLeafFatalNode() = default;
 
   static TreeNodePtr Make() {
@@ -137,7 +144,9 @@ struct TreeLeafFatalNode : TreeNode {
 };
 
 template<typename ConditionNodePtr>
-struct TreeBranchNode : TreeNode {
+struct TreeBranchNode : TreeNode<ConditionNodePtr> {
+  using TreeNodePtr = typename TreeNode<ConditionNodePtr>::pointer;
+
   ConditionNodePtr cond;
   TreeNodePtr then_branch;
   TreeNodePtr else_branch;