From 3f3b3f3bc3dc2a9a106c52df3a709a2d90513abe Mon Sep 17 00:00:00 2001
From: =?utf8?q?=EB=B0=95=EC=A2=85=ED=98=84/=EB=8F=99=EC=9E=91=EC=A0=9C?=
 =?utf8?q?=EC=96=B4Lab=28SR=29/Staff=20Engineer/=EC=82=BC=EC=84=B1?=
 =?utf8?q?=EC=A0=84=EC=9E=90?= <jh1302.park@samsung.com>
Date: Wed, 5 Dec 2018 12:41:26 +0900
Subject: [PATCH] [coco] Introduce OpManager::destroy_all (#2498)

With this commit, OpManager now allows users to destroy a Op tree via
destroy_all method.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>
---
 contrib/coco/core/include/coco/IR/OpManager.h |  5 +++++
 contrib/coco/core/src/IR/OpManager.cpp        | 26 ++++++++++++++++++++++++++
 contrib/coco/core/src/IR/OpManager.test.cpp   | 13 +++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/contrib/coco/core/include/coco/IR/OpManager.h b/contrib/coco/core/include/coco/IR/OpManager.h
index 95ac849..95769f9 100644
--- a/contrib/coco/core/include/coco/IR/OpManager.h
+++ b/contrib/coco/core/include/coco/IR/OpManager.h
@@ -44,6 +44,11 @@ public:
   //
   // NOTE destroy(op) WILL NOT update op->parent(). Client SHOULD detach op before destroy(op) call
   void destroy(Op *);
+
+  // @brief Destroy a Op tree
+  //
+  // @require op->parent() == nullptr && op->up() == nullptr
+  void destroy_all(Op *);
 };
 
 } // namespace coco
diff --git a/contrib/coco/core/src/IR/OpManager.cpp b/contrib/coco/core/src/IR/OpManager.cpp
index 5f8965c..ce94b2d 100644
--- a/contrib/coco/core/src/IR/OpManager.cpp
+++ b/contrib/coco/core/src/IR/OpManager.cpp
@@ -19,6 +19,7 @@
 #include <nncc/foundation/Memory.h>
 
 #include <cassert>
+#include <queue>
 
 using nncc::foundation::make_unique;
 
@@ -44,4 +45,29 @@ void OpManager::destroy(Op *op)
   release(op);
 }
 
+void OpManager::destroy_all(Op *op)
+{
+  assert(op->parent() == nullptr);
+  assert(op->up() == nullptr);
+
+  std::queue<coco::Op *> q;
+
+  q.emplace(op);
+
+  while (q.size() > 0)
+  {
+    auto cur = q.front();
+    q.pop();
+
+    // Insert child op nodes
+    for (uint32_t n = 0; n < cur->arity(); ++n)
+    {
+      q.emplace(cur->arg(n));
+    }
+
+    // Destroy the current op node
+    destroy(cur);
+  }
+}
+
 } // namespace coco
diff --git a/contrib/coco/core/src/IR/OpManager.test.cpp b/contrib/coco/core/src/IR/OpManager.test.cpp
index daaea4c..3afc2da 100644
--- a/contrib/coco/core/src/IR/OpManager.test.cpp
+++ b/contrib/coco/core/src/IR/OpManager.test.cpp
@@ -74,3 +74,16 @@ TEST_F(OpManagerTest, destroy)
   mgr.destroy(op);
   ASSERT_EQ(mgr.size(), 0);
 }
+
+TEST_F(OpManagerTest, destroy_all)
+{
+  // Create a Op tree
+  auto load_op = mgr.create<coco::Load>();
+  auto conv_op = mgr.create<coco::Conv2D>();
+
+  conv_op->arg(load_op);
+
+  mgr.destroy_all(conv_op);
+
+  ASSERT_EQ(mgr.size(), 0);
+}
-- 
2.7.4