Refactor region_stack logic out of Context (#4927)

toolchain/check/BUILD

@@ -71,6 +71,7 @@ cc_library(
         "param_and_arg_refs_stack.h",
         "pattern_match.h",
         "pending_block.h",
+        "region_stack.h",
         "return.h",
         "subst.h",
         "type_completion.h",

toolchain/check/context.cpp

@@ -54,7 +54,9 @@ Context::Context(DiagnosticEmitter* emitter,
       decl_name_stack_(this),
       scope_stack_(sem_ir_->identifiers()),
       vtable_stack_("vtable_stack_", *sem_ir, vlog_stream),
-      global_init_(this) {
+      global_init_(this),
+      region_stack_(
+          [this](SemIRLoc loc, std::string label) { TODO(loc, label); }) {
   // Prepare fields which relate to the number of IRs available for import.
   import_irs().Reserve(imported_ir_count);
   import_ir_constant_values_.reserve(imported_ir_count);
@@ -745,29 +747,14 @@ auto Context::LookupNameInCore(SemIR::LocId loc_id, llvm::StringRef name)
   return constant_values().GetConstantInstId(scope_result.target_inst_id());
 }
 
-auto Context::AddToRegion(SemIR::InstBlockId block_id, SemIR::LocId loc_id)
-    -> void {
-  if (region_stack_.empty()) {
-    TODO(loc_id,
-         "Control flow expressions are currently only supported inside "
-         "functions.");
-    return;
-  }
-  if (block_id == SemIR::InstBlockId::Unreachable) {
-    return;
-  }
-
-  region_stack_.AppendToTop(block_id);
-}
-
 auto Context::BeginSubpattern() -> void {
   inst_block_stack().Push();
-  PushRegion(inst_block_stack().PeekOrAdd());
+  region_stack_.PushRegion(inst_block_stack().PeekOrAdd());
 }
 
 auto Context::EndSubpatternAsExpr(SemIR::InstId result_id)
     -> SemIR::ExprRegionId {
-  if (region_stack_.PeekArray().size() > 1) {
+  if (region_stack_.PeekRegion().size() > 1) {
     // End the exit block with a branch to a successor block, whose contents
     // will be determined later.
     AddInst(SemIR::LocIdAndInst::NoLoc<SemIR::Branch>(
@@ -777,20 +764,20 @@ auto Context::EndSubpatternAsExpr(SemIR::InstId result_id)
     // need control flow out of it.
   }
   auto block_id = inst_block_stack().Pop();
-  CARBON_CHECK(block_id == region_stack_.PeekArray().back());
+  CARBON_CHECK(block_id == region_stack_.PeekRegion().back());
 
   // TODO: Is it possible to validate that this region is genuinely
   // single-entry, single-exit?
   return sem_ir().expr_regions().Add(
-      {.block_ids = PopRegion(), .result_id = result_id});
+      {.block_ids = region_stack_.PopRegion(), .result_id = result_id});
 }
 
 auto Context::EndSubpatternAsEmpty() -> void {
   auto block_id = inst_block_stack().Pop();
-  CARBON_CHECK(block_id == region_stack_.PeekArray().back());
-  CARBON_CHECK(region_stack_.PeekArray().size() == 1);
+  CARBON_CHECK(block_id == region_stack_.PeekRegion().back());
+  CARBON_CHECK(region_stack_.PeekRegion().size() == 1);
   CARBON_CHECK(inst_blocks().Get(block_id).empty());
-  region_stack_.PopArray();
+  region_stack_.PopAndDiscardRegion();
 }
 
 auto Context::InsertHere(SemIR::ExprRegionId region_id) -> SemIR::InstId {
@@ -823,12 +810,12 @@ auto Context::InsertHere(SemIR::ExprRegionId region_id) -> SemIR::InstId {
   inst_block_stack_.Pop();
   // TODO: this will cumulatively cost O(MN) running time for M blocks
   // at the Nth level of the stack. Figure out how to do better.
-  region_stack_.AppendToTop(region.block_ids);
+  region_stack_.AddToRegion(region.block_ids);
   auto resume_with_block_id =
       insts().GetAs<SemIR::Branch>(exit_block.back()).target_id;
   CARBON_CHECK(inst_blocks().GetOrEmpty(resume_with_block_id).empty());
   inst_block_stack_.Push(resume_with_block_id);
-  AddToRegion(resume_with_block_id, loc_id);
+  region_stack_.AddToRegion(resume_with_block_id, loc_id);
   return region.result_id;
 }
 

toolchain/check/context.h

@@ -17,6 +17,7 @@
 #include "toolchain/check/inst_block_stack.h"
 #include "toolchain/check/node_stack.h"
 #include "toolchain/check/param_and_arg_refs_stack.h"
+#include "toolchain/check/region_stack.h"
 #include "toolchain/check/scope_index.h"
 #include "toolchain/check/scope_stack.h"
 #include "toolchain/parse/node_ids.h"
@@ -293,27 +294,6 @@ class Context {
     return scope_stack().GetCurrentScopeAs<InstT>(sem_ir());
   }
 
-  // Mark the start of a new single-entry region with the given entry block.
-  auto PushRegion(SemIR::InstBlockId entry_block_id) -> void {
-    region_stack_.PushArray();
-    region_stack_.AppendToTop(entry_block_id);
-  }
-
-  // Add `block_id` to the most recently pushed single-entry region. To preserve
-  // the single-entry property, `block_id` must not be directly reachable from
-  // any block outside the region. To ensure the region's blocks are in lexical
-  // order, this should be called when the first parse node associated with this
-  // block is handled, or as close as possible.
-  auto AddToRegion(SemIR::InstBlockId block_id, SemIR::LocId loc_id) -> void;
-
-  // Complete creation of the most recently pushed single-entry region, and
-  // return a list of its blocks.
-  auto PopRegion() -> llvm::SmallVector<SemIR::InstBlockId> {
-    llvm::SmallVector<SemIR::InstBlockId> result(region_stack_.PeekArray());
-    region_stack_.PopArray();
-    return result;
-  }
-
   // Returns the type ID for a constant of type `type`.
   auto GetTypeIdForTypeConstant(SemIR::ConstantId constant_id) -> SemIR::TypeId;
 
@@ -602,6 +582,8 @@ class Context {
     return var_storage_map_;
   }
 
+  auto region_stack() -> RegionStack& { return region_stack_; }
+
   auto full_pattern_stack() -> FullPatternStack& {
     return scope_stack_.full_pattern_stack();
   }
@@ -728,7 +710,7 @@ class Context {
   Map<SemIR::InstId, SemIR::InstId> var_storage_map_;
 
   // Stack of single-entry regions being built.
-  ArrayStack<SemIR::InstBlockId> region_stack_;
+  RegionStack region_stack_;
 };
 
 }  // namespace Carbon::Check

toolchain/check/control_flow.cpp

@@ -54,7 +54,7 @@ auto AddConvergenceBlockAndPush(Context& context, Parse::NodeId node_id,
     context.inst_block_stack().Pop();
   }
   context.inst_block_stack().Push(new_block_id);
-  context.AddToRegion(new_block_id, node_id);
+  context.region_stack().AddToRegion(new_block_id, node_id);
 }
 
 auto AddConvergenceBlockWithArgAndPush(
@@ -74,7 +74,7 @@ auto AddConvergenceBlockWithArgAndPush(
     context.inst_block_stack().Pop();
   }
   context.inst_block_stack().Push(new_block_id);
-  context.AddToRegion(new_block_id, node_id);
+  context.region_stack().AddToRegion(new_block_id, node_id);
 
   // Acquire the result value.
   SemIR::TypeId result_type_id =

toolchain/check/handle_function.cpp

@@ -401,7 +401,7 @@ static auto HandleFunctionDefinitionAfterSignature(
   // Create the function scope and the entry block.
   context.return_scope_stack().push_back({.decl_id = decl_id});
   context.inst_block_stack().Push();
-  context.PushRegion(context.inst_block_stack().PeekOrAdd());
+  context.region_stack().PushRegion(context.inst_block_stack().PeekOrAdd());
   context.scope_stack().Push(decl_id);
   StartGenericDefinition(context);
 
@@ -465,7 +465,7 @@ auto HandleParseNode(Context& context, Parse::FunctionDefinitionId node_id)
   context.decl_name_stack().PopScope();
 
   auto& function = context.functions().Get(function_id);
-  function.body_block_ids = context.PopRegion();
+  function.body_block_ids = context.region_stack().PopRegion();
 
   // If this is a generic function, collect information about the definition.
   FinishGenericDefinition(context, function.generic_id);

toolchain/check/handle_if_expr.cpp

@@ -26,7 +26,7 @@ auto HandleParseNode(Context& context, Parse::IfExprIfId node_id) -> bool {
   // Start emitting the `then` block.
   context.inst_block_stack().Pop();
   context.inst_block_stack().Push(then_block_id);
-  context.AddToRegion(then_block_id, node_id);
+  context.region_stack().AddToRegion(then_block_id, node_id);
 
   context.node_stack().Push(if_node, else_block_id);
   return true;
@@ -57,7 +57,7 @@ auto HandleParseNode(Context& context, Parse::IfExprThenId node_id) -> bool {
 
   // Start emitting the `else` block.
   context.inst_block_stack().Push(else_block_id);
-  context.AddToRegion(else_block_id, node_id);
+  context.region_stack().AddToRegion(else_block_id, node_id);
 
   context.node_stack().Push(node_id, then_value_id);
   return true;

toolchain/check/handle_if_statement.cpp

@@ -29,7 +29,7 @@ auto HandleParseNode(Context& context, Parse::IfConditionId node_id) -> bool {
   // Start emitting the `then` block.
   context.inst_block_stack().Pop();
   context.inst_block_stack().Push(then_block_id);
-  context.AddToRegion(then_block_id, node_id);
+  context.region_stack().AddToRegion(then_block_id, node_id);
 
   context.node_stack().Push(node_id, else_block_id);
   return true;
@@ -41,7 +41,7 @@ auto HandleParseNode(Context& context, Parse::IfStatementElseId node_id)
 
   // Switch to emitting the `else` block.
   context.inst_block_stack().Push(else_block_id);
-  context.AddToRegion(else_block_id, node_id);
+  context.region_stack().AddToRegion(else_block_id, node_id);
 
   context.node_stack().Push(node_id);
   return true;
@@ -57,7 +57,7 @@ auto HandleParseNode(Context& context, Parse::IfStatementId node_id) -> bool {
       context.AddInst<SemIR::Branch>(node_id, {.target_id = else_block_id});
       context.inst_block_stack().Pop();
       context.inst_block_stack().Push(else_block_id);
-      context.AddToRegion(else_block_id, node_id);
+      context.region_stack().AddToRegion(else_block_id, node_id);
       break;
     }
 

toolchain/check/handle_loop_statement.cpp

@@ -22,7 +22,7 @@ auto HandleParseNode(Context& context, Parse::WhileConditionStartId node_id)
 
   // Start emitting the loop header block.
   context.inst_block_stack().Push(loop_header_id);
-  context.AddToRegion(loop_header_id, node_id);
+  context.region_stack().AddToRegion(loop_header_id, node_id);
 
   context.node_stack().Push(node_id, loop_header_id);
   return true;
@@ -43,7 +43,7 @@ auto HandleParseNode(Context& context, Parse::WhileConditionId node_id)
 
   // Start emitting the loop body.
   context.inst_block_stack().Push(loop_body_id);
-  context.AddToRegion(loop_body_id, node_id);
+  context.region_stack().AddToRegion(loop_body_id, node_id);
   context.break_continue_stack().push_back(
       {.break_target = loop_exit_id, .continue_target = loop_header_id});
 
@@ -65,7 +65,7 @@ auto HandleParseNode(Context& context, Parse::WhileStatementId node_id)
 
   // Start emitting the loop exit block.
   context.inst_block_stack().Push(loop_exit_id);
-  context.AddToRegion(loop_exit_id, node_id);
+  context.region_stack().AddToRegion(loop_exit_id, node_id);
   return true;
 }
 

toolchain/check/handle_operator.cpp

@@ -364,7 +364,7 @@ static auto HandleShortCircuitOperand(Context& context, Parse::NodeId node_id,
   context.inst_block_stack().Pop();
   context.inst_block_stack().Push(end_block_id);
   context.inst_block_stack().Push(rhs_block_id);
-  context.AddToRegion(rhs_block_id, node_id);
+  context.region_stack().AddToRegion(rhs_block_id, node_id);
 
   // HandleShortCircuitOperator will follow, and doesn't need the operand on the
   // node stack.
@@ -405,7 +405,7 @@ static auto HandleShortCircuitOperator(Context& context, Parse::NodeId node_id)
   context.AddInst<SemIR::BranchWithArg>(
       node_id, {.target_id = resume_block_id, .arg_id = rhs_id});
   context.inst_block_stack().Pop();
-  context.AddToRegion(resume_block_id, node_id);
+  context.region_stack().AddToRegion(resume_block_id, node_id);
 
   // Collect the result from either the first or second operand.
   auto result_id = context.AddInst<SemIR::BlockArg>(

toolchain/check/region_stack.h

@@ -0,0 +1,82 @@
+// 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_REGION_STACK_H_
+#define CARBON_TOOLCHAIN_CHECK_REGION_STACK_H_
+
+#include <utility>
+
+#include "common/array_stack.h"
+#include "toolchain/check/diagnostic_helpers.h"
+#include "toolchain/sem_ir/ids.h"
+
+namespace Carbon::Check {
+
+// Provides a stack of single-entry regions being built.
+class RegionStack {
+ public:
+  // A callback for Context::TODO.
+  using TodoFn = std::function<auto(SemIRLoc, std::string)->void>;
+
+  explicit RegionStack(TodoFn todo_fn) : todo_fn_(std::move(todo_fn)) {}
+
+  // Mark the start of a new single-entry region with the given entry block.
+  auto PushRegion(SemIR::InstBlockId entry_block_id) -> void {
+    stack_.PushArray();
+    stack_.AppendToTop(entry_block_id);
+  }
+
+  // Add `block_id` to the most recently pushed single-entry region. To preserve
+  // the single-entry property, `block_id` must not be directly reachable from
+  // any block outside the region. To ensure the region's blocks are in lexical
+  // order, this should be called when the first parse node associated with this
+  // block is handled, or as close as possible.
+  auto AddToRegion(SemIR::InstBlockId block_id, SemIR::LocId loc_id) -> void {
+    if (stack_.empty()) {
+      todo_fn_(loc_id,
+               "Control flow expressions are currently only supported inside "
+               "functions.");
+      return;
+    }
+    if (block_id == SemIR::InstBlockId::Unreachable) {
+      return;
+    }
+
+    stack_.AppendToTop(block_id);
+  }
+
+  // Adds multiple blocks at once. The caller is responsible for validating that
+  // each block is reachable.
+  auto AddToRegion(llvm::ArrayRef<SemIR::InstBlockId> block_ids) -> void {
+    stack_.AppendToTop(block_ids);
+  }
+
+  // Complete creation of the most recently pushed single-entry region, and
+  // return a list of its blocks.
+  auto PopRegion() -> llvm::SmallVector<SemIR::InstBlockId> {
+    llvm::SmallVector<SemIR::InstBlockId> result(stack_.PeekArray());
+    stack_.PopArray();
+    return result;
+  }
+
+  // Pops a region, and does not return its contents.
+  auto PopAndDiscardRegion() -> void { stack_.PopArray(); }
+
+  // Returns the top-most region.
+  auto PeekRegion() -> llvm::ArrayRef<SemIR::InstBlockId> {
+    return stack_.PeekArray();
+  }
+
+  // Returns true if any regions have been added.
+  auto empty() -> bool { return stack_.empty(); }
+
+ private:
+  TodoFn todo_fn_;
+
+  ArrayStack<SemIR::InstBlockId> stack_;
+};
+
+}  // namespace Carbon::Check
+
+#endif  // CARBON_TOOLCHAIN_CHECK_REGION_STACK_H_