Factor the scope stack and lexical lookups out of Check::Context. (#3688)

toolchain/check/BUILD

@@ -28,6 +28,25 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "scope_stack",
+    srcs = ["scope_stack.cpp"],
+    hdrs = [
+        "lexical_lookup.h",
+        "scope_index.h",
+        "scope_stack.h",
+    ],
+    deps = [
+        "//common:check",
+        "//common:ostream",
+        "//common:vlog",
+        "//toolchain/base:index_base",
+        "//toolchain/sem_ir:file",
+        "//toolchain/sem_ir:ids",
+        "@llvm-project//llvm:Support",
+    ],
+)
+
 cc_library(
     name = "context",
     srcs = [
@@ -46,17 +65,16 @@ cc_library(
         "decl_state.h",
         "eval.h",
         "inst_block_stack.h",
-        "lexical_lookup.h",
         "modifiers.h",
         "pending_block.h",
         "return.h",
-        "scope_index.h",
     ],
     deps = [
         ":node_stack",
         "//common:check",
         "//common:vlog",
         "//toolchain/base:index_base",
+        "//toolchain/check:scope_stack",
         "//toolchain/lex:tokenized_buffer",
         "//toolchain/parse:node_kind",
         "//toolchain/parse:tree",

toolchain/check/check.cpp

@@ -175,13 +175,13 @@ static auto InitPackageScopeAndImports(Context& context, UnitInfo& unit_info)
     if (error_in_import) {
       context.name_scopes().Get(SemIR::NameScopeId::Package).has_error = true;
     }
-    context.PushScope(package_inst_id, SemIR::NameScopeId::Package,
-                      error_in_import);
+    context.scope_stack().Push(package_inst_id, SemIR::NameScopeId::Package,
+                               error_in_import);
   } else {
     // Push the scope; there are no names to add.
-    context.PushScope(package_inst_id, SemIR::NameScopeId::Package);
+    context.scope_stack().Push(package_inst_id, SemIR::NameScopeId::Package);
   }
-  CARBON_CHECK(context.current_scope_index() == ScopeIndex::Package);
+  CARBON_CHECK(context.scope_stack().PeekIndex() == ScopeIndex::Package);
 
   for (auto& [package_id, package_imports] : unit_info.package_imports_map) {
     if (!package_id.is_valid()) {
@@ -255,7 +255,7 @@ static auto CheckParseTree(
 
   // Pop information for the file-level scope.
   sem_ir.set_top_inst_block_id(context.inst_block_stack().Pop());
-  context.PopScope();
+  context.scope_stack().Pop();
   context.FinalizeExports();
 
   context.VerifyOnFinish();

toolchain/check/context.cpp

@@ -38,7 +38,7 @@ Context::Context(const Lex::TokenizedBuffer& tokens, DiagnosticEmitter& emitter,
       params_or_args_stack_("params_or_args_stack_", sem_ir, vlog_stream),
       args_type_info_stack_("args_type_info_stack_", sem_ir, vlog_stream),
       decl_name_stack_(this),
-      lexical_lookup_(sem_ir_->identifiers()) {
+      scope_stack_(sem_ir_->identifiers()) {
   // Map the builtin `<error>` and `type` type constants to their corresponding
   // special `TypeId` values.
   type_ids_for_type_constants_.insert(
@@ -61,7 +61,7 @@ auto Context::VerifyOnFinish() -> void {
   // various pieces of context go out of scope. At this point, nothing should
   // remain.
   // node_stack_ will still contain top-level entities.
-  CARBON_CHECK(scope_stack_.empty()) << scope_stack_.size();
+  scope_stack_.VerifyOnFinish();
   CARBON_CHECK(inst_block_stack_.empty()) << inst_block_stack_.size();
   CARBON_CHECK(params_or_args_stack_.empty()) << params_or_args_stack_.size();
 }
@@ -206,18 +206,9 @@ auto Context::AddPackageImports(Parse::NodeId import_node,
 
 auto Context::AddNameToLookup(SemIR::NameId name_id, SemIR::InstId target_id)
     -> void {
-  if (current_scope().names.insert(name_id).second) {
-    // TODO: Reject if we previously performed a failed lookup for this name in
-    // this scope or a scope nested within it.
-    auto& lexical_results = lexical_lookup_.Get(name_id);
-    CARBON_CHECK(lexical_results.empty() ||
-                 lexical_results.back().scope_index < current_scope_index())
-        << "Failed to clean up after scope nested within the current scope";
-    lexical_results.push_back(
-        {.inst_id = target_id, .scope_index = current_scope_index()});
-  } else {
-    DiagnoseDuplicateName(target_id,
-                          lexical_lookup_.Get(name_id).back().inst_id);
+  if (auto existing = scope_stack().LookupOrAddName(name_id, target_id);
+      existing.is_valid()) {
+    DiagnoseDuplicateName(target_id, existing);
   }
 }
 
@@ -287,15 +278,11 @@ auto Context::LookupNameInDecl(Parse::NodeId /*parse_node*/,
     //    In this case, we're not in the correct scope to define a member of
     //    class A, so we should reject, and we achieve this by not finding the
     //    name A from the outer scope.
-    auto& lexical_results = lexical_lookup_.Get(name_id);
-    if (!lexical_results.empty()) {
-      auto result = lexical_results.back();
-      if (result.scope_index == current_scope_index()) {
-        ResolveIfImportRefUnused(result.inst_id);
-        return result.inst_id;
-      }
+    auto result = scope_stack().LookupInCurrentScope(name_id);
+    if (result.is_valid()) {
+      ResolveIfImportRefUnused(result);
     }
-    return SemIR::InstId::Invalid;
+    return result;
   } else {
     // We do not look into `extend`ed scopes here. A qualified name in a
     // declaration must specify the exact scope in which the name was originally
@@ -316,20 +303,11 @@ auto Context::LookupUnqualifiedName(Parse::NodeId parse_node,
 
   // Find the results from enclosing lexical scopes. These will be combined with
   // results from non-lexical scopes such as namespaces and classes.
-  llvm::ArrayRef<LexicalLookup::Result> lexical_results =
-      lexical_lookup_.Get(name_id);
+  auto [lexical_result, non_lexical_scopes] =
+      scope_stack().LookupInEnclosingScopes(name_id);
 
   // Walk the non-lexical scopes and perform lookups into each of them.
-  for (auto [index, name_scope_id] : llvm::reverse(non_lexical_scope_stack_)) {
-    // If the innermost lexical result is within this non-lexical scope, then
-    // it shadows all further non-lexical results and we're done.
-    if (!lexical_results.empty() &&
-        lexical_results.back().scope_index > index) {
-      auto inst_id = lexical_results.back().inst_id;
-      ResolveIfImportRefUnused(inst_id);
-      return inst_id;
-    }
-
+  for (auto [index, name_scope_id] : llvm::reverse(non_lexical_scopes)) {
     if (auto non_lexical_result =
             LookupQualifiedName(parse_node, name_id, name_scope_id,
                                 /*required=*/false);
@@ -338,16 +316,13 @@ auto Context::LookupUnqualifiedName(Parse::NodeId parse_node,
     }
   }
 
-  if (!lexical_results.empty()) {
-    auto inst_id = lexical_results.back().inst_id;
-    ResolveIfImportRefUnused(inst_id);
-    return inst_id;
+  if (lexical_result.is_valid()) {
+    ResolveIfImportRefUnused(lexical_result);
+    return lexical_result;
   }
 
   // We didn't find anything at all.
-  if (!lexical_lookup_has_load_error_) {
-    DiagnoseNameNotFound(parse_node, name_id);
-  }
+  DiagnoseNameNotFound(parse_node, name_id);
   return SemIR::InstId::BuiltinError;
 }
 
@@ -410,76 +385,6 @@ auto Context::LookupQualifiedName(Parse::NodeId parse_node,
   return result_id;
 }
 
-auto Context::PushScope(SemIR::InstId scope_inst_id,
-                        SemIR::NameScopeId scope_id,
-                        bool lexical_lookup_has_load_error) -> void {
-  scope_stack_.push_back(
-      {.index = next_scope_index_,
-       .scope_inst_id = scope_inst_id,
-       .scope_id = scope_id,
-       .prev_lexical_lookup_has_load_error = lexical_lookup_has_load_error_});
-  if (scope_id.is_valid()) {
-    non_lexical_scope_stack_.push_back({next_scope_index_, scope_id});
-  }
-
-  lexical_lookup_has_load_error_ |= lexical_lookup_has_load_error;
-
-  // TODO: Handle this case more gracefully.
-  CARBON_CHECK(next_scope_index_.index != std::numeric_limits<int32_t>::max())
-      << "Ran out of scopes";
-  ++next_scope_index_.index;
-}
-
-auto Context::PopScope() -> void {
-  auto scope = scope_stack_.pop_back_val();
-
-  lexical_lookup_has_load_error_ = scope.prev_lexical_lookup_has_load_error;
-
-  for (const auto& str_id : scope.names) {
-    auto& lexical_results = lexical_lookup_.Get(str_id);
-    CARBON_CHECK(lexical_results.back().scope_index == scope.index)
-        << "Inconsistent scope index for name " << names().GetFormatted(str_id);
-    lexical_results.pop_back();
-  }
-
-  if (scope.scope_id.is_valid()) {
-    CARBON_CHECK(non_lexical_scope_stack_.back().first == scope.index);
-    non_lexical_scope_stack_.pop_back();
-  }
-
-  if (scope.has_returned_var) {
-    CARBON_CHECK(!return_scope_stack_.empty());
-    CARBON_CHECK(return_scope_stack_.back().returned_var.is_valid());
-    return_scope_stack_.back().returned_var = SemIR::InstId::Invalid;
-  }
-}
-
-auto Context::PopToScope(ScopeIndex index) -> void {
-  while (current_scope_index() > index) {
-    PopScope();
-  }
-  CARBON_CHECK(current_scope_index() == index)
-      << "Scope index " << index << " does not enclose the current scope "
-      << current_scope_index();
-}
-
-auto Context::SetReturnedVarOrGetExisting(SemIR::InstId inst_id)
-    -> SemIR::InstId {
-  CARBON_CHECK(!return_scope_stack_.empty()) << "`returned var` in no function";
-  auto& returned_var = return_scope_stack_.back().returned_var;
-  if (returned_var.is_valid()) {
-    return returned_var;
-  }
-
-  returned_var = inst_id;
-  CARBON_CHECK(!current_scope().has_returned_var)
-      << "Scope has returned var but none is set";
-  if (inst_id.is_valid()) {
-    current_scope().has_returned_var = true;
-  }
-  return SemIR::InstId::Invalid;
-}
-
 template <typename BranchNode, typename... Args>
 static auto AddDominatedBlockAndBranchImpl(Context& context,
                                            Parse::NodeId parse_node,

toolchain/check/context.h

@@ -6,14 +6,13 @@
 #define CARBON_TOOLCHAIN_CHECK_CONTEXT_H_
 
 #include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/FoldingSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "toolchain/check/decl_name_stack.h"
 #include "toolchain/check/decl_state.h"
 #include "toolchain/check/inst_block_stack.h"
-#include "toolchain/check/lexical_lookup.h"
 #include "toolchain/check/node_stack.h"
+#include "toolchain/check/scope_stack.h"
 #include "toolchain/parse/node_ids.h"
 #include "toolchain/parse/tree.h"
 #include "toolchain/parse/tree_node_location_translator.h"
@@ -51,23 +50,6 @@ class Context {
   using DiagnosticEmitter = Carbon::DiagnosticEmitter<SemIRLocation>;
   using DiagnosticBuilder = DiagnosticEmitter::DiagnosticBuilder;
 
-  // A scope in which `break` and `continue` can be used.
-  struct BreakContinueScope {
-    SemIR::InstBlockId break_target;
-    SemIR::InstBlockId continue_target;
-  };
-
-  // A scope in which `return` can be used.
-  struct ReturnScope {
-    // The declaration from which we can return. Inside a function, this will
-    // be a `FunctionDecl`.
-    SemIR::InstId decl_id;
-
-    // The value corresponding to the current `returned var`, if any. Will be
-    // set and unset as `returned var`s are declared and go out of scope.
-    SemIR::InstId returned_var = SemIR::InstId::Invalid;
-  };
-
   // Stores references for work.
   explicit Context(const Lex::TokenizedBuffer& tokens,
                    DiagnosticEmitter& emitter, const Parse::Tree& parse_tree,
@@ -168,68 +150,13 @@ class Context {
   auto NoteIncompleteClass(SemIR::ClassId class_id, DiagnosticBuilder& builder)
       -> void;
 
-  // Pushes a scope onto scope_stack_. NameScopeId::Invalid is used for new
-  // scopes. lexical_lookup_has_load_error is used to limit diagnostics when a
-  // given namespace may contain a mix of both successful and failed name
-  // imports.
-  auto PushScope(SemIR::InstId scope_inst_id = SemIR::InstId::Invalid,
-                 SemIR::NameScopeId scope_id = SemIR::NameScopeId::Invalid,
-                 bool lexical_lookup_has_load_error = false) -> void;
-
-  // Pops the top scope from scope_stack_, cleaning up names from
-  // lexical_lookup_.
-  auto PopScope() -> void;
-
-  // Pops the top scope from scope_stack_ if it contains no names.
-  auto PopScopeIfEmpty() -> void {
-    if (scope_stack_.back().names.empty()) {
-      PopScope();
-    }
-  }
-
-  // Pops scopes until we return to the specified scope index.
-  auto PopToScope(ScopeIndex index) -> void;
-
-  // Returns the scope index associated with the current scope.
-  auto current_scope_index() const -> ScopeIndex {
-    return current_scope().index;
-  }
-
-  // Returns the name scope associated with the current lexical scope, if any.
-  auto current_scope_id() const -> SemIR::NameScopeId {
-    return current_scope().scope_id;
-  }
-
-  // Returns the instruction associated with the current scope, or Invalid if
-  // there is no such instruction, such as for a block scope.
-  auto current_scope_inst_id() const -> SemIR::InstId {
-    return current_scope().scope_inst_id;
-  }
-
-  auto GetCurrentScopeParseNode() const -> Parse::NodeId {
-    auto inst_id = current_scope_inst_id();
-    if (!inst_id.is_valid()) {
-      return Parse::NodeId::Invalid;
-    }
-    return sem_ir_->insts().GetParseNode(inst_id);
-  }
-
   // Returns the current scope, if it is of the specified kind. Otherwise,
   // returns nullopt.
   template <typename InstT>
   auto GetCurrentScopeAs() -> std::optional<InstT> {
-    auto inst_id = current_scope_inst_id();
-    if (!inst_id.is_valid()) {
-      return std::nullopt;
-    }
-    return insts().TryGetAs<InstT>(inst_id);
+    return scope_stack().GetCurrentScopeAs<InstT>(sem_ir());
   }
 
-  // If there is no `returned var` in scope, sets the given instruction to be
-  // the current `returned var` and returns an invalid instruction ID. If there
-  // is already a `returned var`, returns it instead.
-  auto SetReturnedVarOrGetExisting(SemIR::InstId inst_id) -> SemIR::InstId;
-
   // Adds a `Branch` instruction branching to a new instruction block, and
   // returns the ID of the new block. All paths to the branch target must go
   // through the current block, though not necessarily through this branch.
@@ -391,19 +318,20 @@ class Context {
     return args_type_info_stack_;
   }
 
-  auto return_scope_stack() -> llvm::SmallVector<ReturnScope>& {
-    return return_scope_stack_;
-  }
-
-  auto break_continue_stack() -> llvm::SmallVector<BreakContinueScope>& {
-    return break_continue_stack_;
-  }
-
   auto decl_name_stack() -> DeclNameStack& { return decl_name_stack_; }
 
   auto decl_state_stack() -> DeclStateStack& { return decl_state_stack_; }
 
-  auto lexical_lookup() -> LexicalLookup& { return lexical_lookup_; }
+  auto scope_stack() -> ScopeStack& { return scope_stack_; }
+
+  auto return_scope_stack() -> llvm::SmallVector<ScopeStack::ReturnScope>& {
+    return scope_stack().return_scope_stack();
+  }
+
+  auto break_continue_stack()
+      -> llvm::SmallVector<ScopeStack::BreakContinueScope>& {
+    return scope_stack().break_continue_stack();
+  }
 
   // Directly expose SemIR::File data accessors for brevity in calls.
 
@@ -461,45 +389,10 @@ class Context {
     SemIR::TypeId type_id_;
   };
 
-  // An entry in scope_stack_.
-  struct ScopeStackEntry {
-    // The sequential index of this scope entry within the file.
-    ScopeIndex index;
-
-    // The instruction associated with this entry, if any. This can be one of:
-    //
-    // - A `ClassDecl`, for a class definition scope.
-    // - A `FunctionDecl`, for the outermost scope in a function
-    //   definition.
-    // - Invalid, for any other scope.
-    SemIR::InstId scope_inst_id;
-
-    // The name scope associated with this entry, if any.
-    SemIR::NameScopeId scope_id;
-
-    // The previous state of lexical_lookup_has_load_error_, restored on pop.
-    bool prev_lexical_lookup_has_load_error;
-
-    // Names which are registered with lexical_lookup_, and will need to be
-    // unregistered when the scope ends.
-    llvm::DenseSet<SemIR::NameId> names;
-
-    // Whether a `returned var` was introduced in this scope, and needs to be
-    // unregistered when the scope ends.
-    bool has_returned_var = false;
-
-    // TODO: This likely needs to track things which need to be destructed.
-  };
-
   // If the passed in instruction ID is a ImportRefUnused, resolves it for use.
   // Called when name lookup intends to return an inst_id.
   auto ResolveIfImportRefUnused(SemIR::InstId inst_id) -> void;
 
-  auto current_scope() -> ScopeStackEntry& { return scope_stack_.back(); }
-  auto current_scope() const -> const ScopeStackEntry& {
-    return scope_stack_.back();
-  }
-
   // Tokens for getting data on literals.
   const Lex::TokenizedBuffer* tokens_;
 
@@ -533,36 +426,14 @@ class Context {
   // for a type separate from the literal arguments.
   InstBlockStack args_type_info_stack_;
 
-  // A stack of scopes from which we can `return`.
-  llvm::SmallVector<ReturnScope> return_scope_stack_;
-
-  // A stack of `break` and `continue` targets.
-  llvm::SmallVector<BreakContinueScope> break_continue_stack_;
-
-  // A stack for scope context.
-  llvm::SmallVector<ScopeStackEntry> scope_stack_;
-
-  // Information about non-lexical scopes. This is a subset of the entries and
-  // the information in scope_stack_.
-  llvm::SmallVector<std::pair<ScopeIndex, SemIR::NameScopeId>>
-      non_lexical_scope_stack_;
-
-  // The index of the next scope that will be pushed onto scope_stack_. The
-  // first is always the package scope.
-  ScopeIndex next_scope_index_ = ScopeIndex::Package;
-
   // The stack used for qualified declaration name construction.
   DeclNameStack decl_name_stack_;
 
   // The stack of declarations that could have modifiers.
   DeclStateStack decl_state_stack_;
 
-  // Tracks lexical lookup results.
-  LexicalLookup lexical_lookup_;
-
-  // Whether lexical_lookup_ has load errors, updated whenever scope_stack_ is
-  // pushed or popped.
-  bool lexical_lookup_has_load_error_ = false;
+  // The stack of scopes we are currently within.
+  ScopeStack scope_stack_;
 
   // Cache of reverse mapping from type constants to types.
   //

toolchain/check/decl_name_stack.cpp

@@ -10,8 +10,9 @@
 namespace Carbon::Check {
 
 auto DeclNameStack::MakeEmptyNameContext() -> NameContext {
-  return NameContext{.enclosing_scope = context_->current_scope_index(),
-                     .target_scope_id = context_->current_scope_id()};
+  return NameContext{
+      .enclosing_scope = context_->scope_stack().PeekIndex(),
+      .target_scope_id = context_->scope_stack().PeekNameScopeId()};
 }
 
 auto DeclNameStack::MakeUnqualifiedName(Parse::NodeId parse_node,
@@ -25,7 +26,7 @@ auto DeclNameStack::PushScopeAndStartName() -> void {
   decl_name_stack_.push_back(MakeEmptyNameContext());
 
   // Create a scope for any parameters introduced in this name.
-  context_->PushScope();
+  context_->scope_stack().Push();
 }
 
 auto DeclNameStack::FinishName() -> NameContext {
@@ -49,7 +50,7 @@ auto DeclNameStack::FinishName() -> NameContext {
 auto DeclNameStack::PopScope() -> void {
   CARBON_CHECK(decl_name_stack_.back().state == NameContext::State::Finished)
       << "Missing call to FinishName before PopScope";
-  context_->PopToScope(decl_name_stack_.back().enclosing_scope);
+  context_->scope_stack().PopTo(decl_name_stack_.back().enclosing_scope);
   decl_name_stack_.pop_back();
 }
 
@@ -154,12 +155,12 @@ static auto PushNameQualifierScope(Context& context,
                                    bool has_error = false) -> void {
   // If the qualifier has no parameters, we don't need to keep around a
   // parameter scope.
-  context.PopScopeIfEmpty();
+  context.scope_stack().PopIfEmpty();
 
-  context.PushScope(scope_inst_id, scope_id, has_error);
+  context.scope_stack().Push(scope_inst_id, scope_id, has_error);
 
   // Enter a parameter scope in case the qualified name itself has parameters.
-  context.PushScope();
+  context.scope_stack().Push();
 }
 
 auto DeclNameStack::UpdateScopeIfNeeded(NameContext& name_context) -> void {

toolchain/check/handle_binding_pattern.cpp

@@ -26,7 +26,8 @@ auto HandleAnyBindingPattern(Context& context, Parse::NodeId parse_node,
     // TODO: Eventually the name will need to support associations with other
     // scopes, but right now we don't support qualified names here.
     auto bind_name_id = context.bind_names().Add(
-        {.name_id = name_id, .enclosing_scope_id = context.current_scope_id()});
+        {.name_id = name_id,
+         .enclosing_scope_id = context.scope_stack().PeekNameScopeId()});
     if (is_generic) {
       // TODO: Create a `BindTemplateName` instead inside a `template` pattern.
       return {name_node,

toolchain/check/handle_class.cpp

@@ -144,7 +144,7 @@ auto HandleClassDefinitionStart(Context& context,
   }
 
   // Enter the class scope.
-  context.PushScope(class_decl_id, class_info.scope_id);
+  context.scope_stack().Push(class_decl_id, class_info.scope_id);
 
   // Introduce `Self`.
   context.AddNameToLookup(SemIR::NameId::SelfType,
@@ -329,7 +329,7 @@ auto HandleClassDefinition(Context& context,
   auto class_id =
       context.node_stack().Pop<Parse::NodeKind::ClassDefinitionStart>();
   context.inst_block_stack().Pop();
-  context.PopScope();
+  context.scope_stack().Pop();
   context.decl_name_stack().PopScope();
 
   // The class type is now fully defined.

toolchain/check/handle_codeblock.cpp

@@ -9,13 +9,13 @@ namespace Carbon::Check {
 auto HandleCodeBlockStart(Context& context, Parse::CodeBlockStartId parse_node)
     -> bool {
   context.node_stack().Push(parse_node);
-  context.PushScope();
+  context.scope_stack().Push();
   return true;
 }
 
 auto HandleCodeBlock(Context& context, Parse::CodeBlockId /*parse_node*/)
     -> bool {
-  context.PopScope();
+  context.scope_stack().Pop();
   context.node_stack()
       .PopAndDiscardSoloParseNode<Parse::NodeKind::CodeBlockStart>();
   return true;

toolchain/check/handle_function.cpp

@@ -220,7 +220,7 @@ auto HandleFunctionDefinitionStart(Context& context,
   // Create the function scope and the entry block.
   context.return_scope_stack().push_back({.decl_id = decl_id});
   context.inst_block_stack().Push();
-  context.PushScope(decl_id);
+  context.scope_stack().Push(decl_id);
   context.AddCurrentCodeBlockToFunction();
 
   // Bring the implicit and explicit parameters into scope.
@@ -270,7 +270,7 @@ auto HandleFunctionDefinition(Context& context,
     }
   }
 
-  context.PopScope();
+  context.scope_stack().Pop();
   context.inst_block_stack().Pop();
   context.return_scope_stack().pop_back();
   context.decl_name_stack().PopScope();

toolchain/check/handle_impl.cpp

@@ -20,7 +20,7 @@ auto HandleImplIntroducer(Context& context, Parse::ImplIntroducerId parse_node)
 
   // Create a scope for implicit parameters. We may not use it, but it's simpler
   // to create it unconditionally than to track whether it exists.
-  context.PushScope();
+  context.scope_stack().Push();
   return true;
 }
 
@@ -66,7 +66,7 @@ static auto BuildImplDecl(Context& context, Parse::AnyImplDeclId /*parse_node*/)
 
 auto HandleImplDecl(Context& context, Parse::ImplDeclId parse_node) -> bool {
   BuildImplDecl(context, parse_node);
-  context.PopScope();
+  context.scope_stack().Pop();
   return true;
 }
 
@@ -77,14 +77,14 @@ auto HandleImplDefinitionStart(Context& context,
   auto enclosing_scope_id = SemIR::NameScopeId::Invalid;
   auto scope_id = context.name_scopes().Add(
       impl_decl_id, SemIR::NameId::Invalid, enclosing_scope_id);
-  context.PushScope(impl_decl_id, scope_id);
+  context.scope_stack().Push(impl_decl_id, scope_id);
   return true;
 }
 
 auto HandleImplDefinition(Context& context,
                           Parse::ImplDefinitionId /*parse_node*/) -> bool {
-  context.PopScope();
-  context.PopScope();
+  context.scope_stack().Pop();
+  context.scope_stack().Pop();
   return true;
 }
 

toolchain/check/handle_interface.cpp

@@ -125,7 +125,7 @@ auto HandleInterfaceDefinitionStart(
   }
 
   // Enter the interface scope.
-  context.PushScope(interface_decl_id, interface_info.scope_id);
+  context.scope_stack().Push(interface_decl_id, interface_info.scope_id);
 
   // TODO: Introduce `Self`.
 
@@ -152,7 +152,7 @@ auto HandleInterfaceDefinition(Context& context,
   auto interface_id =
       context.node_stack().Pop<Parse::NodeKind::InterfaceDefinitionStart>();
   context.inst_block_stack().Pop();
-  context.PopScope();
+  context.scope_stack().Pop();
   context.decl_name_stack().PopScope();
 
   // The interface type is now fully defined.

toolchain/check/handle_let.cpp

@@ -34,9 +34,9 @@ auto HandleLetDecl(Context& context, Parse::LetDeclId parse_node) -> bool {
   // TODO: For a qualified `let` declaration, this should use the target scope
   // of the name introduced in the declaration. See #2590.
   CheckAccessModifiersOnDecl(context, Lex::TokenKind::Let,
-                             context.current_scope_id());
+                             context.scope_stack().PeekNameScopeId());
   RequireDefaultFinalOnlyInInterfaces(context, Lex::TokenKind::Let,
-                                      context.current_scope_id());
+                                      context.scope_stack().PeekNameScopeId());
   LimitModifiersOnDecl(
       context, KeywordModifierSet::Access | KeywordModifierSet::Interface,
       Lex::TokenKind::Let);

toolchain/check/handle_struct.cpp

@@ -10,7 +10,7 @@ namespace Carbon::Check {
 auto HandleStructLiteralOrStructTypeLiteralStart(
     Context& context, Parse::StructLiteralOrStructTypeLiteralStartId parse_node)
     -> bool {
-  context.PushScope();
+  context.scope_stack().Push();
   context.node_stack().Push(parse_node);
   // At this point we aren't sure whether this will be a value or type literal,
   // so we push onto args irrespective. It just won't be used for a type
@@ -94,7 +94,7 @@ auto HandleStructLiteral(Context& context, Parse::StructLiteralId parse_node)
   auto refs_id = context.ParamOrArgEnd(
       Parse::NodeKind::StructLiteralOrStructTypeLiteralStart);
 
-  context.PopScope();
+  context.scope_stack().Pop();
   context.node_stack()
       .PopAndDiscardSoloParseNode<
           Parse::NodeKind::StructLiteralOrStructTypeLiteralStart>();
@@ -117,7 +117,7 @@ auto HandleStructTypeLiteral(Context& context,
   auto refs_id = context.ParamOrArgEnd(
       Parse::NodeKind::StructLiteralOrStructTypeLiteralStart);
 
-  context.PopScope();
+  context.scope_stack().Pop();
   context.node_stack()
       .PopAndDiscardSoloParseNode<
           Parse::NodeKind::StructLiteralOrStructTypeLiteralStart>();

toolchain/check/handle_variable.cpp

@@ -85,7 +85,7 @@ auto HandleVariableDecl(Context& context, Parse::VariableDeclId parse_node)
   // TODO: For a qualified `var` declaration, this should use the target scope
   // of the name introduced in the declaration. See #2590.
   CheckAccessModifiersOnDecl(context, Lex::TokenKind::Var,
-                             context.current_scope_id());
+                             context.scope_stack().PeekNameScopeId());
   LimitModifiersOnDecl(context, KeywordModifierSet::Access,
                        Lex::TokenKind::Var);
   auto modifiers = context.decl_state_stack().innermost().modifier_set;

toolchain/check/return.cpp

@@ -91,7 +91,7 @@ auto CheckReturnedVar(Context& context, Parse::NodeId returned_node,
 }
 
 auto RegisterReturnedVar(Context& context, SemIR::InstId bind_id) -> void {
-  auto existing_id = context.SetReturnedVarOrGetExisting(bind_id);
+  auto existing_id = context.scope_stack().SetReturnedVarOrGetExisting(bind_id);
   if (existing_id.is_valid()) {
     CARBON_DIAGNOSTIC(ReturnedVarShadowed, Error,
                       "Cannot declare a `returned var` in the scope of "

toolchain/check/scope_stack.cpp

@@ -0,0 +1,145 @@
+// 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 "toolchain/check/scope_stack.h"
+
+#include "common/check.h"
+#include "toolchain/sem_ir/ids.h"
+
+namespace Carbon::Check {
+
+auto ScopeStack::VerifyOnFinish() -> void {
+  CARBON_CHECK(scope_stack_.empty()) << scope_stack_.size();
+}
+
+auto ScopeStack::Push(SemIR::InstId scope_inst_id, SemIR::NameScopeId scope_id,
+                      bool lexical_lookup_has_load_error) -> void {
+  scope_stack_.push_back(
+      {.index = next_scope_index_,
+       .scope_inst_id = scope_inst_id,
+       .scope_id = scope_id,
+       .prev_lexical_lookup_has_load_error = lexical_lookup_has_load_error_});
+  if (scope_id.is_valid()) {
+    non_lexical_scope_stack_.push_back({next_scope_index_, scope_id});
+  }
+
+  lexical_lookup_has_load_error_ |= lexical_lookup_has_load_error;
+
+  // TODO: Handle this case more gracefully.
+  CARBON_CHECK(next_scope_index_.index != std::numeric_limits<int32_t>::max())
+      << "Ran out of scopes";
+  ++next_scope_index_.index;
+}
+
+auto ScopeStack::Pop() -> void {
+  auto scope = scope_stack_.pop_back_val();
+
+  lexical_lookup_has_load_error_ = scope.prev_lexical_lookup_has_load_error;
+
+  for (const auto& str_id : scope.names) {
+    auto& lexical_results = lexical_lookup_.Get(str_id);
+    CARBON_CHECK(lexical_results.back().scope_index == scope.index)
+        << "Inconsistent scope index for name " << str_id;
+    lexical_results.pop_back();
+  }
+
+  if (scope.scope_id.is_valid()) {
+    CARBON_CHECK(non_lexical_scope_stack_.back().scope_index == scope.index);
+    non_lexical_scope_stack_.pop_back();
+  }
+
+  if (scope.has_returned_var) {
+    CARBON_CHECK(!return_scope_stack_.empty());
+    CARBON_CHECK(return_scope_stack_.back().returned_var.is_valid());
+    return_scope_stack_.back().returned_var = SemIR::InstId::Invalid;
+  }
+}
+
+auto ScopeStack::PopTo(ScopeIndex index) -> void {
+  while (PeekIndex() > index) {
+    Pop();
+  }
+  CARBON_CHECK(PeekIndex() == index)
+      << "Scope index " << index << " does not enclose the current scope "
+      << PeekIndex();
+}
+
+auto ScopeStack::LookupInCurrentScope(SemIR::NameId name_id) -> SemIR::InstId {
+  auto& lexical_results = lexical_lookup_.Get(name_id);
+  if (lexical_results.empty()) {
+    return SemIR::InstId::Invalid;
+  }
+
+  auto result = lexical_results.back();
+  if (result.scope_index != PeekIndex()) {
+    return SemIR::InstId::Invalid;
+  }
+
+  return result.inst_id;
+}
+
+auto ScopeStack::LookupInEnclosingScopes(SemIR::NameId name_id)
+    -> std::pair<SemIR::InstId, llvm::ArrayRef<NonLexicalScope>> {
+  // Find the results from enclosing lexical scopes. These will be combined with
+  // results from non-lexical scopes such as namespaces and classes.
+  llvm::ArrayRef<LexicalLookup::Result> lexical_results =
+      lexical_lookup_.Get(name_id);
+
+  // If we have no lexical results, check all non-lexical scopes.
+  if (lexical_results.empty()) {
+    return {lexical_lookup_has_load_error_ ? SemIR::InstId::BuiltinError
+                                           : SemIR::InstId::Invalid,
+            non_lexical_scope_stack_};
+  }
+
+  // Find the first non-lexical scope that is within the scope of the lexical
+  // lookup result.
+  auto* first_non_lexical_scope = std::lower_bound(
+      non_lexical_scope_stack_.begin(), non_lexical_scope_stack_.end(),
+      lexical_results.back().scope_index,
+      [](const NonLexicalScope& scope, ScopeIndex index) {
+        return scope.scope_index < index;
+      });
+  return {
+      lexical_results.back().inst_id,
+      llvm::ArrayRef(first_non_lexical_scope, non_lexical_scope_stack_.end())};
+}
+
+auto ScopeStack::LookupOrAddName(SemIR::NameId name_id, SemIR::InstId target_id)
+    -> SemIR::InstId {
+  if (!scope_stack_.back().names.insert(name_id).second) {
+    auto existing = lexical_lookup_.Get(name_id).back().inst_id;
+    CARBON_CHECK(existing.is_valid())
+        << "Name in scope but not in lexical lookups";
+    return existing;
+  }
+
+  // TODO: Reject if we previously performed a failed lookup for this name
+  // in this scope or a scope nested within it.
+  auto& lexical_results = lexical_lookup_.Get(name_id);
+  CARBON_CHECK(lexical_results.empty() ||
+               lexical_results.back().scope_index < PeekIndex())
+      << "Failed to clean up after scope nested within the current scope";
+  lexical_results.push_back({.inst_id = target_id, .scope_index = PeekIndex()});
+  return SemIR::InstId::Invalid;
+}
+
+auto ScopeStack::SetReturnedVarOrGetExisting(SemIR::InstId inst_id)
+    -> SemIR::InstId {
+  CARBON_CHECK(!return_scope_stack_.empty()) << "`returned var` in no function";
+  auto& returned_var = return_scope_stack_.back().returned_var;
+  if (returned_var.is_valid()) {
+    return returned_var;
+  }
+
+  returned_var = inst_id;
+  CARBON_CHECK(!scope_stack_.back().has_returned_var)
+      << "Scope has returned var but none is set";
+  if (inst_id.is_valid()) {
+    scope_stack_.back().has_returned_var = true;
+  }
+  return SemIR::InstId::Invalid;
+}
+
+}  // namespace Carbon::Check

toolchain/check/scope_stack.h

@@ -0,0 +1,186 @@
+// 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_SCOPE_STACK_H_
+#define CARBON_TOOLCHAIN_CHECK_SCOPE_STACK_H_
+
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "toolchain/check/lexical_lookup.h"
+#include "toolchain/check/scope_index.h"
+#include "toolchain/sem_ir/file.h"
+#include "toolchain/sem_ir/ids.h"
+
+namespace Carbon::Check {
+
+// A stack of lexical and semantic scopes that we are currently performing
+// checking within.
+class ScopeStack {
+ public:
+  explicit ScopeStack(const StringStoreWrapper<IdentifierId>& identifiers)
+      : lexical_lookup_(identifiers) {}
+
+  // A scope in which `break` and `continue` can be used.
+  struct BreakContinueScope {
+    SemIR::InstBlockId break_target;
+    SemIR::InstBlockId continue_target;
+  };
+
+  // A scope in which `return` can be used.
+  struct ReturnScope {
+    // The declaration from which we can return. Inside a function, this will
+    // be a `FunctionDecl`.
+    SemIR::InstId decl_id;
+
+    // The value corresponding to the current `returned var`, if any. Will be
+    // set and unset as `returned var`s are declared and go out of scope.
+    SemIR::InstId returned_var = SemIR::InstId::Invalid;
+  };
+
+  // A non-lexical scope in which unqualified lookup may be required.
+  struct NonLexicalScope {
+    // The index of the scope in the scope stack.
+    ScopeIndex scope_index;
+
+    // The corresponding name scope.
+    SemIR::NameScopeId name_scope_id;
+  };
+
+  // Pushes a scope onto scope_stack_. NameScopeId::Invalid is used for new
+  // scopes. lexical_lookup_has_load_error is used to limit diagnostics when a
+  // given namespace may contain a mix of both successful and failed name
+  // imports.
+  auto Push(SemIR::InstId scope_inst_id = SemIR::InstId::Invalid,
+            SemIR::NameScopeId scope_id = SemIR::NameScopeId::Invalid,
+            bool lexical_lookup_has_load_error = false) -> void;
+
+  // Pops the top scope from scope_stack_, cleaning up names from
+  // lexical_lookup_.
+  auto Pop() -> void;
+
+  // Pops the top scope from scope_stack_ if it contains no names.
+  auto PopIfEmpty() -> void {
+    if (scope_stack_.back().names.empty()) {
+      Pop();
+    }
+  }
+
+  // Pops scopes until we return to the specified scope index.
+  auto PopTo(ScopeIndex index) -> void;
+
+  // Returns the scope index associated with the current scope.
+  auto PeekIndex() const -> ScopeIndex { return Peek().index; }
+
+  // Returns the name scope associated with the current lexical scope, if any.
+  auto PeekNameScopeId() const -> SemIR::NameScopeId { return Peek().scope_id; }
+
+  // Returns the instruction associated with the current scope, or Invalid if
+  // there is no such instruction, such as for a block scope.
+  auto PeekInstId() const -> SemIR::InstId { return Peek().scope_inst_id; }
+
+  // Returns the current scope, if it is of the specified kind. Otherwise,
+  // returns nullopt.
+  template <typename InstT>
+  auto GetCurrentScopeAs(const SemIR::File& sem_ir) -> std::optional<InstT> {
+    auto inst_id = PeekInstId();
+    if (!inst_id.is_valid()) {
+      return std::nullopt;
+    }
+    return sem_ir.insts().TryGetAs<InstT>(inst_id);
+  }
+
+  // If there is no `returned var` in scope, sets the given instruction to be
+  // the current `returned var` and returns an invalid instruction ID. If there
+  // is already a `returned var`, returns it instead.
+  auto SetReturnedVarOrGetExisting(SemIR::InstId inst_id) -> SemIR::InstId;
+
+  // Looks up the name `name_id` in the current scope. Returns the existing
+  // lookup result, if any.
+  auto LookupInCurrentScope(SemIR::NameId name_id) -> SemIR::InstId;
+
+  // Looks up the name `name_id` in the current scope and its enclosing scopes.
+  // Returns the innermost lexical lookup result, if any, along with a list of
+  // non-lexical scopes in which lookup should also be performed, ordered from
+  // outermost to innermost.
+  auto LookupInEnclosingScopes(SemIR::NameId name_id)
+      -> std::pair<SemIR::InstId, llvm::ArrayRef<NonLexicalScope>>;
+
+  // Looks up the name `name_id` in the current scope. Returns the existing
+  // instruction if any, and otherwise adds the name with the value `target_id`
+  // and returns Invalid.
+  auto LookupOrAddName(SemIR::NameId name_id, SemIR::InstId target_id)
+      -> SemIR::InstId;
+
+  // Runs verification that the processing cleanly finished.
+  auto VerifyOnFinish() -> void;
+
+  auto return_scope_stack() -> llvm::SmallVector<ReturnScope>& {
+    return return_scope_stack_;
+  }
+
+  auto break_continue_stack() -> llvm::SmallVector<BreakContinueScope>& {
+    return break_continue_stack_;
+  }
+
+ private:
+  // An entry in scope_stack_.
+  struct ScopeStackEntry {
+    // The sequential index of this scope entry within the file.
+    ScopeIndex index;
+
+    // The instruction associated with this entry, if any. This can be one of:
+    //
+    // - A `ClassDecl`, for a class definition scope.
+    // - A `FunctionDecl`, for the outermost scope in a function
+    //   definition.
+    // - Invalid, for any other scope.
+    SemIR::InstId scope_inst_id;
+
+    // The name scope associated with this entry, if any.
+    SemIR::NameScopeId scope_id;
+
+    // The previous state of lexical_lookup_has_load_error_, restored on pop.
+    bool prev_lexical_lookup_has_load_error;
+
+    // Names which are registered with lexical_lookup_, and will need to be
+    // unregistered when the scope ends.
+    llvm::DenseSet<SemIR::NameId> names;
+
+    // Whether a `returned var` was introduced in this scope, and needs to be
+    // unregistered when the scope ends.
+    bool has_returned_var = false;
+
+    // TODO: This likely needs to track things which need to be destructed.
+  };
+
+  auto Peek() const -> const ScopeStackEntry& { return scope_stack_.back(); }
+
+  // A stack of scopes from which we can `return`.
+  llvm::SmallVector<ReturnScope> return_scope_stack_;
+
+  // A stack of `break` and `continue` targets.
+  llvm::SmallVector<BreakContinueScope> break_continue_stack_;
+
+  // A stack for scope context.
+  llvm::SmallVector<ScopeStackEntry> scope_stack_;
+
+  // Information about non-lexical scopes. This is a subset of the entries and
+  // the information in scope_stack_.
+  llvm::SmallVector<NonLexicalScope> non_lexical_scope_stack_;
+
+  // The index of the next scope that will be pushed onto scope_stack_. The
+  // first is always the package scope.
+  ScopeIndex next_scope_index_ = ScopeIndex::Package;
+
+  // Tracks lexical lookup results.
+  LexicalLookup lexical_lookup_;
+
+  // Whether lexical_lookup_ has load errors, updated whenever scope_stack_ is
+  // pushed or popped.
+  bool lexical_lookup_has_load_error_ = false;
+};
+
+}  // namespace Carbon::Check
+
+#endif  // CARBON_TOOLCHAIN_CHECK_SCOPE_STACK_H_