Create an array stack type for a shared use-case (#4100)

Based on discussion around the region handling in generic_region_stack,
create a generic structure for the stack-of-vectors support. I also want
to add this to InstBlockStack, but that's a little more complex due to
GlobalInit, so cutting a PR here to check with review.

My work here is how I noticed #4099; I want to be sure that I'm correct
about the issue, but it's the difference between being able to use
PeekArray or not.

Note in scope_stack.h, I believe we could remove next_compile_time_index
and make it just based on elements_size(). However, I want to verify
with you before I make further changes there.

---------

Co-authored-by: Chandler Carruth <chandlerc@gmail.com>

common/BUILD

@@ -7,6 +7,26 @@ load("//bazel/version:rules.bzl", "expand_version_build_info")
 
 package(default_visibility = ["//visibility:public"])
 
+cc_library(
+    name = "array_stack",
+    hdrs = ["array_stack.h"],
+    deps = [
+        ":check",
+        "@llvm-project//llvm:Support",
+    ],
+)
+
+cc_test(
+    name = "array_stack_test",
+    size = "small",
+    srcs = ["array_stack_test.cpp"],
+    deps = [
+        ":array_stack",
+        "//testing/base:gtest_main",
+        "@googletest//:gtest",
+    ],
+)
+
 cc_library(
     name = "bazel_working_dir",
     hdrs = ["bazel_working_dir.h"],

common/array_stack.h

@@ -0,0 +1,74 @@
+// 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_COMMON_ARRAY_STACK_H_
+#define CARBON_COMMON_ARRAY_STACK_H_
+
+#include "common/check.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace Carbon {
+
+// Provides a stack of arrays. Only the array at the top of the stack can have
+// elements added.
+//
+// Example usage:
+//   // Push to start.
+//   PushArray();
+//   // Add values.
+//   AppendToTop(3);
+//   // Look at values.
+//   PeekArray();
+//   // Pop when done.
+//   PopArray();
+//
+// By using a single vector for elements, the intent is that as arrays are
+// pushed and popped, the same storage will be reused. This should yield
+// efficiencies for heap allocations. For example, in the toolchain we
+// frequently have an array per scope, and only add to the current scope's
+// array; this allows better reuse when entering and leaving scopes.
+template <typename ValueT>
+class ArrayStack {
+ public:
+  // Pushes a new array onto the stack.
+  auto PushArray() -> void { array_offsets_.push_back(values_.size()); }
+
+  // Pops the top array from the stack.
+  auto PopArray() -> void {
+    auto region = array_offsets_.pop_back_val();
+    values_.truncate(region);
+  }
+
+  // Returns the top array from the stack.
+  auto PeekArray() const -> llvm::ArrayRef<ValueT> {
+    CARBON_CHECK(!array_offsets_.empty());
+    return llvm::ArrayRef(values_).slice(array_offsets_.back());
+  }
+
+  // Returns the full set of values on the stack, regardless of whether any
+  // arrays are pushed.
+  auto PeekAllValues() const -> llvm::ArrayRef<ValueT> { return values_; }
+
+  // Appends a value to the top array on the stack.
+  auto AppendToTop(ValueT value) -> void {
+    CARBON_CHECK(!array_offsets_.empty())
+        << "Must call PushArray before PushValue.";
+    values_.push_back(value);
+  }
+
+  // Returns the current number of values in all arrays.
+  auto all_values_size() const -> size_t { return values_.size(); }
+
+ private:
+  // For each pushed array, the start index in elements_.
+  llvm::SmallVector<int32_t> array_offsets_;
+
+  // The full set of elements in all arrays.
+  llvm::SmallVector<ValueT> values_;
+};
+
+}  // namespace Carbon
+
+#endif  // CARBON_COMMON_ARRAY_STACK_H_

common/array_stack_test.cpp

@@ -0,0 +1,77 @@
+// 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 "common/array_stack.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+namespace Carbon::Testing {
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::IsEmpty;
+
+TEST(ArrayStack, Basics) {
+  ArrayStack<int> stack;
+
+  // PeekAllValues is valid when there are no arrays.
+  EXPECT_THAT(stack.PeekAllValues(), IsEmpty());
+
+  // An array starts empty.
+  stack.PushArray();
+  EXPECT_THAT(stack.PeekArray(), IsEmpty());
+  EXPECT_THAT(stack.PeekAllValues(), IsEmpty());
+
+  // Pushing a couple values works.
+  stack.AppendToTop(1);
+  stack.AppendToTop(2);
+  EXPECT_THAT(stack.PeekArray(), ElementsAre(1, 2));
+  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2));
+
+  // Pushing a new array starts empty, old values are still there.
+  stack.PushArray();
+  EXPECT_THAT(stack.PeekArray(), IsEmpty());
+  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2));
+
+  // The added value goes to the 2nd array.
+  stack.AppendToTop(3);
+  EXPECT_THAT(stack.PeekArray(), ElementsAre(3));
+  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2, 3));
+
+  // Popping goes back to the 1st array.
+  stack.PopArray();
+  EXPECT_THAT(stack.PeekArray(), ElementsAre(1, 2));
+  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2));
+
+  // Push a couple arrays, then a value on the now-3rd array.
+  stack.PushArray();
+  stack.PushArray();
+  stack.AppendToTop(4);
+  EXPECT_THAT(stack.PeekArray(), ElementsAre(4));
+  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2, 4));
+
+  // Popping the 3rd array goes to the 2nd array, which is empty.
+  stack.PopArray();
+  EXPECT_THAT(stack.PeekArray(), IsEmpty());
+  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2));
+
+  // Again back to the 1st array.
+  stack.PopArray();
+  EXPECT_THAT(stack.PeekArray(), ElementsAre(1, 2));
+  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2));
+
+  // Go down to no arrays.
+  stack.PopArray();
+  EXPECT_THAT(stack.PeekAllValues(), IsEmpty());
+
+  // Add a new 1st array.
+  stack.PushArray();
+  stack.AppendToTop(5);
+  EXPECT_THAT(stack.PeekArray(), ElementsAre(5));
+  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(5));
+}
+
+}  // namespace
+}  // namespace Carbon::Testing

toolchain/check/BUILD

@@ -152,6 +152,7 @@ cc_library(
     srcs = ["generic_region_stack.cpp"],
     hdrs = ["generic_region_stack.h"],
     deps = [
+        "//common:array_stack",
         "//common:check",
         "//common:ostream",
         "//common:vlog",
@@ -282,6 +283,7 @@ cc_library(
         "scope_stack.h",
     ],
     deps = [
+        "//common:array_stack",
         "//common:check",
         "//common:ostream",
         "//common:set",

toolchain/check/generic.cpp

@@ -25,7 +25,10 @@ auto StartGenericDefinition(Context& context) -> void {
 
 auto FinishGenericDecl(Context& context, SemIR::InstId decl_id)
     -> SemIR::GenericId {
-  if (context.scope_stack().compile_time_binding_stack().empty()) {
+  auto all_bindings =
+      context.scope_stack().compile_time_bindings_stack().PeekAllValues();
+
+  if (all_bindings.empty()) {
     CARBON_CHECK(context.generic_region_stack().PeekDependentInsts().empty())
         << "Have dependent instructions but no compile time bindings are in "
            "scope.";
@@ -33,8 +36,7 @@ auto FinishGenericDecl(Context& context, SemIR::InstId decl_id)
     return SemIR::GenericId::Invalid;
   }
 
-  auto bindings_id = context.inst_blocks().Add(
-      context.scope_stack().compile_time_binding_stack());
+  auto bindings_id = context.inst_blocks().Add(all_bindings);
   // TODO: Track the list of dependent instructions in this region.
   context.generic_region_stack().Pop();
   return context.generics().Add(

toolchain/check/generic_region_stack.cpp

@@ -6,27 +6,16 @@
 
 namespace Carbon::Check {
 
-auto GenericRegionStack::Push() -> void {
-  regions_.push_back(
-      {.first_dependent_inst = static_cast<int32_t>(dependent_insts_.size())});
-}
+auto GenericRegionStack::Push() -> void { dependent_insts_stack_.PushArray(); }
 
-auto GenericRegionStack::Pop() -> void {
-  auto region = regions_.pop_back_val();
-  dependent_insts_.truncate(region.first_dependent_inst);
-}
+auto GenericRegionStack::Pop() -> void { dependent_insts_stack_.PopArray(); }
 
 auto GenericRegionStack::AddDependentInst(DependentInst inst) -> void {
-  CARBON_CHECK(!regions_.empty())
-      << "Formed a dependent instruction while not in a generic region.";
-  CARBON_CHECK(inst.kind != DependencyKind::None);
-  dependent_insts_.push_back(inst);
+  dependent_insts_stack_.AppendToTop(inst);
 }
 
 auto GenericRegionStack::PeekDependentInsts() -> llvm::ArrayRef<DependentInst> {
-  CARBON_CHECK(!regions_.empty());
-  return llvm::ArrayRef(dependent_insts_)
-      .slice(regions_.back().first_dependent_inst);
+  return dependent_insts_stack_.PeekArray();
 }
 
 }  // namespace Carbon::Check

toolchain/check/generic_region_stack.h

@@ -5,6 +5,7 @@
 #ifndef CARBON_TOOLCHAIN_CHECK_GENERIC_REGION_STACK_H_
 #define CARBON_TOOLCHAIN_CHECK_GENERIC_REGION_STACK_H_
 
+#include "common/array_stack.h"
 #include "llvm/ADT/BitmaskEnum.h"
 #include "toolchain/sem_ir/ids.h"
 
@@ -58,22 +59,8 @@ class GenericRegionStack {
   auto PeekDependentInsts() -> llvm::ArrayRef<DependentInst>;
 
  private:
-  // Information about an enclosing generic region that has been pushed onto the
-  // stack.
-  struct RegionInfo {
-    // The size of `dependent_insts_` at the start of this region. Equivalently,
-    // this is the first index within `dependent_insts_` that belongs to this
-    // region or a region nested within it.
-    int32_t first_dependent_inst;
-  };
-
-  // The current set of enclosing generic regions.
-  llvm::SmallVector<RegionInfo> regions_;
-
-  // List of symbolic constants used in any of the enclosing generic regions. We
-  // keep a single vector rather than one vector per region in order to minimize
-  // heap allocations.
-  llvm::SmallVector<DependentInst> dependent_insts_;
+  // A stack of symbolic constants for enclosing generic regions.
+  ArrayStack<DependentInst> dependent_insts_stack_;
 };
 
 }  // namespace Carbon::Check

toolchain/check/scope_stack.cpp

@@ -15,14 +15,13 @@ auto ScopeStack::VerifyOnFinish() -> void {
 
 auto ScopeStack::Push(SemIR::InstId scope_inst_id, SemIR::NameScopeId scope_id,
                       bool lexical_lookup_has_load_error) -> void {
+  compile_time_binding_stack_.PushArray();
   scope_stack_.push_back(
       {.index = next_scope_index_,
        .scope_inst_id = scope_inst_id,
        .scope_id = scope_id,
-       .next_compile_time_bind_index =
-           scope_stack_.empty()
-               ? SemIR::CompileTimeBindIndex(0)
-               : scope_stack_.back().next_compile_time_bind_index,
+       .next_compile_time_bind_index = SemIR::CompileTimeBindIndex(
+           compile_time_binding_stack_.all_values_size()),
        .lexical_lookup_has_load_error =
            LexicalLookupHasLoadError() || lexical_lookup_has_load_error});
   if (scope_id.is_valid()) {
@@ -57,15 +56,13 @@ auto ScopeStack::Pop() -> void {
     return_scope_stack_.back().returned_var = SemIR::InstId::Invalid;
   }
 
-  CARBON_CHECK(scope.next_compile_time_bind_index.index ==
-               static_cast<int32_t>(compile_time_binding_stack_.size()))
+  CARBON_CHECK(
+      scope.next_compile_time_bind_index.index ==
+      static_cast<int32_t>(compile_time_binding_stack_.all_values_size()))
       << "Wrong number of entries in compile-time binding stack, have "
-      << compile_time_binding_stack_.size() << ", expected "
+      << compile_time_binding_stack_.all_values_size() << ", expected "
       << scope.next_compile_time_bind_index.index;
-  compile_time_binding_stack_.truncate(
-      scope_stack_.empty()
-          ? 0
-          : scope_stack_.back().next_compile_time_bind_index.index);
+  compile_time_binding_stack_.PopArray();
 }
 
 auto ScopeStack::PopTo(ScopeIndex index) -> void {
@@ -163,13 +160,9 @@ auto ScopeStack::Suspend() -> SuspendedScope {
     non_lexical_scope_stack_.pop_back();
   }
 
-  auto remaining_compile_time_bindings =
-      scope_stack_.empty()
-          ? 0
-          : scope_stack_.back().next_compile_time_bind_index.index;
+  auto peek_compile_time_bindings = compile_time_binding_stack_.PeekArray();
   result.suspended_items.reserve(result.entry.num_names +
-                                 compile_time_binding_stack_.size() -
-                                 remaining_compile_time_bindings);
+                                 peek_compile_time_bindings.size());
 
   result.entry.names.ForEach([&](SemIR::NameId name_id) {
     auto [index, inst_id] = lexical_lookup_.Suspend(name_id);
@@ -181,13 +174,12 @@ auto ScopeStack::Suspend() -> SuspendedScope {
                result.entry.num_names);
 
   // Move any compile-time bindings into the suspended scope.
-  for (auto inst_id : llvm::ArrayRef(compile_time_binding_stack_)
-                          .drop_front(remaining_compile_time_bindings)) {
+  for (auto inst_id : peek_compile_time_bindings) {
     result.suspended_items.push_back(
         {.index = SuspendedScope::ScopeItem::IndexForCompileTimeBinding,
          .inst_id = inst_id});
   }
-  compile_time_binding_stack_.truncate(remaining_compile_time_bindings);
+  compile_time_binding_stack_.PopArray();
 
   // This would be easy to support if we had a need, but currently we do not.
   CARBON_CHECK(!result.entry.has_returned_var)
@@ -196,20 +188,22 @@ auto ScopeStack::Suspend() -> SuspendedScope {
 }
 
 auto ScopeStack::Restore(SuspendedScope scope) -> void {
+  compile_time_binding_stack_.PushArray();
   for (auto [index, inst_id] : scope.suspended_items) {
     if (index == SuspendedScope::ScopeItem::IndexForCompileTimeBinding) {
-      compile_time_binding_stack_.push_back(inst_id);
+      compile_time_binding_stack_.AppendToTop(inst_id);
     } else {
       lexical_lookup_.Restore({.index = index, .inst_id = inst_id},
                               scope.entry.index);
     }
   }
 
-  CARBON_CHECK(scope.entry.next_compile_time_bind_index.index ==
-               static_cast<int32_t>(compile_time_binding_stack_.size()))
+  CARBON_CHECK(
+      scope.entry.next_compile_time_bind_index.index ==
+      static_cast<int32_t>(compile_time_binding_stack_.all_values_size()))
       << "Wrong number of entries in compile-time binding stack "
          "when restoring, have "
-      << compile_time_binding_stack_.size() << ", expected "
+      << compile_time_binding_stack_.all_values_size() << ", expected "
       << scope.entry.next_compile_time_bind_index.index;
 
   if (scope.entry.scope_id.is_valid()) {

toolchain/check/scope_stack.h

@@ -5,6 +5,7 @@
 #ifndef CARBON_TOOLCHAIN_CHECK_SCOPE_STACK_H_
 #define CARBON_TOOLCHAIN_CHECK_SCOPE_STACK_H_
 
+#include "common/array_stack.h"
 #include "common/set.h"
 #include "llvm/ADT/SmallVector.h"
 #include "toolchain/check/lexical_lookup.h"
@@ -126,7 +127,7 @@ class ScopeStack {
 
   // Pushes a compile-time binding into the current scope.
   auto PushCompileTimeBinding(SemIR::InstId bind_id) -> void {
-    compile_time_binding_stack_.push_back(bind_id);
+    compile_time_binding_stack_.AppendToTop(bind_id);
   }
 
   // Temporarily removes the top of the stack and its lexical lookup results.
@@ -146,7 +147,7 @@ class ScopeStack {
     return break_continue_stack_;
   }
 
-  auto compile_time_binding_stack() -> llvm::SmallVector<SemIR::InstId>& {
+  auto compile_time_bindings_stack() -> ArrayStack<SemIR::InstId>& {
     return compile_time_binding_stack_;
   }
 
@@ -210,7 +211,7 @@ class ScopeStack {
   llvm::SmallVector<NonLexicalScope> non_lexical_scope_stack_;
 
   // A stack of the current compile time bindings.
-  llvm::SmallVector<SemIR::InstId> compile_time_binding_stack_;
+  ArrayStack<SemIR::InstId> compile_time_binding_stack_;
 
   // The index of the next scope that will be pushed onto scope_stack_. The
   // first is always the package scope.