lang_2023/include/utils.hpp

#pragma once

#include <cstdlib>
#include <unordered_set>
#include <vector>
#include <string>
#include <memory>
#include <optional>
#include <unordered_map>

namespace utils {

using std::size_t;

using IdType = size_t;

const std::string ClassInternalVarName = "self";
const size_t MaxOperatorPrecedence = 4;

struct ValueReturnedMarker {};
struct ValueBroughtMarker {};

enum class ReferenceModifier { Reference = 0, UniqueReference = 1, Dereference = 2 };
enum class IsConstModifier { Const = 0, Var = 1 };
enum class ClassInternalsModifier { Static = 0, Const = 1, Var = 2};
enum class ClassModifier { Struct = 0, Class = 1 };
enum class AssignmentModifier { Assign = 0, Move = 1 };
enum class AliasModifier { Alias = 0, Type = 1, Let = 2 };
enum class AbstractTypeModifier { Basic = 0, Abstract = 1 };
enum class PartitionModifier { Exec = 0, Test = 1 };

enum class ValueType { Const = 0, Var = 1, Tmp = 2 };

inline ValueType IsConstModifierToValueType(IsConstModifier modifier) {
  switch (modifier) {
    case IsConstModifier::Const:
      return ValueType::Const;
    case IsConstModifier::Var:
      return ValueType::Var;
  }

  exit(1); // unreachable
}

inline ValueType ClassInternalsModifierToValueType(ClassInternalsModifier modifier) {
  switch (modifier) {
    case ClassInternalsModifier::Const:
      return ValueType::Const;
    case ClassInternalsModifier::Var:
      return ValueType::Var;
    case ClassInternalsModifier::Static:
      throw std::bad_cast(); // ??
  }

  exit(1); // unreachable
}

inline bool IsBuiltinFunction(const std::string& name) { // optimize ??
  std::unordered_set<std::string> builtin_functions;

  builtin_functions.insert("=");
  builtin_functions.insert("<-");
  builtin_functions.insert("==");
  builtin_functions.insert("!=");
  builtin_functions.insert("<");
  builtin_functions.insert(">");
  builtin_functions.insert("<=");
  builtin_functions.insert(">=");
  builtin_functions.insert("+=");
  builtin_functions.insert("-=");
  builtin_functions.insert("*=");
  builtin_functions.insert("//=");
  builtin_functions.insert("%=");
  builtin_functions.insert("/=");
  builtin_functions.insert("+");
  builtin_functions.insert("-");
  builtin_functions.insert("*");
  builtin_functions.insert("//");
  builtin_functions.insert("%");
  builtin_functions.insert("/");
  builtin_functions.insert("&&");
  builtin_functions.insert("||");
  builtin_functions.insert("size");
  builtin_functions.insert("random");
  builtin_functions.insert("print");
  builtin_functions.insert("scan");

  return builtin_functions.count(name) != 0;
}

template<typename T>
class Storage {
public:
  Storage() = default;

  IdType GetId(const T& value) {
    IdType id = 0;
    auto value_position = value_to_id_.find(value);

    if (value_position == value_to_id_.end()) {
      id = id_to_value_.size();
      value_to_id_[value] = id;
      id_to_value_.push_back(value);
    } else {
      id = value_position->second;
    }

    return id;
  }

  const T& GetValue(IdType id) {
    return id_to_value_[id];
  }

private:
  std::vector<T> id_to_value_;
  std::unordered_map<T, IdType> value_to_id_;
};

template<typename Key, typename Value>
class Trie { // optimize ??
public:
  Trie() {
    verticles_.emplace_back();
  }

  bool Insert(const std::vector<Key>& path, const Value& value) {
    return RecursiveInsert(verticles_[0], path, 0, value);
  }

  std::optional<Value*> Find(const std::vector<Key>& path) {
    return RecursiveFind(verticles_[0], path, 0);
  }

  std::vector<Value*> FindByPrefix(const std::vector<Key>& path) {
    return RecursiveFindByPrefix(verticles_[0], path, 0);
  }
private:
  struct Vertex {
    std::unordered_map<Key, size_t> children_;
    std::optional<Value> value;
  };

  bool RecursiveInsert(Vertex& vertex,
                       const std::vector<Key>& path,
                       size_t path_position,
                       const Value& value) {
    if (path_position == path.size()) {
      if (vertex.value.has_value()) {
        return false;
      }
      vertex.value = value;
      return true;
    }

    auto child_iter = vertex.children_.find(path[path_position]);

    if (child_iter != vertex.children_.end()) {
      return RecursiveInsert(verticles_[child_iter->second], path, path_position + 1, value);
    }

    vertex.children_[path[path_position]] = verticles_.size();
    verticles_.emplace_back();
    return RecursiveInsert(verticles_.back(), path, path_position + 1, value);
  }

  std::optional<Value*> RecursiveFind(Vertex& vertex,
                                      const std::vector<Key>& path,
                                      size_t path_position) {
    if (path_position == path.size()) {
      if (!vertex.value.has_value()) {
        return std::nullopt;
      }
      return &vertex.value.value();
    }

    auto child_iter = vertex.children_.find(path[path_position]);

    if (child_iter != vertex.children_.end()) {
      return RecursiveFind(verticles_[child_iter->second], path, path_position + 1);
    }

    return std::nullopt;
  }

  std::vector<Value*> RecursiveFindByPrefix(Vertex& vertex,
                                            const std::vector<Key>& path,
                                            size_t path_position) {
    if (path_position == path.size()) {
      std::vector<Value*> ans;
      RecursiveGetAll(vertex, ans);
      return ans;
    }

    auto child_iter = vertex.children_.find(path[path_position]);

    if (child_iter != vertex.children_.end()) {
      return RecursiveFindByPrefix(verticles_[child_iter->second], path, path_position + 1);
    }

    return {};
  }

  void RecursiveGetAll(Vertex& vertex, std::vector<Value*>& accumulator) {
    std::vector<Value*> ans;

    if (vertex.value.has_value()) {
      accumulator.push_back(&vertex.value.value());
    }

    for (auto& child : vertex.children_) {
      RecursiveGetAll(verticles_[child.second], accumulator);
    }
  }

private:
  std::vector<Vertex> verticles_;
};

class GroupsManager { // TODO: recall right algorithm name
public:
  GroupsManager() = default;

  explicit GroupsManager(size_t n) {
    edges_.resize(n);
    ranks_.resize(n);

    for (size_t i = 0; i < n; ++i) {
      edges_[i] = i;
      ranks_[i] = 1;
    }
  }

  void Unite(size_t u, size_t v) { // TODO: recall choice algorithm
    u = GetGroupRoot(u);
    v = GetGroupRoot(v);
    if (ranks_[v] >= ranks_[u]) {
      edges_[u] = v;
      ranks_[v] = std::max(ranks_[u] + 1, ranks_[v]);
    } else {
      edges_[v] = u;
      // always ranks_[u] > ranks_[v]
    }
  }

  bool IsInOneGroup(size_t u, size_t v) {
    return GetGroupRoot(u) == GetGroupRoot(v);
  }

  size_t AddElement() {
    size_t id = edges_.size();
    edges_.push_back(id);
    ranks_.push_back(1);
    return id;
  }

  size_t GetGroupRoot(size_t v) {
    if (edges_[v] == v) {
      return v;
    }
    return edges_[v] = GetGroupRoot(edges_[v]);
  }
private:
  std::vector<size_t> edges_;
  std::vector<size_t> ranks_;
};

// move to .cpp ??
inline void BackVisitDfs(size_t id,
                     std::vector<size_t>& verticles,
                     std::vector<size_t>& marks,
                     const std::vector<std::vector<size_t>>& edges,
                     size_t mark) {
  if (marks[id] != 0) {
    return;
  }

  marks[id] = mark;
  verticles.push_back(id);

  for (size_t i = 0; i < edges[id].size(); ++i) {
    BackVisitDfs(id, verticles, marks, edges, mark);
  }
}

// move to .cpp ??
inline std::vector<size_t> BackTopSort(const std::vector<std::vector<size_t>>& edges_) {
  std::vector<size_t> sorted_verticles;
  std::vector<size_t> marks(edges_.size(), 0);

  for (size_t i = 0; i < marks.size(); ++i) {
    BackVisitDfs(i, sorted_verticles, marks, edges_, 1);
  }

  return sorted_verticles;
}

} // namespace utils