Skip to content

Commit df56315

Browse files
arturgontijoarturgontijo
authored
Merge pull request #1018 from singnet/artur-handletrie-remove
Implement `HandleTrie::remove()`
2 parents 7a768a7 + e4dbaed commit df56315

3 files changed

Lines changed: 198 additions & 14 deletions

File tree

‎src/commons/HandleTrie.cc‎

Lines changed: 41 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -145,14 +145,24 @@ HandleTrie::TrieValue* HandleTrie::insert(const string& key, TrieValue* value) {
145145
}
146146
}
147147
if (match) {
148-
tree_cursor->value->merge(value);
149-
this->size--;
150-
delete value;
151-
if (tree_cursor != parent) {
152-
parent->trie_node_mutex.unlock();
148+
if (tree_cursor->value != NULL) {
149+
tree_cursor->value->merge(value);
150+
this->size--;
151+
delete value;
152+
if (tree_cursor != parent) {
153+
parent->trie_node_mutex.unlock();
154+
}
155+
tree_cursor->trie_node_mutex.unlock();
156+
return tree_cursor->value;
157+
} else {
158+
// Value was removed, set it to the new value
159+
tree_cursor->value = value;
160+
if (tree_cursor != parent) {
161+
parent->trie_node_mutex.unlock();
162+
}
163+
tree_cursor->trie_node_mutex.unlock();
164+
return tree_cursor->value;
153165
}
154-
tree_cursor->trie_node_mutex.unlock();
155-
return tree_cursor->value;
156166
}
157167
}
158168
key_cursor++;
@@ -161,12 +171,16 @@ HandleTrie::TrieValue* HandleTrie::insert(const string& key, TrieValue* value) {
161171
}
162172

163173
HandleTrie::TrieValue* HandleTrie::lookup(const string& key) {
174+
TrieNode* node = lookup_node(key);
175+
return node != NULL ? node->value : NULL;
176+
}
177+
178+
HandleTrie::TrieNode* HandleTrie::lookup_node(const string& key) {
164179
if (key.size() != key_size) {
165180
Utils::error("Invalid key size: " + to_string(key.size()) + " != " + to_string(key_size));
166181
}
167182

168183
TrieNode* tree_cursor = root;
169-
TrieValue* value;
170184
unsigned char key_cursor = 0;
171185
tree_cursor->trie_node_mutex.lock();
172186
while (tree_cursor != NULL) {
@@ -179,13 +193,9 @@ HandleTrie::TrieValue* HandleTrie::lookup(const string& key) {
179193
break;
180194
}
181195
}
182-
if (match) {
183-
value = tree_cursor->value;
184-
} else {
185-
value = NULL;
186-
}
196+
TrieNode* node = match ? tree_cursor : NULL;
187197
tree_cursor->trie_node_mutex.unlock();
188-
return value;
198+
return node;
189199
} else {
190200
unsigned char c = TLB[(unsigned char) key[key_cursor]];
191201
TrieNode* child = tree_cursor->children[c];
@@ -200,6 +210,23 @@ HandleTrie::TrieValue* HandleTrie::lookup(const string& key) {
200210
return NULL;
201211
}
202212

213+
bool HandleTrie::remove(const string& key) {
214+
TrieNode* node = lookup_node(key);
215+
if (node == NULL) {
216+
return false;
217+
}
218+
node->trie_node_mutex.lock();
219+
if (node->value == NULL) {
220+
node->trie_node_mutex.unlock();
221+
return false;
222+
}
223+
delete node->value;
224+
node->value = NULL;
225+
this->size--;
226+
node->trie_node_mutex.unlock();
227+
return true;
228+
}
229+
203230
void HandleTrie::traverse(bool keep_root_locked,
204231
bool (*visit_function)(TrieNode* node, void* data),
205232
void* data) {

‎src/commons/HandleTrie.h‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,15 @@ class HandleTrie {
7474
*/
7575
TrieValue* lookup(const string& key);
7676

77+
/**
78+
* Remove a key from this HandleTrie and its associated value.
79+
*
80+
* @param key Handle being removed.
81+
*
82+
* @return true if the key was found and removed, false otherwise.
83+
*/
84+
bool remove(const string& key);
85+
7786
/**
7887
* Traverse all keys (in-order) calling the passed visit_function once per stored value.
7988
*
@@ -88,6 +97,16 @@ class HandleTrie {
8897
unsigned int size;
8998

9099
private:
100+
/**
101+
* Lookup for a node containing a given handle.
102+
* Similar to lookup() but returns the node pointer instead of the value.
103+
*
104+
* @param key Handle being searched.
105+
*
106+
* @return The HandleTrie::TrieNode containing the key or NULL if none.
107+
*/
108+
TrieNode* lookup_node(const string& key);
109+
91110
static unsigned char TLB[256];
92111
static bool TLB_INITIALIZED;
93112
static void TLB_INIT() {

‎src/tests/cpp/handle_trie_test.cc‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -537,3 +537,141 @@ TEST(HandleTrieTest, multithread_limited_handle_set) {
537537
}
538538
}
539539
}
540+
541+
TEST(HandleTrieTest, remove_basic) {
542+
HandleTrie trie(4);
543+
TestValue* value;
544+
545+
unsigned int initial_size = trie.size;
546+
EXPECT_EQ(initial_size, 0);
547+
548+
// Insert some values
549+
trie.insert("ABCD", new TestValue(3));
550+
trie.insert("ABCF", new TestValue(4));
551+
trie.insert("ABFD", new TestValue(5));
552+
trie.insert("FBCD", new TestValue(6));
553+
trie.insert("AFCD", new TestValue(7));
554+
555+
EXPECT_EQ(trie.size, initial_size + 5);
556+
557+
// Verify they exist
558+
value = (TestValue*) trie.lookup("ABCD");
559+
EXPECT_TRUE(value != NULL);
560+
EXPECT_TRUE(value->count == 3);
561+
value = (TestValue*) trie.lookup("ABCF");
562+
EXPECT_TRUE(value != NULL);
563+
EXPECT_TRUE(value->count == 4);
564+
565+
// Remove one
566+
bool removed = trie.remove("ABCD");
567+
EXPECT_TRUE(removed);
568+
569+
// Verify it's gone
570+
value = (TestValue*) trie.lookup("ABCD");
571+
EXPECT_TRUE(value == NULL);
572+
573+
EXPECT_EQ(trie.size, initial_size + 4);
574+
575+
// Verify others still exist
576+
value = (TestValue*) trie.lookup("ABCF");
577+
EXPECT_TRUE(value != NULL);
578+
EXPECT_TRUE(value->count == 4);
579+
value = (TestValue*) trie.lookup("ABFD");
580+
EXPECT_TRUE(value != NULL);
581+
EXPECT_TRUE(value->count == 5);
582+
value = (TestValue*) trie.lookup("FBCD");
583+
EXPECT_TRUE(value != NULL);
584+
EXPECT_TRUE(value->count == 6);
585+
value = (TestValue*) trie.lookup("AFCD");
586+
EXPECT_TRUE(value != NULL);
587+
EXPECT_TRUE(value->count == 7);
588+
589+
// Try to remove non-existent key
590+
removed = trie.remove("XXXX");
591+
EXPECT_FALSE(removed);
592+
593+
EXPECT_EQ(trie.size, initial_size + 4);
594+
}
595+
596+
TEST(HandleTrieTest, remove_and_reinsert) {
597+
HandleTrie trie(4);
598+
TestValue* value;
599+
600+
// Insert
601+
trie.insert("ABCD", new TestValue(3));
602+
value = (TestValue*) trie.lookup("ABCD");
603+
EXPECT_TRUE(value != NULL);
604+
EXPECT_TRUE(value->count == 3);
605+
606+
// Remove
607+
bool removed = trie.remove("ABCD");
608+
EXPECT_TRUE(removed);
609+
value = (TestValue*) trie.lookup("ABCD");
610+
EXPECT_TRUE(value == NULL);
611+
612+
// Re-insert
613+
trie.insert("ABCD", new TestValue(5));
614+
value = (TestValue*) trie.lookup("ABCD");
615+
EXPECT_TRUE(value != NULL);
616+
EXPECT_TRUE(value->count == 5);
617+
}
618+
619+
TEST(HandleTrieTest, remove_multiple) {
620+
HandleTrie trie(4);
621+
622+
// Insert multiple values
623+
trie.insert("ABCD", new TestValue(3));
624+
trie.insert("ABCF", new TestValue(4));
625+
trie.insert("ABFD", new TestValue(5));
626+
trie.insert("FBCD", new TestValue(6));
627+
trie.insert("AFCD", new TestValue(7));
628+
629+
// Remove all
630+
EXPECT_TRUE(trie.remove("ABCD"));
631+
EXPECT_TRUE(trie.remove("ABCF"));
632+
EXPECT_TRUE(trie.remove("ABFD"));
633+
EXPECT_TRUE(trie.remove("FBCD"));
634+
EXPECT_TRUE(trie.remove("AFCD"));
635+
636+
// Verify all are gone
637+
EXPECT_TRUE(trie.lookup("ABCD") == NULL);
638+
EXPECT_TRUE(trie.lookup("ABCF") == NULL);
639+
EXPECT_TRUE(trie.lookup("ABFD") == NULL);
640+
EXPECT_TRUE(trie.lookup("FBCD") == NULL);
641+
EXPECT_TRUE(trie.lookup("AFCD") == NULL);
642+
}
643+
644+
TEST(HandleTrieTest, remove_already_removed) {
645+
HandleTrie trie(4);
646+
647+
// Insert
648+
trie.insert("ABCD", new TestValue(3));
649+
650+
// Remove once
651+
bool removed = trie.remove("ABCD");
652+
EXPECT_TRUE(removed);
653+
654+
// Try to remove again
655+
removed = trie.remove("ABCD");
656+
EXPECT_FALSE(removed);
657+
}
658+
659+
TEST(HandleTrieTest, remove_after_merge) {
660+
HandleTrie trie(4);
661+
AccumulatorValue* value;
662+
663+
// Insert and merge
664+
trie.insert("ABCD", new AccumulatorValue());
665+
trie.insert("ABCD", new AccumulatorValue());
666+
value = (AccumulatorValue*) trie.lookup("ABCD");
667+
EXPECT_TRUE(value != NULL);
668+
EXPECT_TRUE(value->count == 2);
669+
670+
// Remove
671+
bool removed = trie.remove("ABCD");
672+
EXPECT_TRUE(removed);
673+
674+
// Verify it's gone
675+
value = (AccumulatorValue*) trie.lookup("ABCD");
676+
EXPECT_TRUE(value == NULL);
677+
}

0 commit comments

Comments
 (0)