diff --git a/llvm/include/llvm/ADT/MapVector.h b/llvm/include/llvm/ADT/MapVector.h --- a/llvm/include/llvm/ADT/MapVector.h +++ b/llvm/include/llvm/ADT/MapVector.h @@ -46,6 +46,7 @@ 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; @@ -198,6 +199,19 @@ 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 diff --git a/llvm/unittests/ADT/MapVectorTest.cpp b/llvm/unittests/ADT/MapVectorTest.cpp --- a/llvm/unittests/ADT/MapVectorTest.cpp +++ b/llvm/unittests/ADT/MapVectorTest.cpp @@ -100,6 +100,25 @@ 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) { @@ -257,6 +276,25 @@ 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) { @@ -375,6 +413,25 @@ 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) {