better partitions system in global_info, trie in utils

include/global_info.hpp

@@ -67,8 +67,7 @@ public:
                                  definition::Constructor&& constructor_info,
                                  const interpreter::tokens::BaseNode& base_node);
 
-    utils::IdType AddPartition(const std::vector<std::string>& path,
-                               const std::string& name,
+    utils::IdType AddPartition(const std::vector<std::string>& path, // including name
                                interpreter::tokens::PartitionStatement* node);
 
     std::optional<definition::Namespace*> FindNamespace(const std::optional<std::vector<std::string>>& path);
@@ -164,6 +163,23 @@ public:
     return partitions_.at(id);
   }
 
+  std::optional<utils::IdType> FindPartition(const std::vector<std::string>& path) {
+    auto trie_ans = partitions_trie_.Find(path);
+    return trie_ans.has_value() ? std::optional<utils::IdType>(*trie_ans.value()) : std::nullopt;
+  }
+
+  std::vector<utils::IdType> FindPartitionsByPrefix(const std::vector<std::string>& path) { // optimize ??
+    auto trie_ans = partitions_trie_.FindByPrefix(path);
+
+    std::vector<utils::IdType> ans(trie_ans.size());
+
+    for (size_t i = 0; i < ans.size(); ++i) {
+      ans[i] = *trie_ans[i];
+    }
+
+    return ans;
+  }
+
 private:
 
   std::vector<definition::Function> functions_;
@@ -176,6 +192,7 @@ private:
   std::unordered_map<std::string, utils::IdType> name_to_abstract_type_;
 
   std::vector<PartitionInfo> partitions_;
+  utils::Trie<std::string, utils::IdType> partitions_trie_;
 
   definition::Namespace global_namespace_;
   std::vector<definition::Import> imports_;

include/utils.hpp

@@ -3,6 +3,8 @@
 #include <cstdlib>
 #include <vector>
 #include <string>
+#include <memory>
+#include <optional>
 #include <unordered_map>
 
 namespace utils {
@@ -64,6 +66,106 @@ private:
   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;