From 18589e82b5e780de303afd5540ed88b44a366158 Mon Sep 17 00:00:00 2001
From: Priyank Warkhede <priyankw@meta.com>
Date: Fri, 13 Dec 2024 16:51:12 -0800
Subject: [PATCH] CowStorage: handle get/set for hybrid Thrift object

Summary: In CowStorage get/set implementation, handle Thrift object under HybridNode.

Reviewed By: wilsonwinhi

Differential Revision: D67162094

fbshipit-source-id: d8b250fc1ef54172e4b41986f6646400e0a4f037
---
 fboss/thrift_cow/storage/CowStorage.h | 36 ++++++++++++++++++++-------
 1 file changed, 27 insertions(+), 9 deletions(-)

diff --git a/fboss/thrift_cow/storage/CowStorage.h b/fboss/thrift_cow/storage/CowStorage.h
index 998bda778a2a4..3d5fad4ac719c 100644
--- a/fboss/thrift_cow/storage/CowStorage.h
+++ b/fboss/thrift_cow/storage/CowStorage.h
@@ -67,11 +67,20 @@ class CowStorage : public Storage<Root, CowStorage<Root, Node>> {
       const {
     T out;
     auto op = thrift_cow::pvlambda([&](auto& node) {
-      auto val = node.toThrift();
-      if constexpr (std::is_assignable_v<decltype(out)&, decltype(val)>) {
-        out = std::move(val);
+      if constexpr (!thrift_cow::is_cow_type_v<decltype(node)>) {
+        // Thrift object under HybridNode
+        if constexpr (std::is_assignable_v<decltype(out)&, decltype(node)>) {
+          out = std::move(node);
+        } else {
+          throw std::runtime_error("Type mismatch");
+        }
       } else {
-        throw std::runtime_error("Type mismatch");
+        auto val = node.toThrift();
+        if constexpr (std::is_assignable_v<decltype(out)&, decltype(val)>) {
+          out = std::move(val);
+        } else {
+          throw std::runtime_error("Type mismatch");
+        }
       }
     });
     const auto& rootNode = *root_;
@@ -130,14 +139,23 @@ class CowStorage : public Storage<Root, CowStorage<Root, Node>> {
   std::optional<StorageError>
   set_impl(PathIter begin, PathIter end, T&& value) {
     StorageImpl::modifyPath(&root_, begin, end);
+    using ValueT = typename folly::remove_cvref_t<decltype(value)>;
     auto op = thrift_cow::pvlambda([&](auto& node) {
       using NodeT = typename folly::remove_cvref_t<decltype(node)>;
-      using TType = typename NodeT::ThriftType;
-      using ValueT = typename folly::remove_cvref_t<decltype(value)>;
-      if constexpr (std::is_same_v<ValueT, TType>) {
-        node.fromThrift(std::forward<T>(value));
+      if constexpr (!thrift_cow::is_cow_type_v<NodeT>) {
+        // Thrift object under HybridNode
+        if constexpr (std::is_same_v<ValueT, NodeT>) {
+          node = std::forward<T>(value);
+        } else {
+          throw std::runtime_error("set: type mismatch for passed in path");
+        }
       } else {
-        throw std::runtime_error("set: type mismatch for passed in path");
+        using TType = typename NodeT::ThriftType;
+        if constexpr (std::is_same_v<ValueT, TType>) {
+          node.fromThrift(std::forward<T>(value));
+        } else {
+          throw std::runtime_error("set: type mismatch for passed in path");
+        }
       }
     });
     auto traverseResult = thrift_cow::RootPathVisitor::visit(