Index: include/llvm/ADT/SetVector.h =================================================================== --- include/llvm/ADT/SetVector.h +++ include/llvm/ADT/SetVector.h @@ -151,6 +151,20 @@ return false; } + /// Erase a single element from the set vector. + /// \returns an iterator pointing to the next element that followed the + /// element erased. This is the end of the SetVector if the last element is + /// erased. + iterator erase(iterator I) { + const key_type &V = *I; + typename vector_type::iterator VI = + std::find(vector_.begin(), vector_.end(), V); + assert(VI != vector_.end() && "Iterator to erase is out of bounds."); + assert(set_.count(V) && "Corrupted SetVector instances!"); + set_.erase(V); + return vector_.erase(VI); + } + /// \brief Remove items from the set vector based on a predicate function. /// /// This is intended to be equivalent to the following code, if we could Index: unittests/ADT/CMakeLists.txt =================================================================== --- unittests/ADT/CMakeLists.txt +++ unittests/ADT/CMakeLists.txt @@ -32,6 +32,7 @@ PostOrderIteratorTest.cpp RangeAdapterTest.cpp SCCIteratorTest.cpp + SetVectorTest.cpp SmallPtrSetTest.cpp SmallStringTest.cpp SmallVectorTest.cpp Index: unittests/ADT/SetVectorTest.cpp =================================================================== --- /dev/null +++ unittests/ADT/SetVectorTest.cpp @@ -0,0 +1,100 @@ +//===- llvm/unittest/ADT/SetVector.cpp ------------------------------===// +// +// The LLVM Compiler Infrastructure +// +// This file is distributed under the University of Illinois Open Source +// License. See LICENSE.TXT for details. +// +//===----------------------------------------------------------------------===// +// +// SetVector unit tests. +// +//===----------------------------------------------------------------------===// + +#include "llvm/ADT/SetVector.h" +#include "gtest/gtest.h" + +using namespace llvm; + +void checkEraseAndIteratorsForSetVector() { + SetVector S; + + S.insert(0); + S.insert(1); + S.insert(2); + + auto B = S.begin(); + auto M = std::next(B); + + EXPECT_EQ(B, S.begin()); + EXPECT_EQ(M, std::next(B)); + EXPECT_TRUE(3 == S.size()); + EXPECT_TRUE(*B == 0); + EXPECT_TRUE(*M == 1); + + // Erase will return the iterator pointing to the next element after the + // erased one. + M = S.erase(M); + EXPECT_EQ(B, S.begin()); + EXPECT_EQ(M, std::next(B)); + EXPECT_TRUE(2 == S.size()); + EXPECT_TRUE(*B == 0); + EXPECT_TRUE(*M == 2); + + M = S.erase(M); + EXPECT_EQ(B, S.begin()); + EXPECT_EQ(M, std::next(B)); + EXPECT_TRUE(1 == S.size()); + EXPECT_EQ(M, S.end()); + EXPECT_TRUE(*B == 0); + + M = S.erase(B); + EXPECT_EQ(S.begin(), S.end()); + EXPECT_EQ(M, S.end()); + EXPECT_TRUE(0 == S.size()); +} + +void checkEraseAndIteratorsForSmallSetVector() { + SmallSetVector S; + + S.insert(0); + S.insert(1); + S.insert(2); + + auto B = S.begin(); + auto M = std::next(B); + + EXPECT_EQ(B, S.begin()); + EXPECT_EQ(M, std::next(B)); + EXPECT_TRUE(3 == S.size()); + EXPECT_TRUE(*B == 0); + EXPECT_TRUE(*M == 1); + + // Erase will return the iterator pointing to the next element after the + // erased one. + M = S.erase(M); + EXPECT_EQ(B, S.begin()); + EXPECT_EQ(M, std::next(B)); + EXPECT_TRUE(2 == S.size()); + EXPECT_TRUE(*B == 0); + EXPECT_TRUE(*M == 2); + + M = S.erase(M); + EXPECT_EQ(B, S.begin()); + EXPECT_EQ(M, std::next(B)); + EXPECT_TRUE(1 == S.size()); + EXPECT_EQ(M, S.end()); + EXPECT_TRUE(*B == 0); + + M = S.erase(B); + EXPECT_EQ(S.begin(), S.end()); + EXPECT_EQ(M, S.end()); + EXPECT_TRUE(0 == S.size()); +} + +TEST(SetVector, EraseTest) { + // Test for SetVector. + checkEraseAndIteratorsForSetVector(); + // Test for SmallSetVector. + checkEraseAndIteratorsForSmallSetVector(); +}