Split global init out from InstBlockStack. (#4101)

Creates a `GlobalInit` class for storing relevant values, pulling
functions off `InstBlockStack` and `Context`. Adds a `Context` pointer
just so that it doesn't need to be passed in on each call (`Finalize` in
particular uses several members).

Note we have several different `InstBlockStack` instances, so several
copies of the relevant members were simply unused.

toolchain/check/BUILD

@@ -21,6 +21,7 @@ cc_library(
         "eval.cpp",
         "function.cpp",
         "generic.cpp",
+        "global_init.cpp",
         "import_ref.cpp",
         "inst_block_stack.cpp",
         "merge.cpp",
@@ -38,6 +39,7 @@ cc_library(
         "eval.h",
         "function.h",
         "generic.h",
+        "global_init.h",
         "import_ref.h",
         "inst_block_stack.h",
         "keyword_modifier_set.h",

toolchain/check/check.cpp

@@ -915,7 +915,7 @@ static auto CheckParseTree(
   sem_ir.set_top_inst_block_id(context.inst_block_stack().Pop());
   context.scope_stack().Pop();
   context.FinalizeExports();
-  context.FinalizeGlobalInit();
+  context.global_init().Finalize();
 
   DiagnoseMissingDefinitions(context, emitter);
 

toolchain/check/context.cpp

@@ -45,7 +45,8 @@ Context::Context(const Lex::TokenizedBuffer& tokens, DiagnosticEmitter& emitter,
       param_and_arg_refs_stack_(sem_ir, vlog_stream, node_stack_),
       args_type_info_stack_("args_type_info_stack_", sem_ir, vlog_stream),
       decl_name_stack_(this),
-      scope_stack_(sem_ir_->identifiers()) {
+      scope_stack_(sem_ir_->identifiers()),
+      global_init_(this) {
   // Map the builtin `<error>` and `type` type constants to their corresponding
   // special `TypeId` values.
   type_ids_for_type_constants_.Insert(
@@ -632,31 +633,6 @@ auto Context::is_current_position_reachable() -> bool {
          SemIR::TerminatorKind::Terminator;
 }
 
-auto Context::FinalizeGlobalInit() -> void {
-  inst_block_stack().PushGlobalInit();
-  if (!inst_block_stack().PeekCurrentBlockContents().empty()) {
-    AddInst<SemIR::Return>(Parse::NodeId::Invalid, {});
-    // Pop the GlobalInit block here to finalize it.
-    inst_block_stack().Pop();
-
-    // __global_init is only added if there are initialization instructions.
-    auto name_id = sem_ir().identifiers().Add("__global_init");
-    sem_ir().functions().Add(
-        {.name_id = SemIR::NameId::ForIdentifier(name_id),
-         .parent_scope_id = SemIR::NameScopeId::Package,
-         .decl_id = SemIR::InstId::Invalid,
-         .generic_id = SemIR::GenericId::Invalid,
-         .implicit_param_refs_id = SemIR::InstBlockId::Invalid,
-         .param_refs_id = SemIR::InstBlockId::Empty,
-         .return_storage_id = SemIR::InstId::Invalid,
-         .is_extern = false,
-         .return_slot = SemIR::Function::ReturnSlot::Absent,
-         .body_block_ids = {SemIR::InstBlockId::GlobalInit}});
-  } else {
-    inst_block_stack().PopGlobalInit();
-  }
-}
-
 namespace {
 // Worklist-based type completion mechanism.
 //

toolchain/check/context.h

@@ -12,6 +12,7 @@
 #include "toolchain/check/decl_name_stack.h"
 #include "toolchain/check/diagnostic_helpers.h"
 #include "toolchain/check/generic_region_stack.h"
+#include "toolchain/check/global_init.h"
 #include "toolchain/check/inst_block_stack.h"
 #include "toolchain/check/node_stack.h"
 #include "toolchain/check/param_and_arg_refs_stack.h"
@@ -296,9 +297,6 @@ class Context {
     inst_blocks().Set(SemIR::InstBlockId::Exports, exports_);
   }
 
-  // Finalizes the initialization function (__global_init).
-  auto FinalizeGlobalInit() -> void;
-
   // Sets the total number of IRs which exist. This is used to prepare a map
   // from IR to imported IR.
   auto SetTotalIRCount(int num_irs) -> void {
@@ -427,6 +425,8 @@ class Context {
     return definitions_required_;
   }
 
+  auto global_init() -> GlobalInit& { return global_init_; }
+
  private:
   // A FoldingSet node for a type.
   class TypeNode : public llvm::FastFoldingSetNode {
@@ -513,6 +513,9 @@ class Context {
   // Declaration instructions of entities that should have definitions by the
   // end of the current source file.
   llvm::SmallVector<SemIR::InstId> definitions_required_;
+
+  // State for global initialization.
+  GlobalInit global_init_;
 };
 
 }  // namespace Carbon::Check

toolchain/check/global_init.cpp

@@ -0,0 +1,50 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "toolchain/check/global_init.h"
+
+#include "toolchain/check/context.h"
+
+namespace Carbon::Check {
+
+auto GlobalInit::Resume() -> void {
+  context_->inst_block_stack().Push(block_id_, block_);
+}
+
+auto GlobalInit::Suspend() -> void {
+  // TODO: Consider splicing together blocks in order to avoid sizable copies
+  // here.
+  auto contents = context_->inst_block_stack().PeekCurrentBlockContents();
+  block_.assign(contents.begin(), contents.end());
+
+  block_id_ = context_->inst_block_stack().PeekOrAdd();
+  context_->inst_block_stack().PopAndDiscard();
+}
+
+auto GlobalInit::Finalize() -> void {
+  // __global_init is only added if there are initialization instructions.
+  if (block_.empty() && block_id_ == SemIR::InstBlockId::GlobalInit) {
+    return;
+  }
+
+  Resume();
+  context_->AddInst<SemIR::Return>(Parse::NodeId::Invalid, {});
+  // Pop the GlobalInit block here to finalize it.
+  context_->inst_block_stack().Pop();
+
+  auto name_id = context_->sem_ir().identifiers().Add("__global_init");
+  context_->sem_ir().functions().Add(
+      {.name_id = SemIR::NameId::ForIdentifier(name_id),
+       .parent_scope_id = SemIR::NameScopeId::Package,
+       .decl_id = SemIR::InstId::Invalid,
+       .generic_id = SemIR::GenericId::Invalid,
+       .implicit_param_refs_id = SemIR::InstBlockId::Invalid,
+       .param_refs_id = SemIR::InstBlockId::Empty,
+       .return_storage_id = SemIR::InstId::Invalid,
+       .is_extern = false,
+       .return_slot = SemIR::Function::ReturnSlot::Absent,
+       .body_block_ids = {SemIR::InstBlockId::GlobalInit}});
+}
+
+}  // namespace Carbon::Check

toolchain/check/global_init.h

@@ -0,0 +1,47 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef CARBON_TOOLCHAIN_CHECK_GLOBAL_INIT_H_
+#define CARBON_TOOLCHAIN_CHECK_GLOBAL_INIT_H_
+
+#include "llvm/ADT/SmallVector.h"
+#include "toolchain/sem_ir/ids.h"
+
+namespace Carbon::Check {
+
+class Context;
+
+// Tracks state for global initialization. Handles should `Resume` when entering
+// an expression that's used for global init, and `Suspend` when the expression
+// is finished. Instructions in the middle will be tracked for the
+// `__global_init` function.
+class GlobalInit {
+ public:
+  explicit GlobalInit(Context* context) : context_(context) {}
+
+  // Resumes adding instructions to global init.
+  auto Resume() -> void;
+
+  // Suspends adding instructions to global init.
+  auto Suspend() -> void;
+
+  // Finalizes the global initialization state, creating `__global_init` if
+  // needed. Only called once at the end of checking.
+  auto Finalize() -> void;
+
+ private:
+  // The associated context. Stored for convenience.
+  Context* context_;
+
+  // The currently suspended global init block. The value may change as a result
+  // of control flow in initialization.
+  SemIR::InstBlockId block_id_ = SemIR::InstBlockId::GlobalInit;
+
+  // The contents for the currently suspended global init block.
+  llvm::SmallVector<SemIR::InstId> block_;
+};
+
+}  // namespace Carbon::Check
+
+#endif  // CARBON_TOOLCHAIN_CHECK_GLOBAL_INIT_H_

toolchain/check/handle_let_and_var.cpp

@@ -43,7 +43,7 @@ auto HandleReturnedModifier(Context& context, Parse::ReturnedModifierId node_id)
 
 static auto HandleInitializer(Context& context, Parse::NodeId node_id) -> bool {
   if (context.scope_stack().PeekIndex() == ScopeIndex::Package) {
-    context.inst_block_stack().PushGlobalInit();
+    context.global_init().Resume();
   }
   context.node_stack().Push(node_id);
   return true;
@@ -155,7 +155,7 @@ static auto HandleDecl(Context& context, NodeT node_id)
   if (context.node_stack().PeekIs<Parse::NodeKind::TuplePattern>()) {
     if (decl_info->init_id &&
         context.scope_stack().PeekIndex() == ScopeIndex::Package) {
-      context.inst_block_stack().PopGlobalInit();
+      context.global_init().Suspend();
     }
     context.TODO(node_id, "tuple pattern in let/var");
     decl_info = std::nullopt;
@@ -250,7 +250,7 @@ auto HandleLetDecl(Context& context, Parse::LetDeclId node_id) -> bool {
 
   if (decl_info->init_id &&
       context.scope_stack().PeekIndex() == ScopeIndex::Package) {
-    context.inst_block_stack().PopGlobalInit();
+    context.global_init().Suspend();
   }
 
   return true;
@@ -288,7 +288,7 @@ auto HandleVariableDecl(Context& context, Parse::VariableDeclId node_id)
     }
 
     if (context.scope_stack().PeekIndex() == ScopeIndex::Package) {
-      context.inst_block_stack().PopGlobalInit();
+      context.global_init().Suspend();
     }
   }
 

toolchain/check/inst_block_stack.cpp

@@ -21,9 +21,10 @@ auto InstBlockStack::Push(SemIR::InstBlockId id) -> void {
   ++size_;
 }
 
-auto InstBlockStack::PushGlobalInit() -> void {
-  Push(init_block_id_);
-  stack_[size_ - 1].content = std::move(init_block_);
+auto InstBlockStack::Push(SemIR::InstBlockId id,
+                          llvm::ArrayRef<SemIR::InstId> inst_ids) -> void {
+  Push(id);
+  stack_[size_ - 1].content = inst_ids;
 }
 
 auto InstBlockStack::PeekOrAdd(int depth) -> SemIR::InstBlockId {
@@ -57,12 +58,6 @@ auto InstBlockStack::Pop() -> SemIR::InstBlockId {
   return back.id;
 }
 
-auto InstBlockStack::PopGlobalInit() -> void {
-  init_block_ = std::move(stack_[size_ - 1].content);
-  init_block_id_ = stack_[size_ - 1].id;
-  PopAndDiscard();
-}
-
 auto InstBlockStack::PopAndDiscard() -> void {
   CARBON_CHECK(!empty()) << "no current block";
   --size_;

toolchain/check/inst_block_stack.h

@@ -25,17 +25,14 @@ class InstBlockStack {
   // Pushes an existing instruction block.
   auto Push(SemIR::InstBlockId id) -> void;
 
+  // Pushes an existing instruction block with a set of instructions.
+  auto Push(SemIR::InstBlockId id, llvm::ArrayRef<SemIR::InstId> inst_ids)
+      -> void;
+
   // Pushes a new instruction block. It will be invalid unless PeekOrAdd is
   // called in order to support lazy allocation.
   auto Push() -> void { Push(SemIR::InstBlockId::Invalid); }
 
-  // Pushes the `GlobalInit` inst block onto the stack, this block is handled
-  // separately from the rest.
-  // This method shall be used in conjunction with `PopGlobalInit` method to
-  // allow emitting initialization instructions to `GlobalInit` block from
-  // separate parts of the tree, accumulating them all in one block.
-  auto PushGlobalInit() -> void;
-
   // Pushes a new unreachable code block.
   auto PushUnreachable() -> void { Push(SemIR::InstBlockId::Unreachable); }
 
@@ -53,11 +50,6 @@ class InstBlockStack {
   // allocated.
   auto PopAndDiscard() -> void;
 
-  // Pops the `GlobalInit` inst block from the stack without finalizing it.
-  // `Pop` should be called at the end of the check phase, while `GlobalInit`
-  // is pushed, to finalize the block.
-  auto PopGlobalInit() -> void;
-
   // Adds the given instruction ID to the block at the top of the stack.
   auto AddInstId(SemIR::InstId inst_id) -> void {
     CARBON_CHECK(!empty()) << "no current block";
@@ -122,11 +114,6 @@ class InstBlockStack {
   // Whether to print verbose output.
   llvm::raw_ostream* vlog_stream_;
 
-  std::vector<SemIR::InstId> init_block_;
-
-  // Current global init block to push.
-  SemIR::InstBlockId init_block_id_ = SemIR::InstBlockId::GlobalInit;
-
   // The actual stack.
   llvm::SmallVector<StackEntry> stack_;