ast: Remove IfStatement::set_else_statements()

Move it to the constructor.

Bug: tint:390
Change-Id: Ib4ac1a1c83aa59963472ac7c14c9e0cbcf2734e6
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/35007
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: dan sinclair <dsinclair@chromium.org>
diff --git a/src/ast/block_statement_test.cc b/src/ast/block_statement_test.cc
index f0941b7..fc2cb92 100644
--- a/src/ast/block_statement_test.cc
+++ b/src/ast/block_statement_test.cc
@@ -88,7 +88,8 @@
 
 TEST_F(BlockStatementTest, IsValid_InvalidBodyStatement) {
   BlockStatement b;
-  b.append(create<IfStatement>(nullptr, create<BlockStatement>()));
+  b.append(create<IfStatement>(Source{}, nullptr, create<BlockStatement>(),
+                               ElseStatementList{}));
   EXPECT_FALSE(b.IsValid());
 }
 
diff --git a/src/ast/case_statement_test.cc b/src/ast/case_statement_test.cc
index ba62191..37617a8 100644
--- a/src/ast/case_statement_test.cc
+++ b/src/ast/case_statement_test.cc
@@ -124,7 +124,8 @@
   b.push_back(create<SintLiteral>(&i32, 2));
 
   auto* body = create<BlockStatement>();
-  body->append(create<IfStatement>(nullptr, create<BlockStatement>()));
+  body->append(create<IfStatement>(Source{}, nullptr, create<BlockStatement>(),
+                                   ElseStatementList{}));
 
   CaseStatement c({b}, body);
   EXPECT_FALSE(c.IsValid());
diff --git a/src/ast/else_statement_test.cc b/src/ast/else_statement_test.cc
index 166d805..98dfd71 100644
--- a/src/ast/else_statement_test.cc
+++ b/src/ast/else_statement_test.cc
@@ -97,7 +97,8 @@
 
 TEST_F(ElseStatementTest, IsValid_InvalidBodyStatement) {
   auto* body = create<BlockStatement>();
-  body->append(create<IfStatement>(nullptr, create<BlockStatement>()));
+  body->append(create<IfStatement>(Source{}, nullptr, create<BlockStatement>(),
+                                   ElseStatementList{}));
 
   ElseStatement e(body);
   EXPECT_FALSE(e.IsValid());
diff --git a/src/ast/if_statement.cc b/src/ast/if_statement.cc
index 451c123..aac56a8 100644
--- a/src/ast/if_statement.cc
+++ b/src/ast/if_statement.cc
@@ -23,13 +23,14 @@
 namespace tint {
 namespace ast {
 
-IfStatement::IfStatement(Expression* condition, BlockStatement* body)
-    : Base(), condition_(condition), body_(body) {}
-
 IfStatement::IfStatement(const Source& source,
                          Expression* condition,
-                         BlockStatement* body)
-    : Base(source), condition_(condition), body_(body) {}
+                         BlockStatement* body,
+                         ElseStatementList else_stmts)
+    : Base(source),
+      condition_(condition),
+      body_(body),
+      else_statements_(std::move(else_stmts)) {}
 
 IfStatement::IfStatement(IfStatement&&) = default;
 
@@ -37,8 +38,8 @@
 
 IfStatement* IfStatement::Clone(CloneContext* ctx) const {
   auto* cloned = ctx->mod->create<IfStatement>(
-      ctx->Clone(source()), ctx->Clone(condition_), ctx->Clone(body_));
-  cloned->else_statements_ = ctx->Clone(else_statements_);
+      ctx->Clone(source()), ctx->Clone(condition_), ctx->Clone(body_),
+      ctx->Clone(else_statements_));
   return cloned;
 }
 
diff --git a/src/ast/if_statement.h b/src/ast/if_statement.h
index 32a4d31..f489c4a 100644
--- a/src/ast/if_statement.h
+++ b/src/ast/if_statement.h
@@ -30,16 +30,14 @@
 class IfStatement : public Castable<IfStatement, Statement> {
  public:
   /// Constructor
-  /// @param condition the if condition
-  /// @param body the if body
-  IfStatement(Expression* condition, BlockStatement* body);
-  /// Constructor
   /// @param source the source information
   /// @param condition the if condition
   /// @param body the if body
+  /// @param else_stmts the else statements
   IfStatement(const Source& source,
               Expression* condition,
-              BlockStatement* body);
+              BlockStatement* body,
+              ElseStatementList else_stmts);
   /// Move constructor
   IfStatement(IfStatement&&);
   ~IfStatement() override;
@@ -51,11 +49,6 @@
   /// @returns the if body
   BlockStatement* body() { return body_; }
 
-  /// Sets the else statements
-  /// @param else_statements the else statements to set
-  void set_else_statements(ElseStatementList else_statements) {
-    else_statements_ = std::move(else_statements);
-  }
   /// @returns the else statements
   const ElseStatementList& else_statements() const { return else_statements_; }
   /// @returns the else statements
diff --git a/src/ast/if_statement_test.cc b/src/ast/if_statement_test.cc
index a87930a..e3117cd 100644
--- a/src/ast/if_statement_test.cc
+++ b/src/ast/if_statement_test.cc
@@ -27,28 +27,18 @@
 TEST_F(IfStatementTest, Creation) {
   auto* cond = create<IdentifierExpression>("cond");
   auto* body = create<BlockStatement>();
-  auto* discard = create<DiscardStatement>();
-  body->append(discard);
-
-  IfStatement stmt(cond, body);
-  EXPECT_EQ(stmt.condition(), cond);
-  ASSERT_EQ(stmt.body()->size(), 1u);
-  EXPECT_EQ(stmt.body()->get(0), discard);
-}
-
-TEST_F(IfStatementTest, Creation_WithSource) {
-  auto* cond = create<IdentifierExpression>("cond");
-  auto* body = create<BlockStatement>();
   body->append(create<DiscardStatement>());
 
-  IfStatement stmt(Source{Source::Location{20, 2}}, cond, body);
+  IfStatement stmt(Source{Source::Location{20, 2}}, cond, body,
+                   ElseStatementList{});
   auto src = stmt.source();
   EXPECT_EQ(src.range.begin.line, 20u);
   EXPECT_EQ(src.range.begin.column, 2u);
 }
 
 TEST_F(IfStatementTest, IsIf) {
-  IfStatement stmt(nullptr, create<BlockStatement>());
+  IfStatement stmt(Source{}, nullptr, create<BlockStatement>(),
+                   ElseStatementList{});
   EXPECT_TRUE(stmt.Is<IfStatement>());
 }
 
@@ -57,7 +47,7 @@
   auto* body = create<BlockStatement>();
   body->append(create<DiscardStatement>());
 
-  IfStatement stmt(cond, body);
+  IfStatement stmt(Source{}, cond, body, ElseStatementList{});
   EXPECT_TRUE(stmt.IsValid());
 }
 
@@ -66,13 +56,13 @@
   auto* body = create<BlockStatement>();
   body->append(create<DiscardStatement>());
 
-  ElseStatementList else_stmts;
-  else_stmts.push_back(create<ElseStatement>(
-      create<IdentifierExpression>("Ident"), create<BlockStatement>()));
-  else_stmts.push_back(create<ElseStatement>(create<BlockStatement>()));
-
-  IfStatement stmt(cond, body);
-  stmt.set_else_statements(else_stmts);
+  IfStatement stmt(
+      Source{}, cond, body,
+      {
+          create<ElseStatement>(create<IdentifierExpression>("Ident"),
+                                create<BlockStatement>()),
+          create<ElseStatement>(create<BlockStatement>()),
+      });
   EXPECT_TRUE(stmt.IsValid());
 }
 
@@ -80,7 +70,7 @@
   auto* body = create<BlockStatement>();
   body->append(create<DiscardStatement>());
 
-  IfStatement stmt(nullptr, body);
+  IfStatement stmt(Source{}, nullptr, body, ElseStatementList{});
   EXPECT_FALSE(stmt.IsValid());
 }
 
@@ -89,7 +79,7 @@
   auto* body = create<BlockStatement>();
   body->append(create<DiscardStatement>());
 
-  IfStatement stmt(cond, body);
+  IfStatement stmt(Source{}, cond, body, ElseStatementList{});
   EXPECT_FALSE(stmt.IsValid());
 }
 
@@ -99,7 +89,7 @@
   body->append(create<DiscardStatement>());
   body->append(nullptr);
 
-  IfStatement stmt(cond, body);
+  IfStatement stmt(Source{}, cond, body, ElseStatementList{});
   EXPECT_FALSE(stmt.IsValid());
 }
 
@@ -107,9 +97,10 @@
   auto* cond = create<IdentifierExpression>("cond");
   auto* body = create<BlockStatement>();
   body->append(create<DiscardStatement>());
-  body->append(create<IfStatement>(nullptr, create<BlockStatement>()));
+  body->append(create<IfStatement>(Source{}, nullptr, create<BlockStatement>(),
+                                   ast::ElseStatementList{}));
 
-  IfStatement stmt(cond, body);
+  IfStatement stmt(Source{}, cond, body, ElseStatementList{});
   EXPECT_FALSE(stmt.IsValid());
 }
 
@@ -118,14 +109,14 @@
   auto* body = create<BlockStatement>();
   body->append(create<DiscardStatement>());
 
-  ElseStatementList else_stmts;
-  else_stmts.push_back(create<ElseStatement>(
-      create<IdentifierExpression>("Ident"), create<BlockStatement>()));
-  else_stmts.push_back(create<ElseStatement>(create<BlockStatement>()));
-  else_stmts.push_back(nullptr);
-
-  IfStatement stmt(cond, body);
-  stmt.set_else_statements(else_stmts);
+  IfStatement stmt(
+      Source{}, cond, body,
+      {
+          create<ElseStatement>(create<IdentifierExpression>("Ident"),
+                                create<BlockStatement>()),
+          create<ElseStatement>(create<BlockStatement>()),
+          nullptr,
+      });
   EXPECT_FALSE(stmt.IsValid());
 }
 
@@ -134,12 +125,11 @@
   auto* body = create<BlockStatement>();
   body->append(create<DiscardStatement>());
 
-  ElseStatementList else_stmts;
-  else_stmts.push_back(create<ElseStatement>(create<IdentifierExpression>(""),
-                                             create<BlockStatement>()));
-
-  IfStatement stmt(cond, body);
-  stmt.set_else_statements(else_stmts);
+  IfStatement stmt(Source{}, cond, body,
+                   {
+                       create<ElseStatement>(create<IdentifierExpression>(""),
+                                             create<BlockStatement>()),
+                   });
   EXPECT_FALSE(stmt.IsValid());
 }
 
@@ -148,12 +138,11 @@
   auto* body = create<BlockStatement>();
   body->append(create<DiscardStatement>());
 
-  ElseStatementList else_stmts;
-  else_stmts.push_back(create<ElseStatement>(create<BlockStatement>()));
-  else_stmts.push_back(create<ElseStatement>(create<BlockStatement>()));
-
-  IfStatement stmt(cond, body);
-  stmt.set_else_statements(else_stmts);
+  IfStatement stmt(Source{}, cond, body,
+                   {
+                       create<ElseStatement>(create<BlockStatement>()),
+                       create<ElseStatement>(create<BlockStatement>()),
+                   });
   EXPECT_FALSE(stmt.IsValid());
 }
 
@@ -162,13 +151,13 @@
   auto* body = create<BlockStatement>();
   body->append(create<DiscardStatement>());
 
-  ElseStatementList else_stmts;
-  else_stmts.push_back(create<ElseStatement>(create<BlockStatement>()));
-  else_stmts.push_back(create<ElseStatement>(
-      create<IdentifierExpression>("ident"), create<BlockStatement>()));
-
-  IfStatement stmt(cond, body);
-  stmt.set_else_statements(else_stmts);
+  IfStatement stmt(
+      Source{}, cond, body,
+      {
+          create<ElseStatement>(create<BlockStatement>()),
+          create<ElseStatement>(create<IdentifierExpression>("ident"),
+                                create<BlockStatement>()),
+      });
   EXPECT_FALSE(stmt.IsValid());
 }
 
@@ -177,7 +166,7 @@
   auto* body = create<BlockStatement>();
   body->append(create<DiscardStatement>());
 
-  IfStatement stmt(cond, body);
+  IfStatement stmt(Source{}, cond, body, ElseStatementList{});
 
   std::ostringstream out;
   stmt.to_str(out, 2);
@@ -204,13 +193,12 @@
   else_body->append(create<DiscardStatement>());
   else_body->append(create<DiscardStatement>());
 
-  ElseStatementList else_stmts;
-  else_stmts.push_back(create<ElseStatement>(
-      create<IdentifierExpression>("ident"), else_if_body));
-  else_stmts.push_back(create<ElseStatement>(else_body));
-
-  IfStatement stmt(cond, body);
-  stmt.set_else_statements(else_stmts);
+  IfStatement stmt(Source{}, cond, body,
+                   {
+                       create<ElseStatement>(
+                           create<IdentifierExpression>("ident"), else_if_body),
+                       create<ElseStatement>(else_body),
+                   });
 
   std::ostringstream out;
   stmt.to_str(out, 2);
diff --git a/src/ast/loop_statement_test.cc b/src/ast/loop_statement_test.cc
index 12001de..63d4e13 100644
--- a/src/ast/loop_statement_test.cc
+++ b/src/ast/loop_statement_test.cc
@@ -118,7 +118,8 @@
 TEST_F(LoopStatementTest, IsValid_InvalidBodyStatement) {
   auto* body = create<BlockStatement>();
   body->append(create<DiscardStatement>());
-  body->append(create<IfStatement>(nullptr, create<BlockStatement>()));
+  body->append(create<IfStatement>(Source{}, nullptr, create<BlockStatement>(),
+                                   ElseStatementList{}));
 
   auto* continuing = create<BlockStatement>();
   continuing->append(create<DiscardStatement>());
@@ -145,7 +146,8 @@
 
   auto* continuing = create<BlockStatement>();
   continuing->append(create<DiscardStatement>());
-  continuing->append(create<IfStatement>(nullptr, create<BlockStatement>()));
+  continuing->append(create<IfStatement>(
+      Source{}, nullptr, create<BlockStatement>(), ElseStatementList{}));
 
   LoopStatement l(body, continuing);
   EXPECT_FALSE(l.IsValid());
diff --git a/src/reader/spirv/function.cc b/src/reader/spirv/function.cc
index 5d25d78..26e6580 100644
--- a/src/reader/spirv/function.cc
+++ b/src/reader/spirv/function.cc
@@ -572,7 +572,8 @@
 
   auto* cond = create<ast::IdentifierExpression>(guard_name);
   auto* body = create<ast::BlockStatement>();
-  AddStatement(create<ast::IfStatement>(cond, body));
+  AddStatement(
+      create<ast::IfStatement>(Source{}, cond, body, ast::ElseStatementList{}));
   PushNewStatementBlock(top.construct_, end_id, body, nullptr, nullptr);
 }
 
@@ -582,7 +583,8 @@
 
   auto* cond = MakeTrue();
   auto* body = create<ast::BlockStatement>();
-  AddStatement(create<ast::IfStatement>(cond, body));
+  AddStatement(
+      create<ast::IfStatement>(Source{}, cond, body, ast::ElseStatementList{}));
   PushNewStatementBlock(top.construct_, end_id, body, nullptr, nullptr);
 }
 
@@ -2028,10 +2030,41 @@
       block_info.basic_block->terminator()->GetSingleWordInOperand(0);
   auto* cond = MakeExpression(condition_id).expr;
   auto* body = create<ast::BlockStatement>();
-  auto* const if_stmt = AddStatement(create<ast::IfStatement>(cond, body))
-                            ->As<ast::IfStatement>();
 
   // Generate the code for the condition.
+  // Use the IfBuilder to create the if-statement. The IfBuilder is constructed
+  // as a std::shared_ptr and is captured by the then and else clause
+  // CompletionAction lambdas, and so will only be destructed when the last
+  // block is completed. The IfBuilder destructor constructs the IfStatement,
+  // inserting it at the current insertion point in the current
+  // ast::BlockStatement.
+  struct IfBuilder {
+    IfBuilder(ast::Module* mod,
+              StatementBlock& statement_block,
+              tint::ast::Expression* cond,
+              ast::BlockStatement* body)
+        : mod_(mod),
+          dst_block_(statement_block.statements_),
+          dst_block_insertion_point_(statement_block.statements_->size()),
+          cond_(cond),
+          body_(body) {}
+
+    ~IfBuilder() {
+      dst_block_->insert(
+          dst_block_insertion_point_,
+          mod_->create<ast::IfStatement>(Source{}, cond_, body_, else_stmts_));
+    }
+
+    ast::Module* mod_;
+    ast::BlockStatement* dst_block_;
+    size_t dst_block_insertion_point_;
+    tint::ast::Expression* cond_;
+    ast::BlockStatement* body_;
+    ast::ElseStatementList else_stmts_;
+  };
+
+  auto if_builder = std::make_shared<IfBuilder>(
+      &ast_module_, statements_stack_.back(), cond, body);
 
   // Compute the block IDs that should end the then-clause and the else-clause.
 
@@ -2069,20 +2102,18 @@
 
   // Push statement blocks for the then-clause and the else-clause.
   // But make sure we do it in the right order.
-
-  auto push_else = [this, if_stmt, else_end, construct]() {
+  auto push_else = [this, if_builder, else_end, construct]() {
     // Push the else clause onto the stack first.
     auto* else_body = create<ast::BlockStatement>();
     PushNewStatementBlock(
-        construct, else_end, else_body, nullptr, [this, if_stmt, else_body]() {
+        construct, else_end, else_body, nullptr,
+        [this, if_builder, else_body]() {
           // Only set the else-clause if there are statements to fill it.
           if (!else_body->empty()) {
             // The "else" consists of the statement list from the top of
-            // statments stack, without an elseif condition.
-            ast::ElseStatementList else_stmts;
-            else_stmts.emplace_back(
+            // statements stack, without an elseif condition.
+            if_builder->else_stmts_.emplace_back(
                 create<ast::ElseStatement>(nullptr, else_body));
-            if_stmt->set_else_statements(else_stmts);
           }
         });
   };
@@ -2121,7 +2152,7 @@
     }
 
     // Push the then clause onto the stack.
-    PushNewStatementBlock(construct, then_end, body, nullptr, nullptr);
+    PushNewStatementBlock(construct, then_end, body, nullptr, [if_builder] {});
   }
 
   return success();
@@ -2439,18 +2470,17 @@
   if ((then_stmt == nullptr) && (else_stmt == nullptr)) {
     return nullptr;
   }
-  auto* if_block = create<ast::BlockStatement>();
-  auto* if_stmt = create<ast::IfStatement>(condition, if_block);
-  if (then_stmt != nullptr) {
-    if_block->append(then_stmt);
-  }
+  ast::ElseStatementList else_stmts;
   if (else_stmt != nullptr) {
     auto* stmts = create<ast::BlockStatement>();
     stmts->append(else_stmt);
-
-    ast::ElseStatementList else_stmts;
     else_stmts.emplace_back(create<ast::ElseStatement>(nullptr, stmts));
-    if_stmt->set_else_statements(else_stmts);
+  }
+  auto* if_block = create<ast::BlockStatement>();
+  auto* if_stmt =
+      create<ast::IfStatement>(Source{}, condition, if_block, else_stmts);
+  if (then_stmt != nullptr) {
+    if_block->append(then_stmt);
   }
   return if_stmt;
 }
diff --git a/src/reader/wgsl/parser_impl.cc b/src/reader/wgsl/parser_impl.cc
index 0e7e9d7..ec4a748 100644
--- a/src/reader/wgsl/parser_impl.cc
+++ b/src/reader/wgsl/parser_impl.cc
@@ -1647,14 +1647,11 @@
   auto el = else_stmt();
   if (el.errored)
     return Failure::kErrored;
-
-  auto* stmt = create<ast::IfStatement>(source, condition.value, body.value);
-  if (el.matched) {
+  if (el.matched)
     elseif.value.push_back(el.value);
-  }
-  stmt->set_else_statements(elseif.value);
 
-  return stmt;
+  return create<ast::IfStatement>(source, condition.value, body.value,
+                                  elseif.value);
 }
 
 // elseif_stmt
@@ -1945,8 +1942,9 @@
     auto* break_body = create<ast::BlockStatement>(not_condition->source());
     break_body->append(break_stmt);
     // if (!condition) { break; }
-    auto* break_if_not_condition = create<ast::IfStatement>(
-        not_condition->source(), not_condition, break_body);
+    auto* break_if_not_condition =
+        create<ast::IfStatement>(not_condition->source(), not_condition,
+                                 break_body, ast::ElseStatementList{});
     body->insert(0, break_if_not_condition);
   }
 
diff --git a/src/type_determiner_test.cc b/src/type_determiner_test.cc
index af76cee..fc1df28 100644
--- a/src/type_determiner_test.cc
+++ b/src/type_determiner_test.cc
@@ -225,9 +225,6 @@
                                      create<ast::SintLiteral>(&i32, 3)),
                                  else_body);
 
-  ast::ElseStatementList else_stmts;
-  else_stmts.push_back(else_stmt);
-
   auto* lhs = create<ast::ScalarConstructorExpression>(
       create<ast::SintLiteral>(&i32, 2));
   auto* rhs = create<ast::ScalarConstructorExpression>(
@@ -236,10 +233,10 @@
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(lhs, rhs));
 
-  ast::IfStatement stmt(create<ast::ScalarConstructorExpression>(
+  ast::IfStatement stmt(Source{},
+                        create<ast::ScalarConstructorExpression>(
                             create<ast::SintLiteral>(&i32, 3)),
-                        body);
-  stmt.set_else_statements(else_stmts);
+                        body, ast::ElseStatementList{else_stmt});
 
   EXPECT_TRUE(td()->DetermineResultType(&stmt));
   ASSERT_NE(stmt.condition()->result_type(), nullptr);
diff --git a/src/validator/validator_test.cc b/src/validator/validator_test.cc
index e15a5e0..1978a18 100644
--- a/src/validator/validator_test.cc
+++ b/src/validator/validator_test.cc
@@ -353,7 +353,8 @@
       create<ast::FloatLiteral>(&f32, 3.14f));
 
   auto* outer_body = create<ast::BlockStatement>();
-  outer_body->append(create<ast::IfStatement>(cond, body));
+  outer_body->append(
+      create<ast::IfStatement>(Source{}, cond, body, ast::ElseStatementList{}));
   outer_body->append(create<ast::AssignmentStatement>(
       Source{Source::Location{12, 34}}, lhs, rhs));
 
@@ -388,7 +389,8 @@
 
   auto* outer_body = create<ast::BlockStatement>();
   outer_body->append(create<ast::VariableDeclStatement>(var));
-  outer_body->append(create<ast::IfStatement>(cond, body));
+  outer_body->append(
+      create<ast::IfStatement>(Source{}, cond, body, ast::ElseStatementList{}));
   EXPECT_TRUE(td()->DetermineStatements(outer_body)) << td()->error();
   ASSERT_NE(lhs->result_type(), nullptr);
   ASSERT_NE(rhs->result_type(), nullptr);
@@ -554,7 +556,8 @@
       create<ast::FloatLiteral>(&f32, 3.14)));
 
   auto* outer_body = create<ast::BlockStatement>();
-  outer_body->append(create<ast::IfStatement>(cond, body));
+  outer_body->append(
+      create<ast::IfStatement>(Source{}, cond, body, ast::ElseStatementList{}));
   outer_body->append(create<ast::VariableDeclStatement>(
       Source{Source::Location{12, 34}}, var_a_float));
 
@@ -588,7 +591,8 @@
 
   auto* outer_body = create<ast::BlockStatement>();
   outer_body->append(create<ast::VariableDeclStatement>(var_a_float));
-  outer_body->append(create<ast::IfStatement>(cond, body));
+  outer_body->append(
+      create<ast::IfStatement>(Source{}, cond, body, ast::ElseStatementList{}));
 
   EXPECT_TRUE(td()->DetermineStatements(outer_body)) << td()->error();
   EXPECT_FALSE(v()->ValidateStatements(outer_body));
diff --git a/src/writer/hlsl/generator_impl_binary_test.cc b/src/writer/hlsl/generator_impl_binary_test.cc
index 18c46f3..52e5fff 100644
--- a/src/writer/hlsl/generator_impl_binary_test.cc
+++ b/src/writer/hlsl/generator_impl_binary_test.cc
@@ -388,21 +388,21 @@
                                     create<ast::IdentifierExpression>("c")),
       body);
 
-  ast::ElseStatementList else_stmts;
-  else_stmts.push_back(else_if_stmt);
-  else_stmts.push_back(else_stmt);
-
   body = create<ast::BlockStatement>();
   body->append(
       create<ast::ReturnStatement>(create<ast::ScalarConstructorExpression>(
           create<ast::SintLiteral>(&i32, 1))));
 
   ast::IfStatement expr(
+      Source{},
       create<ast::BinaryExpression>(ast::BinaryOp::kLogicalAnd,
                                     create<ast::IdentifierExpression>("a"),
                                     create<ast::IdentifierExpression>("b")),
-      body);
-  expr.set_else_statements(else_stmts);
+      body,
+      {
+          else_if_stmt,
+          else_stmt,
+      });
 
   ASSERT_TRUE(gen.EmitStatement(out, &expr)) << gen.error();
   EXPECT_EQ(result(), R"(bool _tint_tmp = a;
diff --git a/src/writer/hlsl/generator_impl_function_test.cc b/src/writer/hlsl/generator_impl_function_test.cc
index 476f720..b3bb0e7 100644
--- a/src/writer/hlsl/generator_impl_function_test.cc
+++ b/src/writer/hlsl/generator_impl_function_test.cc
@@ -889,12 +889,13 @@
   list->append(create<ast::ReturnStatement>());
 
   body->append(create<ast::IfStatement>(
+      Source{},
       create<ast::BinaryExpression>(ast::BinaryOp::kEqual,
                                     create<ast::ScalarConstructorExpression>(
                                         create<ast::SintLiteral>(&i32, 1)),
                                     create<ast::ScalarConstructorExpression>(
                                         create<ast::SintLiteral>(&i32, 1))),
-      list));
+      list, ast::ElseStatementList{}));
 
   body->append(create<ast::ReturnStatement>());
   auto* func_1 = create<ast::Function>("ep_1", params, &void_type, body);
diff --git a/src/writer/hlsl/generator_impl_if_test.cc b/src/writer/hlsl/generator_impl_if_test.cc
index 6e9bd7a..0ce0fed 100644
--- a/src/writer/hlsl/generator_impl_if_test.cc
+++ b/src/writer/hlsl/generator_impl_if_test.cc
@@ -31,7 +31,7 @@
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
 
-  ast::IfStatement i(cond, body);
+  ast::IfStatement i(Source{}, cond, body, ast::ElseStatementList{});
   gen.increment_indent();
 
   ASSERT_TRUE(gen.EmitStatement(out, &i)) << gen.error();
@@ -46,15 +46,12 @@
   auto* else_body = create<ast::BlockStatement>();
   else_body->append(create<ast::ReturnStatement>());
 
-  ast::ElseStatementList elses;
-  elses.push_back(create<ast::ElseStatement>(else_cond, else_body));
-
   auto* cond = create<ast::IdentifierExpression>("cond");
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
 
-  ast::IfStatement i(cond, body);
-  i.set_else_statements(elses);
+  ast::IfStatement i(Source{}, cond, body,
+                     {create<ast::ElseStatement>(else_cond, else_body)});
 
   gen.increment_indent();
 
@@ -73,15 +70,12 @@
   auto* else_body = create<ast::BlockStatement>();
   else_body->append(create<ast::ReturnStatement>());
 
-  ast::ElseStatementList elses;
-  elses.push_back(create<ast::ElseStatement>(else_body));
-
   auto* cond = create<ast::IdentifierExpression>("cond");
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
 
-  ast::IfStatement i(cond, body);
-  i.set_else_statements(elses);
+  ast::IfStatement i(Source{}, cond, body,
+                     {create<ast::ElseStatement>(else_body)});
 
   gen.increment_indent();
 
@@ -103,16 +97,15 @@
   auto* else_body_2 = create<ast::BlockStatement>();
   else_body_2->append(create<ast::ReturnStatement>());
 
-  ast::ElseStatementList elses;
-  elses.push_back(create<ast::ElseStatement>(else_cond, else_body));
-  elses.push_back(create<ast::ElseStatement>(else_body_2));
-
   auto* cond = create<ast::IdentifierExpression>("cond");
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
 
-  ast::IfStatement i(cond, body);
-  i.set_else_statements(elses);
+  ast::IfStatement i(Source{}, cond, body,
+                     {
+                         create<ast::ElseStatement>(else_cond, else_body),
+                         create<ast::ElseStatement>(else_body_2),
+                     });
 
   gen.increment_indent();
 
diff --git a/src/writer/msl/generator_impl_function_test.cc b/src/writer/msl/generator_impl_function_test.cc
index 53378d4..088455a 100644
--- a/src/writer/msl/generator_impl_function_test.cc
+++ b/src/writer/msl/generator_impl_function_test.cc
@@ -911,12 +911,13 @@
   list->append(create<ast::ReturnStatement>());
 
   body->append(create<ast::IfStatement>(
+      Source{},
       create<ast::BinaryExpression>(ast::BinaryOp::kEqual,
                                     create<ast::ScalarConstructorExpression>(
                                         create<ast::SintLiteral>(&i32, 1)),
                                     create<ast::ScalarConstructorExpression>(
                                         create<ast::SintLiteral>(&i32, 1))),
-      list));
+      list, ast::ElseStatementList{}));
 
   body->append(create<ast::ReturnStatement>());
 
diff --git a/src/writer/msl/generator_impl_if_test.cc b/src/writer/msl/generator_impl_if_test.cc
index 71428a8..dd34be7 100644
--- a/src/writer/msl/generator_impl_if_test.cc
+++ b/src/writer/msl/generator_impl_if_test.cc
@@ -33,7 +33,7 @@
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
 
-  ast::IfStatement i(cond, body);
+  ast::IfStatement i(Source{}, cond, body, ast::ElseStatementList{});
 
   gen.increment_indent();
 
@@ -49,15 +49,12 @@
   auto* else_body = create<ast::BlockStatement>();
   else_body->append(create<ast::ReturnStatement>());
 
-  ast::ElseStatementList elses;
-  elses.push_back(create<ast::ElseStatement>(else_cond, else_body));
-
   auto* cond = create<ast::IdentifierExpression>("cond");
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
 
-  ast::IfStatement i(cond, body);
-  i.set_else_statements(elses);
+  ast::IfStatement i(Source{}, cond, body,
+                     {create<ast::ElseStatement>(else_cond, else_body)});
 
   gen.increment_indent();
 
@@ -74,15 +71,12 @@
   auto* else_body = create<ast::BlockStatement>();
   else_body->append(create<ast::ReturnStatement>());
 
-  ast::ElseStatementList elses;
-  elses.push_back(create<ast::ElseStatement>(else_body));
-
   auto* cond = create<ast::IdentifierExpression>("cond");
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
 
-  ast::IfStatement i(cond, body);
-  i.set_else_statements(elses);
+  ast::IfStatement i(Source{}, cond, body,
+                     {create<ast::ElseStatement>(else_body)});
 
   gen.increment_indent();
 
@@ -104,16 +98,15 @@
   auto* else_body_2 = create<ast::BlockStatement>();
   else_body_2->append(create<ast::ReturnStatement>());
 
-  ast::ElseStatementList elses;
-  elses.push_back(create<ast::ElseStatement>(else_cond, else_body));
-  elses.push_back(create<ast::ElseStatement>(else_body_2));
-
   auto* cond = create<ast::IdentifierExpression>("cond");
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
 
-  ast::IfStatement i(cond, body);
-  i.set_else_statements(elses);
+  ast::IfStatement i(Source{}, cond, body,
+                     {
+                         create<ast::ElseStatement>(else_cond, else_body),
+                         create<ast::ElseStatement>(else_body_2),
+                     });
 
   gen.increment_indent();
 
diff --git a/src/writer/spirv/builder_if_test.cc b/src/writer/spirv/builder_if_test.cc
index 71115c3..768604a 100644
--- a/src/writer/spirv/builder_if_test.cc
+++ b/src/writer/spirv/builder_if_test.cc
@@ -48,7 +48,8 @@
   auto* cond = create<ast::ScalarConstructorExpression>(
       create<ast::BoolLiteral>(&bool_type, true));
 
-  ast::IfStatement expr(cond, create<ast::BlockStatement>());
+  ast::IfStatement expr(Source{}, cond, create<ast::BlockStatement>(),
+                        ast::ElseStatementList{});
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
@@ -85,7 +86,7 @@
   auto* cond = create<ast::ScalarConstructorExpression>(
       create<ast::BoolLiteral>(&bool_type, true));
 
-  ast::IfStatement expr(cond, body);
+  ast::IfStatement expr(Source{}, cond, body, ast::ElseStatementList{});
 
   td.RegisterVariableForTesting(var);
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
@@ -135,14 +136,11 @@
                                        create<ast::ScalarConstructorExpression>(
                                            create<ast::SintLiteral>(&i32, 3))));
 
-  ast::ElseStatementList else_stmts;
-  else_stmts.push_back(create<ast::ElseStatement>(else_body));
-
   auto* cond = create<ast::ScalarConstructorExpression>(
       create<ast::BoolLiteral>(&bool_type, true));
 
-  ast::IfStatement expr(cond, body);
-  expr.set_else_statements(else_stmts);
+  ast::IfStatement expr(Source{}, cond, body,
+                        {create<ast::ElseStatement>(else_body)});
 
   td.RegisterVariableForTesting(var);
 
@@ -200,14 +198,11 @@
   auto* else_cond = create<ast::ScalarConstructorExpression>(
       create<ast::BoolLiteral>(&bool_type, true));
 
-  ast::ElseStatementList else_stmts;
-  else_stmts.push_back(create<ast::ElseStatement>(else_cond, else_body));
-
   auto* cond = create<ast::ScalarConstructorExpression>(
       create<ast::BoolLiteral>(&bool_type, true));
 
-  ast::IfStatement expr(cond, body);
-  expr.set_else_statements(else_stmts);
+  ast::IfStatement expr(Source{}, cond, body,
+                        {create<ast::ElseStatement>(else_cond, else_body)});
 
   td.RegisterVariableForTesting(var);
 
@@ -285,18 +280,16 @@
   auto* elseif_2_cond = create<ast::ScalarConstructorExpression>(
       create<ast::BoolLiteral>(&bool_type, false));
 
-  ast::ElseStatementList else_stmts;
-  else_stmts.push_back(
-      create<ast::ElseStatement>(elseif_1_cond, elseif_1_body));
-  else_stmts.push_back(
-      create<ast::ElseStatement>(elseif_2_cond, elseif_2_body));
-  else_stmts.push_back(create<ast::ElseStatement>(else_body));
-
   auto* cond = create<ast::ScalarConstructorExpression>(
       create<ast::BoolLiteral>(&bool_type, true));
 
-  ast::IfStatement expr(cond, body);
-  expr.set_else_statements(else_stmts);
+  ast::IfStatement expr(
+      Source{}, cond, body,
+      {
+          create<ast::ElseStatement>(elseif_1_cond, elseif_1_body),
+          create<ast::ElseStatement>(elseif_2_cond, elseif_2_body),
+          create<ast::ElseStatement>(else_body),
+      });
 
   td.RegisterVariableForTesting(var);
 
@@ -360,7 +353,8 @@
   auto* if_body = create<ast::BlockStatement>();
   if_body->append(create<ast::BreakStatement>());
 
-  auto* if_stmt = create<ast::IfStatement>(cond, if_body);
+  auto* if_stmt = create<ast::IfStatement>(Source{}, cond, if_body,
+                                           ast::ElseStatementList{});
 
   auto* loop_body = create<ast::BlockStatement>();
   loop_body->append(if_stmt);
@@ -407,11 +401,9 @@
   auto* else_body = create<ast::BlockStatement>();
   else_body->append(create<ast::BreakStatement>());
 
-  ast::ElseStatementList else_stmts;
-  else_stmts.push_back(create<ast::ElseStatement>(else_body));
-
-  auto* if_stmt = create<ast::IfStatement>(cond, create<ast::BlockStatement>());
-  if_stmt->set_else_statements(else_stmts);
+  auto* if_stmt = create<ast::IfStatement>(
+      Source{}, cond, create<ast::BlockStatement>(),
+      ast::ElseStatementList{create<ast::ElseStatement>(else_body)});
 
   auto* loop_body = create<ast::BlockStatement>();
   loop_body->append(if_stmt);
@@ -459,7 +451,8 @@
   auto* if_body = create<ast::BlockStatement>();
   if_body->append(create<ast::ContinueStatement>());
 
-  auto* if_stmt = create<ast::IfStatement>(cond, if_body);
+  auto* if_stmt = create<ast::IfStatement>(Source{}, cond, if_body,
+                                           ast::ElseStatementList{});
 
   auto* loop_body = create<ast::BlockStatement>();
   loop_body->append(if_stmt);
@@ -506,11 +499,9 @@
   auto* else_body = create<ast::BlockStatement>();
   else_body->append(create<ast::ContinueStatement>());
 
-  ast::ElseStatementList else_stmts;
-  else_stmts.push_back(create<ast::ElseStatement>(else_body));
-
-  auto* if_stmt = create<ast::IfStatement>(cond, create<ast::BlockStatement>());
-  if_stmt->set_else_statements(else_stmts);
+  auto* if_stmt = create<ast::IfStatement>(
+      Source{}, cond, create<ast::BlockStatement>(),
+      ast::ElseStatementList{create<ast::ElseStatement>(else_body)});
 
   auto* loop_body = create<ast::BlockStatement>();
   loop_body->append(if_stmt);
@@ -556,7 +547,7 @@
   auto* if_body = create<ast::BlockStatement>();
   if_body->append(create<ast::ReturnStatement>());
 
-  ast::IfStatement expr(cond, if_body);
+  ast::IfStatement expr(Source{}, cond, if_body, ast::ElseStatementList{});
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
@@ -588,7 +579,7 @@
   auto* if_body = create<ast::BlockStatement>();
   if_body->append(create<ast::ReturnStatement>(cond2));
 
-  ast::IfStatement expr(cond, if_body);
+  ast::IfStatement expr(Source{}, cond, if_body, ast::ElseStatementList{});
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
@@ -618,8 +609,9 @@
       create<ast::Variable>("a", ast::StorageClass::kFunction, &bool_type);
   td.RegisterVariableForTesting(var);
 
-  ast::IfStatement expr(create<ast::IdentifierExpression>("a"),
-                        create<ast::BlockStatement>());
+  ast::IfStatement expr(Source{}, create<ast::IdentifierExpression>("a"),
+                        create<ast::BlockStatement>(),
+                        ast::ElseStatementList{});
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
diff --git a/src/writer/spirv/builder_switch_test.cc b/src/writer/spirv/builder_switch_test.cc
index add53b8..aeb8724 100644
--- a/src/writer/spirv/builder_switch_test.cc
+++ b/src/writer/spirv/builder_switch_test.cc
@@ -447,9 +447,10 @@
 
   auto* case_1_body = create<ast::BlockStatement>();
   case_1_body->append(
-      create<ast::IfStatement>(create<ast::ScalarConstructorExpression>(
+      create<ast::IfStatement>(Source{},
+                               create<ast::ScalarConstructorExpression>(
                                    create<ast::BoolLiteral>(&bool_type, true)),
-                               if_body));
+                               if_body, ast::ElseStatementList{}));
 
   case_1_body->append(
       create<ast::AssignmentStatement>(create<ast::IdentifierExpression>("v"),
diff --git a/src/writer/wgsl/generator_impl_if_test.cc b/src/writer/wgsl/generator_impl_if_test.cc
index b0c68b8..b2e2205 100644
--- a/src/writer/wgsl/generator_impl_if_test.cc
+++ b/src/writer/wgsl/generator_impl_if_test.cc
@@ -32,7 +32,7 @@
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::DiscardStatement>());
 
-  ast::IfStatement i(cond, body);
+  ast::IfStatement i(Source{}, cond, body, ast::ElseStatementList{});
 
   gen.increment_indent();
 
@@ -48,15 +48,12 @@
   auto* else_body = create<ast::BlockStatement>();
   else_body->append(create<ast::DiscardStatement>());
 
-  ast::ElseStatementList elses;
-  elses.push_back(create<ast::ElseStatement>(else_cond, else_body));
-
   auto* cond = create<ast::IdentifierExpression>("cond");
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::DiscardStatement>());
 
-  ast::IfStatement i(cond, body);
-  i.set_else_statements(elses);
+  ast::IfStatement i(Source{}, cond, body,
+                     {create<ast::ElseStatement>(else_cond, else_body)});
 
   gen.increment_indent();
 
@@ -73,15 +70,12 @@
   auto* else_body = create<ast::BlockStatement>();
   else_body->append(create<ast::DiscardStatement>());
 
-  ast::ElseStatementList elses;
-  elses.push_back(create<ast::ElseStatement>(else_body));
-
   auto* cond = create<ast::IdentifierExpression>("cond");
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::DiscardStatement>());
 
-  ast::IfStatement i(cond, body);
-  i.set_else_statements(elses);
+  ast::IfStatement i(Source{}, cond, body,
+                     {create<ast::ElseStatement>(else_body)});
 
   gen.increment_indent();
 
@@ -103,16 +97,15 @@
   auto* else_body_2 = create<ast::BlockStatement>();
   else_body_2->append(create<ast::DiscardStatement>());
 
-  ast::ElseStatementList elses;
-  elses.push_back(create<ast::ElseStatement>(else_cond, else_body));
-  elses.push_back(create<ast::ElseStatement>(else_body_2));
-
   auto* cond = create<ast::IdentifierExpression>("cond");
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::DiscardStatement>());
 
-  ast::IfStatement i(cond, body);
-  i.set_else_statements(elses);
+  ast::IfStatement i(Source{}, cond, body,
+                     {
+                         create<ast::ElseStatement>(else_cond, else_body),
+                         create<ast::ElseStatement>(else_body_2),
+                     });
 
   gen.increment_indent();