From: Duncan P. N. Exon Smith <dexonsmith@apple.com>
Date: Tue, 15 Jul 2014 20:24:56 +0000 (+0000)
Subject: ADT: Add MapVector::remove_if
X-Git-Url: http://plrg.eecs.uci.edu/git/?p=oota-llvm.git;a=commitdiff_plain;h=c12ce2b66186404177a3a9f854df0272beb52d81

ADT: Add MapVector::remove_if

Add a `MapVector::remove_if()` that erases items in bulk in linear time,
as opposed to quadratic time for repeated calls to `MapVector::erase()`.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@213090 91177308-0d34-0410-b5e6-96231b3b80d8
---

diff --git a/docs/ProgrammersManual.rst b/docs/ProgrammersManual.rst
index e828e6bf501..40eabbe442d 100644
--- a/docs/ProgrammersManual.rst
+++ b/docs/ProgrammersManual.rst
@@ -1441,8 +1441,10 @@ order, making it an easy (but somewhat expensive) solution for non-deterministic
 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 removing elements takes linear time.
+pairs.  This provides fast lookup and iteration, but has two main drawbacks:
+the key is stored twice and removing elements takes linear time.  If it is
+necessary to remove elements, it's best to remove them in bulk using
+``remove_if()``.
 
 .. _dss_inteqclasses:
 
diff --git a/include/llvm/ADT/MapVector.h b/include/llvm/ADT/MapVector.h
index 0fa68245ff4..4e1fc152727 100644
--- a/include/llvm/ADT/MapVector.h
+++ b/include/llvm/ADT/MapVector.h
@@ -146,8 +146,36 @@ public:
     }
     return Next;
   }
+
+  /// \brief Remove the elements that match the predicate.
+  ///
+  /// Erase all elements that match \c Pred in a single pass.  Takes linear
+  /// time.
+  template <class Predicate> void remove_if(Predicate Pred);
 };
 
+template <typename KeyT, typename ValueT, typename MapType, typename VectorType>
+template <class Function>
+void MapVector<KeyT, ValueT, MapType, VectorType>::remove_if(Function Pred) {
+  auto O = Vector.begin();
+  for (auto I = O, E = Vector.end(); I != E; ++I) {
+    if (Pred(*I)) {
+      // Erase from the map.
+      Map.erase(I->first);
+      continue;
+    }
+
+    if (I != O) {
+      // Move the value and update the index in the map.
+      *O = std::move(*I);
+      Map[O->first] = O - Vector.begin();
+    }
+    ++O;
+  }
+  // Erase trailing entries in the vector.
+  Vector.erase(O, Vector.end());
+}
+
 } // end namespace llvm
 
 #endif
diff --git a/unittests/ADT/MapVectorTest.cpp b/unittests/ADT/MapVectorTest.cpp
index dd84e7ebd13..92f0dc41910 100644
--- a/unittests/ADT/MapVectorTest.cpp
+++ b/unittests/ADT/MapVectorTest.cpp
@@ -68,3 +68,24 @@ TEST(MapVectorTest, erase) {
   ASSERT_EQ(MV[3], 4);
   ASSERT_EQ(MV[5], 6);
 }
+
+TEST(MapVectorTest, remove_if) {
+  MapVector<int, int> MV;
+
+  MV.insert(std::make_pair(1, 11));
+  MV.insert(std::make_pair(2, 12));
+  MV.insert(std::make_pair(3, 13));
+  MV.insert(std::make_pair(4, 14));
+  MV.insert(std::make_pair(5, 15));
+  MV.insert(std::make_pair(6, 16));
+  ASSERT_EQ(MV.size(), 6u);
+
+  MV.remove_if([](const std::pair<int, int> &Val) { return Val.second % 2; });
+  ASSERT_EQ(MV.size(), 3u);
+  ASSERT_EQ(MV.find(1), MV.end());
+  ASSERT_EQ(MV.find(3), MV.end());
+  ASSERT_EQ(MV.find(5), MV.end());
+  ASSERT_EQ(MV[2], 12);
+  ASSERT_EQ(MV[4], 14);
+  ASSERT_EQ(MV[6], 16);
+}