Merge pull request #19214 from mhaidrygoog/fix_map_erase

Fixed erase() method by changing RemoveRecursive()

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<iterator, Entry*> 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<iterator, Entry*> 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<Key, T, Compare>::iterator Map<Key, T, Compare>::erase(
     iterator iter) {
   if (iter == end()) return iter;
   key_type& del_key = iter->first;
-  iter++;
-  root_ = RemoveRecursive(root_, del_key);
+  Pair<iterator, Entry*> ret = RemoveRecursive(root_, del_key);
+  root_ = ret.second;
   size_--;
-  return iter;
+  return ret.first;
 }
 
 template <class Key, class T, class Compare>
@@ -373,34 +376,38 @@ Map<Key, T, Compare>::RebalanceTreeAfterDeletion(Entry* root) {
 }
 
 template <class Key, class T, class Compare>
-typename Map<Key, T, Compare>::Entry* Map<Key, T, Compare>::RemoveRecursive(
-    Entry* root, const key_type& k) {
-  if (root == nullptr) return root;
+typename ::grpc_core::Pair<typename Map<Key, T, Compare>::iterator,
+                           typename Map<Key, T, Compare>::Entry*>
+Map<Key, T, Compare>::RemoveRecursive(Entry* root, const key_type& k) {
+  Pair<iterator, Entry*> 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 <class Key, class T, class Compare>

test/core/gprpp/map_test.cc

@@ -17,7 +17,9 @@
  */
 
 #include "src/core/lib/gprpp/map.h"
+
 #include <gtest/gtest.h>
+
 #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<const char*, Payload, StringLess> 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<const char*, UniquePtr<Payload>, StringLess> test_map;
-  for (int i = 0; i < 5; i++) {
+  for (int i = 4; i >= 0; --i) {
     test_map.emplace(kKeys[i], MakeUnique<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 to char key
 TEST_F(MapTest, IterationWithUniquePtrKey) {
   Map<UniquePtr<char>, 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,