2 files changed, 71 insertions, 34 deletions

diff --git a/src/trie.h b/src/trie.h
index 873a375..46fcb4b 100644
--- a/src/trie.h
+++ b/src/trie.h
@@ -5,31 +5,57 @@
 #include <map>
 
 template <
-	typename Key
+	typename Key,
+	typename Value = std::nullptr_t
 >
 class Trie {
 	public:
 		typedef std::vector<Key> key_list;
 
 		Trie():
-			children_() { }
+			value_(),
+			children_() {
+			this->value_.first = false;
+		}
 
 		inline void add(key_list path) {
 			this->add(path, path.begin());
 		}
 
+		inline void add(key_list path, Value value) {
+			Trie* tmp = this->add(path, path.begin());
+
+			tmp->value_.first  = true;
+			tmp->value_.second = value;
+		}
+
 		inline void remove(key_list path) {
 			this->remove(path, path.begin());
 		}
 
-		inline std::pair<bool, const Trie*> resolve(key_list path) const {
-			return this->resolve(path, path.begin());
+		inline bool check(key_list path) {
+			return this->resolve(path).first;
+		}
+
+		inline std::pair<bool, Value*> get(key_list path) {
+			std::pair<bool, Trie*> tmp(this->resolve(path));
+
+			if ( tmp.first ) {
+				if ( tmp.second->value_.first ) {
+					return std::make_pair(true, &tmp.second->value_.second);
+				} else {
+					return std::make_pair(false, nullptr);
+				}
+			} else {
+				return std::make_pair(false, nullptr);
+			}
 		}
 
 	private:
+		std::pair<bool, Value> value_;
 		std::map<Key, Trie> children_;
 
-		inline void add(
+		inline Trie* add(
 			key_list& path,
 			typename key_list::const_iterator currStep
 		) {
@@ -38,7 +64,9 @@ class Trie {
 					this->children_[*currStep]
 				);
 
-				trie.add(path, ++currStep);
+				return trie.add(path, ++currStep);
+			} else {
+				return this;
 			}
 		}
 
@@ -66,11 +94,15 @@ class Trie {
 			}
 		}
 
-		inline std::pair<bool, const Trie*> resolve(
+		inline std::pair<bool, Trie*> resolve(key_list path) {
+			return this->resolve(path, path.begin());
+		}
+
+		inline std::pair<bool, Trie*> resolve(
 			key_list& path,
 			typename key_list::const_iterator currStep
-		) const {
-			typename std::map<Key, Trie>::const_iterator matchingTrie(
+		) {
+			typename std::map<Key, Trie>::iterator matchingTrie(
 				this->children_.find(*currStep)
 			);
 
diff --git a/test.cc b/test.cc
index 2c38a3f..ff0b04f 100644
--- a/test.cc
+++ b/test.cc
@@ -14,20 +14,12 @@ TEST_F(TrieTest, Basic) {
 	trie.add({2, 3, 4, 5});
 	trie.add({2, 3, 1, 2});
 
-	EXPECT_EQ(trie.resolve({1, 1, 1, 1}).first,  true);
-	EXPECT_NE(trie.resolve({1, 1, 1, 1}).second, nullptr);
+	EXPECT_EQ(trie.check({1, 1, 1, 1}),  true);
+	EXPECT_EQ(trie.check({1, 2, 1, 2}),  true);
+	EXPECT_EQ(trie.check({1, 2}),        true);
 
-	EXPECT_EQ(trie.resolve({1, 2, 1, 2}).first,  true);
-	EXPECT_NE(trie.resolve({1, 2, 1, 2}).second, nullptr);
-
-	EXPECT_EQ(trie.resolve({1, 2}).first,        true);
-	EXPECT_NE(trie.resolve({1, 2}).second,       nullptr);
-
-	EXPECT_EQ(trie.resolve({1, 1, 2, 3}).first,  false);
-	EXPECT_EQ(trie.resolve({1, 1, 2, 3}).second, nullptr);
-
-	EXPECT_EQ(trie.resolve({2, 1, 4}).first,     false);
-	EXPECT_EQ(trie.resolve({2, 1, 4}).second,    nullptr);
+	EXPECT_EQ(trie.check({1, 1, 2, 3}),  false);
+	EXPECT_EQ(trie.check({2, 1, 4}),     false);
 }
 
 TEST_F(TrieTest, Remove) {
@@ -40,22 +32,35 @@ TEST_F(TrieTest, Remove) {
 
 	trie.remove({1, 1, 1, 1});
 
-	EXPECT_EQ(trie.resolve({1, 1, 1, 1}).first,  false);
-	EXPECT_EQ(trie.resolve({1, 1, 1, 1}).second, nullptr);
-
-	EXPECT_EQ(trie.resolve({1, 2, 1, 2}).first,  true);
-	EXPECT_NE(trie.resolve({1, 2, 1, 2}).second, nullptr);
+	EXPECT_EQ(trie.check({1, 1, 1, 1}),  false);
+	EXPECT_EQ(trie.check({1, 2, 1, 2}),  true);
 
 	trie.remove({2, 3});
 
-	EXPECT_EQ(trie.resolve({2, 3, 4, 5}).first,  false);
-	EXPECT_EQ(trie.resolve({2, 3, 4, 5}).second, nullptr);
-
-	EXPECT_EQ(trie.resolve({2, 3}).first,        false);
-	EXPECT_EQ(trie.resolve({2, 3}).second,       nullptr);
+	EXPECT_EQ(trie.check({2, 3, 4, 5}),  false);
+	EXPECT_EQ(trie.check({2, 3}),        false);
+	EXPECT_EQ(trie.check({2}),           true);
+}
 
-	EXPECT_EQ(trie.resolve({2}).first,           true);
-	EXPECT_NE(trie.resolve({2}).second,          nullptr);
+TEST_F(TrieTest, Value) {
+	Trie<uint8_t, uint8_t> trie;
+
+	trie.add({1, 1, 1, 1}, 42);
+	trie.add({1, 2, 1, 2}, 43);
+	trie.add({2, 3, 4, 5}, 44);
+	trie.add({2, 3, 1, 2}, 45);
+
+	EXPECT_EQ(trie.get({1, 1, 1, 1}).first,  true);
+	EXPECT_EQ(*trie.get({1, 1, 1, 1}).second, 42);
+	EXPECT_EQ(trie.get({1, 2, 1, 2}).first,  true);
+	EXPECT_EQ(*trie.get({1, 2, 1, 2}).second, 43);
+	EXPECT_EQ(trie.get({2, 3, 4, 5}).first,  true);
+	EXPECT_EQ(*trie.get({2, 3, 4, 5}).second, 44);
+	EXPECT_EQ(trie.get({2, 3, 4, 5}).first,  true);
+	EXPECT_EQ(*trie.get({2, 3, 1, 2}).second, 45);
+
+	EXPECT_EQ(trie.get({1, 2}).first,        false);
+	EXPECT_EQ(trie.get({1, 2}).second,       nullptr);
 }
 
 int main(int argc, char **argv) {