Detect control flow in entities nested inside functions (#5336)

Right now, return_scope_stack is being used to determine whether logic
is in a function scope. However, we need to handle nested entities
inside function scopes. For example where this crashes right now:

```
base class C(B:! bool) {}

fn F() {
  class B {
    extend base: C(true or false);
  }
}
```

This is doing a few things to make this kind of code not crash:

- Split `scope_stack().Push` into `PushForDeclName`, `PushForEntity`,
`PushForExpr`, and `PushForFunction` so that better decisions can be
made about behaviors.
- Hide `return_scope_stack` in the API, instead using interfaces to get
at the underlying data.
- Also using `PushForFunction` to update it similar to the other stacks
that `ScopeStack` manages.
- Add `IsInFunctionScope` as the best way to determine presence in
function scope.
- Remove `PeekIsLexicalScope` since destruction really wants function
scope information anyways.
- Clean up `destroy_id_stack` handling to be for function scopes rather
than lexical scopes.
- Return after related `context.TODO`s in a couple more spots, so that
code doesn't proceed to add control flow in spite of the lack of
support.

---------

Co-authored-by: Geoff Romer <gromer@google.com>

toolchain/check/scope_stack.h

@@ -31,17 +31,6 @@ class ScopeStack {
     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::None;
-  };
-
   // A non-lexical scope in which unqualified lookup may be required.
   struct NonLexicalScope {
     // The index of the scope in the scope stack.
@@ -57,17 +46,24 @@ class ScopeStack {
   // Information about a scope that has been temporarily removed from the stack.
   struct SuspendedScope;
 
-  // Pushes a scope onto scope_stack_. NameScopeId::None 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::None,
-            SemIR::NameScopeId scope_id = SemIR::NameScopeId::None,
-            SemIR::SpecificId specific_id = SemIR::SpecificId::None,
-            bool lexical_lookup_has_load_error = false) -> void;
+  // Pushes a scope for a declaration name's parameters.
+  auto PushForDeclName() -> void;
 
-  // Pops the top scope from scope_stack_, cleaning up names from
-  // lexical_lookup_.
+  // Pushes a non-function entity scope. Functions must use
+  // `PushForFunctionBody` instead.
+  auto PushForEntity(SemIR::InstId scope_inst_id, SemIR::NameScopeId scope_id,
+                     SemIR::SpecificId specific_id,
+                     bool lexical_lookup_has_load_error = false) -> void;
+
+  // Pushes a scope which should be in the same region as the current scope.
+  // These can be in a function without breaking `return` scoping. For example,
+  // this is used by struct literals and code blocks.
+  auto PushForSameRegion() -> void;
+
+  // Pushes a function scope.
+  auto PushForFunctionBody(SemIR::InstId scope_inst_id) -> void;
+
+  // Pops the top scope from scope_stack_. Removes names from lexical_lookup_.
   auto Pop() -> void;
 
   // Pops the top scope from scope_stack_ if it contains no names.
@@ -98,8 +94,12 @@ class ScopeStack {
     return Peek().specific_id;
   }
 
-  // Returns true if current scope is lexical.
-  auto PeekIsLexicalScope() const -> bool { return Peek().is_lexical_scope(); }
+  // Returns true if the current scope is inside a function scope (either the
+  // scope itself, or a lexical scope), without an intervening entity scope.
+  auto IsInFunctionScope() const -> bool {
+    return !return_scope_stack_.empty() &&
+           !return_scope_stack_.back().nested_scope_index.has_value();
+  }
 
   // Returns the current scope, if it is of the specified kind. Otherwise,
   // returns nullopt.
@@ -117,6 +117,19 @@ class ScopeStack {
   // is already a `returned var`, returns it instead.
   auto SetReturnedVarOrGetExisting(SemIR::InstId inst_id) -> SemIR::InstId;
 
+  // Returns the `returned var` instruction that's currently in scope, or `None`
+  // if there isn't one.
+  auto GetReturnedVar() -> SemIR::InstId {
+    CARBON_CHECK(IsInFunctionScope(), "Handling return but not in a function");
+    return return_scope_stack_.back().returned_var;
+  }
+
+  // Returns the decl ID for the current return scope.
+  auto GetReturnScopeDeclId() -> SemIR::InstId {
+    CARBON_CHECK(IsInFunctionScope(), "Handling return but not in a function");
+    return return_scope_stack_.back().decl_id;
+  }
+
   // Looks up the name `name_id` in the current scope and enclosing scopes, but
   // do not look past `scope_index`. Returns the existing lookup result, if any.
   auto LookupInLexicalScopesWithin(SemIR::NameId name_id,
@@ -160,10 +173,6 @@ class ScopeStack {
   // Runs verification that the processing cleanly finished.
   auto VerifyOnFinish() const -> void;
 
-  auto return_scope_stack() -> llvm::SmallVector<ReturnScope>& {
-    return return_scope_stack_;
-  }
-
   auto break_continue_stack() -> llvm::SmallVector<BreakContinueScope>& {
     return break_continue_stack_;
   }
@@ -219,6 +228,35 @@ class ScopeStack {
     Set<SemIR::NameId> names = {};
   };
 
+  // A scope in which `return` can be used.
+  struct ReturnScope {
+    // The `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::None;
+
+    // When a nested scope interrupts a return scope, this is the index of the
+    // outermost interrupting scope (the one closest to the function scope).
+    // This can then be used to determine whether we're actually inside the most
+    // recent `ReturnScope`, or inside a different entity scope.
+    //
+    // This won't be set for functions directly inside functions, because they
+    // will have their own `ReturnScope`.
+
+    // For example, when a `class` is inside a `fn`, it interrupts the function
+    // body by setting this on `PushEntity`; `Pop` will set it back to `None`.
+    ScopeIndex nested_scope_index = ScopeIndex::None;
+  };
+
+  // Pushes a scope onto scope_stack_. NameScopeId::None 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::NameScopeId scope_id,
+            SemIR::SpecificId specific_id, bool lexical_lookup_has_load_error)
+      -> void;
+
   auto Peek() const -> const ScopeStackEntry& { return scope_stack_.back(); }
 
   // Returns whether lexical lookup currently has any load errors.
@@ -227,6 +265,15 @@ class ScopeStack {
            scope_stack_.back().lexical_lookup_has_load_error;
   }
 
+  // If inside a return scope, marks a nested scope (see `nested_scope_index`).
+  // Called after pushing the new scope.
+  auto MarkNestingIfInReturnScope() -> void {
+    if (!return_scope_stack_.empty() &&
+        !return_scope_stack_.back().nested_scope_index.has_value()) {
+      return_scope_stack_.back().nested_scope_index = scope_stack_.back().index;
+    }
+  }
+
   // Checks that the provided scope's `next_compile_time_bind_index` matches the
   // full size of the current `compile_time_binding_stack_`. The values should
   // always match, and this is used to validate the correspondence during
@@ -246,8 +293,8 @@ class ScopeStack {
   // A stack for scope context.
   llvm::SmallVector<ScopeStackEntry> scope_stack_;
 
-  // A stack of `destroy` functions to call. This only has entries for lexical
-  // scopes, because non-lexical scopes don't have destruction on scope exit.
+  // A stack of `destroy` functions to call. This only has entries inside of
+  // function bodies, where destruction on scope exit is required.
   ArrayStack<SemIR::InstId> destroy_id_stack_;
 
   // Information about non-lexical scopes. This is a subset of the entries and