[llvm] r213084 - ADT: Fix MapVector::erase()

[Duncan P. N. Exon Smith]

Author: dexonsmith
Date: Tue Jul 15 13:32:30 2014
New Revision: 213084

URL: http://llvm.org/viewvc/llvm-project?rev=213084&view=rev
Log:
ADT: Fix MapVector::erase()

Actually update the changed indexes in the map portion of `MapVector`
when erasing from the middle.  Add a unit test that checks for this.

Note that `MapVector::erase()` is a linear time operation (it was and
still is).  I'll commit a new method in a moment called
`MapVector::remove_if()` that deletes multiple entries in linear time,
which should be slightly less painful.

Modified:
    llvm/trunk/docs/ProgrammersManual.rst
    llvm/trunk/include/llvm/ADT/MapVector.h
    llvm/trunk/unittests/ADT/MapVectorTest.cpp

Modified: llvm/trunk/docs/ProgrammersManual.rst
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/docs/ProgrammersManual.rst?rev=213084&r1=213083&r2=213084&view=diff
==============================================================================
--- llvm/trunk/docs/ProgrammersManual.rst (original)
+++ llvm/trunk/docs/ProgrammersManual.rst Tue Jul 15 13:32:30 2014
@@ -1442,7 +1442,7 @@ iteration over maps of pointers.
 
 It is implemented by mapping from key to an index in a vector of key,value
 pairs.  This provides fast lookup and iteration, but has two main drawbacks: The
-key is stored twice and it doesn't support removing elements.
+key is stored twice and removing elements takes linear time.
 
 .. _dss_inteqclasses:
 

Modified: llvm/trunk/include/llvm/ADT/MapVector.h
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/include/llvm/ADT/MapVector.h?rev=213084&r1=213083&r2=213084&view=diff
==============================================================================
--- llvm/trunk/include/llvm/ADT/MapVector.h (original)
+++ llvm/trunk/include/llvm/ADT/MapVector.h Tue Jul 15 13:32:30 2014
@@ -125,12 +125,26 @@ public:
   }
 
   /// \brief Remove the element given by Iterator.
+  ///
   /// Returns an iterator to the element following the one which was removed,
   /// which may be end().
+  ///
+  /// \note This is a deceivingly expensive operation (linear time).  It's
+  /// usually better to use \a remove_if() if possible.
   typename VectorType::iterator erase(typename VectorType::iterator Iterator) {
-    typename MapType::iterator MapIterator = Map.find(Iterator->first);
-    Map.erase(MapIterator);
-    return Vector.erase(Iterator);
+    Map.erase(Iterator->first);
+    auto Next = Vector.erase(Iterator);
+    if (Next == Vector.end())
+      return Next;
+
+    // Update indices in the map.
+    size_t Index = Next - Vector.begin();
+    for (auto &I : Map) {
+      assert(I.second != Index && "Index was already erased!");
+      if (I.second > Index)
+        --I.second;
+    }
+    return Next;
   }
 };
 

Modified: llvm/trunk/unittests/ADT/MapVectorTest.cpp
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/unittests/ADT/MapVectorTest.cpp?rev=213084&r1=213083&r2=213084&view=diff
==============================================================================
--- llvm/trunk/unittests/ADT/MapVectorTest.cpp (original)
+++ llvm/trunk/unittests/ADT/MapVectorTest.cpp Tue Jul 15 13:32:30 2014
@@ -53,3 +53,18 @@ TEST(MapVectorTest, insert_pop) {
   EXPECT_EQ(MV[1], 2);
   EXPECT_EQ(MV[4], 7);
 }
+
+TEST(MapVectorTest, erase) {
+  MapVector<int, int> MV;
+
+  MV.insert(std::make_pair(1, 2));
+  MV.insert(std::make_pair(3, 4));
+  MV.insert(std::make_pair(5, 6));
+  ASSERT_EQ(MV.size(), 3u);
+
+  MV.erase(MV.find(1));
+  ASSERT_EQ(MV.size(), 2u);
+  ASSERT_EQ(MV.find(1), MV.end());
+  ASSERT_EQ(MV[3], 4);
+  ASSERT_EQ(MV[5], 6);
+}