// 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_SEMANTICS_SEMANTICS_NODE_STACK_H_
#define CARBON_TOOLCHAIN_SEMANTICS_SEMANTICS_NODE_STACK_H_

#include <type_traits>

#include "common/vlog.h"
#include "llvm/ADT/SmallVector.h"
#include "toolchain/parser/parse_node_kind.h"
#include "toolchain/parser/parse_tree.h"
#include "toolchain/semantics/semantics_node.h"

namespace Carbon {

// Wraps the stack of nodes for SemanticsParseTreeHandler.
//
// All pushes and pops will be vlogged.
//
// Pop APIs will run basic verification:
//
// - If receiving a pop_parse_kind, verify that the parse_node being popped is
//   of pop_parse_kind.
// - Validates presence of node_id based on whether it's a solo
//   parse_node.
//
// These should be assumed API constraints unless otherwise mentioned on a
// method. The main exception is PopAndIgnore, which doesn't do verification.
class SemanticsNodeStack {
 public:
  explicit SemanticsNodeStack(const ParseTree& parse_tree,
                              llvm::raw_ostream* vlog_stream)
      : parse_tree_(&parse_tree), vlog_stream_(vlog_stream) {}

  // Pushes a solo parse tree node onto the stack. Used when there is no
  // IR generated by the node.
  auto Push(ParseTree::Node parse_node) -> void {
    CARBON_VLOG() << "Node Push " << stack_.size() << ": "
                  << parse_tree_->node_kind(parse_node) << " -> <none>\n";
    CARBON_CHECK(stack_.size() < (1 << 20))
        << "Excessive stack size: likely infinite loop";
    stack_.push_back(Entry(parse_node, SemanticsNodeId::Invalid));
  }

  // Pushes a parse tree node onto the stack with an ID.
  template <typename IdT>
  auto Push(ParseTree::Node parse_node, IdT id) -> void {
    CARBON_VLOG() << "Node Push " << stack_.size() << ": "
                  << parse_tree_->node_kind(parse_node) << " -> " << id << "\n";
    CARBON_CHECK(stack_.size() < (1 << 20))
        << "Excessive stack size: likely infinite loop";
    stack_.push_back(Entry(parse_node, id));
  }

  // Pops the top of the stack without any verification.
  auto PopAndIgnore() -> void { PopEntry<SemanticsNodeId>(); }

  // Pops the top of the stack and returns the parse_node.
  auto PopForSoloParseNode() -> ParseTree::Node {
    Entry back = PopEntry<SemanticsNodeId>();
    RequireSoloParseNode(back);
    return back.parse_node;
  }

  // Pops the top of the stack and returns the parse_node.
  auto PopForSoloParseNode(ParseNodeKind pop_parse_kind) -> ParseTree::Node {
    auto parse_node = PopForSoloParseNode();
    RequireParseKind(parse_node, pop_parse_kind);
    return parse_node;
  }

  // Pops the top of the stack.
  auto PopAndDiscardSoloParseNode(ParseNodeKind pop_parse_kind) -> void {
    PopForSoloParseNode(pop_parse_kind);
  }

  // Pops the top of the stack and returns the parse_node and the ID.
  template <typename IdT>
  auto PopWithParseNode() -> std::pair<ParseTree::Node, IdT> {
    Entry back = PopEntry<IdT>();
    RequireValidId(back);
    return {back.parse_node, back.id<IdT>()};
  }

  // Pops the top of the stack and returns the parse_node and the ID.
  template <typename IdT>
  auto PopWithParseNode(ParseNodeKind pop_parse_kind)
      -> std::pair<ParseTree::Node, IdT> {
    auto back = PopWithParseNode<IdT>();
    RequireParseKind(back.first, pop_parse_kind);
    return back;
  }

  // Pops the top of the stack and returns the ID.
  template <typename IdT>
  auto Pop() -> IdT {
    return PopWithParseNode<IdT>().second;
  }

  // Pops the top of the stack and returns the ID.
  template <typename IdT>
  auto Pop(ParseNodeKind pop_parse_kind) -> IdT {
    return PopWithParseNode<IdT>(pop_parse_kind).second;
  }

  // Pops the top of the stack, and discards the ID.
  auto PopAndDiscardId() -> void { PopWithParseNode<SemanticsNodeId>(); }

  // Pops the top of the stack, and discards the ID.
  auto PopAndDiscardId(ParseNodeKind pop_parse_kind) -> void {
    PopWithParseNode<SemanticsNodeId>(pop_parse_kind);
  }

  // Peeks at the parse_node of the top of the stack.
  auto PeekParseNode() -> ParseTree::Node { return stack_.back().parse_node; }

  // Peeks at the ID of the top of the stack.
  template <typename IdT>
  auto Peek(ParseNodeKind parse_kind) -> IdT {
    Entry back = stack_.back();
    RequireParseKind(back.parse_node, parse_kind);
    RequireValidId(back);
    return back.id<IdT>();
  }

  // Prints the stack for a stack dump.
  auto PrintForStackDump(llvm::raw_ostream& output) const -> void;

  auto empty() const -> bool { return stack_.empty(); }
  auto size() const -> size_t { return stack_.size(); }

 private:
  // An entry in stack_.
  struct Entry {
    explicit Entry(ParseTree::Node parse_node, SemanticsNodeId node_id)
        : parse_node(parse_node), node_id(node_id) {}
    explicit Entry(ParseTree::Node parse_node,
                   SemanticsNodeBlockId node_block_id)
        : parse_node(parse_node), node_block_id(node_block_id) {}
    explicit Entry(ParseTree::Node parse_node, SemanticsFunctionId function_id)
        : parse_node(parse_node), function_id(function_id) {}
    explicit Entry(ParseTree::Node parse_node, SemanticsStringId name_id)
        : parse_node(parse_node), name_id(name_id) {}
    explicit Entry(ParseTree::Node parse_node, SemanticsTypeId type_id)
        : parse_node(parse_node), type_id(type_id) {}

    // Returns the appropriate ID basaed on type.
    template <typename T>
    auto id() -> T& {
      if constexpr (std::is_same<T, SemanticsNodeId>()) {
        return node_id;
      }
      if constexpr (std::is_same<T, SemanticsNodeBlockId>()) {
        return node_block_id;
      }
      if constexpr (std::is_same<T, SemanticsFunctionId>()) {
        return function_id;
      }
      if constexpr (std::is_same<T, SemanticsStringId>()) {
        return name_id;
      }
      if constexpr (std::is_same<T, SemanticsTypeId>()) {
        return type_id;
      }
    }

    // The node associated with the stack entry.
    ParseTree::Node parse_node;

    // The entries will evaluate as invalid if and only if they're a solo
    // parse_node. Invalid is used instead of optional to save space.
    //
    // A discriminator isn't needed because the caller can determine which field
    // is used based on the ParseNodeKind.
    union {
      SemanticsNodeId node_id;
      SemanticsNodeBlockId node_block_id;
      SemanticsFunctionId function_id;
      SemanticsStringId name_id;
      SemanticsTypeId type_id;
    };

    // APIs rely on type punning. They read node_id.is_valid, even though that
    // may not be the active union member. These asserts enforce standard layout
    // in order to help ensure that works.
    // TODO: Use is_layout_compatible in C++20.
    static_assert(std::is_standard_layout_v<SemanticsNodeId>,
                  "Need standard layout for type punning");
    static_assert(std::is_standard_layout_v<SemanticsNodeBlockId>,
                  "Need standard layout for type punning");
    static_assert(std::is_standard_layout_v<SemanticsFunctionId>,
                  "Need standard layout for type punning");
    static_assert(std::is_standard_layout_v<SemanticsStringId>,
                  "Need standard layout for type punning");
    static_assert(std::is_standard_layout_v<SemanticsTypeId>,
                  "Need standard layout for type punning");
  };
  static_assert(sizeof(Entry) == 8, "Unexpected Entry size");

  // Pops an entry.
  template <typename IdT>
  auto PopEntry() -> Entry {
    Entry back = stack_.pop_back_val();
    CARBON_VLOG() << "Node Pop " << stack_.size() << ": "
                  << parse_tree_->node_kind(back.parse_node) << " -> "
                  << back.id<IdT>() << "\n";
    return back;
  }

  // Require an entry to have the given ParseNodeKind.
  auto RequireParseKind(ParseTree::Node parse_node, ParseNodeKind require_kind)
      -> void {
    auto actual_kind = parse_tree_->node_kind(parse_node);
    CARBON_CHECK(require_kind == actual_kind)
        << "Expected " << require_kind << ", found " << actual_kind;
  }

  // Requires an entry to have a invalid node_id.
  auto RequireSoloParseNode(Entry entry) -> void {
    // See above comment on type punning.
    CARBON_CHECK(!entry.node_id.is_valid())
        << "Expected invalid id on " << parse_tree_->node_kind(entry.parse_node)
        << ", was " << entry.node_id << " (may not be node)";
  }

  // Requires an entry to have a valid id.
  auto RequireValidId(Entry entry) -> void {
    // See above comment on type punning.
    CARBON_CHECK(entry.node_id.is_valid())
        << "Expected valid id on " << parse_tree_->node_kind(entry.parse_node);
  }

  // The file's parse tree.
  const ParseTree* parse_tree_;

  // Whether to print verbose output.
  llvm::raw_ostream* vlog_stream_;

  // The actual stack.
  // PushEntry and PopEntry control modification in order to centralize
  // vlogging.
  llvm::SmallVector<Entry> stack_;
};

}  // namespace Carbon

#endif  // CARBON_TOOLCHAIN_SEMANTICS_SEMANTICS_NODE_STACK_H_