[ir] [refactor] BasicStmtVisitor includes Frontend sstatements (#732)

taichi-dev · Apr 9, 2020 · b5283a9 · b5283a9
1 parent 69bf35f
commit b5283a9
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 14 deletions.
diff --git a/taichi/analysis/detect_for_with_break.cpp b/taichi/analysis/detect_for_with_break.cpp
@@ -15,6 +15,13 @@ class DetectForWithBreak : public BasicStmtVisitor {
   DetectForWithBreak(IRNode *root) : root(root) {
   }
 
+  void visit(FrontendBreakStmt *stmt) override {
+    TI_ASSERT_INFO(loop_stack.size() != 0, "break statement out of loop scope");
+    auto loop = loop_stack.back();
+    if (loop->is<FrontendForStmt>())
+      fors_with_break.insert(loop);
+  }
+
   void visit(FrontendWhileStmt *stmt) override {
     loop_stack.push_back(stmt);
     stmt->body->accept(this);
@@ -27,20 +34,6 @@ class DetectForWithBreak : public BasicStmtVisitor {
     loop_stack.pop_back();
   }
 
-  void visit(FrontendIfStmt *stmt) override {
-    if (stmt->true_statements)
-      stmt->true_statements->accept(this);
-    if (stmt->false_statements)
-      stmt->false_statements->accept(this);
-  }
-
-  void visit(FrontendBreakStmt *stmt) override {
-    TI_ASSERT_INFO(loop_stack.size() != 0, "break statement out of loop scope");
-    auto loop = loop_stack.back();
-    if (loop->is<FrontendForStmt>())
-      fors_with_break.insert(loop);
-  }
-
   std::unordered_set<Stmt *> run() {
     root->accept(this);
     return fors_with_break;

diff --git a/taichi/ir/basic_stmt_visitor.cpp b/taichi/ir/basic_stmt_visitor.cpp
@@ -43,4 +43,19 @@ void BasicStmtVisitor::visit(OffloadedStmt *stmt) {
     stmt->body->accept(this);
 }
 
+void BasicStmtVisitor::visit(FrontendWhileStmt *stmt) {
+  stmt->body->accept(this);
+}
+
+void BasicStmtVisitor::visit(FrontendForStmt *stmt) {
+  stmt->body->accept(this);
+}
+
+void BasicStmtVisitor::visit(FrontendIfStmt *stmt) {
+  if (stmt->true_statements)
+    stmt->true_statements->accept(this);
+  if (stmt->false_statements)
+    stmt->false_statements->accept(this);
+}
+
 TLANG_NAMESPACE_END
diff --git a/taichi/ir/visitors.h b/taichi/ir/visitors.h
@@ -22,6 +22,12 @@ class BasicStmtVisitor : public IRVisitor {
   void visit(StructForStmt *for_stmt) override;
 
   void visit(OffloadedStmt *stmt) override;
+
+  void visit(FrontendWhileStmt *stmt) override;
+
+  void visit(FrontendForStmt *stmt) override;
+
+  void visit(FrontendIfStmt *stmt) override;
 };
 
 TLANG_NAMESPACE_END