diff --git a/include/global_info.hpp b/include/global_info.hpp index 8c8971a..27a8b93 100644 --- a/include/global_info.hpp +++ b/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& path, - const std::string& name, + utils::IdType AddPartition(const std::vector& path, // including name interpreter::tokens::PartitionStatement* node); std::optional FindNamespace(const std::optional>& path); @@ -164,6 +163,23 @@ public: return partitions_.at(id); } + std::optional FindPartition(const std::vector& path) { + auto trie_ans = partitions_trie_.Find(path); + return trie_ans.has_value() ? std::optional(*trie_ans.value()) : std::nullopt; + } + + std::vector FindPartitionsByPrefix(const std::vector& path) { // optimize ?? + auto trie_ans = partitions_trie_.FindByPrefix(path); + + std::vector ans(trie_ans.size()); + + for (size_t i = 0; i < ans.size(); ++i) { + ans[i] = *trie_ans[i]; + } + + return ans; + } + private: std::vector functions_; @@ -176,6 +192,7 @@ private: std::unordered_map name_to_abstract_type_; std::vector partitions_; + utils::Trie partitions_trie_; definition::Namespace global_namespace_; std::vector imports_; diff --git a/include/utils.hpp b/include/utils.hpp index 34155f9..8ff214f 100644 --- a/include/utils.hpp +++ b/include/utils.hpp @@ -3,6 +3,8 @@ #include #include #include +#include +#include #include namespace utils { @@ -64,6 +66,106 @@ private: std::unordered_map value_to_id_; }; +template +class Trie { // optimize ?? +public: + Trie() { + verticles_.emplace_back(); + } + + bool Insert(const std::vector& path, const Value& value) { + return RecursiveInsert(verticles_[0], path, 0, value); + } + + std::optional Find(const std::vector& path) { + return RecursiveFind(verticles_[0], path, 0); + } + + std::vector FindByPrefix(const std::vector& path) { + return RecursiveFindByPrefix(verticles_[0], path, 0); + } +private: + struct Vertex { + std::unordered_map children_; + std::optional value; + }; + + bool RecursiveInsert(Vertex& vertex, + const std::vector& 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 RecursiveFind(Vertex& vertex, + const std::vector& 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 RecursiveFindByPrefix(Vertex& vertex, + const std::vector& path, + size_t path_position) { + if (path_position == path.size()) { + std::vector 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& accumulator) { + std::vector 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 verticles_; +}; + class GroupsManager { // TODO: recall right algorithm name public: GroupsManager() = default; diff --git a/src/global_info.cpp b/src/global_info.cpp index de4b937..8f5054a 100644 --- a/src/global_info.cpp +++ b/src/global_info.cpp @@ -248,22 +248,25 @@ utils::IdType GlobalInfo::NamespaceVisitor::AddConstructor(const std::string& co } utils::IdType GlobalInfo::NamespaceVisitor::AddPartition(const std::vector& path, - const std::string& name, interpreter::tokens::PartitionStatement* node) { PartitionInfo partition; - partition.path.reserve(current_path_.size() + path.size()); + partition.path.reserve(current_path_.size() + path.size() - 1); partition.path = current_path_; - for (auto& path_namespace : path) { - partition.path.push_back(path_namespace); + for (size_t i = 0; i + 1 < path.size(); ++i) { + partition.path.push_back(path[i]); } - partition.name = name; + partition.name = path.back(); partition.node = node; + utils::IdType id = global_info_.partitions_.size(); + global_info_.partitions_.push_back(partition); - return global_info_.partitions_.size() - 1; + global_info_.partitions_trie_.Insert(partition.path, id); + + return id; } std::optional GlobalInfo::NamespaceVisitor::FindNamespace(const std::optional>& path) {