Store contents of incomplete node blocks on the node block stack. (#3223)

`SemIR` now only stores complete blocks, and placeholders for blocks
where we need a block ID before we know the block contents.

This is a preparation step towards the new node block allocation design.

toolchain/check/context.cpp

@@ -153,7 +153,7 @@ static auto AddDominatedBlockAndBranchImpl(Context& context,
   if (!context.node_block_stack().is_current_block_reachable()) {
     return SemIR::NodeBlockId::Unreachable;
   }
-  auto block_id = context.semantics_ir().AddNodeBlock();
+  auto block_id = context.semantics_ir().AddNodeBlockId();
   context.AddNode(BranchNode::Make(parse_node, block_id, args...));
   return block_id;
 }
@@ -185,7 +185,7 @@ auto Context::AddConvergenceBlockAndPush(Parse::Node parse_node, int num_blocks)
   for ([[maybe_unused]] auto _ : llvm::seq(num_blocks)) {
     if (node_block_stack().is_current_block_reachable()) {
       if (new_block_id == SemIR::NodeBlockId::Unreachable) {
-        new_block_id = semantics_ir().AddNodeBlock();
+        new_block_id = semantics_ir().AddNodeBlockId();
       }
       AddNode(SemIR::Node::Branch::Make(parse_node, new_block_id));
     }
@@ -203,7 +203,7 @@ auto Context::AddConvergenceBlockWithArgAndPush(
   for (auto arg_id : block_args) {
     if (node_block_stack().is_current_block_reachable()) {
       if (new_block_id == SemIR::NodeBlockId::Unreachable) {
-        new_block_id = semantics_ir().AddNodeBlock();
+        new_block_id = semantics_ir().AddNodeBlockId();
       }
       AddNode(
           SemIR::Node::BranchWithArg::Make(parse_node, new_block_id, arg_id));

toolchain/check/node_block_stack.cpp

@@ -6,43 +6,64 @@
 
 #include "common/vlog.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringExtras.h"
 #include "toolchain/sem_ir/node.h"
 
 namespace Carbon::Check {
 
 auto NodeBlockStack::Push(SemIR::NodeBlockId id) -> void {
-  CARBON_VLOG() << name_ << " Push " << stack_.size() << "\n";
-  CARBON_CHECK(stack_.size() < (1 << 20))
+  CARBON_VLOG() << name_ << " Push " << size_ << "\n";
+  CARBON_CHECK(size_ < (1 << 20))
       << "Excessive stack size: likely infinite loop";
-  stack_.push_back(id);
+  if (size_ == static_cast<int>(stack_.size())) {
+    stack_.emplace_back();
+  }
+  stack_[size_].Reset(id);
+  ++size_;
 }
 
 auto NodeBlockStack::PeekForAdd(int depth) -> SemIR::NodeBlockId {
-  CARBON_CHECK(static_cast<int>(stack_.size()) > depth) << "no such block";
-  int index = stack_.size() - depth - 1;
+  CARBON_CHECK(size() > depth) << "no such block";
+  int index = size() - depth - 1;
   auto& slot = stack_[index];
-  if (!slot.is_valid()) {
-    slot = semantics_ir_->AddNodeBlock();
-    CARBON_VLOG() << name_ << " Add " << index << ": " << slot << "\n";
+  if (!slot.id.is_valid()) {
+    slot.id = semantics_ir_->AddNodeBlockId();
+    CARBON_VLOG() << name_ << " Add " << index << ": " << slot.id << "\n";
   }
-  return slot;
+  return slot.id;
 }
 
 auto NodeBlockStack::Pop() -> SemIR::NodeBlockId {
-  CARBON_CHECK(!stack_.empty()) << "no current block";
-  auto back = stack_.pop_back_val();
-  CARBON_VLOG() << name_ << " Pop " << stack_.size() << ": " << back << "\n";
-  if (!back.is_valid()) {
+  CARBON_CHECK(!empty()) << "no current block";
+  --size_;
+  auto& back = stack_[size_];
+
+  // Finalize the block.
+  if (!back.content.empty() && back.id != SemIR::NodeBlockId::Unreachable) {
+    if (back.id.is_valid()) {
+      semantics_ir_->SetNodeBlock(back.id, back.content);
+    } else {
+      back.id = semantics_ir_->AddNodeBlock(back.content);
+    }
+  }
+
+  CARBON_VLOG() << name_ << " Pop " << size_ << ": " << back.id << "\n";
+  if (!back.id.is_valid()) {
     return SemIR::NodeBlockId::Empty;
   }
-  return back;
+  return back.id;
 }
 
 auto NodeBlockStack::PrintForStackDump(llvm::raw_ostream& output) const
     -> void {
   output << name_ << ":\n";
-  for (auto [i, entry] : llvm::enumerate(stack_)) {
-    output << "\t" << i << ".\t" << entry << "\n";
+  for (const auto& [i, entry] : llvm::enumerate(stack_)) {
+    output << "\t" << i << ".\t" << entry.id << "\t{";
+    llvm::ListSeparator sep;
+    for (auto id : entry.content) {
+      output << sep << id;
+    }
+    output << "}\n";
   }
 }
 

toolchain/check/node_block_stack.h

@@ -39,8 +39,8 @@ class NodeBlockStack {
   // Peeks at the top node block. This does not trigger lazy allocation, so the
   // returned node block may be invalid.
   auto Peek() -> SemIR::NodeBlockId {
-    CARBON_CHECK(!stack_.empty()) << "no current block";
-    return stack_.back();
+    CARBON_CHECK(!empty()) << "no current block";
+    return stack_[size() - 1].id;
   }
 
   // Returns the top node block, allocating one if it's still invalid. If
@@ -52,16 +52,10 @@ class NodeBlockStack {
   // SemIR::NodeBlockId::Empty is returned if one wasn't allocated.
   auto Pop() -> SemIR::NodeBlockId;
 
-  // Pops the top node block, ensuring that it is lazily allocated if it's
-  // empty. For use when more nodes will be added to the block later.
-  auto PopForAdd() -> SemIR::NodeBlockId {
-    PeekForAdd();
-    return Pop();
-  }
-
   // Adds the given node ID to the block at the top of the stack.
   auto AddNodeId(SemIR::NodeId node_id) -> void {
-    semantics_ir_->AddNodeIdForNodeBlockStack(PeekForAdd(), node_id);
+    CARBON_CHECK(!empty()) << "no current block";
+    stack_[size_ - 1].content.push_back(node_id);
   }
 
   // Returns whether the current block is statically reachable.
@@ -72,10 +66,30 @@ class NodeBlockStack {
   // Prints the stack for a stack dump.
   auto PrintForStackDump(llvm::raw_ostream& output) const -> void;
 
-  auto empty() const -> bool { return stack_.empty(); }
-  auto size() const -> size_t { return stack_.size(); }
+  auto empty() const -> bool { return size() == 0; }
+  auto size() const -> int { return size_; }
 
  private:
+  struct StackEntry {
+    // Preallocate an arbitrary size for the stack entries.
+    // TODO: Perform measurements to pick a good starting size to avoid
+    // reallocation.
+    StackEntry() { content.reserve(32); }
+
+    auto Reset(SemIR::NodeBlockId new_id) {
+      id = new_id;
+      content.clear();
+    }
+
+    // The block ID, if one has been allocated, Invalid if no block has been
+    // allocated, or Unreachable if this block is known to be unreachable.
+    SemIR::NodeBlockId id = SemIR::NodeBlockId::Invalid;
+
+    // The content of the block. Stored as a vector rather than as a SmallVector
+    // to reduce the cost of resizing `stack_` and performing swaps.
+    std::vector<SemIR::NodeId> content;
+  };
+
   // A name for debugging.
   llvm::StringLiteral name_;
 
@@ -86,9 +100,11 @@ class NodeBlockStack {
   llvm::raw_ostream* vlog_stream_;
 
   // The actual stack.
-  // PushEntry and PopEntry control modification in order to centralize
-  // vlogging.
-  llvm::SmallVector<SemIR::NodeBlockId> stack_;
+  llvm::SmallVector<StackEntry> stack_;
+
+  // The size of the stack. Entries after this in `stack_` are kept around so
+  // that we can reuse the allocated buffer for their content.
+  int size_ = 0;
 };
 
 }  // namespace Carbon::Check

toolchain/check/testdata/basics/multifile_raw_and_textual_ir.carbon

@@ -12,7 +12,7 @@
 // CHECK:STDOUT:   sem_ir:
 // CHECK:STDOUT:   - cross_reference_irs_size: 1
 // CHECK:STDOUT:     functions: [
-// CHECK:STDOUT:       {name: str0, param_refs: block0, body: [block2]},
+// CHECK:STDOUT:       {name: str0, param_refs: block0, body: [block1]},
 // CHECK:STDOUT:     ]
 // CHECK:STDOUT:     integer_literals: [
 // CHECK:STDOUT:     ]
@@ -33,10 +33,10 @@
 // CHECK:STDOUT:       [
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
-// CHECK:STDOUT:         node+0,
+// CHECK:STDOUT:         node+1,
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
-// CHECK:STDOUT:         node+1,
+// CHECK:STDOUT:         node+0,
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:     ]
 // CHECK:STDOUT:
@@ -52,7 +52,7 @@
 // CHECK:STDOUT:   sem_ir:
 // CHECK:STDOUT:   - cross_reference_irs_size: 1
 // CHECK:STDOUT:     functions: [
-// CHECK:STDOUT:       {name: str0, param_refs: block0, body: [block2]},
+// CHECK:STDOUT:       {name: str0, param_refs: block0, body: [block1]},
 // CHECK:STDOUT:     ]
 // CHECK:STDOUT:     integer_literals: [
 // CHECK:STDOUT:     ]
@@ -73,10 +73,10 @@
 // CHECK:STDOUT:       [
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
-// CHECK:STDOUT:         node+0,
+// CHECK:STDOUT:         node+1,
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
-// CHECK:STDOUT:         node+1,
+// CHECK:STDOUT:         node+0,
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:     ]
 // CHECK:STDOUT:

toolchain/check/testdata/basics/multifile_raw_ir.carbon

@@ -12,7 +12,7 @@
 // CHECK:STDOUT:   sem_ir:
 // CHECK:STDOUT:   - cross_reference_irs_size: 1
 // CHECK:STDOUT:     functions: [
-// CHECK:STDOUT:       {name: str0, param_refs: block0, body: [block2]},
+// CHECK:STDOUT:       {name: str0, param_refs: block0, body: [block1]},
 // CHECK:STDOUT:     ]
 // CHECK:STDOUT:     integer_literals: [
 // CHECK:STDOUT:     ]
@@ -33,17 +33,17 @@
 // CHECK:STDOUT:       [
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
-// CHECK:STDOUT:         node+0,
+// CHECK:STDOUT:         node+1,
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
-// CHECK:STDOUT:         node+1,
+// CHECK:STDOUT:         node+0,
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:     ]
 // CHECK:STDOUT: - filename: b.carbon
 // CHECK:STDOUT:   sem_ir:
 // CHECK:STDOUT:   - cross_reference_irs_size: 1
 // CHECK:STDOUT:     functions: [
-// CHECK:STDOUT:       {name: str0, param_refs: block0, body: [block2]},
+// CHECK:STDOUT:       {name: str0, param_refs: block0, body: [block1]},
 // CHECK:STDOUT:     ]
 // CHECK:STDOUT:     integer_literals: [
 // CHECK:STDOUT:     ]
@@ -64,10 +64,10 @@
 // CHECK:STDOUT:       [
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
-// CHECK:STDOUT:         node+0,
+// CHECK:STDOUT:         node+1,
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
-// CHECK:STDOUT:         node+1,
+// CHECK:STDOUT:         node+0,
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:     ]
 // --- a.carbon

toolchain/check/testdata/basics/raw_and_textual_ir.carbon

@@ -16,7 +16,7 @@ fn Foo(n: i32) -> (i32, f64) {
 // CHECK:STDOUT:   sem_ir:
 // CHECK:STDOUT:   - cross_reference_irs_size: 1
 // CHECK:STDOUT:     functions: [
-// CHECK:STDOUT:       {name: str0, param_refs: block2, return_type: type3, return_slot: node+6, body: [block5]},
+// CHECK:STDOUT:       {name: str0, param_refs: block1, return_type: type3, return_slot: node+6, body: [block4]},
 // CHECK:STDOUT:     ]
 // CHECK:STDOUT:     integer_literals: [
 // CHECK:STDOUT:       2,
@@ -50,7 +50,7 @@ fn Foo(n: i32) -> (i32, f64) {
 // CHECK:STDOUT:       {kind: StubReference, arg0: nodeIntegerType, type: typeTypeType},
 // CHECK:STDOUT:       {kind: StubReference, arg0: nodeFloatingPointType, type: typeTypeType},
 // CHECK:STDOUT:       {kind: TupleType, arg0: typeBlock0, type: typeTypeType},
-// CHECK:STDOUT:       {kind: TupleLiteral, arg0: block3, type: type1},
+// CHECK:STDOUT:       {kind: TupleLiteral, arg0: block2, type: type1},
 // CHECK:STDOUT:       {kind: TupleType, arg0: typeBlock1, type: typeTypeType},
 // CHECK:STDOUT:       {kind: VarStorage, arg0: str2, type: type3},
 // CHECK:STDOUT:       {kind: FunctionDeclaration, arg0: function0},
@@ -59,7 +59,7 @@ fn Foo(n: i32) -> (i32, f64) {
 // CHECK:STDOUT:       {kind: StubReference, arg0: node+9, type: type0},
 // CHECK:STDOUT:       {kind: RealLiteral, arg0: real0, type: type2},
 // CHECK:STDOUT:       {kind: StubReference, arg0: node+11, type: type2},
-// CHECK:STDOUT:       {kind: TupleLiteral, arg0: block6, type: type3},
+// CHECK:STDOUT:       {kind: TupleLiteral, arg0: block5, type: type3},
 // CHECK:STDOUT:       {kind: ReturnExpression, arg0: node+13},
 // CHECK:STDOUT:     ]
 // CHECK:STDOUT:     node_blocks: [
@@ -67,22 +67,19 @@ fn Foo(n: i32) -> (i32, f64) {
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
 // CHECK:STDOUT:         node+0,
+// CHECK:STDOUT:       ],
+// CHECK:STDOUT:       [
 // CHECK:STDOUT:         node+1,
 // CHECK:STDOUT:         node+2,
-// CHECK:STDOUT:         node+3,
-// CHECK:STDOUT:         node+4,
-// CHECK:STDOUT:         node+5,
-// CHECK:STDOUT:         node+6,
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
 // CHECK:STDOUT:         node+0,
-// CHECK:STDOUT:       ],
-// CHECK:STDOUT:       [
 // CHECK:STDOUT:         node+1,
 // CHECK:STDOUT:         node+2,
-// CHECK:STDOUT:       ],
-// CHECK:STDOUT:       [
-// CHECK:STDOUT:         node+7,
+// CHECK:STDOUT:         node+3,
+// CHECK:STDOUT:         node+4,
+// CHECK:STDOUT:         node+5,
+// CHECK:STDOUT:         node+6,
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
 // CHECK:STDOUT:         node+8,
@@ -97,6 +94,9 @@ fn Foo(n: i32) -> (i32, f64) {
 // CHECK:STDOUT:         node+10,
 // CHECK:STDOUT:         node+12,
 // CHECK:STDOUT:       ],
+// CHECK:STDOUT:       [
+// CHECK:STDOUT:         node+7,
+// CHECK:STDOUT:       ],
 // CHECK:STDOUT:     ]
 // CHECK:STDOUT:
 // CHECK:STDOUT: file "raw_and_textual_ir.carbon" {

toolchain/check/testdata/basics/raw_ir.carbon

@@ -16,7 +16,7 @@ fn Foo(n: i32) -> (i32, f64) {
 // CHECK:STDOUT:   sem_ir:
 // CHECK:STDOUT:   - cross_reference_irs_size: 1
 // CHECK:STDOUT:     functions: [
-// CHECK:STDOUT:       {name: str0, param_refs: block2, return_type: type3, return_slot: node+6, body: [block5]},
+// CHECK:STDOUT:       {name: str0, param_refs: block1, return_type: type3, return_slot: node+6, body: [block4]},
 // CHECK:STDOUT:     ]
 // CHECK:STDOUT:     integer_literals: [
 // CHECK:STDOUT:       2,
@@ -50,7 +50,7 @@ fn Foo(n: i32) -> (i32, f64) {
 // CHECK:STDOUT:       {kind: StubReference, arg0: nodeIntegerType, type: typeTypeType},
 // CHECK:STDOUT:       {kind: StubReference, arg0: nodeFloatingPointType, type: typeTypeType},
 // CHECK:STDOUT:       {kind: TupleType, arg0: typeBlock0, type: typeTypeType},
-// CHECK:STDOUT:       {kind: TupleLiteral, arg0: block3, type: type1},
+// CHECK:STDOUT:       {kind: TupleLiteral, arg0: block2, type: type1},
 // CHECK:STDOUT:       {kind: TupleType, arg0: typeBlock1, type: typeTypeType},
 // CHECK:STDOUT:       {kind: VarStorage, arg0: str2, type: type3},
 // CHECK:STDOUT:       {kind: FunctionDeclaration, arg0: function0},
@@ -59,7 +59,7 @@ fn Foo(n: i32) -> (i32, f64) {
 // CHECK:STDOUT:       {kind: StubReference, arg0: node+9, type: type0},
 // CHECK:STDOUT:       {kind: RealLiteral, arg0: real0, type: type2},
 // CHECK:STDOUT:       {kind: StubReference, arg0: node+11, type: type2},
-// CHECK:STDOUT:       {kind: TupleLiteral, arg0: block6, type: type3},
+// CHECK:STDOUT:       {kind: TupleLiteral, arg0: block5, type: type3},
 // CHECK:STDOUT:       {kind: ReturnExpression, arg0: node+13},
 // CHECK:STDOUT:     ]
 // CHECK:STDOUT:     node_blocks: [
@@ -67,22 +67,19 @@ fn Foo(n: i32) -> (i32, f64) {
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
 // CHECK:STDOUT:         node+0,
+// CHECK:STDOUT:       ],
+// CHECK:STDOUT:       [
 // CHECK:STDOUT:         node+1,
 // CHECK:STDOUT:         node+2,
-// CHECK:STDOUT:         node+3,
-// CHECK:STDOUT:         node+4,
-// CHECK:STDOUT:         node+5,
-// CHECK:STDOUT:         node+6,
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
 // CHECK:STDOUT:         node+0,
-// CHECK:STDOUT:       ],
-// CHECK:STDOUT:       [
 // CHECK:STDOUT:         node+1,
 // CHECK:STDOUT:         node+2,
-// CHECK:STDOUT:       ],
-// CHECK:STDOUT:       [
-// CHECK:STDOUT:         node+7,
+// CHECK:STDOUT:         node+3,
+// CHECK:STDOUT:         node+4,
+// CHECK:STDOUT:         node+5,
+// CHECK:STDOUT:         node+6,
 // CHECK:STDOUT:       ],
 // CHECK:STDOUT:       [
 // CHECK:STDOUT:         node+8,
@@ -97,4 +94,7 @@ fn Foo(n: i32) -> (i32, f64) {
 // CHECK:STDOUT:         node+10,
 // CHECK:STDOUT:         node+12,
 // CHECK:STDOUT:       ],
+// CHECK:STDOUT:       [
+// CHECK:STDOUT:         node+7,
+// CHECK:STDOUT:       ],
 // CHECK:STDOUT:     ]

toolchain/check/testdata/basics/verbose.carbon

@@ -8,7 +8,7 @@
 // NOAUTOUPDATE
 // SET-CHECK-SUBSET
 // CHECK:STDERR: Node Push 0: FunctionIntroducer -> <none>
-// CHECK:STDERR: AddNode block{{[0-9]+}}: {kind: FunctionDeclaration, arg0: function{{[0-9]+}}}
+// CHECK:STDERR: AddNode block{{[0-9]+}}: {kind: Return}
 
 fn Foo() {
   return;

toolchain/sem_ir/file.h

@@ -157,15 +157,6 @@ class File : public Printable<File> {
     return node_id;
   }
 
-  // Adds the ID of an existing node to the specified block. Temporary, should
-  // only be called by Check::NodeBlockStack.
-  auto AddNodeIdForNodeBlockStack(NodeBlockId block_id, NodeId node_id)
-      -> void {
-    if (block_id != NodeBlockId::Unreachable) {
-      node_blocks_[block_id.index].push_back(node_id);
-    }
-  }
-
   // Overwrites a given node with a new value.
   auto ReplaceNode(NodeId node_id, Node node) -> void {
     nodes_[node_id.index] = node;
@@ -174,13 +165,24 @@ class File : public Printable<File> {
   // Returns the requested node.
   auto GetNode(NodeId node_id) const -> Node { return nodes_[node_id.index]; }
 
-  // Adds an empty node block, returning an ID to reference it.
-  auto AddNodeBlock() -> NodeBlockId {
+  // Reserves and returns a node block ID. The contents of the node block
+  // should be specified by calling SetNodeBlock, or by pushing the ID onto the
+  // NodeBlockStack.
+  auto AddNodeBlockId() -> NodeBlockId {
     NodeBlockId id(node_blocks_.size());
     node_blocks_.push_back({});
     return id;
   }
 
+  // Sets the contents of an empty node block to the given content.
+  auto SetNodeBlock(NodeBlockId block_id, llvm::ArrayRef<NodeId> content)
+      -> void {
+    CARBON_CHECK(block_id != NodeBlockId::Unreachable);
+    CARBON_CHECK(node_blocks_[block_id.index].empty())
+        << "node block content set more than once";
+    node_blocks_[block_id.index].assign(content.begin(), content.end());
+  }
+
   // Adds a node block with the given content, returning an ID to reference it.
   auto AddNodeBlock(llvm::ArrayRef<NodeId> content) -> NodeBlockId {
     NodeBlockId id(node_blocks_.size());