Stop using ArrayStack for the cleanup stack. (#7505)

Because we merge cleanups across scopes in various cases, and want to
use linear indexes into the complete stack, the ArrayStack abstraction
is getting in the way more than it's helping. Switch to just a
SmallVector.

This loses the unit testing of the MergeIntoGrandparent logic. This is
covered indirectly by check tests still, but direct testing of it is a
bit tricky given that ScopeStack isn't set up for use without a Context.

common/array_stack.h

@@ -86,24 +86,6 @@ class ArrayStack {
     llvm::append_range(values_, values);
   }
 
-  // Truncates the top array on the stack to the given size.
-  auto TruncateTopArray(size_t size) -> void {
-    CARBON_CHECK(!array_offsets_.empty());
-    values_.truncate(array_offsets_.back() + size);
-  }
-
-  // Merges the top array into the the array two before it, leaving the top
-  // array empty.
-  auto MergeTopArrayIntoGrandparent() -> void {
-    CARBON_CHECK(array_offsets_.size() >= 3);
-    auto new_mid =
-        std::rotate(values_.begin() + array_offsets_[array_offsets_.size() - 2],
-                    values_.begin() + array_offsets_[array_offsets_.size() - 1],
-                    values_.end());
-    array_offsets_[array_offsets_.size() - 2] = new_mid - values_.begin();
-    array_offsets_[array_offsets_.size() - 1] = values_.size();
-  }
-
   // Returns the current number of values in all arrays.
   auto all_values_size() const -> size_t { return values_.size(); }
 

common/array_stack_test.cpp

@@ -103,123 +103,5 @@ TEST(ArrayStack, PeekArrayAt) {
   EXPECT_THAT(stack.PeekArrayAt(2), ElementsAre(3));
 }
 
-TEST(ArrayStack, MergeTopArrayIntoGrandparent) {
-  ArrayStack<int> stack;
-
-  // Basic case with 3 arrays.
-  stack.PushArray();
-  stack.AppendToTop(1);
-  stack.AppendToTop(2);
-
-  stack.PushArray();
-  stack.AppendToTop(3);
-  stack.AppendToTop(4);
-  stack.AppendToTop(5);
-
-  stack.PushArray();
-  stack.AppendToTop(6);
-  stack.AppendToTop(7);
-
-  stack.MergeTopArrayIntoGrandparent();
-
-  EXPECT_THAT(stack.PeekArrayAt(0), ElementsAre(1, 2, 6, 7));
-  EXPECT_THAT(stack.PeekArrayAt(1), ElementsAre(3, 4, 5));
-  EXPECT_THAT(stack.PeekArrayAt(2), IsEmpty());
-  EXPECT_THAT(stack.PeekArray(), IsEmpty());
-  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2, 6, 7, 3, 4, 5));
-
-  // Appending to the now-empty top array and popping works as expected.
-  stack.AppendToTop(8);
-  EXPECT_THAT(stack.PeekArray(), ElementsAre(8));
-  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2, 6, 7, 3, 4, 5, 8));
-
-  stack.PopArray();
-  EXPECT_THAT(stack.PeekArray(), ElementsAre(3, 4, 5));
-  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2, 6, 7, 3, 4, 5));
-
-  stack.PopArray();
-  EXPECT_THAT(stack.PeekArray(), ElementsAre(1, 2, 6, 7));
-  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2, 6, 7));
-}
-
-TEST(ArrayStack, MergeTopArrayIntoGrandparentDeeperStack) {
-  ArrayStack<int> stack;
-
-  // Verify behavior when there are more than 3 arrays on the stack.
-  stack.PushArray();
-  stack.AppendToTop(10);
-
-  stack.PushArray();
-  stack.AppendToTop(20);
-
-  stack.PushArray();
-  stack.AppendToTop(30);
-
-  stack.PushArray();
-  stack.AppendToTop(40);
-  stack.AppendToTop(50);
-
-  stack.MergeTopArrayIntoGrandparent();
-
-  EXPECT_THAT(stack.PeekArrayAt(0), ElementsAre(10));
-  EXPECT_THAT(stack.PeekArrayAt(1), ElementsAre(20, 40, 50));
-  EXPECT_THAT(stack.PeekArrayAt(2), ElementsAre(30));
-  EXPECT_THAT(stack.PeekArrayAt(3), IsEmpty());
-  EXPECT_THAT(stack.PeekAllValues(), ElementsAre(10, 20, 40, 50, 30));
-}
-
-TEST(ArrayStack, MergeTopArrayIntoGrandparentEmptyArrays) {
-  // Test when the parent array is initially empty.
-  {
-    ArrayStack<int> stack;
-    stack.PushArray();
-    stack.AppendToTop(1);
-    stack.PushArray();
-    stack.PushArray();
-    stack.AppendToTop(2);
-    stack.AppendToTop(3);
-
-    stack.MergeTopArrayIntoGrandparent();
-    EXPECT_THAT(stack.PeekArrayAt(0), ElementsAre(1, 2, 3));
-    EXPECT_THAT(stack.PeekArrayAt(1), IsEmpty());
-    EXPECT_THAT(stack.PeekArrayAt(2), IsEmpty());
-    EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2, 3));
-  }
-
-  // Test when the top array is initially empty.
-  {
-    ArrayStack<int> stack;
-    stack.PushArray();
-    stack.AppendToTop(1);
-    stack.PushArray();
-    stack.AppendToTop(2);
-    stack.PushArray();
-
-    stack.MergeTopArrayIntoGrandparent();
-    EXPECT_THAT(stack.PeekArrayAt(0), ElementsAre(1));
-    EXPECT_THAT(stack.PeekArrayAt(1), ElementsAre(2));
-    EXPECT_THAT(stack.PeekArrayAt(2), IsEmpty());
-    EXPECT_THAT(stack.PeekAllValues(), ElementsAre(1, 2));
-  }
-
-  // Test when the grandparent array is initially empty.
-  {
-    ArrayStack<int> stack;
-    stack.PushArray();
-    stack.PushArray();
-    stack.AppendToTop(1);
-    stack.AppendToTop(2);
-    stack.PushArray();
-    stack.AppendToTop(3);
-    stack.AppendToTop(4);
-
-    stack.MergeTopArrayIntoGrandparent();
-    EXPECT_THAT(stack.PeekArrayAt(0), ElementsAre(3, 4));
-    EXPECT_THAT(stack.PeekArrayAt(1), ElementsAre(1, 2));
-    EXPECT_THAT(stack.PeekArrayAt(2), IsEmpty());
-    EXPECT_THAT(stack.PeekAllValues(), ElementsAre(3, 4, 1, 2));
-  }
-}
-
 }  // namespace
 }  // namespace Carbon::Testing

toolchain/check/control_flow.cpp

@@ -144,7 +144,7 @@ auto MaybeAddCleanupForInst(Context& context, SemIR::InstId inst_id) -> void {
     return;
   }
 
-  context.scope_stack().destroy_id_stack().AppendToTop(inst_id);
+  context.scope_stack().PushCleanupFor(inst_id);
 }
 
 // Adds cleanups for variables added after `depth`.

toolchain/check/scope_stack.cpp

@@ -29,8 +29,7 @@ auto ScopeStack::VerifyOnFinish() const -> void {
   CARBON_CHECK(break_continue_stack_.empty(), "{0}",
                break_continue_stack_.size());
   CARBON_CHECK(scope_stack_.empty(), "{0}", scope_stack_.size());
-  CARBON_CHECK(destroy_id_stack_.empty(), "{0}",
-               destroy_id_stack_.all_values_size());
+  CARBON_CHECK(destroy_id_stack_.empty(), "{0}", destroy_id_stack_.size());
   CARBON_CHECK(non_lexical_scope_stack_.empty(), "{0}",
                non_lexical_scope_stack_.size());
   CARBON_CHECK(compile_time_binding_stack_.empty(), "{0}",
@@ -71,7 +70,8 @@ auto ScopeStack::Push(SemIR::InstId scope_inst_id, SemIR::NameScopeId scope_id,
            compile_time_binding_stack_.all_values_size()),
        .lexical_lookup_has_load_error =
            LexicalLookupHasLoadError() || lexical_lookup_has_load_error,
-       .cleanup_scope_kind = cleanup_scope_kind});
+       .cleanup_scope_kind = cleanup_scope_kind,
+       .cleanup_scope_depth = CleanupScopeDepth(destroy_id_stack_.size())});
   if (scope_stack_.back().is_lexical_scope()) {
     // For lexical lookups, unqualified lookup doesn't know how to find the
     // associated specific, so if we start adding lexical scopes associated with
@@ -96,10 +96,6 @@ auto ScopeStack::Push(SemIR::InstId scope_inst_id, SemIR::NameScopeId scope_id,
   ++next_scope_index_.index;
 
   VerifyNextCompileTimeBindIndex("Push", scope_stack_.back());
-
-  if (cleanup_scope_kind == CleanupScopeKind::Owned) {
-    destroy_id_stack_.PushArray();
-  }
 }
 
 auto ScopeStack::PushForDeclName() -> void {
@@ -162,11 +158,11 @@ auto ScopeStack::Pop(bool check_unused) -> void {
     non_lexical_scope_stack_.pop_back();
   }
 
-  if (scope.cleanup_scope_kind == CleanupScopeKind::Owned) {
-    CARBON_CHECK(destroy_id_stack_.PeekArray().empty(),
-                 "Popping scope with cleanups");
-    destroy_id_stack_.PopArray();
-  }
+  CARBON_CHECK(scope.cleanup_scope_kind == CleanupScopeKind::Inherited ||
+                   static_cast<size_t>(scope.cleanup_scope_depth.index) ==
+                       destroy_id_stack_.size(),
+               "Popping scope with cleanups: have {0} but expected {1}",
+               destroy_id_stack_.size(), scope.cleanup_scope_depth.index);
 
   if (!return_scope_stack_.empty()) {
     if (scope.has_returned_var) {
@@ -208,7 +204,14 @@ auto ScopeStack::MergeTopScopeIntoGrandparentAndPop() -> void {
   CARBON_CHECK(grandparent.cleanup_scope_kind != CleanupScopeKind::None &&
                parent.cleanup_scope_kind == CleanupScopeKind::Owned &&
                current.cleanup_scope_kind == CleanupScopeKind::Owned);
-  destroy_id_stack_.MergeTopArrayIntoGrandparent();
+
+  // NOLINTNEXTLINE(readability-qualified-auto)
+  auto new_mid =
+      std::rotate(destroy_id_stack_.begin() + parent.cleanup_scope_depth.index,
+                  destroy_id_stack_.begin() + current.cleanup_scope_depth.index,
+                  destroy_id_stack_.end());
+  parent.cleanup_scope_depth.index = new_mid - destroy_id_stack_.begin();
+  current.cleanup_scope_depth.index = destroy_id_stack_.size();
 
   Pop();
 }

toolchain/check/scope_stack.h

@@ -244,10 +244,16 @@ class ScopeStack {
     return Peek().cleanup_scope_kind != CleanupScopeKind::None;
   }
 
+  // Registers a cleanup for `inst_id` within the current cleanup scope.
+  auto PushCleanupFor(SemIR::InstId inst_id) -> void {
+    CARBON_CHECK(IsCleanupScope());
+    destroy_id_stack_.push_back(inst_id);
+  }
+
   // Returns all values on `destroy_id_stack_` added since `depth`.
   auto GetCleanupsSince(CleanupScopeDepth depth) const
       -> llvm::ArrayRef<SemIR::InstId> {
-    return destroy_id_stack_.PeekAllValues().slice(depth.index);
+    return llvm::ArrayRef(destroy_id_stack_).slice(depth.index);
   }
 
   // Discards cleanups after the given depth, which must be within the current
@@ -255,20 +261,17 @@ class ScopeStack {
   auto DiscardCleanupsSince(CleanupScopeDepth depth) -> void {
     auto enclosing = enclosing_cleanup_scope_depth();
     CARBON_CHECK(depth >= enclosing);
-    // TODO: Consider switching `destroy_id_stack_` from `ArrayStack` to
-    // `SmallVector` - `ArrayStack` gives us nothing here.
-    destroy_id_stack_.TruncateTopArray(depth.index - enclosing.index);
+    destroy_id_stack_.truncate(depth.index);
   }
 
   // Returns the current depth of the cleanup stack.
   auto cleanup_scope_depth() const -> CleanupScopeDepth {
-    return CleanupScopeDepth(destroy_id_stack_.all_values_size());
+    return CleanupScopeDepth(destroy_id_stack_.size());
   }
 
   // Returns the depth of the cleanup stack enclosing the current scope.
   auto enclosing_cleanup_scope_depth() const -> CleanupScopeDepth {
-    return CleanupScopeDepth(destroy_id_stack_.all_values_size() -
-                             destroy_id_stack_.PeekArray().size());
+    return Peek().cleanup_scope_depth;
   }
 
   // Returns the depth of the cleanup stack enclosing this function scope.
@@ -280,10 +283,6 @@ class ScopeStack {
     return break_continue_stack_;
   }
 
-  auto destroy_id_stack() -> ArrayStack<SemIR::InstId>& {
-    return destroy_id_stack_;
-  }
-
   auto compile_time_binding_stack() -> ArrayStack<SemIR::InstId>& {
     return compile_time_binding_stack_;
   }
@@ -330,6 +329,9 @@ class ScopeStack {
     // The kind of cleanup scope that is associated with this scope.
     CleanupScopeKind cleanup_scope_kind = CleanupScopeKind::None;
 
+    // The cleanup scope depth on entry to this scope.
+    CleanupScopeDepth cleanup_scope_depth;
+
     // Whether there are any ids in the `names` set.
     int num_names = 0;
 
@@ -415,7 +417,7 @@ class ScopeStack {
 
   // A stack of instances to destroy. This only has entries inside of function
   // bodies, where destruction on scope exit is required.
-  ArrayStack<SemIR::InstId> destroy_id_stack_;
+  llvm::SmallVector<SemIR::InstId> destroy_id_stack_;
 
   // Information about non-lexical scopes. This is a subset of the entries and
   // the information in scope_stack_.