From 5d6122d1d8c1bd6ee44ffccfd53ea8592541c7c2 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: Thu, 19 Jul 2018 19:43:13 +0900
Subject: [PATCH] [nest] Support 'ret' (#722)

This commit introduces 'ret' methods which allows users to set/get the
return specification.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>
---
 contrib/nest/include/nest/Module.h | 12 ++++++++++++
 contrib/nest/src/Module.cpp        | 16 ++++++++++++++++
 contrib/nest/src/Module.test.cpp   | 16 ++++++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/contrib/nest/include/nest/Module.h b/contrib/nest/include/nest/Module.h
index e8af83b..7e4f1f3 100644
--- a/contrib/nest/include/nest/Module.h
+++ b/contrib/nest/include/nest/Module.h
@@ -3,6 +3,7 @@
 
 #include "nest/VarContext.h"
 #include "nest/DomainContext.h"
+#include "nest/Ret.h"
 #include "nest/Block.h"
 
 namespace nest
@@ -39,6 +40,17 @@ public:
 
 public:
   void push(const Expr &expr);
+
+private:
+  std::shared_ptr<Ret> _ret;
+
+public:
+  // NOTE Do NOT invoke ret() before ret(expr) call
+  const Ret &ret(void) const;
+
+public:
+  // NOTE Only one ret(expr) call is allowed for each module
+  void ret(const Closure &closure);
 };
 
 } // namespace nest
diff --git a/contrib/nest/src/Module.cpp b/contrib/nest/src/Module.cpp
index 18c953b..acfb6aa 100644
--- a/contrib/nest/src/Module.cpp
+++ b/contrib/nest/src/Module.cpp
@@ -1,5 +1,7 @@
 #include "nest/Module.h"
 
+#include <cassert>
+
 namespace nest
 {
 
@@ -9,4 +11,18 @@ void Module::push(const Expr &expr)
   _block.append(stmt);
 }
 
+void Module::ret(const Closure &clo)
+{
+  // Only one RET is allowed for each module
+  assert(_ret == nullptr);
+  _ret = std::make_shared<Ret>(clo.id(), clo.sub());
+}
+
+const Ret &Module::ret(void) const
+{
+  // Caller should NOT invoke this method before setting ret
+  assert(_ret != nullptr);
+  return *_ret;
+}
+
 } // namespace nest
diff --git a/contrib/nest/src/Module.test.cpp b/contrib/nest/src/Module.test.cpp
index 5def2c7..975ee12 100644
--- a/contrib/nest/src/Module.test.cpp
+++ b/contrib/nest/src/Module.test.cpp
@@ -53,3 +53,19 @@ TEST(MODULE, push)
   ASSERT_EQ(m.block().size(), 1);
   ASSERT_NE(m.block().at(0)->asPush(), nullptr);
 }
+
+TEST(MODULE, ret)
+{
+  nest::Module m;
+
+  auto ifm = m.domain().make({1});
+  auto ofm = m.domain().make({1});
+
+  auto ind = m.var().make();
+
+  m.push(ifm(ind));
+  m.ret(ofm(ind));
+
+  ASSERT_EQ(m.ret().id(), ofm.id());
+  ASSERT_EQ(m.ret().sub().rank(), 1);
+}
-- 
2.7.4