[llvm] [STLForwardCompat] Improve category handling in transformOptional (PR #149539)

[Krzysztof Parzyszek via llvm-commits]

https://github.com/kparzysz updated https://github.com/llvm/llvm-project/pull/149539

>From 69a16974093e1c1d84abae6653589b0569771a6b Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Fri, 18 Jul 2025 09:15:42 -0500
Subject: [PATCH 1/3] [STLForwardCompat] Improve category handling in
 transformOptional

The old version would prefer the "const &" overload over the "&&" one
unless the former was not allowed in the given situation. In particular,
if the function passed was "[](auto &&)" the argument would be "const &"
even if the value passed to transformOptional was an rvalue reference.

This version improves the handling of expression categories, and the
lambda argument category will reflect the argument category in the
above scenario.
---
 llvm/include/llvm/ADT/STLForwardCompat.h    | 22 +++++-----------
 llvm/unittests/ADT/STLForwardCompatTest.cpp | 29 +++++++++++++++++++++
 2 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/llvm/include/llvm/ADT/STLForwardCompat.h b/llvm/include/llvm/ADT/STLForwardCompat.h
index 7bd2c8705f393..c4bb5cd713139 100644
--- a/llvm/include/llvm/ADT/STLForwardCompat.h
+++ b/llvm/include/llvm/ADT/STLForwardCompat.h
@@ -55,21 +55,13 @@ using type_identity_t // NOLINT(readability-identifier-naming)
 
 // TODO: Remove this in favor of std::optional<T>::transform once we switch to
 // C++23.
-template <typename T, typename Function>
-auto transformOptional(const std::optional<T> &O, const Function &F)
-    -> std::optional<decltype(F(*O))> {
-  if (O)
-    return F(*O);
-  return std::nullopt;
-}
-
-// TODO: Remove this in favor of std::optional<T>::transform once we switch to
-// C++23.
-template <typename T, typename Function>
-auto transformOptional(std::optional<T> &&O, const Function &F)
-    -> std::optional<decltype(F(*std::move(O)))> {
-  if (O)
-    return F(*std::move(O));
+template <typename Optional, typename Function,
+          typename Value = typename llvm::remove_cvref_t<Optional>::value_type>
+auto transformOptional(Optional &&O, Function &&F)
+    -> std::optional<std::invoke_result_t<Function, Value>> {
+  if (O) {
+    return F(*std::forward<Optional>(O));
+  }
   return std::nullopt;
 }
 
diff --git a/llvm/unittests/ADT/STLForwardCompatTest.cpp b/llvm/unittests/ADT/STLForwardCompatTest.cpp
index e3d500aa7b55a..b8b2ca1c85054 100644
--- a/llvm/unittests/ADT/STLForwardCompatTest.cpp
+++ b/llvm/unittests/ADT/STLForwardCompatTest.cpp
@@ -10,6 +10,11 @@
 #include "CountCopyAndMove.h"
 #include "gtest/gtest.h"
 
+#include <optional>
+#include <tuple>
+#include <type_traits>
+#include <utility>
+
 namespace {
 
 template <typename T>
@@ -142,6 +147,30 @@ TEST(TransformTest, MoveTransformLlvm) {
   EXPECT_EQ(0, CountCopyAndMove::Destructions);
 }
 
+template <typename T> constexpr bool IsRvalueReference(T &&arg) {
+  return std::is_rvalue_reference_v<decltype(arg)>;
+}
+
+TEST(TransformTest, TransformCategory) {
+  struct StructA {
+    int x;
+  };
+  struct StructB : StructA {
+    StructB(StructA &&A) : StructA(std::move(A)) {}
+  };
+
+  std::optional<StructA> A{StructA{}};
+  llvm::transformOptional(A, [](auto &&s) {
+    EXPECT_FALSE(std::is_rvalue_reference_v<decltype(s)>);
+    return StructB{std::move(s)};
+  });
+
+  llvm::transformOptional(std::move(A), [](auto &&s) {
+    EXPECT_TRUE(std::is_rvalue_reference_v<decltype(s)>);
+    return StructB{std::move(s)};
+  });
+}
+
 TEST(TransformTest, ToUnderlying) {
   enum E { A1 = 0, B1 = -1 };
   static_assert(llvm::to_underlying(A1) == 0);

>From 377b12dcfb56dd213d566932cdb507eb7676dc7c Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Fri, 18 Jul 2025 11:28:40 -0500
Subject: [PATCH 2/3] Delete accidentally left over function

---
 llvm/unittests/ADT/STLForwardCompatTest.cpp | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/llvm/unittests/ADT/STLForwardCompatTest.cpp b/llvm/unittests/ADT/STLForwardCompatTest.cpp
index b8b2ca1c85054..4a8f53cf72f94 100644
--- a/llvm/unittests/ADT/STLForwardCompatTest.cpp
+++ b/llvm/unittests/ADT/STLForwardCompatTest.cpp
@@ -147,10 +147,6 @@ TEST(TransformTest, MoveTransformLlvm) {
   EXPECT_EQ(0, CountCopyAndMove::Destructions);
 }
 
-template <typename T> constexpr bool IsRvalueReference(T &&arg) {
-  return std::is_rvalue_reference_v<decltype(arg)>;
-}
-
 TEST(TransformTest, TransformCategory) {
   struct StructA {
     int x;

>From 648354fc01d85a5b4688fbb070e401d04c98fb8e Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Fri, 18 Jul 2025 12:05:20 -0500
Subject: [PATCH 3/3] Untrail the return type

---
 llvm/include/llvm/ADT/STLForwardCompat.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/ADT/STLForwardCompat.h b/llvm/include/llvm/ADT/STLForwardCompat.h
index c4bb5cd713139..81b9a685e11d2 100644
--- a/llvm/include/llvm/ADT/STLForwardCompat.h
+++ b/llvm/include/llvm/ADT/STLForwardCompat.h
@@ -57,8 +57,8 @@ using type_identity_t // NOLINT(readability-identifier-naming)
 // C++23.
 template <typename Optional, typename Function,
           typename Value = typename llvm::remove_cvref_t<Optional>::value_type>
-auto transformOptional(Optional &&O, Function &&F)
-    -> std::optional<std::invoke_result_t<Function, Value>> {
+std::optional<std::invoke_result_t<Function, Value>>
+transformOptional(Optional &&O, Function &&F) {
   if (O) {
     return F(*std::forward<Optional>(O));
   }