tint/transform/utils: HoistToDeclBefore polish

A new HoistToDeclBefore::InsertBefore() overload that takes a statement
builder. Required for supporting multiple clones.

Remove Apply() - it was API smell that wasn't needed.

Spring-clean the implementation

Change-Id: If448d2e1945ad6d988d1bdb30487d89efced2f0e
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/104043
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: Antonio Maiorano <amaiorano@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
diff --git a/src/tint/transform/expand_compound_assignment.cc b/src/tint/transform/expand_compound_assignment.cc
index 85b9bf6..f15e28c 100644
--- a/src/tint/transform/expand_compound_assignment.cc
+++ b/src/tint/transform/expand_compound_assignment.cc
@@ -159,10 +159,7 @@
     }
 
     /// Finalize the transformation and clone the module.
-    void Finalize() {
-        hoist_to_decl_before.Apply();
-        ctx.Clone();
-    }
+    void Finalize() { ctx.Clone(); }
 };
 
 }  // namespace
diff --git a/src/tint/transform/promote_initializers_to_let.cc b/src/tint/transform/promote_initializers_to_let.cc
index 22cdc20..315a4ce 100644
--- a/src/tint/transform/promote_initializers_to_let.cc
+++ b/src/tint/transform/promote_initializers_to_let.cc
@@ -102,7 +102,6 @@
         }
     }
 
-    hoist_to_decl_before.Apply();
     ctx.Clone();
 }
 
diff --git a/src/tint/transform/promote_side_effects_to_decl.cc b/src/tint/transform/promote_side_effects_to_decl.cc
index e36c416..2bcda04 100644
--- a/src/tint/transform/promote_side_effects_to_decl.cc
+++ b/src/tint/transform/promote_side_effects_to_decl.cc
@@ -74,8 +74,6 @@
                 hoist_to_decl_before.Prepare(sem_expr);
             }
         }
-
-        hoist_to_decl_before.Apply();
         ctx.Clone();
     }
 };
diff --git a/src/tint/transform/utils/hoist_to_decl_before.cc b/src/tint/transform/utils/hoist_to_decl_before.cc
index 85bd2b2..0aed996 100644
--- a/src/tint/transform/utils/hoist_to_decl_before.cc
+++ b/src/tint/transform/utils/hoist_to_decl_before.cc
@@ -14,7 +14,7 @@
 
 #include "src/tint/transform/utils/hoist_to_decl_before.h"
 
-#include <unordered_map>
+#include <utility>
 
 #include "src/tint/program_builder.h"
 #include "src/tint/sem/block_statement.h"
@@ -23,7 +23,9 @@
 #include "src/tint/sem/reference.h"
 #include "src/tint/sem/variable.h"
 #include "src/tint/sem/while_statement.h"
+#include "src/tint/utils/hashmap.h"
 #include "src/tint/utils/reverse.h"
+#include "src/tint/utils/transform.h"
 
 namespace tint::transform {
 
@@ -36,46 +38,78 @@
     /// loop, so that declaration statements can be inserted before the
     /// condition expression or continuing statement.
     struct LoopInfo {
-        utils::Vector<const ast::Statement*, 8> cond_decls;
-        utils::Vector<const ast::Statement*, 8> cont_decls;
+        utils::Vector<StmtBuilder, 8> cond_decls;
+        utils::Vector<StmtBuilder, 8> cont_decls;
     };
 
     /// Info for each else-if that needs decomposing
     struct ElseIfInfo {
         /// Decls to insert before condition
-        utils::Vector<const ast::Statement*, 8> cond_decls;
+        utils::Vector<StmtBuilder, 8> cond_decls;
     };
 
     /// For-loops that need to be decomposed to loops.
-    std::unordered_map<const sem::ForLoopStatement*, LoopInfo> for_loops;
+    utils::Hashmap<const sem::ForLoopStatement*, LoopInfo, 4> for_loops;
 
     /// Whiles that need to be decomposed to loops.
-    std::unordered_map<const sem::WhileStatement*, LoopInfo> while_loops;
+    utils::Hashmap<const sem::WhileStatement*, LoopInfo, 4> while_loops;
 
     /// 'else if' statements that need to be decomposed to 'else {if}'
-    std::unordered_map<const ast::IfStatement*, ElseIfInfo> else_ifs;
+    utils::Hashmap<const ast::IfStatement*, ElseIfInfo, 4> else_ifs;
 
-    // Converts any for-loops marked for conversion to loops, inserting
-    // registered declaration statements before the condition or continuing
-    // statement.
-    void ForLoopsToLoops() {
-        if (for_loops.empty()) {
-            return;
+    template <size_t N>
+    static auto Build(const utils::Vector<StmtBuilder, N>& builders) {
+        return utils::Transform(builders, [&](auto& builder) { return builder(); });
+    }
+
+    /// @returns a new LoopInfo reference for the given @p for_loop.
+    /// @note if this is the first call to this method, then RegisterForLoopTransform() is
+    /// automatically called.
+    /// @warning the returned reference is invalid if this is called a second time, or the
+    /// #for_loops map is mutated.
+    LoopInfo& ForLoop(const sem::ForLoopStatement* for_loop) {
+        if (for_loops.IsEmpty()) {
+            RegisterForLoopTransform();
         }
+        return for_loops.GetOrZero(for_loop);
+    }
 
-        // At least one for-loop needs to be transformed into a loop.
+    /// @returns a new LoopInfo reference for the given @p while_loop.
+    /// @note if this is the first call to this method, then RegisterWhileLoopTransform() is
+    /// automatically called.
+    /// @warning the returned reference is invalid if this is called a second time, or the
+    /// #for_loops map is mutated.
+    LoopInfo& WhileLoop(const sem::WhileStatement* while_loop) {
+        if (while_loops.IsEmpty()) {
+            RegisterWhileLoopTransform();
+        }
+        return while_loops.GetOrZero(while_loop);
+    }
+
+    /// @returns a new ElseIfInfo reference for the given @p else_if.
+    /// @note if this is the first call to this method, then RegisterElseIfTransform() is
+    /// automatically called.
+    /// @warning the returned reference is invalid if this is called a second time, or the
+    /// #else_ifs map is mutated.
+    ElseIfInfo& ElseIf(const ast::IfStatement* else_if) {
+        if (else_ifs.IsEmpty()) {
+            RegisterElseIfTransform();
+        }
+        return else_ifs.GetOrZero(else_if);
+    }
+
+    /// Registers the handler for transforming for-loops based on the content of the #for_loops map.
+    void RegisterForLoopTransform() const {
         ctx.ReplaceAll([&](const ast::ForLoopStatement* stmt) -> const ast::Statement* {
             auto& sem = ctx.src->Sem();
 
             if (auto* fl = sem.Get(stmt)) {
-                if (auto it = for_loops.find(fl); it != for_loops.end()) {
-                    auto& info = it->second;
+                if (auto* info = for_loops.Find(fl)) {
                     auto* for_loop = fl->Declaration();
                     // For-loop needs to be decomposed to a loop.
                     // Build the loop body's statements.
-                    // Start with any let declarations for the conditional
-                    // expression.
-                    auto body_stmts = info.cond_decls;
+                    // Start with any let declarations for the conditional expression.
+                    auto body_stmts = Build(info->cond_decls);
                     // If the for-loop has a condition, emit this next as:
                     //   if (!cond) { break; }
                     if (auto* cond = for_loop->condition) {
@@ -95,7 +129,7 @@
                     if (auto* cont = for_loop->continuing) {
                         // Continuing block starts with any let declarations used by
                         // the continuing.
-                        auto cont_stmts = info.cont_decls;
+                        auto cont_stmts = Build(info->cont_decls);
                         cont_stmts.Push(ctx.Clone(cont));
                         continuing = b.Block(cont_stmts);
                     }
@@ -112,34 +146,29 @@
         });
     }
 
-    // Converts any while-loops marked for conversion to loops, inserting
-    // registered declaration statements before the condition.
-    void WhilesToLoops() {
-        if (while_loops.empty()) {
-            return;
-        }
-
+    /// Registers the handler for transforming while-loops based on the content of the #while_loops
+    /// map.
+    void RegisterWhileLoopTransform() const {
         // At least one while needs to be transformed into a loop.
         ctx.ReplaceAll([&](const ast::WhileStatement* stmt) -> const ast::Statement* {
             auto& sem = ctx.src->Sem();
 
             if (auto* w = sem.Get(stmt)) {
-                if (auto it = while_loops.find(w); it != while_loops.end()) {
-                    auto& info = it->second;
+                if (auto* info = while_loops.Find(w)) {
                     auto* while_loop = w->Declaration();
                     // While needs to be decomposed to a loop.
                     // Build the loop body's statements.
                     // Start with any let declarations for the conditional
                     // expression.
-                    auto body_stmts = info.cond_decls;
+                    auto body_stmts = utils::Transform(info->cond_decls,
+                                                       [&](auto& builder) { return builder(); });
                     // Emit the condition as:
                     //   if (!cond) { break; }
                     auto* cond = while_loop->condition;
                     // !condition
-                    auto* not_cond =
-                        b.create<ast::UnaryOpExpression>(ast::UnaryOp::kNot, ctx.Clone(cond));
+                    auto* not_cond = b.Not(ctx.Clone(cond));
                     // { break; }
-                    auto* break_body = b.Block(b.create<ast::BreakStatement>());
+                    auto* break_body = b.Block(b.Break());
                     // if (!condition) { break; }
                     body_stmts.Push(b.If(not_cond, break_body));
 
@@ -157,81 +186,47 @@
         });
     }
 
-    void ElseIfsToElseWithNestedIfs() {
+    /// Registers the handler for transforming if-statements based on the content of the #else_ifs
+    /// map.
+    void RegisterElseIfTransform() const {
         // Decompose 'else-if' statements into 'else { if }' blocks.
-        ctx.ReplaceAll([&](const ast::IfStatement* else_if) -> const ast::Statement* {
-            if (!else_ifs.count(else_if)) {
-                return nullptr;
+        ctx.ReplaceAll([&](const ast::IfStatement* stmt) -> const ast::Statement* {
+            if (auto* info = else_ifs.Find(stmt)) {
+                // Build the else block's body statements, starting with let decls for the
+                // conditional expression.
+                auto body_stmts = Build(info->cond_decls);
+
+                // Move the 'else-if' into the new `else` block as a plain 'if'.
+                auto* cond = ctx.Clone(stmt->condition);
+                auto* body = ctx.Clone(stmt->body);
+                auto* new_if = b.If(cond, body, b.Else(ctx.Clone(stmt->else_statement)));
+                body_stmts.Push(new_if);
+
+                // Replace the 'else-if' with the new 'else' block.
+                return b.Block(body_stmts);
             }
-            auto& else_if_info = else_ifs[else_if];
-
-            // Build the else block's body statements, starting with let decls for
-            // the conditional expression.
-            auto& body_stmts = else_if_info.cond_decls;
-
-            // Move the 'else-if' into the new `else` block as a plain 'if'.
-            auto* cond = ctx.Clone(else_if->condition);
-            auto* body = ctx.Clone(else_if->body);
-            auto* new_if = b.If(cond, body, b.Else(ctx.Clone(else_if->else_statement)));
-            body_stmts.Push(new_if);
-
-            // Replace the 'else-if' with the new 'else' block.
-            return b.Block(body_stmts);
+            return nullptr;
         });
     }
 
-  public:
-    /// Constructor
-    /// @param ctx_in the clone context
-    explicit State(CloneContext& ctx_in) : ctx(ctx_in), b(*ctx_in.dst) {}
+    /// A type used to signal to InsertBeforeImpl that no insertion should take place - instead flow
+    /// control statements should just be decomposed.
+    struct Decompose {};
 
-    /// Hoists `expr` to a `let` or `var` with optional `decl_name`, inserting it
-    /// before `before_expr`.
-    /// @param before_expr expression to insert `expr` before
-    /// @param expr expression to hoist
-    /// @param as_let hoist to `let` if true, otherwise to `var`
-    /// @param decl_name optional name to use for the variable/constant name
-    /// @return true on success
-    bool Add(const sem::Expression* before_expr,
-             const ast::Expression* expr,
-             bool as_let,
-             const char* decl_name) {
-        auto name = b.Symbols().New(decl_name);
-
-        // Construct the let/var that holds the hoisted expr
-        auto* v = as_let ? static_cast<const ast::Variable*>(b.Let(name, ctx.Clone(expr)))
-                         : static_cast<const ast::Variable*>(b.Var(name, ctx.Clone(expr)));
-        auto* decl = b.Decl(v);
-
-        if (!InsertBefore(before_expr->Stmt(), decl)) {
-            return false;
-        }
-
-        // Replace the initializer expression with a reference to the let
-        ctx.Replace(expr, b.Expr(name));
-        return true;
-    }
-
-    /// Inserts `stmt` before `before_stmt`, possibly marking a for-loop to be
-    /// converted to a loop, or an else-if to an else { if }. If `decl` is
-    /// nullptr, for-loop and else-if conversions are marked, but no hoisting
-    /// takes place.
-    /// @param before_stmt statement to insert `stmt` before
-    /// @param stmt statement to insert
-    /// @return true on success
-    bool InsertBefore(const sem::Statement* before_stmt, const ast::Statement* stmt) {
+    template <typename BUILDER>
+    bool InsertBeforeImpl(const sem::Statement* before_stmt, BUILDER&& builder) {
         auto* ip = before_stmt->Declaration();
 
         auto* else_if = before_stmt->As<sem::IfStatement>();
         if (else_if && else_if->Parent()->Is<sem::IfStatement>()) {
             // Insertion point is an 'else if' condition.
             // Need to convert 'else if' to 'else { if }'.
-            auto& else_if_info = else_ifs[else_if->Declaration()];
+            auto& else_if_info = ElseIf(else_if->Declaration());
 
             // Index the map to convert this else if, even if `stmt` is nullptr.
             auto& decls = else_if_info.cond_decls;
-            if (stmt) {
-                decls.Push(stmt);
+            if constexpr (!std::is_same_v<BUILDER, Decompose>) {
+                decls.Push(std::forward<BUILDER>(builder));
             }
             return true;
         }
@@ -241,9 +236,9 @@
             // For-loop needs to be decomposed to a loop.
 
             // Index the map to convert this for-loop, even if `stmt` is nullptr.
-            auto& decls = for_loops[fl].cond_decls;
-            if (stmt) {
-                decls.Push(stmt);
+            auto& decls = ForLoop(fl).cond_decls;
+            if constexpr (!std::is_same_v<BUILDER, Decompose>) {
+                decls.Push(std::forward<BUILDER>(builder));
             }
             return true;
         }
@@ -253,9 +248,9 @@
             // While needs to be decomposed to a loop.
 
             // Index the map to convert this while, even if `stmt` is nullptr.
-            auto& decls = while_loops[w].cond_decls;
-            if (stmt) {
-                decls.Push(stmt);
+            auto& decls = WhileLoop(w).cond_decls;
+            if constexpr (!std::is_same_v<BUILDER, Decompose>) {
+                decls.Push(std::forward<BUILDER>(builder));
             }
             return true;
         }
@@ -264,8 +259,9 @@
         if (auto* block = parent->As<sem::BlockStatement>()) {
             // Insert point sits in a block. Simple case.
             // Insert the stmt before the parent statement.
-            if (stmt) {
-                ctx.InsertBefore(block->Declaration()->statements, ip, stmt);
+            if constexpr (!std::is_same_v<BUILDER, Decompose>) {
+                ctx.InsertBefore(block->Declaration()->statements, ip,
+                                 std::forward<BUILDER>(builder));
             }
             return true;
         }
@@ -276,9 +272,9 @@
             if (fl->Declaration()->initializer == ip) {
                 // Insertion point is a for-loop initializer.
                 // Insert the new statement above the for-loop.
-                if (stmt) {
+                if constexpr (!std::is_same_v<BUILDER, Decompose>) {
                     ctx.InsertBefore(fl->Block()->Declaration()->statements, fl->Declaration(),
-                                     stmt);
+                                     std::forward<BUILDER>(builder));
                 }
                 return true;
             }
@@ -288,9 +284,9 @@
                 // For-loop needs to be decomposed to a loop.
 
                 // Index the map to convert this for-loop, even if `stmt` is nullptr.
-                auto& decls = for_loops[fl].cont_decls;
-                if (stmt) {
-                    decls.Push(stmt);
+                auto& decls = ForLoop(fl).cont_decls;
+                if constexpr (!std::is_same_v<BUILDER, Decompose>) {
+                    decls.Push(std::forward<BUILDER>(builder));
                 }
                 return true;
             }
@@ -304,23 +300,56 @@
         return false;
     }
 
-    /// Use to signal that we plan on hoisting a decl before `before_expr`. This
-    /// will convert 'for-loop's to 'loop's and 'else-if's to 'else {if}'s if
-    /// needed.
-    /// @param before_expr expression we would hoist a decl before
-    /// @return true on success
-    bool Prepare(const sem::Expression* before_expr) {
-        return InsertBefore(before_expr->Stmt(), nullptr);
+  public:
+    /// Constructor
+    /// @param ctx_in the clone context
+    explicit State(CloneContext& ctx_in) : ctx(ctx_in), b(*ctx_in.dst) {}
+
+    /// @copydoc HoistToDeclBefore::Add()
+    bool Add(const sem::Expression* before_expr,
+             const ast::Expression* expr,
+             bool as_let,
+             const char* decl_name) {
+        auto name = b.Symbols().New(decl_name);
+
+        if (as_let) {
+            auto builder = [this, expr, name] {
+                return b.Decl(b.Let(name, ctx.CloneWithoutTransform(expr)));
+            };
+            if (!InsertBeforeImpl(before_expr->Stmt(), std::move(builder))) {
+                return false;
+            }
+        } else {
+            auto builder = [this, expr, name] {
+                return b.Decl(b.Var(name, ctx.CloneWithoutTransform(expr)));
+            };
+            if (!InsertBeforeImpl(before_expr->Stmt(), std::move(builder))) {
+                return false;
+            }
+        }
+
+        // Replace the initializer expression with a reference to the let
+        ctx.Replace(expr, b.Expr(name));
+        return true;
     }
 
-    /// Applies any scheduled insertions from previous calls to Add() to
-    /// CloneContext. Call this once before ctx.Clone().
-    /// @return true on success
-    bool Apply() {
-        ForLoopsToLoops();
-        WhilesToLoops();
-        ElseIfsToElseWithNestedIfs();
-        return true;
+    /// @copydoc HoistToDeclBefore::InsertBefore(const sem::Statement*, const ast::Statement*)
+    bool InsertBefore(const sem::Statement* before_stmt, const ast::Statement* stmt) {
+        if (stmt) {
+            auto builder = [stmt] { return stmt; };
+            return InsertBeforeImpl(before_stmt, std::move(builder));
+        }
+        return InsertBeforeImpl(before_stmt, Decompose{});
+    }
+
+    /// @copydoc HoistToDeclBefore::InsertBefore(const sem::Statement*, const StmtBuilder&)
+    bool InsertBefore(const sem::Statement* before_stmt, const StmtBuilder& builder) {
+        return InsertBeforeImpl(before_stmt, std::move(builder));
+    }
+
+    /// @copydoc HoistToDeclBefore::Prepare()
+    bool Prepare(const sem::Expression* before_expr) {
+        return InsertBefore(before_expr->Stmt(), nullptr);
     }
 };
 
@@ -340,12 +369,13 @@
     return state_->InsertBefore(before_stmt, stmt);
 }
 
-bool HoistToDeclBefore::Prepare(const sem::Expression* before_expr) {
-    return state_->Prepare(before_expr);
+bool HoistToDeclBefore::InsertBefore(const sem::Statement* before_stmt,
+                                     const StmtBuilder& builder) {
+    return state_->InsertBefore(before_stmt, builder);
 }
 
-bool HoistToDeclBefore::Apply() {
-    return state_->Apply();
+bool HoistToDeclBefore::Prepare(const sem::Expression* before_expr) {
+    return state_->Prepare(before_expr);
 }
 
 }  // namespace tint::transform
diff --git a/src/tint/transform/utils/hoist_to_decl_before.h b/src/tint/transform/utils/hoist_to_decl_before.h
index d0b96e0..b2e993e 100644
--- a/src/tint/transform/utils/hoist_to_decl_before.h
+++ b/src/tint/transform/utils/hoist_to_decl_before.h
@@ -15,6 +15,7 @@
 #ifndef SRC_TINT_TRANSFORM_UTILS_HOIST_TO_DECL_BEFORE_H_
 #define SRC_TINT_TRANSFORM_UTILS_HOIST_TO_DECL_BEFORE_H_
 
+#include <functional>
 #include <memory>
 
 #include "src/tint/sem/expression.h"
@@ -34,25 +35,40 @@
     /// Destructor
     ~HoistToDeclBefore();
 
-    /// Hoists `expr` to a `let` or `var` with optional `decl_name`, inserting it
-    /// before `before_expr`.
+    /// StmtBuilder is a builder of an AST statement
+    using StmtBuilder = std::function<const ast::Statement*()>;
+
+    /// Hoists @p expr to a `let` or `var` with optional `decl_name`, inserting it
+    /// before @p before_expr.
     /// @param before_expr expression to insert `expr` before
     /// @param expr expression to hoist
-    /// @param as_const hoist to `let` if true, otherwise to `var`
+    /// @param as_let hoist to `let` if true, otherwise to `var`
     /// @param decl_name optional name to use for the variable/constant name
     /// @return true on success
     bool Add(const sem::Expression* before_expr,
              const ast::Expression* expr,
-             bool as_const,
+             bool as_let,
              const char* decl_name = "");
 
-    /// Inserts `stmt` before `before_stmt`, possibly converting 'for-loop's to
-    /// 'loop's if necessary.
-    /// @param before_stmt statement to insert `stmt` before
+    /// Inserts @p stmt before @p before_stmt, possibly converting 'for-loop's to 'loop's if
+    /// necessary.
+    /// @warning If the container of @p before_stmt is cloned multiple times, then the resolver will
+    /// ICE as the same statement cannot be shared.
+    /// @param before_stmt statement to insert @p stmt before
     /// @param stmt statement to insert
     /// @return true on success
     bool InsertBefore(const sem::Statement* before_stmt, const ast::Statement* stmt);
 
+    /// Inserts the returned statement of @p builder before @p before_stmt, possibly converting
+    /// 'for-loop's to 'loop's if necessary.
+    /// @note If the container of @p before_stmt is cloned multiple times, then @p builder will be
+    /// called for each clone.
+    /// @param before_stmt the preceding statement that the statement of @p builder will be inserted
+    /// before
+    /// @param builder the statement builder used to create the new statement
+    /// @return true on success
+    bool InsertBefore(const sem::Statement* before_stmt, const StmtBuilder& builder);
+
     /// Use to signal that we plan on hoisting a decl before `before_expr`. This
     /// will convert 'for-loop's to 'loop's and 'else-if's to 'else {if}'s if
     /// needed.
@@ -60,11 +76,6 @@
     /// @return true on success
     bool Prepare(const sem::Expression* before_expr);
 
-    /// Applies any scheduled insertions from previous calls to Add() to
-    /// CloneContext. Call this once before ctx.Clone().
-    /// @return true on success
-    bool Apply();
-
   private:
     class State;
     std::unique_ptr<State> state_;
diff --git a/src/tint/transform/utils/hoist_to_decl_before_test.cc b/src/tint/transform/utils/hoist_to_decl_before_test.cc
index f0d3b36..7b45e01 100644
--- a/src/tint/transform/utils/hoist_to_decl_before_test.cc
+++ b/src/tint/transform/utils/hoist_to_decl_before_test.cc
@@ -45,7 +45,6 @@
     HoistToDeclBefore hoistToDeclBefore(ctx);
     auto* sem_expr = ctx.src->Sem().Get(expr);
     hoistToDeclBefore.Add(sem_expr, expr, true);
-    hoistToDeclBefore.Apply();
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -77,7 +76,6 @@
     HoistToDeclBefore hoistToDeclBefore(ctx);
     auto* sem_expr = ctx.src->Sem().Get(expr);
     hoistToDeclBefore.Add(sem_expr, expr, true);
-    hoistToDeclBefore.Apply();
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -112,7 +110,6 @@
     HoistToDeclBefore hoistToDeclBefore(ctx);
     auto* sem_expr = ctx.src->Sem().Get(expr);
     hoistToDeclBefore.Add(sem_expr, expr, true);
-    hoistToDeclBefore.Apply();
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -151,7 +148,6 @@
     HoistToDeclBefore hoistToDeclBefore(ctx);
     auto* sem_expr = ctx.src->Sem().Get(expr);
     hoistToDeclBefore.Add(sem_expr, expr, true);
-    hoistToDeclBefore.Apply();
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -195,7 +191,6 @@
     HoistToDeclBefore hoistToDeclBefore(ctx);
     auto* sem_expr = ctx.src->Sem().Get(expr);
     hoistToDeclBefore.Add(sem_expr, expr, true);
-    hoistToDeclBefore.Apply();
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -240,7 +235,6 @@
     HoistToDeclBefore hoistToDeclBefore(ctx);
     auto* sem_expr = ctx.src->Sem().Get(expr);
     hoistToDeclBefore.Add(sem_expr, expr, true);
-    hoistToDeclBefore.Apply();
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -279,7 +273,6 @@
     HoistToDeclBefore hoistToDeclBefore(ctx);
     auto* sem_expr = ctx.src->Sem().Get(expr);
     hoistToDeclBefore.Add(sem_expr, expr, true);
-    hoistToDeclBefore.Apply();
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -314,7 +307,6 @@
     HoistToDeclBefore hoistToDeclBefore(ctx);
     auto* sem_expr = ctx.src->Sem().Get(expr);
     hoistToDeclBefore.Add(sem_expr, expr, true);
-    hoistToDeclBefore.Apply();
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -349,7 +341,6 @@
     HoistToDeclBefore hoistToDeclBefore(ctx);
     auto* sem_expr = ctx.src->Sem().Get(expr);
     hoistToDeclBefore.Prepare(sem_expr);
-    hoistToDeclBefore.Apply();
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -387,7 +378,6 @@
     HoistToDeclBefore hoistToDeclBefore(ctx);
     auto* sem_expr = ctx.src->Sem().Get(expr);
     hoistToDeclBefore.Prepare(sem_expr);
-    hoistToDeclBefore.Apply();
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -434,7 +424,6 @@
     HoistToDeclBefore hoistToDeclBefore(ctx);
     auto* sem_expr = ctx.src->Sem().Get(expr);
     hoistToDeclBefore.Prepare(sem_expr);
-    hoistToDeclBefore.Apply();
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -473,7 +462,42 @@
     auto* before_stmt = ctx.src->Sem().Get(var);
     auto* new_stmt = ctx.dst->CallStmt(ctx.dst->Call("foo"));
     hoistToDeclBefore.InsertBefore(before_stmt, new_stmt);
-    hoistToDeclBefore.Apply();
+
+    ctx.Clone();
+    Program cloned(std::move(cloned_b));
+
+    auto* expect = R"(
+fn foo() {
+}
+
+fn f() {
+  foo();
+  var a = 1i;
+}
+)";
+
+    EXPECT_EQ(expect, str(cloned));
+}
+
+TEST_F(HoistToDeclBeforeTest, InsertBefore_Block_Function) {
+    // fn foo() {
+    // }
+    // fn f() {
+    //     var a = 1i;
+    // }
+    ProgramBuilder b;
+    b.Func("foo", utils::Empty, b.ty.void_(), utils::Empty);
+    auto* var = b.Decl(b.Var("a", b.Expr(1_i)));
+    b.Func("f", utils::Empty, b.ty.void_(), utils::Vector{var});
+
+    Program original(std::move(b));
+    ProgramBuilder cloned_b;
+    CloneContext ctx(&cloned_b, &original);
+
+    HoistToDeclBefore hoistToDeclBefore(ctx);
+    auto* before_stmt = ctx.src->Sem().Get(var);
+    hoistToDeclBefore.InsertBefore(before_stmt,
+                                   [&] { return ctx.dst->CallStmt(ctx.dst->Call("foo")); });
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -512,7 +536,45 @@
     auto* before_stmt = ctx.src->Sem().Get(var);
     auto* new_stmt = ctx.dst->CallStmt(ctx.dst->Call("foo"));
     hoistToDeclBefore.InsertBefore(before_stmt, new_stmt);
-    hoistToDeclBefore.Apply();
+
+    ctx.Clone();
+    Program cloned(std::move(cloned_b));
+
+    auto* expect = R"(
+fn foo() {
+}
+
+fn f() {
+  foo();
+  for(var a = 1i; true; ) {
+  }
+}
+)";
+
+    EXPECT_EQ(expect, str(cloned));
+}
+
+TEST_F(HoistToDeclBeforeTest, InsertBefore_ForLoopInit_Function) {
+    // fn foo() {
+    // }
+    // fn f() {
+    //     for(var a = 1i; true;) {
+    //     }
+    // }
+    ProgramBuilder b;
+    b.Func("foo", utils::Empty, b.ty.void_(), utils::Empty);
+    auto* var = b.Decl(b.Var("a", b.Expr(1_i)));
+    auto* s = b.For(var, b.Expr(true), nullptr, b.Block());
+    b.Func("f", utils::Empty, b.ty.void_(), utils::Vector{s});
+
+    Program original(std::move(b));
+    ProgramBuilder cloned_b;
+    CloneContext ctx(&cloned_b, &original);
+
+    HoistToDeclBefore hoistToDeclBefore(ctx);
+    auto* before_stmt = ctx.src->Sem().Get(var);
+    hoistToDeclBefore.InsertBefore(before_stmt,
+                                   [&] { return ctx.dst->CallStmt(ctx.dst->Call("foo")); });
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -554,7 +616,57 @@
     auto* before_stmt = ctx.src->Sem().Get(cont->As<ast::Statement>());
     auto* new_stmt = ctx.dst->CallStmt(ctx.dst->Call("foo"));
     hoistToDeclBefore.InsertBefore(before_stmt, new_stmt);
-    hoistToDeclBefore.Apply();
+
+    ctx.Clone();
+    Program cloned(std::move(cloned_b));
+
+    auto* expect = R"(
+fn foo() {
+}
+
+fn f() {
+  var a = 1i;
+  loop {
+    if (!(true)) {
+      break;
+    }
+    {
+    }
+
+    continuing {
+      foo();
+      a += 1i;
+    }
+  }
+}
+)";
+
+    EXPECT_EQ(expect, str(cloned));
+}
+
+TEST_F(HoistToDeclBeforeTest, InsertBefore_ForLoopCont_Function) {
+    // fn foo() {
+    // }
+    // fn f() {
+    //     var a = 1i;
+    //     for(; true; a+=1i) {
+    //     }
+    // }
+    ProgramBuilder b;
+    b.Func("foo", utils::Empty, b.ty.void_(), utils::Empty);
+    auto* var = b.Decl(b.Var("a", b.Expr(1_i)));
+    auto* cont = b.CompoundAssign("a", b.Expr(1_i), ast::BinaryOp::kAdd);
+    auto* s = b.For(nullptr, b.Expr(true), cont, b.Block());
+    b.Func("f", utils::Empty, b.ty.void_(), utils::Vector{var, s});
+
+    Program original(std::move(b));
+    ProgramBuilder cloned_b;
+    CloneContext ctx(&cloned_b, &original);
+
+    HoistToDeclBefore hoistToDeclBefore(ctx);
+    auto* before_stmt = ctx.src->Sem().Get(cont->As<ast::Statement>());
+    hoistToDeclBefore.InsertBefore(before_stmt,
+                                   [&] { return ctx.dst->CallStmt(ctx.dst->Call("foo")); });
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
@@ -609,7 +721,55 @@
     auto* before_stmt = ctx.src->Sem().Get(elseif);
     auto* new_stmt = ctx.dst->CallStmt(ctx.dst->Call("foo"));
     hoistToDeclBefore.InsertBefore(before_stmt, new_stmt);
-    hoistToDeclBefore.Apply();
+
+    ctx.Clone();
+    Program cloned(std::move(cloned_b));
+
+    auto* expect = R"(
+fn foo() {
+}
+
+fn f() {
+  var a : bool;
+  if (true) {
+  } else {
+    foo();
+    if (a) {
+    } else {
+    }
+  }
+}
+)";
+
+    EXPECT_EQ(expect, str(cloned));
+}
+
+TEST_F(HoistToDeclBeforeTest, InsertBefore_ElseIf_Function) {
+    // fn foo() {
+    // }
+    // fn f() {
+    //     var a : bool;
+    //     if (true) {
+    //     } else if (a) {
+    //     } else {
+    //     }
+    // }
+    ProgramBuilder b;
+    b.Func("foo", utils::Empty, b.ty.void_(), utils::Empty);
+    auto* var = b.Decl(b.Var("a", b.ty.bool_()));
+    auto* elseif = b.If(b.Expr("a"), b.Block(), b.Else(b.Block()));
+    auto* s = b.If(b.Expr(true), b.Block(),  //
+                   b.Else(elseif));
+    b.Func("f", utils::Empty, b.ty.void_(), utils::Vector{var, s});
+
+    Program original(std::move(b));
+    ProgramBuilder cloned_b;
+    CloneContext ctx(&cloned_b, &original);
+
+    HoistToDeclBefore hoistToDeclBefore(ctx);
+    auto* before_stmt = ctx.src->Sem().Get(elseif);
+    hoistToDeclBefore.InsertBefore(before_stmt,
+                                   [&] { return ctx.dst->CallStmt(ctx.dst->Call("foo")); });
 
     ctx.Clone();
     Program cloned(std::move(cloned_b));
diff --git a/src/tint/transform/var_for_dynamic_index.cc b/src/tint/transform/var_for_dynamic_index.cc
index aaebdc7..d30831e 100644
--- a/src/tint/transform/var_for_dynamic_index.cc
+++ b/src/tint/transform/var_for_dynamic_index.cc
@@ -57,7 +57,6 @@
         }
     }
 
-    hoist_to_decl_before.Apply();
     ctx.Clone();
 }