[llvm] [llvm][ADT] Add `getSingleElement` helper (PR #131508)

[Matthias Springer via llvm-commits]

https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/131508

>From 78eca92aea88e855c081f658e4743ed6397c6c4e Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Sun, 16 Mar 2025 11:21:49 +0100
Subject: [PATCH 1/2] [llvm][ADT] Add `getSingleElement` helper

---
 llvm/include/llvm/ADT/STLExtras.h    |  8 ++++++++
 llvm/unittests/ADT/STLExtrasTest.cpp | 29 ++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+)

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 <typename ContainerTy> 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 <typename ContainerTy>
+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 <typename T> 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..6800cd6eb03d3 100644
--- a/llvm/unittests/ADT/STLExtrasTest.cpp
+++ b/llvm/unittests/ADT/STLExtrasTest.cpp
@@ -1016,6 +1016,35 @@ TEST(STLExtrasTest, hasSingleElement) {
   EXPECT_FALSE(hasSingleElement(S));
 }
 
+TEST(STLExtrasTest, getSingleElement) {
+  // Note: Asserting behavior of getSingleElement cannot be tested because the
+  // program would crash.
+
+  // Test const and non-const containers.
+  const std::vector<int> V1 = {7};
+  EXPECT_EQ(getSingleElement(V1), 7);
+  std::vector<int> V2 = {8};
+  EXPECT_EQ(getSingleElement(V2), 8);
+
+  // Test LLVM container.
+  SmallVector<int> 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<int> 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);
+}
+
 TEST(STLExtrasTest, hasNItems) {
   const std::list<int> V0 = {}, V1 = {1}, V2 = {1, 2};
   const std::list<int> V3 = {1, 3, 5};

>From 4b8b49a2d4ec79e4be34448475f5a9d0cc9fa59a Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Sun, 16 Mar 2025 20:26:02 +0100
Subject: [PATCH 2/2] add negative test

---
 llvm/unittests/ADT/STLExtrasTest.cpp | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/llvm/unittests/ADT/STLExtrasTest.cpp b/llvm/unittests/ADT/STLExtrasTest.cpp
index 6800cd6eb03d3..38c196b00fddb 100644
--- a/llvm/unittests/ADT/STLExtrasTest.cpp
+++ b/llvm/unittests/ADT/STLExtrasTest.cpp
@@ -1017,9 +1017,6 @@ TEST(STLExtrasTest, hasSingleElement) {
 }
 
 TEST(STLExtrasTest, getSingleElement) {
-  // Note: Asserting behavior of getSingleElement cannot be tested because the
-  // program would crash.
-
   // Test const and non-const containers.
   const std::vector<int> V1 = {7};
   EXPECT_EQ(getSingleElement(V1), 7);
@@ -1043,6 +1040,14 @@ TEST(STLExtrasTest, getSingleElement) {
   some_namespace::some_struct S;
   S.data = V2;
   EXPECT_EQ(getSingleElement(S), 8);
+
+  // Make sure that we crash on empty or too many elements.
+  SmallVector<int> V4;
+  EXPECT_DEATH(getSingleElement(V4), "expected container with single element");
+  SmallVector<int> V5{12, 13, 14};
+  EXPECT_DEATH(getSingleElement(V5), "expected container with single element");
+  std::list<int> L2;
+  EXPECT_DEATH(getSingleElement(L2), "expected container with single element");
 }
 
 TEST(STLExtrasTest, hasNItems) {