1 files changed, 56 insertions, 7 deletions

diff --git a/include/llvm/ADT/MapVector.h b/include/llvm/ADT/MapVector.h
index f6fcb0888de3..4e1fc1527270 100644
--- a/include/llvm/ADT/MapVector.h
+++ b/include/llvm/ADT/MapVector.h
@@ -1,4 +1,4 @@
-//===- llvm/ADT/MapVector.h - Map with deterministic value order *- C++ -*-===//
+//===- llvm/ADT/MapVector.h - Map w/ deterministic value order --*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
@@ -17,9 +17,7 @@
 #ifndef LLVM_ADT_MAPVECTOR_H
 #define LLVM_ADT_MAPVECTOR_H
 
-#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/STLExtras.h"
 #include <vector>
 
 namespace llvm {
@@ -31,7 +29,7 @@ template<typename KeyT, typename ValueT,
          typename MapType = llvm::DenseMap<KeyT, unsigned>,
          typename VectorType = std::vector<std::pair<KeyT, ValueT> > >
 class MapVector {
-  typedef typename VectorType::size_type SizeType;
+  typedef typename VectorType::size_type size_type;
 
   MapType Map;
   VectorType Vector;
@@ -40,7 +38,7 @@ public:
   typedef typename VectorType::iterator iterator;
   typedef typename VectorType::const_iterator const_iterator;
 
-  SizeType size() const {
+  size_type size() const {
     return Vector.size();
   }
 
@@ -97,12 +95,12 @@ public:
     if (Result.second) {
       Vector.push_back(std::make_pair(KV.first, KV.second));
       I = Vector.size() - 1;
-      return std::make_pair(llvm::prior(end()), true);
+      return std::make_pair(std::prev(end()), true);
     }
     return std::make_pair(begin() + I, false);
   }
 
-  unsigned count(const KeyT &Key) const {
+  size_type count(const KeyT &Key) const {
     typename MapType::const_iterator Pos = Map.find(Key);
     return Pos == Map.end()? 0 : 1;
   }
@@ -125,8 +123,59 @@ public:
     Map.erase(Pos);
     Vector.pop_back();
   }
+
+  /// \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) {
+    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;
+  }
+
+  /// \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