Remove Expression::serialize and add ExpressionHumanFriendlyPrinter

taichi-dev · Mar 29, 2022 · 2a84d1a · 2a84d1a
1 parent 9b887cb
commit 2a84d1a
Show file tree

Hide file tree

Showing 12 changed files with 411 additions and 278 deletions.
diff --git a/taichi/inc/expressions.inc.h b/taichi/inc/expressions.inc.h
@@ -0,0 +1,21 @@
+PER_EXPRESSION(ArgLoadExpression)
+PER_EXPRESSION(RandExpression)
+PER_EXPRESSION(UnaryOpExpression)
+PER_EXPRESSION(BinaryOpExpression)
+PER_EXPRESSION(TernaryOpExpression)
+PER_EXPRESSION(InternalFuncCallExpression)
+PER_EXPRESSION(ExternalTensorExpression)
+PER_EXPRESSION(GlobalVariableExpression)
+PER_EXPRESSION(GlobalPtrExpression)
+PER_EXPRESSION(TensorElementExpression)
+PER_EXPRESSION(RangeAssumptionExpression)
+PER_EXPRESSION(LoopUniqueExpression)
+PER_EXPRESSION(IdExpression)
+PER_EXPRESSION(AtomicOpExpression)
+PER_EXPRESSION(SNodeOpExpression)
+PER_EXPRESSION(ConstExpression)
+PER_EXPRESSION(ExternalTensorShapeAlongAxisExpression)
+PER_EXPRESSION(FuncCallExpression)
+PER_EXPRESSION(MeshPatchIndexExpression)
+PER_EXPRESSION(MeshRelationAccessExpression)
+PER_EXPRESSION(MeshIndexConversionExpression)
diff --git a/taichi/ir/expr.cpp b/taichi/ir/expr.cpp
@@ -6,17 +6,6 @@
 
 TLANG_NAMESPACE_BEGIN
 
-void Expr::serialize(std::ostream &ss) const {
-  TI_ASSERT(expr);
-  expr->serialize(ss);
-}
-
-std::string Expr::serialize() const {
-  std::stringstream ss;
-  serialize(ss);
-  return ss.str();
-}
-
 void Expr::set_tb(const std::string &tb) {
   expr->tb = tb;
 }

diff --git a/taichi/ir/expr.h b/taichi/ir/expr.h
@@ -83,9 +83,6 @@ class Expr {
 
   Expr operator[](const ExprGroup &indices) const;
 
-  std::string serialize() const;
-  void serialize(std::ostream &ss) const;
-
   Expr operator!();
 
   template <typename T, typename... Args>

diff --git a/taichi/ir/expression.cpp b/taichi/ir/expression.cpp
@@ -11,20 +11,5 @@ std::string Expression::get_attribute(const std::string &key) const {
   }
 }
 
-void ExprGroup::serialize(std::ostream &ss) const {
-  for (int i = 0; i < (int)exprs.size(); i++) {
-    exprs[i].serialize(ss);
-    if (i + 1 < (int)exprs.size()) {
-      ss << ", ";
-    }
-  }
-}
-
-std::string ExprGroup::serialize() const {
-  std::stringstream ss;
-  serialize(ss);
-  return ss.str();
-}
-
 }  // namespace lang
 }  // namespace taichi
diff --git a/taichi/ir/expression.h b/taichi/ir/expression.h
@@ -7,6 +7,8 @@
 
 TLANG_NAMESPACE_BEGIN
 
+class ExpressionVisitor;
+
 // always a tree - used as rvalues
 class Expression {
  public:
@@ -42,7 +44,7 @@ class Expression {
     // implemented
   }
 
-  virtual void serialize(std::ostream &ss) = 0;
+  virtual void accept(ExpressionVisitor *visitor) = 0;
 
   virtual void flatten(FlattenContext *ctx) {
     TI_NOT_IMPLEMENTED;
@@ -110,10 +112,6 @@ class ExprGroup {
   Expr &operator[](int i) {
     return exprs[i];
   }
-
-  void serialize(std::ostream &ss) const;
-
-  std::string serialize() const;
 };
 
 inline ExprGroup operator,(const Expr &a, const Expr &b) {
@@ -124,4 +122,53 @@ inline ExprGroup operator,(const ExprGroup &a, const Expr &b) {
   return ExprGroup(a, b);
 }
 
+#define PER_EXPRESSION(x) class x;
+#include "taichi/inc/expressions.inc.h"
+#undef PER_EXPRESSION
+
+class ExpressionVisitor {
+ public:
+  ExpressionVisitor(bool allow_undefined_visitor = false,
+                    bool invoke_default_visitor = false)
+      : allow_undefined_visitor_(allow_undefined_visitor),
+        invoke_default_visitor_(invoke_default_visitor) {
+  }
+
+  virtual ~ExpressionVisitor() = default;
+
+  virtual void visit(ExprGroup &expr_group) = 0;
+
+  void visit(Expr &expr) {
+    expr.expr->accept(this);
+  }
+
+  virtual void visit(Expression *expr) {
+    if (!allow_undefined_visitor_) {
+      TI_ERROR("missing visitor function");
+    }
+  }
+
+#define DEFINE_VISIT(T)             \
+  virtual void visit(T *expr) {     \
+    if (allow_undefined_visitor_) { \
+      if (invoke_default_visitor_)  \
+        visit((Expression *)expr);  \
+    } else                          \
+      TI_NOT_IMPLEMENTED;           \
+  }
+
+#define PER_EXPRESSION(x) DEFINE_VISIT(x)
+#include "taichi/inc/expressions.inc.h"
+#undef PER_EXPRESSION
+#undef DEFINE_VISIT
+ private:
+  bool allow_undefined_visitor_{false};
+  bool invoke_default_visitor_{false};
+};
+
+#define TI_DEFINE_ACCEPT_FOR_EXPRESSION              \
+  void accept(ExpressionVisitor *visitor) override { \
+    visitor->visit(this);                            \
+  }
+
 TLANG_NAMESPACE_END