From 2fa918827fa5d1eac2290b251324708fa2e1936a Mon Sep 17 00:00:00 2001
From: TJ Yin <ytj@meta.com>
Date: Tue, 29 Oct 2024 18:43:48 -0700
Subject: [PATCH] Add DynamicStructPatch::patchIfSet<Tag>(id) and replace the
 existing API.

Summary:
This diff also changed `patchIfSet(badge, id, patch)` to private, just like the static patch:

https://www.internalfb.com/code/fbsource/[39b979388a33844ecef35fedbb3c53bf67e97912]/fbcode/thrift/lib/cpp2/op/detail/StructPatch.h?lines=583-587

Reviewed By: thedavekwon

Differential Revision: D65082475

fbshipit-source-id: 3df7839b51051668a3b44db9788036db93c51c34
---
 .../lib/cpp2/patch/test/DynamicPatchTest.cpp  |  6 ++--
 .../thrift/lib/thrift/detail/DynamicPatch.cpp |  6 ++--
 .../thrift/lib/thrift/detail/DynamicPatch.h   | 29 ++++++++++++++-----
 3 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/third-party/thrift/src/thrift/lib/cpp2/patch/test/DynamicPatchTest.cpp b/third-party/thrift/src/thrift/lib/cpp2/patch/test/DynamicPatchTest.cpp
index c201ea779ebed..6d5267bd36d62 100644
--- a/third-party/thrift/src/thrift/lib/cpp2/patch/test/DynamicPatchTest.cpp
+++ b/third-party/thrift/src/thrift/lib/cpp2/patch/test/DynamicPatchTest.cpp
@@ -802,7 +802,7 @@ TEST(PatchMergeTest, DynamicStructPatch) {
 
   op::I32Patch foo;
   foo += 1;
-  p.patchIfSet(badge, FieldId{1}, DynamicPatch{foo});
+  p.patchIfSet<type::i32_t>(badge, FieldId{1}).merge(foo);
 
   Object obj;
   obj[FieldId(1)].emplace_i32(3);
@@ -819,7 +819,7 @@ TEST(PatchMergeTest, DynamicStructPatch) {
   obj[FieldId(1)].emplace_i32(3);
 
   // patch becomes += 2
-  p.patchIfSet(badge, FieldId{1}, DynamicPatch{foo});
+  p.patchIfSet<type::i32_t>(badge, FieldId{1}).merge(foo);
 
   p.apply(badge, obj);
   EXPECT_EQ(obj[FieldId(1)].as_i32(), 5);
@@ -829,7 +829,7 @@ TEST(PatchMergeTest, DynamicStructPatch) {
   // In dynamic patch, we only use Remove operation to remove field
   // Clear operation will just set field to intrinsic default
   foo.clear();
-  p.patchIfSet(badge, FieldId{1}, DynamicPatch{foo});
+  p.patchIfSet<type::i32_t>(badge, FieldId{1}).merge(foo);
   p.apply(badge, obj);
   EXPECT_EQ(obj[FieldId(1)].as_i32(), 0);
 
diff --git a/third-party/thrift/src/thrift/lib/thrift/detail/DynamicPatch.cpp b/third-party/thrift/src/thrift/lib/thrift/detail/DynamicPatch.cpp
index d5e1aeccf17fd..af1d858db2a3b 100644
--- a/third-party/thrift/src/thrift/lib/thrift/detail/DynamicPatch.cpp
+++ b/third-party/thrift/src/thrift/lib/thrift/detail/DynamicPatch.cpp
@@ -929,7 +929,7 @@ DynamicUnionPatch DiffVisitorBase::diffUnion(
   auto guard = folly::makeGuard([&] { pop(); });
   auto subPatch = diff(badge, src.at(id), dst.at(id));
   if (!subPatch.empty(badge)) {
-    patch.patchIfSet(badge, id, DynamicPatch{std::move(subPatch)});
+    patch.patchIfSet(badge, id).merge(badge, DynamicPatch{std::move(subPatch)});
   }
 
   return patch;
@@ -1002,7 +1002,7 @@ void DiffVisitorBase::diffField(
     auto empty = emptyValue(field.getType());
     auto subPatch = diff(badge, empty, field);
     patch.ensure(badge, id, std::move(empty));
-    patch.patchIfSet(badge, id, DynamicPatch{std::move(subPatch)});
+    patch.patchIfSet(badge, id).merge(badge, DynamicPatch{std::move(subPatch)});
     return;
   }
 
@@ -1010,7 +1010,7 @@ void DiffVisitorBase::diffField(
   auto guard = folly::makeGuard([&] { pop(); });
   auto subPatch = diff(badge, src.at(id), dst.at(id));
   if (!subPatch.empty(badge)) {
-    patch.patchIfSet(badge, id, DynamicPatch{std::move(subPatch)});
+    patch.patchIfSet(badge, id).merge(badge, DynamicPatch{std::move(subPatch)});
   }
 }
 
diff --git a/third-party/thrift/src/thrift/lib/thrift/detail/DynamicPatch.h b/third-party/thrift/src/thrift/lib/thrift/detail/DynamicPatch.h
index 37487e7d11f2f..916e9fdb0e8e9 100644
--- a/third-party/thrift/src/thrift/lib/thrift/detail/DynamicPatch.h
+++ b/third-party/thrift/src/thrift/lib/thrift/detail/DynamicPatch.h
@@ -448,18 +448,27 @@ class DynamicStructurePatch {
                    : ensureStruct(id, std::move(v));
   }
 
+  // patchIfSet
+  template <class Tag>
+  op::patch_type<Tag>& patchIfSet(detail::Badge badge, FieldId id) {
+    using Patch = op::patch_type<Tag>;
+    auto& patch = patchIfSet(badge, id);
+
+    // patch already has the correct type, return it directly.
+    if (auto p = patch.template get_if<Patch>(badge)) {
+      return *p;
+    }
+
+    // Use merge to change patch's type.
+    patch.merge(badge, DynamicPatch{Patch{}});
+    return *patch.template get_if<Patch>(badge);
+  }
+
   DynamicPatch& patchIfSet(detail::Badge, FieldId id) {
     ensurePatchable();
     return ensure_.contains(id) ? patchAfter_[id] : patchPrior_[id];
   }
 
-  template <class SubPatch>
-  DynamicPatch& patchIfSet(detail::Badge badge, FieldId id, SubPatch&& patch) {
-    auto& ret = patchIfSet(badge, id);
-    ret.merge(badge, std::forward<SubPatch>(patch));
-    return ret;
-  }
-
   [[nodiscard]] bool empty(detail::Badge) const {
     return !assign_ && !clear_ && remove_.empty() && ensure_.empty() &&
         patchPrior_.empty() && patchAfter_.empty();
@@ -478,6 +487,12 @@ class DynamicStructurePatch {
   template <class Self, class Visitor>
   static void customVisitImpl(Self&& self, detail::Badge, Visitor&& v);
 
+  // Needed for merge(...). We can consider making this a public API.
+  template <class SubPatch>
+  void patchIfSet(detail::Badge badge, FieldId id, SubPatch&& patch) {
+    patchIfSet(badge, id).merge(badge, std::forward<SubPatch>(patch));
+  }
+
  public:
   FOLLY_FOR_EACH_THIS_OVERLOAD_IN_CLASS_BODY_DELEGATE(
       customVisit, customVisitImpl);