ADT: Add MapVector::remove_if

[oota-llvm.git] / include / llvm / ADT / MapVector.h

diff --git a/include/llvm/ADT/MapVector.h b/include/llvm/ADT/MapVector.h
index 0fa68245ff4fa9f495ebe110b8c7b385232fa1a4..4e1fc1527270f27e06a7094602a887bbc4041d7b 100644 (file)
--- 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