From 82eb816a4636471c10bc380a38d8097bc2b6234f Mon Sep 17 00:00:00 2001
From: Francesco Zoffoli <fzoffoli@meta.com>
Date: Wed, 15 Jan 2025 16:48:37 -0800
Subject: [PATCH] Make `visit_union` take the callable by universal reference

Summary:
`visit_union` is taking the callable by value, forcing a copy.

This is a problem because
1. it's inconsistent with `std::visit`, which instead uses universal references. This results in surprises when using it
2. when the callable returns a coroutine, this guarantees that the coroutine state is destoryed before the coroutine finish execution

```
int a = 0
co_await thrift::visit_union(my_union, [&](auto, auto) -> Task<void> { a += 1; co_return; });
```
This results in a dangling reference in the old implementation, and it's correct in the new one.
This is consistent with the coro wiki: https://www.internalfb.com/wiki/Coro/#lambdas

Note: To be consistent the same should  probably be done to also `visit_by_thrift_field_metadata` and the other visit functions

Triggering problem: https://fb.workplace.com/groups/192968587972839/posts/1656198051649878/?comment_id=1657068398229510&reply_comment_id=1657106148225735

Reviewed By: Mizuchi

Differential Revision: D68151611

fbshipit-source-id: a05c53043642b4274dbab63316e6dcb7d21de9f4
---
 thrift/lib/cpp2/visitation/visit_union.h               | 4 ++--
 thrift/test/reflection/visitation_visit_union_test.cpp | 6 ++++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/thrift/lib/cpp2/visitation/visit_union.h b/thrift/lib/cpp2/visitation/visit_union.h
index 4c5139a0021..c80f5696c29 100644
--- a/thrift/lib/cpp2/visitation/visit_union.h
+++ b/thrift/lib/cpp2/visitation/visit_union.h
@@ -44,8 +44,8 @@ struct VisitUnion {
  * @param f a callable that accepts all member types from union
  */
 template <typename T, typename F>
-decltype(auto) visit_union(T&& t, F f) {
+decltype(auto) visit_union(T&& t, F&& f) {
   return apache::thrift::detail::VisitUnion<folly::remove_cvref_t<T>>()(
-      detail::MetadataForwarder<T, F>{std::move(f)}, static_cast<T&&>(t));
+      detail::MetadataForwarder<T, F>{std::forward<F>(f)}, static_cast<T&&>(t));
 }
 } // namespace apache::thrift
diff --git a/thrift/test/reflection/visitation_visit_union_test.cpp b/thrift/test/reflection/visitation_visit_union_test.cpp
index a209c1532aa..ac3382e4c1e 100644
--- a/thrift/test/reflection/visitation_visit_union_test.cpp
+++ b/thrift/test/reflection/visitation_visit_union_test.cpp
@@ -37,7 +37,7 @@ struct VisitUnionAdapter {
 
 struct ForEachFieldAdapter {
   template <class T, class F>
-  void operator()(T&& t, F f) const {
+  void operator()(T&& t, F&& f) const {
     apache::thrift::for_each_field(
         std::forward<T>(t), [&](auto&& meta, auto&& ref) {
           if (folly::to_underlying(t.getType()) == meta.id_ref()) {
@@ -148,10 +148,12 @@ TYPED_TEST(VisitUnionTest, PassCallableByReference) {
   TestPassCallableByValue f;
   Basic a;
   a.int64_ref() = 42;
-  TestFixture::adapter(a, f);
+  TestFixture::adapter(a, folly::copy(f));
   EXPECT_EQ(f.i, 0);
   TestFixture::adapter(a, std::ref(f));
   EXPECT_EQ(f.i, 1);
+  TestFixture::adapter(a, f);
+  EXPECT_EQ(f.i, 2);
 }
 
 template <class T>