diff --git a/src/core/lib/gprpp/map.h b/src/core/lib/gprpp/map.h index b210c26cbff..525c25347f4 100644 --- a/src/core/lib/gprpp/map.h +++ b/src/core/lib/gprpp/map.h @@ -112,7 +112,10 @@ class Map { // inserted entry and the second value being the new root of the subtree // after a rebalance Pair InsertRecursive(Entry* root, value_type&& p); - static Entry* RemoveRecursive(Entry* root, const key_type& k); + // Returns a pair with the first value being an iterator pointing to the + // successor of the deleted entry and the second value being the new root of + // the subtree after a rebalance + Pair RemoveRecursive(Entry* root, const key_type& k); // Return 0 if lhs = rhs // 1 if lhs > rhs // -1 if lhs < rhs @@ -233,10 +236,10 @@ typename Map::iterator Map::erase( iterator iter) { if (iter == end()) return iter; key_type& del_key = iter->first; - iter++; - root_ = RemoveRecursive(root_, del_key); + Pair ret = RemoveRecursive(root_, del_key); + root_ = ret.second; size_--; - return iter; + return ret.first; } template @@ -373,34 +376,38 @@ Map::RebalanceTreeAfterDeletion(Entry* root) { } template -typename Map::Entry* Map::RemoveRecursive( - Entry* root, const key_type& k) { - if (root == nullptr) return root; +typename ::grpc_core::Pair::iterator, + typename Map::Entry*> +Map::RemoveRecursive(Entry* root, const key_type& k) { + Pair ret = MakePair(end(), root); + if (root == nullptr) return ret; int comp = CompareKeys(root->pair.first, k); if (comp > 0) { - root->left = RemoveRecursive(root->left, k); + ret = RemoveRecursive(root->left, k); + root->left = ret.second; } else if (comp < 0) { - root->right = RemoveRecursive(root->right, k); + ret = RemoveRecursive(root->right, k); + root->right = ret.second; } else { - Entry* ret; + Entry* entry; + Entry* successor = InOrderSuccessor(root); if (root->left == nullptr) { - ret = root->right; + entry = root->right; Delete(root); - return ret; + return MakePair(iterator(this, successor), entry); } else if (root->right == nullptr) { - ret = root->left; + entry = root->left; Delete(root); - return ret; + return MakePair(iterator(this, successor), entry); } else { - ret = root->right; - while (ret->left != nullptr) { - ret = ret->left; - } - root->pair.swap(ret->pair); - root->right = RemoveRecursive(root->right, ret->pair.first); + entry = successor; + root->pair.swap(entry->pair); + ret = RemoveRecursive(root->right, entry->pair.first); + root->right = ret.second; + ret.first = iterator(this, root); } } - return RebalanceTreeAfterDeletion(root); + return MakePair(ret.first, RebalanceTreeAfterDeletion(root)); } template diff --git a/test/core/gprpp/map_test.cc b/test/core/gprpp/map_test.cc index e70bf38f111..6e70a2a2ac2 100644 --- a/test/core/gprpp/map_test.cc +++ b/test/core/gprpp/map_test.cc @@ -17,7 +17,9 @@ */ #include "src/core/lib/gprpp/map.h" + #include + #include "include/grpc/support/string_util.h" #include "src/core/lib/gprpp/inlined_vector.h" #include "src/core/lib/gprpp/memory.h" @@ -319,43 +321,49 @@ TEST_F(MapTest, MapRandomInsertions) { // Test Map iterator TEST_F(MapTest, Iteration) { Map test_map; - for (int i = 0; i < 5; i++) { + for (int i = 4; i >= 0; --i) { test_map.emplace(kKeys[i], Payload(i)); } - int count = 0; - for (auto iter = test_map.begin(); iter != test_map.end(); iter++) { - EXPECT_EQ(iter->second.data(), count); - count++; + auto it = test_map.begin(); + for (int i = 0; i < 5; ++i) { + ASSERT_NE(it, test_map.end()); + EXPECT_STREQ(kKeys[i], it->first); + EXPECT_EQ(i, it->second.data()); + ++it; } - EXPECT_EQ(count, 5); + EXPECT_EQ(it, test_map.end()); } // Test Map iterator with unique ptr payload TEST_F(MapTest, IterationWithUniquePtrValue) { Map, StringLess> test_map; - for (int i = 0; i < 5; i++) { + for (int i = 4; i >= 0; --i) { test_map.emplace(kKeys[i], MakeUnique(i)); } - int count = 0; - for (auto iter = test_map.begin(); iter != test_map.end(); iter++) { - EXPECT_EQ(iter->second->data(), count); - count++; + auto it = test_map.begin(); + for (int i = 0; i < 5; ++i) { + ASSERT_NE(it, test_map.end()); + EXPECT_STREQ(kKeys[i], it->first); + EXPECT_EQ(i, it->second->data()); + ++it; } - EXPECT_EQ(count, 5); + EXPECT_EQ(it, test_map.end()); } // Test Map iterator with unique ptr to char key TEST_F(MapTest, IterationWithUniquePtrKey) { Map, Payload, StringLess> test_map; - for (int i = 0; i < 5; i++) { + for (int i = 4; i >= 0; --i) { test_map.emplace(CopyString(kKeys[i]), Payload(i)); } - int count = 0; - for (auto iter = test_map.begin(); iter != test_map.end(); iter++) { - EXPECT_EQ(iter->second.data(), count); - count++; + auto it = test_map.begin(); + for (int i = 0; i < 5; ++i) { + ASSERT_NE(it, test_map.end()); + EXPECT_STREQ(kKeys[i], it->first.get()); + EXPECT_EQ(i, it->second.data()); + ++it; } - EXPECT_EQ(count, 5); + EXPECT_EQ(it, test_map.end()); } // Test removing entries while iterating the map @@ -367,11 +375,23 @@ TEST_F(MapTest, EraseUsingIterator) { int count = 0; for (auto iter = test_map.begin(); iter != test_map.end();) { EXPECT_EQ(iter->second.data(), count); - iter = test_map.erase(iter); - count++; + if (count % 2 == 1) { + iter = test_map.erase(iter); + } else { + ++iter; + } + ++count; } EXPECT_EQ(count, 5); - EXPECT_TRUE(test_map.empty()); + auto it = test_map.begin(); + for (int i = 0; i < 5; ++i) { + if (i % 2 == 0) { + EXPECT_STREQ(kKeys[i], it->first); + EXPECT_EQ(i, it->second.data()); + ++it; + } + } + EXPECT_EQ(it, test_map.end()); } // Random ops on a Map with Integer key of Payload value,