[llvm][ADT] Add a method to MapVector for erasing a range of elements.

[labrinea]

Summary

We currently support two ways to erase elements in MapVector: either erase a single element in linear time, which is rather expensive, or erase multiple elements based on a predicate. However, we may want to erase a range of contiguous elements. Here's a motivating example, where the current API is suboptimal:

MapVector.remove_if([&MapVector, N](const auto &Entry) {
  auto Iter = MapVector.find(Entry.first);
  return std::distance(Iter, MapVector.end()) > N;
});

Instead we could do:

N = MapVector.size() - N;
MapVector.erase(MapVector.begin(), MapVector.begin() + N);

[efriedma]

This reduces the total cost of the erase from O(N*M), where N is the total number of elements in the map and M is the number of elements you're removing, to O(N). That still seems like a significant performance hazard. I think we should discourage this sort of usage, in favor of a container where the removal is O(M), like std::map.

That said, depending on the usage, that might be okay, so I guess I'm not strongly against this.

llvm/include/llvm/ADT/MapVector.h

207

Asserting that you can't erase an empty range seems weird to me; standard library containers treat it as a no-op.

[ChuanqiXu]

I agree with the complexity analysis here. I don't find a better solution if we don't add new data structures like a map from VectorType::iterator to MapType::iterator. It might benefit for erase but it is redundant for other methods. I think this patch is beneficial. And I feel it is good to comment that MapVector is not good at erase().

llvm/include/llvm/ADT/MapVector.h

49

This looks not necessary.

207

We could also assert begin() < S && S < E && E < end().

[labrinea]

I didn't end up using this patch in https://reviews.llvm.org/D119880, but if you still find it useful I can address the review comments, otherwise I'll abandon it.

llvm/include/llvm/ADT/MapVector.h

class MapVector {
   static_assert(
       std::is_integral<typename MapType::mapped_type>::value,
       "The mapped_type of the specified Map must be an integral type");
 
 public:
   using key_type = KeyT;
   using value_type = typename VectorType::value_type;
   using size_type = typename VectorType::size_type;
+  using difference_type = typename VectorType::difference_type;
 
   using iterator = typename VectorType::iterator;
   using const_iterator = typename VectorType::const_iterator;
   using reverse_iterator = typename VectorType::reverse_iterator;
   using const_reverse_iterator = typename VectorType::const_reverse_iterator;
 
   /// Clear the MapVector and return the underlying vector.
   VectorType takeVector() {
public:
   size_type erase(const KeyT &Key) {
     auto Iterator = find(Key);
     if (Iterator == end())
       return 0;
     erase(Iterator);
     return 1;
   }
 
+  /// Remove elements in range.
+  ///
+  /// Returns an iterator to the element following the range which got removed.
+  iterator erase(iterator S, iterator E) {
+    difference_type NumElts = std::distance(S, E);
+    assert(NumElts > 0 && "Range end is not greater than range start!");
+    for (auto I = S; I != E; ++I)
+      Map.erase(I->first);
+    for (auto I = E; I != Vector.end(); ++I)
+      Map[I->first] -= (unsigned)NumElts;
+    return Vector.erase(S, E);
+  }
+
   /// 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>

llvm/unittests/ADT/MapVectorTest.cpp

TEST(MapVectorTest, erase) {
 
   ASSERT_EQ(MV.erase(3), 1u);
   ASSERT_EQ(MV.size(), 1u);
   ASSERT_EQ(MV.find(3), MV.end());
   ASSERT_EQ(MV[5], 6);
 
   ASSERT_EQ(MV.erase(79), 0u);
   ASSERT_EQ(MV.size(), 1u);
+
+  // Test range removal.
+  MV.insert(std::make_pair(7, 8));
+  MV.insert(std::make_pair(9, 10));
+  MV.insert(std::make_pair(11, 12));
+  auto I = MV.find(7);
+  I = MV.erase(I, I+2);
+  // The iterator points to the element following the range which was removed.
+  ASSERT_EQ(I->first, 11);
+  ASSERT_EQ(I->second, 12);
+  // The entries have been removed.
+  ASSERT_EQ(MV.find(7), MV.end());
+  ASSERT_EQ(MV.find(9), MV.end());
+  ASSERT_EQ(MV.size(), 2u);
+  // The indices have been updated.
+  ASSERT_EQ(MV.find(5), MV.begin());
+  ASSERT_EQ(MV.find(11), MV.begin()+1);
+  ASSERT_EQ(MV[5], 6);
+  ASSERT_EQ(MV[11], 12);
 }
 
 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));
TEST(SmallMapVectorSmallTest, erase) {
 
   ASSERT_EQ(MV.erase(3), 1u);
   ASSERT_EQ(MV.size(), 1u);
   ASSERT_EQ(MV.find(3), MV.end());
   ASSERT_EQ(MV[5], 6);
 
   ASSERT_EQ(MV.erase(79), 0u);
   ASSERT_EQ(MV.size(), 1u);
+
+  // Test range removal.
+  MV.insert(std::make_pair(7, 8));
+  MV.insert(std::make_pair(9, 10));
+  MV.insert(std::make_pair(11, 12));
+  auto I = MV.find(7);
+  I = MV.erase(I, I+2);
+  // The iterator points to the element following the range which was removed.
+  ASSERT_EQ(I->first, 11);
+  ASSERT_EQ(I->second, 12);
+  // The entries have been removed.
+  ASSERT_EQ(MV.find(7), MV.end());
+  ASSERT_EQ(MV.find(9), MV.end());
+  ASSERT_EQ(MV.size(), 2u);
+  // The indices have been updated.
+  ASSERT_EQ(MV.find(5), MV.begin());
+  ASSERT_EQ(MV.find(11), MV.begin()+1);
+  ASSERT_EQ(MV[5], 6);
+  ASSERT_EQ(MV[11], 12);
 }
 
 TEST(SmallMapVectorSmallTest, remove_if) {
   SmallMapVector<int, int, 32> MV;
 
   MV.insert(std::make_pair(1, 11));
   MV.insert(std::make_pair(2, 12));
   MV.insert(std::make_pair(3, 13));
TEST(SmallMapVectorLargeTest, erase) {
 
   ASSERT_EQ(MV.erase(3), 1u);
   ASSERT_EQ(MV.size(), 1u);
   ASSERT_EQ(MV.find(3), MV.end());
   ASSERT_EQ(MV[5], 6);
 
   ASSERT_EQ(MV.erase(79), 0u);
   ASSERT_EQ(MV.size(), 1u);
+
+  // Test range removal.
+  MV.insert(std::make_pair(7, 8));
+  MV.insert(std::make_pair(9, 10));
+  MV.insert(std::make_pair(11, 12));
+  auto I = MV.find(7);
+  I = MV.erase(I, I+2);
+  // The iterator points to the element following the range which was removed.
+  ASSERT_EQ(I->first, 11);
+  ASSERT_EQ(I->second, 12);
+  // The entries have been removed.
+  ASSERT_EQ(MV.find(7), MV.end());
+  ASSERT_EQ(MV.find(9), MV.end());
+  ASSERT_EQ(MV.size(), 2u);
+  // The indices have been updated.
+  ASSERT_EQ(MV.find(5), MV.begin());
+  ASSERT_EQ(MV.find(11), MV.begin()+1);
+  ASSERT_EQ(MV[5], 6);
+  ASSERT_EQ(MV[11], 12);
 }
 
 TEST(SmallMapVectorLargeTest, remove_if) {
   SmallMapVector<int, int, 1> MV;
 
   MV.insert(std::make_pair(1, 11));
   MV.insert(std::make_pair(2, 12));
   MV.insert(std::make_pair(3, 13));