diff --git a/llvm/include/llvm/ADT/STLExtras.h b/llvm/include/llvm/ADT/STLExtras.h index 78b7e94c2b3a1..dc0443c9244be 100644 --- a/llvm/include/llvm/ADT/STLExtras.h +++ b/llvm/include/llvm/ADT/STLExtras.h @@ -325,6 +325,14 @@ template bool hasSingleElement(ContainerTy &&C) { return B != E && std::next(B) == E; } +/// Asserts that the given container has a single element and returns that +/// element. +template +decltype(auto) getSingleElement(ContainerTy &&C) { + assert(hasSingleElement(C) && "expected container with single element"); + return *adl_begin(C); +} + /// Return a range covering \p RangeOrContainer with the first N elements /// excluded. template auto drop_begin(T &&RangeOrContainer, size_t N = 1) { diff --git a/llvm/unittests/ADT/STLExtrasTest.cpp b/llvm/unittests/ADT/STLExtrasTest.cpp index dbb094b0a3088..0101be47a6869 100644 --- a/llvm/unittests/ADT/STLExtrasTest.cpp +++ b/llvm/unittests/ADT/STLExtrasTest.cpp @@ -1016,6 +1016,42 @@ TEST(STLExtrasTest, hasSingleElement) { EXPECT_FALSE(hasSingleElement(S)); } +TEST(STLExtrasTest, getSingleElement) { + // Test const and non-const containers. + const std::vector V1 = {7}; + EXPECT_EQ(getSingleElement(V1), 7); + std::vector V2 = {8}; + EXPECT_EQ(getSingleElement(V2), 8); + + // Test LLVM container. + SmallVector V3{9}; + EXPECT_EQ(getSingleElement(V3), 9); + + // Test that the returned element is a reference. + getSingleElement(V3) = 11; + EXPECT_EQ(V3[0], 11); + + // Test non-random access container. + std::list L1 = {10}; + EXPECT_EQ(getSingleElement(L1), 10); + + // Make sure that we use the `begin`/`end` functions from `some_namespace`, + // using ADL. + some_namespace::some_struct S; + S.data = V2; + EXPECT_EQ(getSingleElement(S), 8); + +#if defined(GTEST_HAS_DEATH_TEST) && !defined(NDEBUG) + // Make sure that we crash on empty or too many elements. + SmallVector V4; + EXPECT_DEATH(getSingleElement(V4), "expected container with single element"); + SmallVector V5{12, 13, 14}; + EXPECT_DEATH(getSingleElement(V5), "expected container with single element"); + std::list L2; + EXPECT_DEATH(getSingleElement(L2), "expected container with single element"); +#endif +} + TEST(STLExtrasTest, hasNItems) { const std::list V0 = {}, V1 = {1}, V2 = {1, 2}; const std::list V3 = {1, 3, 5};