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

[Krzysztof Parzyszek via llvm-commits]

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

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.

>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] [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);