[spirv-reader] use stack of statement lists

This is preparation for emitting nested control flow.

Bug: tint:3
Change-Id: I90fc7edba8cb9937f722e6f5e94c7f222d34c403
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/21801
Reviewed-by: dan sinclair <dsinclair@google.com>

src/reader/spirv/function.cc

@@ -367,10 +367,22 @@ FunctionEmitter::FunctionEmitter(ParserImpl* pi,
       type_mgr_(ir_context_.get_type_mgr()),
       fail_stream_(pi->fail_stream()),
       namer_(pi->namer()),
-      function_(function) {}
+      function_(function) {
+  statements_stack_.emplace_back(ast::StatementList{});
+}
 
 FunctionEmitter::~FunctionEmitter() = default;
 
+const ast::StatementList& FunctionEmitter::ast_body() {
+  assert(!statements_stack_.empty());
+  return statements_stack_[0];
+}
+
+void FunctionEmitter::AddStatement(std::unique_ptr<ast::Statement> statement) {
+  assert(!statements_stack_.empty());
+  statements_stack_.back().emplace_back(std::move(statement));
+}
+
 bool FunctionEmitter::Emit() {
   if (failed()) {
     return false;
@@ -389,7 +401,15 @@ bool FunctionEmitter::Emit() {
   }
 
   // Set the body of the AST function node.
-  parser_impl_.get_module().functions().back()->set_body(std::move(ast_body_));
+  if (statements_stack_.size() != 1) {
+    return Fail() << "internal error: statement-list stack should have 1 "
+                     "element but has "
+                  << statements_stack_.size();
+  }
+  ast::StatementList body(std::move(statements_stack_[0]));
+  parser_impl_.get_module().functions().back()->set_body(std::move(body));
+  // Maintain the invariant by repopulating the one and only element.
+  statements_stack_[0] = ast::StatementList{};
 
   return success();
 }
@@ -1363,7 +1383,7 @@ bool FunctionEmitter::EmitFunctionVariables() {
     // TODO(dneto): Add the initializer via Variable::set_constructor.
     auto var_decl_stmt =
         std::make_unique<ast::VariableDeclStatement>(std::move(var));
-    ast_body_.emplace_back(std::move(var_decl_stmt));
+    AddStatement(std::move(var_decl_stmt));
     // Save this as an already-named value.
     identifier_values_.insert(inst.result_id());
   }
@@ -1445,7 +1465,7 @@ bool FunctionEmitter::EmitConstDefinition(
   }
   ast_const->set_constructor(std::move(ast_expr.expr));
   ast_const->set_is_const(true);
-  ast_body_.emplace_back(
+  AddStatement(
       std::make_unique<ast::VariableDeclStatement>(std::move(ast_const)));
   // Save this as an already-named value.
   identifier_values_.insert(inst.result_id());
@@ -1476,7 +1496,7 @@ bool FunctionEmitter::EmitStatement(const spvtools::opt::Instruction& inst) {
       // TODO(dneto): Order of evaluation?
       auto lhs = MakeExpression(inst.GetSingleWordInOperand(0));
       auto rhs = MakeExpression(inst.GetSingleWordInOperand(1));
-      ast_body_.emplace_back(std::make_unique<ast::AssignmentStatement>(
+      AddStatement(std::make_unique<ast::AssignmentStatement>(
           std::move(lhs.expr), std::move(rhs.expr)));
       return success();
     }

src/reader/spirv/function.h

@@ -172,8 +172,10 @@ class FunctionEmitter {
   /// @returns true if emission has failed.
   bool failed() const { return !success(); }
 
+  /// Returns the body of the function.  It is the bottom of the statement
+  /// stack.
   /// @returns the body of the function.
-  const ast::StatementList& ast_body() { return ast_body_; }
+  const ast::StatementList& ast_body();
 
   /// Records failure.
   /// @returns a FailStream on which to emit diagnostics.
@@ -188,7 +190,8 @@ class FunctionEmitter {
   /// @returns true if emission has not yet failed.
   bool EmitFunctionDeclaration();
 
-  /// Emits the function body, populating |ast_body_|
+  /// Emits the function body, populating the bottom entry of the statements
+  /// stack.
   /// @returns false if emission failed.
   bool EmitBody();
 
@@ -343,6 +346,9 @@ class FunctionEmitter {
   /// or nullptr
   BlockInfo* HeaderIfBreakable(const Construct* c);
 
+  /// Appends a new statement to the top of the statement stack.
+  void AddStatement(std::unique_ptr<ast::Statement> statement);
+
   ParserImpl& parser_impl_;
   ast::Module& ast_module_;
   spvtools::opt::IRContext& ir_context_;
@@ -352,7 +358,12 @@ class FunctionEmitter {
   FailStream& fail_stream_;
   Namer& namer_;
   const spvtools::opt::Function& function_;
-  ast::StatementList ast_body_;
+
+  // A stack of statement lists. Each list is contained in a construct in
+  // the next deeper element of stack. The 0th entry represents the statements
+  // for the entire function.  This stack is never empty.
+  std::vector<ast::StatementList> statements_stack_;
+
   // The set of IDs that have already had an identifier name generated for it.
   std::unordered_set<uint32_t> identifier_values_;
   // Mapping from SPIR-V ID that is used at most once, to its AST expression.