lang/include/name_tree.hpp

#pragma once

#include "basic_printers.hpp"
#include "statement_nodes.hpp"
#include "tree_sitter_wrapper.hpp"

#include <cstddef>
#include <optional>
#include <unordered_map>
#include <vector>

namespace names {

class NameTree;

class AnyStatementProxy {
  friend NameTree;

public:
  nodes::Statement *get();

  const nodes::Statement *get() const;

  //

  bool operator==(const AnyStatementProxy &other_any_statement_proxy) const {
    return name_tree_ == other_any_statement_proxy.name_tree_ &&
           id_ == other_any_statement_proxy.id_;
  }

  bool operator!=(const AnyStatementProxy &other_any_statement_proxy) const {
    return !(*this == other_any_statement_proxy);
  }

private:
  AnyStatementProxy(NameTree &name_tree, size_t id)
      : name_tree_(&name_tree), id_(id) {}

private:
  NameTree *name_tree_;
  size_t id_;
};

template <typename T> class StatementProxy {
public:
  StatementProxy(AnyStatementProxy proxy) : proxy_(proxy) {}

  T *get() { return proxy_.get()->get<T>().value(); }

  const T *get() const { return proxy_.get()->get<T>().value(); }

  //

  bool operator==(const StatementProxy &other_statement_proxy) const {
    return proxy_ == other_statement_proxy.proxy_;
  }

  bool operator!=(const StatementProxy &other_statement_proxy) const {
    return !(*this == other_statement_proxy);
  }

private:
  AnyStatementProxy proxy_;
};

class NameTree {
  friend AnyStatementProxy;

public:
  NameTree() {
    nodes_.emplace_back(); // root
  }

  std::optional<AnyStatementProxy> insert(const std::string &path,
                                          nodes::Statement &&statement);

  std::pair<std::optional<AnyStatementProxy>, nodes::CombineResult>
  insert_combine(const std::string &path, nodes::Statement &&statement);

  std::optional<AnyStatementProxy> find(const std::string &path);

  void print(printers::Printer &printer) const;

  // // TODO
  // void add_statement_children_to_tree();

private:
  struct Node {
  public:
    std::optional<nodes::Statement *> get_statement() {
      if (statement_.has_value()) {
        return &statement_.value();
      }
      return std::nullopt;
    }

    std::optional<const nodes::Statement *> get_statement() const {
      if (statement_.has_value()) {
        return &statement_.value();
      }
      return std::nullopt;
    }

    bool set_statement(nodes::Statement &&statement) {
      if (statement_.has_value()) {
        return false;
      }
      statement_ = std::move(statement);
      return true;
    }

    bool set_statement(const nodes::Statement &statement) {
      if (statement_.has_value()) {
        return false;
      }
      statement_ = statement;
      return true;
    }

    std::optional<size_t> find(size_t name_id) const {
      auto name_iter = children_.find(name_id);

      if (name_iter == children_.end()) {
        return std::nullopt;
      }

      return name_iter->second;
    }

    bool insert(size_t name_id, size_t node_id) {
      if (children_.count(name_id) != 0) {
        return false;
      }
      children_[name_id] = node_id;
      return true;
    }

    void print(const std::vector<Node> &nodes,
               printers::Printer &printer) const {
      printer.print("Node - has statement: ");
      printer.print(statement_.has_value() ? "true" : "false");
      printer.indent();
      for (auto &iter : children_) {
        printer.new_indent_line();
        printer.print(std::to_string(iter.first) + " -> ");
        nodes[iter.second].print(nodes, printer);
      }
      printer.deindent();
      printer.new_indent_line();
    }

  private:
    std::unordered_map<size_t, size_t> children_;
    std::optional<nodes::Statement> statement_;
  };

  size_t get_root() const { return 0; }

  size_t add_node(size_t current_node, const std::vector<size_t> &name_ids,
                  size_t current_name = 0);

  std::optional<size_t> find_node(size_t current_node,
                                  const std::vector<size_t> &name_ids,
                                  size_t current_name = 0) const;

  std::vector<size_t> slice_path_to_name_ids(const std::string &path);

  std::optional<std::vector<size_t>>
  slice_path_to_existing_name_ids(const std::string &path) const;

  size_t add_name_id(const std::string &name);

  std::optional<size_t> find_name_id(const std::string &name) const;

private:
  std::unordered_map<std::string, size_t> name_to_id_;
  std::vector<Node> nodes_;
};

} // namespace names