mirror of
https://github.com/open-goal/jak-project.git
synced 2024-10-20 21:27:52 -04:00
185 lines
4.3 KiB
C++
185 lines
4.3 KiB
C++
#pragma once
|
|
|
|
#include "common/util/assert.h"
|
|
#include <vector>
|
|
#include <string>
|
|
|
|
/*!
|
|
* A simple prefix tree. It works similarly to a map, but also supports fast lookups by prefix with
|
|
* the ability to lookup all entries with keys that begin with a given prefix.
|
|
*
|
|
* It owns the memory for the objects it stores.
|
|
* Doing an insert will create a copy of your object.
|
|
*
|
|
* Other that deleting the whole thing, there is no support for removing a node.
|
|
*/
|
|
template <typename T>
|
|
class Trie {
|
|
public:
|
|
Trie() = default;
|
|
Trie(const Trie&) = delete;
|
|
Trie& operator=(const Trie&) = delete;
|
|
|
|
// Insert an object, replacing an existing one if it exists
|
|
void insert(const std::string& str, const T& obj);
|
|
|
|
// Get the object at the string. Default construct a new one if none exists.
|
|
T* operator[](const std::string& str);
|
|
|
|
// Lookup an existing object. If none exists, return nullptr.
|
|
T* lookup(const std::string& str);
|
|
|
|
// return the number of entries.
|
|
int size() const { return m_size; }
|
|
|
|
// Get all objects starting with the given prefix.
|
|
std::vector<T*> lookup_prefix(const std::string& str);
|
|
~Trie();
|
|
|
|
private:
|
|
static constexpr int CHAR_START = ' ';
|
|
static constexpr int CHAR_END = '~' + 1;
|
|
static constexpr int CHAR_SIZE = CHAR_END - CHAR_START;
|
|
|
|
static int idx(char c) {
|
|
assert(c >= CHAR_START);
|
|
assert(c < CHAR_END);
|
|
return c - CHAR_START;
|
|
}
|
|
|
|
struct Node {
|
|
T* value = nullptr;
|
|
Node* children[CHAR_SIZE] = {0};
|
|
|
|
void delete_children() {
|
|
if (value) {
|
|
delete value;
|
|
value = nullptr;
|
|
}
|
|
for (auto child : children) {
|
|
if (child) {
|
|
child->delete_children();
|
|
delete child;
|
|
child = nullptr;
|
|
}
|
|
}
|
|
}
|
|
|
|
/*!
|
|
* Return true if a new object was inserted.
|
|
*/
|
|
bool insert(const char* str, const T* obj) {
|
|
if (!*str) {
|
|
// we are the child!
|
|
if (value) {
|
|
delete value;
|
|
value = new T(*obj);
|
|
return false; // didn't change the count.
|
|
} else {
|
|
value = new T(*obj);
|
|
return true; // did change the count
|
|
}
|
|
|
|
} else {
|
|
// still more to go
|
|
char first = *str;
|
|
if (!children[idx(first)]) {
|
|
children[idx(first)] = new Node();
|
|
}
|
|
return children[idx(first)]->insert(str + 1, obj);
|
|
}
|
|
}
|
|
|
|
T* bracket_operator(const char* str, bool* inserted) {
|
|
if (!*str) {
|
|
// we are the child!
|
|
if (!value) {
|
|
value = new T();
|
|
*inserted = true;
|
|
} else {
|
|
*inserted = false;
|
|
}
|
|
|
|
return value;
|
|
} else {
|
|
// still more to go
|
|
char first = *str;
|
|
if (!children[idx(first)]) {
|
|
children[idx(first)] = new Node();
|
|
}
|
|
return children[idx(first)]->bracket_operator(str + 1, inserted);
|
|
}
|
|
}
|
|
|
|
T* lookup(const char* str) {
|
|
if (!*str) {
|
|
return value;
|
|
}
|
|
if (children[idx(*str)]) {
|
|
return children[idx(*str)]->lookup(str + 1);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
void get_all_children(std::vector<T*>& result) {
|
|
if (value) {
|
|
result.push_back(value);
|
|
}
|
|
for (auto child : children) {
|
|
if (child) {
|
|
child->get_all_children(result);
|
|
}
|
|
}
|
|
}
|
|
|
|
std::vector<T*> lookup_prefix(const char* str) {
|
|
if (!*str) {
|
|
std::vector<T*> result;
|
|
get_all_children(result);
|
|
return result;
|
|
} else {
|
|
if (children[idx(*str)]) {
|
|
return children[idx(*str)]->lookup_prefix(str + 1);
|
|
} else {
|
|
return {};
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
Node m_root;
|
|
int m_size = 0;
|
|
};
|
|
|
|
template <typename T>
|
|
Trie<T>::~Trie<T>() {
|
|
m_root.delete_children();
|
|
m_size = 0;
|
|
}
|
|
|
|
template <typename T>
|
|
void Trie<T>::insert(const std::string& str, const T& obj) {
|
|
if (m_root.insert(str.c_str(), &obj)) {
|
|
m_size++;
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
T* Trie<T>::lookup(const std::string& str) {
|
|
return m_root.lookup(str.c_str());
|
|
}
|
|
|
|
template <typename T>
|
|
T* Trie<T>::operator[](const std::string& str) {
|
|
bool added = false;
|
|
auto result = m_root.bracket_operator(str.c_str(), &added);
|
|
if (added) {
|
|
m_size++;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
template <typename T>
|
|
std::vector<T*> Trie<T>::lookup_prefix(const std::string& str) {
|
|
return m_root.lookup_prefix(str.c_str());
|
|
} |