Refactor InstBlockStack to use ArrayStack. (#4104)

The use of ArrayStack here is intended to simplify the logic, and also
make better use of the inst heap allocations. Prior changes #4101 and
#4103 removed the less related logic from InstBlockStack, although #4103
is the actual part that blocked using ArrayStack.

BTW, note the PrintForStackDump implementation was incorrect because it
didn't apply size_. This simplification fixes the issue.

common/array_stack.h

@@ -47,6 +47,15 @@ class ArrayStack {
     return llvm::ArrayRef(values_).slice(array_offsets_.back());
   }
 
+  // Returns the array at a specific index.
+  auto PeekArrayAt(int index) const -> llvm::ArrayRef<ValueT> {
+    auto ref = llvm::ArrayRef(values_).slice(array_offsets_[index]);
+    if (index + 1 < static_cast<int>(array_offsets_.size())) {
+      ref = ref.take_front(array_offsets_[index + 1] - array_offsets_[index]);
+    }
+    return ref;
+  }
+
   // Returns the full set of values on the stack, regardless of whether any
   // arrays are pushed.
   auto PeekAllValues() const -> llvm::ArrayRef<ValueT> { return values_; }
@@ -58,6 +67,13 @@ class ArrayStack {
     values_.push_back(value);
   }
 
+  // Adds multiple values to the top array on the stack.
+  auto AppendToTop(llvm::ArrayRef<ValueT> values) -> void {
+    CARBON_CHECK(!array_offsets_.empty())
+        << "Must call PushArray before PushValues.";
+    values_.append(values.begin(), values.end());
+  }
+
   // Returns the current number of values in all arrays.
   auto all_values_size() const -> size_t { return values_.size(); }
 

common/array_stack_test.cpp

@@ -73,5 +73,35 @@ TEST(ArrayStack, Basics) {
   EXPECT_THAT(stack.PeekAllValues(), ElementsAre(5));
 }
 
+TEST(ArrayStack, AppendArray) {
+  ArrayStack<int> stack;
+
+  stack.PushArray();
+  stack.AppendToTop(llvm::ArrayRef<int>());
+  EXPECT_THAT(stack.PeekArray(), IsEmpty());
+  stack.AppendToTop({1, 2});
+  EXPECT_THAT(stack.PeekArray(), ElementsAre(1, 2));
+}
+
+TEST(ArrayStack, PeekArrayAt) {
+  ArrayStack<int> stack;
+
+  // Verify behavior with a single array.
+  stack.PushArray();
+  stack.AppendToTop(1);
+  stack.AppendToTop(2);
+
+  EXPECT_THAT(stack.PeekArrayAt(0), ElementsAre(1, 2));
+
+  // Verify behavior with a couple more arrays.
+  stack.PushArray();
+  stack.PushArray();
+  stack.AppendToTop(3);
+
+  EXPECT_THAT(stack.PeekArrayAt(0), ElementsAre(1, 2));
+  EXPECT_THAT(stack.PeekArrayAt(1), IsEmpty());
+  EXPECT_THAT(stack.PeekArrayAt(2), ElementsAre(3));
+}
+
 }  // namespace
 }  // namespace Carbon::Testing

toolchain/check/BUILD

@@ -53,6 +53,7 @@ cc_library(
     ],
     deps = [
         ":node_stack",
+        "//common:array_stack",
         "//common:check",
         "//common:map",
         "//common:vlog",

toolchain/check/inst_block_stack.cpp

@@ -11,66 +11,63 @@
 namespace Carbon::Check {
 
 auto InstBlockStack::Push(SemIR::InstBlockId id) -> void {
-  CARBON_VLOG() << name_ << " Push " << size_ << "\n";
-  CARBON_CHECK(size_ < (1 << 20))
+  CARBON_VLOG() << name_ << " Push " << id_stack_.size() << "\n";
+  CARBON_CHECK(id_stack_.size() < (1 << 20))
       << "Excessive stack size: likely infinite loop";
-  if (size_ == static_cast<int>(stack_.size())) {
-    stack_.emplace_back();
-  }
-  stack_[size_].Reset(id);
-  ++size_;
+  id_stack_.push_back(id);
+  insts_stack_.PushArray();
 }
 
 auto InstBlockStack::Push(SemIR::InstBlockId id,
                           llvm::ArrayRef<SemIR::InstId> inst_ids) -> void {
   Push(id);
-  stack_[size_ - 1].content = inst_ids;
+  insts_stack_.AppendToTop(inst_ids);
 }
 
 auto InstBlockStack::PeekOrAdd(int depth) -> SemIR::InstBlockId {
-  CARBON_CHECK(size_ > depth) << "no such block";
-  int index = size_ - depth - 1;
-  auto& slot = stack_[index];
-  if (!slot.id.is_valid()) {
-    slot.id = sem_ir_->inst_blocks().AddDefaultValue();
+  CARBON_CHECK(static_cast<int>(id_stack_.size()) > depth) << "no such block";
+  int index = id_stack_.size() - depth - 1;
+  auto& slot = id_stack_[index];
+  if (!slot.is_valid()) {
+    slot = sem_ir_->inst_blocks().AddDefaultValue();
   }
-  return slot.id;
+  return slot;
 }
 
 auto InstBlockStack::Pop() -> SemIR::InstBlockId {
   CARBON_CHECK(!empty()) << "no current block";
-  --size_;
-  auto& back = stack_[size_];
+  auto id = id_stack_.pop_back_val();
+  auto insts = insts_stack_.PeekArray();
 
   // Finalize the block.
-  if (!back.content.empty() && back.id != SemIR::InstBlockId::Unreachable) {
-    if (back.id.is_valid()) {
-      sem_ir_->inst_blocks().Set(back.id, back.content);
+  if (!insts.empty() && id != SemIR::InstBlockId::Unreachable) {
+    if (id.is_valid()) {
+      sem_ir_->inst_blocks().Set(id, insts);
     } else {
-      back.id = sem_ir_->inst_blocks().Add(back.content);
+      id = sem_ir_->inst_blocks().Add(insts);
     }
   }
 
-  CARBON_VLOG() << name_ << " Pop " << size_ << ": " << back.id << "\n";
-  if (!back.id.is_valid()) {
-    return SemIR::InstBlockId::Empty;
-  }
-  return back.id;
+  insts_stack_.PopArray();
+
+  CARBON_VLOG() << name_ << " Pop " << id_stack_.size() << ": " << id << "\n";
+  return id.is_valid() ? id : SemIR::InstBlockId::Empty;
 }
 
 auto InstBlockStack::PopAndDiscard() -> void {
   CARBON_CHECK(!empty()) << "no current block";
-  --size_;
-  CARBON_VLOG() << name_ << " PopAndDiscard " << size_ << "\n";
+  id_stack_.pop_back();
+  insts_stack_.PopArray();
+  CARBON_VLOG() << name_ << " PopAndDiscard " << id_stack_.size() << "\n";
 }
 
 auto InstBlockStack::PrintForStackDump(llvm::raw_ostream& output) const
     -> void {
   output << name_ << ":\n";
-  for (const auto& [i, entry] : llvm::enumerate(stack_)) {
-    output << "\t" << i << ".\t" << entry.id << "\t{";
+  for (const auto& [i, id] : llvm::enumerate(id_stack_)) {
+    output << "\t" << i << ".\t" << id << "\t{";
     llvm::ListSeparator sep;
-    for (auto id : entry.content) {
+    for (auto id : insts_stack_.PeekArrayAt(i)) {
       output << sep << id;
     }
     output << "}\n";

toolchain/check/inst_block_stack.h

@@ -5,6 +5,7 @@
 #ifndef CARBON_TOOLCHAIN_CHECK_INST_BLOCK_STACK_H_
 #define CARBON_TOOLCHAIN_CHECK_INST_BLOCK_STACK_H_
 
+#include "common/array_stack.h"
 #include "llvm/ADT/SmallVector.h"
 #include "toolchain/sem_ir/file.h"
 
@@ -53,50 +54,31 @@ class InstBlockStack {
   // 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";
-    stack_[size_ - 1].content.push_back(inst_id);
+    insts_stack_.AppendToTop(inst_id);
   }
 
   // Returns whether the current block is statically reachable.
   auto is_current_block_reachable() -> bool {
-    return size_ != 0 &&
-           stack_[size_ - 1].id != SemIR::InstBlockId::Unreachable;
+    return id_stack_.back() != SemIR::InstBlockId::Unreachable;
   }
 
   // Returns a view of the contents of the top instruction block on the stack.
-  auto PeekCurrentBlockContents() -> llvm::ArrayRef<SemIR::InstId> {
+  auto PeekCurrentBlockContents() const -> llvm::ArrayRef<SemIR::InstId> {
     CARBON_CHECK(!empty()) << "no current block";
-    return stack_[size_ - 1].content;
+    return insts_stack_.PeekArray();
   }
 
   // Prints the stack for a stack dump.
   auto PrintForStackDump(llvm::raw_ostream& output) const -> void;
 
   // Runs verification that the processing cleanly finished.
-  auto VerifyOnFinish() const -> void { CARBON_CHECK(empty()) << size_; }
+  auto VerifyOnFinish() const -> void {
+    CARBON_CHECK(empty()) << id_stack_.size();
+  }
 
-  auto empty() const -> bool { return size_ == 0; }
+  auto empty() const -> bool { return id_stack_.empty(); }
 
  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::InstBlockId 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::InstBlockId id = SemIR::InstBlockId::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::InstId> content;
-  };
-
   // A name for debugging.
   llvm::StringLiteral name_;
 
@@ -106,12 +88,12 @@ class InstBlockStack {
   // Whether to print verbose output.
   llvm::raw_ostream* vlog_stream_;
 
-  // The actual stack.
-  llvm::SmallVector<StackEntry> stack_;
+  // The stack of block IDs. Valid if allocated, Invalid if no block has been
+  // allocated, or Unreachable if this block is known to be unreachable.
+  llvm::SmallVector<SemIR::InstBlockId> id_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;
+  // The stack of insts in each block.
+  ArrayStack<SemIR::InstId> insts_stack_;
 };
 
 }  // namespace Carbon::Check