[ir] Add RAII guards to IR Builder

taichi/ir/ir_builder.cpp

@@ -44,6 +44,38 @@ void IRBuilder::set_insertion_point_to_false_branch(IfStmt *if_stmt) {
   set_insertion_point({if_stmt->false_statements.get(), 0});
 }
 
+IRBuilder::LoopGuard::~LoopGuard() {
+  if (location_ >= 0 && location_ < loop_->parent->size() &&
+      loop_->parent->statements[location_].get() == loop_) {
+    // faster than set_insertion_point_to_after()
+    builder_.set_insertion_point({loop_->parent, location_ + 1});
+  } else {
+    builder_.set_insertion_point_to_after(loop_);
+  }
+}
+
+IRBuilder::IfGuard::IfGuard(IRBuilder &builder,
+                            IfStmt *if_stmt,
+                            bool true_branch)
+    : builder_(builder), if_stmt_(if_stmt) {
+  location_ = (int)if_stmt_->parent->size() - 1;
+  if (true_branch) {
+    builder_.set_insertion_point_to_true_branch(if_stmt_);
+  } else {
+    builder_.set_insertion_point_to_false_branch(if_stmt_);
+  }
+}
+
+IRBuilder::IfGuard::~IfGuard() {
+  if (location_ >= 0 && location_ < if_stmt_->parent->size() &&
+      if_stmt_->parent->statements[location_].get() == if_stmt_) {
+    // faster than set_insertion_point_to_after()
+    builder_.set_insertion_point({if_stmt_->parent, location_ + 1});
+  } else {
+    builder_.set_insertion_point_to_after(if_stmt_);
+  }
+}
+
 RangeForStmt *IRBuilder::create_range_for(Stmt *begin,
                                           Stmt *end,
                                           int vectorize,

taichi/ir/ir_builder.h

@@ -25,28 +25,28 @@ class IRBuilder {
   std::unique_ptr<IRNode> extract_ir();
 
   // General inserter. Returns stmt.get().
-  template <typename XxxStmt>
-  XxxStmt *insert(std::unique_ptr<XxxStmt> &&stmt) {
+  template <typename XStmt>
+  XStmt *insert(std::unique_ptr<XStmt> &&stmt) {
     return insert(std::move(stmt), &insert_point_);
   }
 
   // Insert to a specific insertion point.
-  template <typename XxxStmt>
-  static XxxStmt *insert(std::unique_ptr<XxxStmt> &&stmt,
-                         InsertPoint *insert_point) {
+  template <typename XStmt>
+  static XStmt *insert(std::unique_ptr<XStmt> &&stmt,
+                       InsertPoint *insert_point) {
     return insert_point->block
         ->insert(std::move(stmt), insert_point->position++)
-        ->template as<XxxStmt>();
+        ->template as<XStmt>();
   }
 
   void set_insertion_point(InsertPoint new_insert_point);
   void set_insertion_point_to_after(Stmt *stmt);
   void set_insertion_point_to_before(Stmt *stmt);
   void set_insertion_point_to_true_branch(IfStmt *if_stmt);
   void set_insertion_point_to_false_branch(IfStmt *if_stmt);
-  template <typename XxxStmt>
-  void set_insertion_point_to_loop_begin(XxxStmt *loop) {
-    using DecayedType = typename std::decay_t<XxxStmt>;
+  template <typename XStmt>
+  void set_insertion_point_to_loop_begin(XStmt *loop) {
+    using DecayedType = typename std::decay_t<XStmt>;
     if constexpr (!std::is_base_of_v<Stmt, DecayedType>) {
       TI_ERROR("The argument is not a statement.");
     }
@@ -59,6 +59,41 @@ class IRBuilder {
     }
   }
 
+  // RAII handles insertion points automatically.
+  class LoopGuard {
+   private:
+    IRBuilder &builder_;
+    Stmt *loop_;
+    int location_;
+
+   public:
+    template <typename XStmt>
+    explicit LoopGuard(IRBuilder &builder, XStmt *loop)
+        : builder_(builder), loop_(loop) {
+      location_ = (int)loop->parent->size() - 1;
+      builder_.set_insertion_point_to_loop_begin(loop);
+    }
+    ~LoopGuard();
+  };
+  class IfGuard {
+   private:
+    IRBuilder &builder_;
+    IfStmt *if_stmt_;
+    int location_;
+
+   public:
+    explicit IfGuard(IRBuilder &builder, IfStmt *if_stmt, bool true_branch);
+    ~IfGuard();
+  };
+
+  template <typename XStmt>
+  LoopGuard get_loop_guard(XStmt *loop) {
+    return LoopGuard(*this, loop);
+  }
+  IfGuard get_if_guard(IfStmt *if_stmt, bool true_branch) {
+    return IfGuard(*this, if_stmt, true_branch);
+  }
+
   // Control flows.
   RangeForStmt *create_range_for(Stmt *begin,
                                  Stmt *end,
@@ -164,9 +199,9 @@ class IRBuilder {
                                    const std::vector<Stmt *> &indices);
   ExternalPtrStmt *create_external_ptr(ArgLoadStmt *ptr,
                                        const std::vector<Stmt *> &indices);
-  template <typename XxxStmt>
-  GlobalLoadStmt *create_global_load(XxxStmt *ptr) {
-    using DecayedType = typename std::decay_t<XxxStmt>;
+  template <typename XStmt>
+  GlobalLoadStmt *create_global_load(XStmt *ptr) {
+    using DecayedType = typename std::decay_t<XStmt>;
     if constexpr (!std::is_base_of_v<Stmt, DecayedType>) {
       TI_ERROR("The argument is not a statement.");
     }
@@ -177,9 +212,9 @@ class IRBuilder {
       TI_ERROR("Statement {} is not a global pointer.", ptr->name());
     }
   }
-  template <typename XxxStmt>
-  void create_global_store(XxxStmt *ptr, Stmt *data) {
-    using DecayedType = typename std::decay_t<XxxStmt>;
+  template <typename XStmt>
+  void create_global_store(XStmt *ptr, Stmt *data) {
+    using DecayedType = typename std::decay_t<XStmt>;
     if constexpr (!std::is_base_of_v<Stmt, DecayedType>) {
       TI_ERROR("The argument is not a statement.");
     }

tests/cpp_new/ir/ir_builder_test.cpp

@@ -44,10 +44,12 @@ TEST(IRBuilder, RangeFor) {
   auto *zero = builder.get_int32(0);
   auto *ten = builder.get_int32(10);
   auto *loop = builder.create_range_for(zero, ten);
-  builder.set_insertion_point_to_loop_begin(loop);
-  auto *index = builder.get_loop_index(loop, 0);
-  builder.set_insertion_point_to_after(loop);
-  auto *ret = builder.create_return(zero);
+  Stmt *index;
+  {
+    auto _ = builder.get_loop_guard(loop);
+    index = builder.get_loop_index(loop, 0);
+  }
+  [[maybe_unused]] auto *ret = builder.create_return(zero);
   EXPECT_EQ(zero->parent->size(), 4);
   ASSERT_TRUE(loop->is<RangeForStmt>());
   auto *loopc = loop->cast<RangeForStmt>();