// 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 "common/array_stack.h"
#include "common/move_only.h"
#include "common/set.h"
#include "llvm/ADT/SmallVector.h"
#include "toolchain/check/full_pattern_stack.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 {

class Context;

// A stack of lexical and semantic scopes that we are currently performing
// checking within.
class ScopeStack {
 public:
  explicit ScopeStack(Context& context);

  // A scope in which `break` and `continue` can be used.
  struct BreakContinueScope {
    SemIR::InstBlockId break_target;
    SemIR::InstBlockId continue_target;
  };

  // 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;

    // The corresponding specific.
    SemIR::SpecificId specific_id;
  };

  // Information about a scope that has been temporarily removed from the stack.
  // This type is large, so moves of this type should be avoided.
  struct SuspendedScope;

  // Pushes a scope for a declaration name's parameters.
  auto PushForDeclName() -> void;

  // 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_.
  // If `check_unused` is set, checks and emits diagnostics for unused names.
  auto Pop(bool check_unused = false) -> void;

  // Pops the top scope from scope_stack_ if it contains no names.
  auto PopIfEmpty(bool check_unused = false) -> void {
    if (scope_stack_.back().num_names == 0) {
      Pop(check_unused);
    }
  }

  // Pops scopes until we return to the specified scope index.
  auto PopTo(ScopeIndex index, bool check_unused = false) -> 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 `None` if
  // there is no such instruction, such as for a block scope.
  auto PeekInstId() const -> SemIR::InstId { return Peek().scope_inst_id; }

  // Returns the specific associated with the innermost enclosing scope that is
  // associated with a specific. This will generally be the self specific of the
  // innermost enclosing generic, as there is no way to enter any other specific
  // scope.
  auto PeekSpecificId() const -> SemIR::SpecificId {
    return Peek().specific_id;
  }

  // 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.
  template <typename InstT>
  auto TryGetCurrentScopeAs() -> std::optional<InstT> {
    auto inst_id = PeekInstId();
    if (!inst_id.has_value()) {
      return std::nullopt;
    }
    return sem_ir().insts().TryGetAs<InstT>(inst_id);
  }

  // Returns the current scope, assuming it is of the specified kind.
  // Check-fails if there is no instruction for a current scope, or the scope is
  // of a different kind.
  template <typename InstT>
  auto GetCurrentScopeAs() -> InstT {
    auto inst_id = PeekInstId();
    CARBON_CHECK(inst_id.has_value());
    return sem_ir().insts().GetAs<InstT>(inst_id);
  }

  // If there is no `returned var` in scope, sets the given instruction to be
  // the current `returned var` and returns an `None`. If there
  // is already a `returned var`, returns it instead.
  auto SetReturnedVarOrGetExisting(SemIR::InstId inst_id, SemIR::NameId name_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.
  // If `use_loc_id` is specified, the name is marked as used at that location.
  auto LookupInLexicalScopesWithin(SemIR::NameId name_id,
                                   ScopeIndex scope_index,
                                   SemIR::LocId use_loc_id, bool is_reachable)
      -> SemIR::InstId;

  // Looks up the name `name_id` in the current scope and related lexical
  // 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. If `use_loc_id` is specified, the
  // name is marked as used at that location.
  auto LookupInLexicalScopes(SemIR::NameId name_id, SemIR::LocId use_loc_id,
                             bool is_reachable)
      -> std::pair<SemIR::InstId, llvm::ArrayRef<NonLexicalScope>>;

  // Looks up the name `name_id` in the current scope, or in `scope_index` if
  // specified. Returns the existing instruction if the name is already declared
  // in that scope or any unfinished scope within it, and otherwise adds the
  // name with the value `target_id` and returns `None`. `is_decl_reachable`
  // indicates whether the name was declared in a reachable position.
  auto LookupOrAddName(SemIR::NameId name_id, SemIR::InstId target_id,
                       ScopeIndex scope_index = ScopeIndex::None,
                       bool is_decl_reachable = true) -> SemIR::InstId;

  // Prepares to add a compile-time binding in the current scope, and returns
  // its index. The added binding must then be pushed using
  // `PushCompileTimeBinding`.
  auto AddCompileTimeBinding() -> SemIR::CompileTimeBindIndex {
    auto index = scope_stack_.back().next_compile_time_bind_index;
    ++scope_stack_.back().next_compile_time_bind_index.index;
    return index;
  }

  // Pushes a compile-time binding into the current scope.
  auto PushCompileTimeBinding(SemIR::InstId bind_id) -> void {
    compile_time_binding_stack_.AppendToTop(bind_id);
  }

  // Temporarily removes the top of the stack and its lexical lookup results.
  auto Suspend() -> SuspendedScope;

  // Restores a suspended scope stack entry.
  auto Restore(SuspendedScope&& scope) -> void;

  // Runs verification that the processing cleanly finished.
  auto VerifyOnFinish() const -> void;

  auto break_continue_stack() -> llvm::SmallVector<BreakContinueScope>& {
    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_;
  }

  auto full_pattern_stack() -> FullPatternStack& { return full_pattern_stack_; }

 private:
  auto sem_ir() const -> const SemIR::File&;

  auto lexical_lookup() -> LexicalLookup& { return lexical_lookup_; }

  // An entry in scope_stack_.
  struct ScopeStackEntry : public MoveOnly<ScopeStackEntry> {
    auto is_lexical_scope() const -> bool { return !scope_id.has_value(); }

    // 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 specific associated with this entry, if any.
    SemIR::SpecificId specific_id;

    // The next compile-time binding index to allocate in this scope.
    SemIR::CompileTimeBindIndex next_compile_time_bind_index;

    // Whether lexical_lookup_ has load errors from this scope or an ancestor
    // scope.
    bool lexical_lookup_has_load_error;

    // Whether a `returned var` was introduced in this scope, and needs to be
    // unregistered when the scope ends.
    bool has_returned_var = false;

    // Whether there are any ids in the `names` set.
    int num_names = 0;

    // Names which are registered with lexical_lookup_, and will need to be
    // unregistered when the scope ends.
    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.
  auto LexicalLookupHasLoadError() const -> bool {
    return !scope_stack_.empty() &&
           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;
    }
  }

  // Marks the name `name_id` as used at the given location.
  auto MarkUsed(SemIR::NameId name_id, SemIR::LocId loc_id, bool is_reachable)
      -> void;

  // 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
  // significant changes.
  auto VerifyNextCompileTimeBindIndex(llvm::StringLiteral label,
                                      const ScopeStackEntry& scope) -> void;

  // Context, used only for checks and emitting diagnostics.
  Context* context_;

  // 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_;

  // 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_;

  // 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_;

  // A stack of the current compile time bindings.
  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.
  ScopeIndex next_scope_index_ = ScopeIndex::Package;

  // Tracks lexical lookup results.
  LexicalLookup lexical_lookup_;

  // Stack of full-patterns currently being checked.
  FullPatternStack full_pattern_stack_;
};

struct ScopeStack::SuspendedScope : public MoveOnly<SuspendedScope> {
  // An item that was suspended within this scope. This represents either a
  // lexical lookup entry in this scope, or a compile time binding entry in this
  // scope.
  //
  // TODO: For compile-time bindings, the common case is that they will both
  // have a suspended lexical lookup entry and a suspended compile time binding
  // entry. We should be able to store that as a single ScopeItem rather than
  // two.
  struct ScopeItem {
    static constexpr uint32_t IndexForCompileTimeBinding = -1;

    // The scope index for a LexicalLookup::SuspendedResult, or
    // CompileTimeBindingIndex for a suspended compile time binding.
    uint32_t index;
    // The instruction within the scope.
    SemIR::InstId inst_id;
    // Whether the name was declared in a reachable position.
    bool is_decl_reachable;
    // The location of the first use of the name, if any.
    SemIR::LocId use_loc_id;
  };

  // The suspended scope stack entry.
  ScopeStackEntry entry;
  // The list of items that were within this scope when it was suspended. The
  // inline size is an attempt to keep the size of a `SuspendedFunction`
  // reasonable while avoiding heap allocations most of the time.
  llvm::SmallVector<ScopeItem, 8> suspended_items;
};

}  // namespace Carbon::Check

#endif  // CARBON_TOOLCHAIN_CHECK_SCOPE_STACK_H_