[llvm] [ADT] Fix llvm::concat_iterator for `ValueT == common_base_class *` (PR #144744)

[Javier Lopez-Gomez via llvm-commits]

https://github.com/jalopezg-git updated https://github.com/llvm/llvm-project/pull/144744

>From 8a774b8b631ee42764fb57a34c1349b5d6d32c39 Mon Sep 17 00:00:00 2001
From: Javier Lopez-Gomez <javier.lopez.gomez at proton.me>
Date: Thu, 19 Jun 2025 23:31:01 +0200
Subject: [PATCH] [ADT] Fix llvm::concat_iterator for `ValueT ==
 common_base_class *`

Fix llvm::concat_iterator for the case of `ValueT` being a pointer
to a common base class to which the result of dereferencing any
iterator in `ItersT` can be casted to.
---
 llvm/include/llvm/ADT/STLExtras.h    | 22 +++++++++++++++-------
 llvm/unittests/ADT/STLExtrasTest.cpp | 14 ++++++++++++++
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/llvm/include/llvm/ADT/STLExtras.h b/llvm/include/llvm/ADT/STLExtras.h
index eea06cfb99ba2..37f895d2a5690 100644
--- a/llvm/include/llvm/ADT/STLExtras.h
+++ b/llvm/include/llvm/ADT/STLExtras.h
@@ -1032,13 +1032,15 @@ class concat_iterator
 
   static constexpr bool ReturnsByValue =
       !(std::is_reference_v<decltype(*std::declval<IterTs>())> && ...);
+  static constexpr bool NoTypeAdapterNeeded =
+      (std::is_pointer_v<IterTs> && ...) && std::is_pointer_v<ValueT>;
 
   using reference_type =
-      typename std::conditional_t<ReturnsByValue, ValueT, ValueT &>;
+      typename std::conditional_t<std::is_pointer_v<ValueT>, ValueT, ValueT &>;
 
-  using handle_type =
-      typename std::conditional_t<ReturnsByValue, std::optional<ValueT>,
-                                  ValueT *>;
+  using handle_type = typename std::conditional_t<
+      NoTypeAdapterNeeded, ValueT,
+      std::conditional_t<ReturnsByValue, std::optional<ValueT>, ValueT *>>;
 
   /// We store both the current and end iterators for each concatenated
   /// sequence in a tuple of pairs.
@@ -1088,7 +1090,9 @@ class concat_iterator
     if (Begin == End)
       return {};
 
-    if constexpr (ReturnsByValue)
+    if constexpr (NoTypeAdapterNeeded)
+      return Begin;
+    else if constexpr (ReturnsByValue)
       return *Begin;
     else
       return &*Begin;
@@ -1105,8 +1109,12 @@ class concat_iterator
 
     // Loop over them, and return the first result we find.
     for (auto &GetHelperFn : GetHelperFns)
-      if (auto P = (this->*GetHelperFn)())
-        return *P;
+      if (auto P = (this->*GetHelperFn)()) {
+        if constexpr (NoTypeAdapterNeeded)
+          return P;
+        else
+          return *P;
+      }
 
     llvm_unreachable("Attempted to get a pointer from an end concat iterator!");
   }
diff --git a/llvm/unittests/ADT/STLExtrasTest.cpp b/llvm/unittests/ADT/STLExtrasTest.cpp
index 286cfa745fd14..7e9b6f19a3d32 100644
--- a/llvm/unittests/ADT/STLExtrasTest.cpp
+++ b/llvm/unittests/ADT/STLExtrasTest.cpp
@@ -398,6 +398,8 @@ struct some_struct {
   std::string swap_val;
 };
 
+struct derives_from_some_struct : some_struct {};
+
 std::vector<int>::const_iterator begin(const some_struct &s) {
   return s.data.begin();
 }
@@ -532,6 +534,18 @@ TEST(STLExtrasTest, ConcatRangeADL) {
   EXPECT_THAT(concat<const int>(S0, S1), ElementsAre(1, 2, 3, 4));
 }
 
+TEST(STLExtrasTest, ConcatRangePtrToDerivedClass) {
+  some_namespace::some_struct S0{};
+  some_namespace::derives_from_some_struct S1{};
+  SmallVector<some_namespace::some_struct *> V0{&S0};
+  SmallVector<some_namespace::derives_from_some_struct *> V1{&S1, &S1};
+
+  // Use concat over ranges of pointers to different (but related) types.
+  EXPECT_THAT(concat<some_namespace::some_struct *>(V0, V1),
+              ElementsAre(&S0, static_cast<some_namespace::some_struct *>(&S1),
+                          static_cast<some_namespace::some_struct *>(&S1)));
+}
+
 TEST(STLExtrasTest, MakeFirstSecondRangeADL) {
   // Make sure that we use the `begin`/`end` functions from `some_namespace`,
   // using ADL.