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,33 @@ +//===- 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; + +TEST(SetVector, EraseTest) { + SetVector S; + S.insert(0); + S.insert(1); + S.insert(2); + + auto I = S.erase(std::next(S.begin())); + + // Test that the returned iterator is the expected one-after-erase + // and the size/contents is the expected sequence {0, 2}. + EXPECT_EQ(std::next(S.begin()), I); + EXPECT_EQ(2u, S.size()); + EXPECT_EQ(0, *S.begin()); + EXPECT_EQ(2, *std::next(S.begin())); +}