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