lang/include/name_tree.hpp

#pragma once

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

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

namespace names {

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

  bool insert(const std::string &path, nodes::Statement &&statement);

  std::optional<nodes::Statement *> find(const std::string &path);

  std::optional<const nodes::Statement *> find(const std::string &path) const;

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

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

  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