From 49c893a41423171d59cba3579c6b652310333018 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BD=AD=E4=BA=8E=E6=96=8C?= <1931127624@qq.com>
Date: Sat, 28 Mar 2020 23:08:42 +0800
Subject: [PATCH] [ir] Basic function definition/call instructions (#612)

* [skip ci] add FuncBodyStmt

* add FrontendFuncDefStmt & begin_func/end_func

* function using scope_stack & real func_call

* [skip ci] enforce code format

* Update TaichiCore.cmake

Co-authored-by: Taichi Gardener <taichigardener@gmail.com>
Co-authored-by: Yuanming Hu <yuanming-hu@users.noreply.github.com>
---
 taichi/inc/statements.inc.h      |  3 +++
 taichi/ir/ir.h                   | 45 ++++++++++++++++++++++++++++++++
 taichi/python/export_lang.cpp    | 15 +++++++++++
 taichi/transforms/ir_printer.cpp | 16 ++++++++++++
 4 files changed, 79 insertions(+)

diff --git a/taichi/inc/statements.inc.h b/taichi/inc/statements.inc.h
index 8d704689d181d2..b31bececc37d75 100644
--- a/taichi/inc/statements.inc.h
+++ b/taichi/inc/statements.inc.h
@@ -20,6 +20,9 @@ PER_STATEMENT(StructForStmt)
 PER_STATEMENT(IfStmt)
 PER_STATEMENT(WhileStmt)
 PER_STATEMENT(WhileControlStmt)
+PER_STATEMENT(FuncBodyStmt)
+PER_STATEMENT(FrontendFuncDefStmt)
+PER_STATEMENT(FuncCallStmt)
 
 PER_STATEMENT(ArgLoadStmt)
 PER_STATEMENT(ExternalPtrStmt)
diff --git a/taichi/ir/ir.h b/taichi/ir/ir.h
index 181ed7aa836c63..5bbfc607c77b68 100644
--- a/taichi/ir/ir.h
+++ b/taichi/ir/ir.h
@@ -1865,6 +1865,51 @@ class StructForStmt : public Stmt {
   DEFINE_ACCEPT
 };
 
+class FuncBodyStmt : public Stmt {
+ public:
+  std::string funcid;
+  std::unique_ptr<Block> body;
+
+  FuncBodyStmt(const std::string &funcid, std::unique_ptr<Block> &&body)
+      : funcid(funcid), body(std::move(body)) {
+  }
+
+  bool is_container_statement() const override {
+    return true;
+  }
+
+  DEFINE_ACCEPT
+};
+
+class FrontendFuncDefStmt : public Stmt {
+ public:
+  std::string funcid;
+  std::unique_ptr<Block> body;
+
+  FrontendFuncDefStmt(const std::string &funcid) : funcid(funcid) {
+  }
+
+  bool is_container_statement() const override {
+    return true;
+  }
+
+  DEFINE_ACCEPT
+};
+
+class FuncCallStmt : public Stmt {
+ public:
+  std::string funcid;
+
+  FuncCallStmt(const std::string &funcid) : funcid(funcid) {
+  }
+
+  bool is_container_statement() const override {
+    return false;
+  }
+
+  DEFINE_ACCEPT
+};
+
 class WhileStmt : public Stmt {
  public:
   Stmt *mask;
diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp
index 35b40b08172a70..608080b16c880f 100644
--- a/taichi/python/export_lang.cpp
+++ b/taichi/python/export_lang.cpp
@@ -297,6 +297,21 @@ void export_lang(py::module &m) {
     current_ast_builder().insert(Stmt::make<FrontendBreakStmt>());
   });
 
+  m.def("begin_func", [&](const std::string &funcid) {
+    auto stmt_unique = std::make_unique<FrontendFuncDefStmt>(funcid);
+    auto stmt = stmt_unique.get();
+    current_ast_builder().insert(std::move(stmt_unique));
+    scope_stack.push_back(current_ast_builder().create_scope(stmt->body));
+  });
+
+  m.def("end_func", [&](const std::string &funcid) { scope_stack.pop_back(); });
+
+  m.def("func_call", [&](const std::string &funcid) {
+    auto func = Stmt::make<FuncCallStmt>(
+        funcid);  // TODO: use FuncCallExpr with return values & args
+    current_ast_builder().insert(std::move(func));
+  });
+
   m.def("layout", layout);
 
   m.def("value_cast", static_cast<Expr (*)(const Expr &expr, DataType)>(cast));
diff --git a/taichi/transforms/ir_printer.cpp b/taichi/transforms/ir_printer.cpp
index 941d93cc2b477f..4212fef231d494 100644
--- a/taichi/transforms/ir_printer.cpp
+++ b/taichi/transforms/ir_printer.cpp
@@ -195,6 +195,22 @@ class IRPrinter : public IRVisitor {
     print("while control {}, {}", stmt->mask->name(), stmt->cond->name());
   }
 
+  void visit(FuncCallStmt *stmt) override {
+    print("{}{} = call \"{}\"", stmt->type_hint(), stmt->name(), stmt->funcid);
+  }
+
+  void visit(FrontendFuncDefStmt *stmt) override {
+    print("function \"{}\" {{", stmt->funcid);
+    stmt->body->accept(this);
+    print("}}");
+  }
+
+  void visit(FuncBodyStmt *stmt) override {
+    print("func \"{}\" {{");
+    stmt->body->accept(this);
+    print("}}");
+  }
+
   void visit(WhileStmt *stmt) override {
     print("while true {{");
     stmt->body->accept(this);