lang_2023/include/utils.hpp

#pragma once

#include <cstdlib>
#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";

enum class ReferenceModifier { Reference = 0, UniqueReference = 1 };
enum class IsConstModifier { Const = 0, Var = 1 };
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 FunctionTypeModifier { Function = 0, Operator = 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;
  }
  // unreachable
  exit(1);
}

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

// static 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);
//   }
// }
//
// static 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