[STLForwardCompat] Improve category handling in transformOptional (#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.

2 files changed, 32 insertions, 15 deletions

diff --git a/llvm/include/llvm/ADT/STLForwardCompat.h b/llvm/include/llvm/ADT/STLForwardCompat.h
index 7bd2c87..81b9a68 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>
+std::optional<std::invoke_result_t<Function, Value>>
+transformOptional(Optional &&O, Function &&F) {
+  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 e3d500a..4a8f53c 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,26 @@ TEST(TransformTest, MoveTransformLlvm) {
   EXPECT_EQ(0, CountCopyAndMove::Destructions);
 }
 
+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);