From 3b49db3bfec019fb52c5b35d2dafffa29b3faacb Mon Sep 17 00:00:00 2001
From: Samuel Susla <samuelsusla@meta.com>
Date: Fri, 12 Apr 2024 02:47:01 -0700
Subject: [PATCH] add option to set traits when cloning up to the root (#44017)

Summary:
Pull Request resolved: https://github.com/facebook/react-native/pull/44017

changelog: [internal]

Add an option to mark all nodes clone indirectly by `ShadowNode::cloneTree`.

This is a pre-requisite for new state reconciliation algorithm. It will be used to mark part of shadow tree that was affected by native state update.

Reviewed By: rubennorte

Differential Revision: D55745323

fbshipit-source-id: 5e2a2e8a572cc5077d907608f83992a43625d58e
---
 .../react/renderer/core/ShadowNode.cpp        |  6 ++--
 .../react/renderer/core/ShadowNode.h          |  4 +--
 .../renderer/core/tests/ShadowNodeTest.cpp    | 30 +++++++++++++++++++
 3 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp b/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp
index c0640e89de4..99e5798206c 100644
--- a/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp
+++ b/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.cpp
@@ -304,7 +304,8 @@ const ShadowNodeFamily& ShadowNode::getFamily() const {
 ShadowNode::Unshared ShadowNode::cloneTree(
     const ShadowNodeFamily& shadowNodeFamily,
     const std::function<ShadowNode::Unshared(ShadowNode const& oldShadowNode)>&
-        callback) const {
+        callback,
+    ShadowNodeTraits traits) const {
   auto ancestors = shadowNodeFamily.getAncestors(*this);
 
   if (ancestors.empty()) {
@@ -332,7 +333,8 @@ ShadowNode::Unshared ShadowNode::cloneTree(
     children[childIndex] = childNode;
 
     childNode = parentNode.clone(
-        {.children = std::make_shared<ShadowNode::ListOfShared>(children)});
+        {.children = std::make_shared<ShadowNode::ListOfShared>(children),
+         .traits = traits});
   }
 
   return std::const_pointer_cast<ShadowNode>(childNode);
diff --git a/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.h b/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.h
index 74f08707b12..23b6a9d2b5f 100644
--- a/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.h
+++ b/packages/react-native/ReactCommon/react/renderer/core/ShadowNode.h
@@ -100,8 +100,8 @@ class ShadowNode : public Sealable,
    */
   Unshared cloneTree(
       const ShadowNodeFamily& shadowNodeFamily,
-      const std::function<Unshared(ShadowNode const& oldShadowNode)>& callback)
-      const;
+      const std::function<Unshared(ShadowNode const& oldShadowNode)>& callback,
+      ShadowNodeTraits traits = {}) const;
 
 #pragma mark - Getters
 
diff --git a/packages/react-native/ReactCommon/react/renderer/core/tests/ShadowNodeTest.cpp b/packages/react-native/ReactCommon/react/renderer/core/tests/ShadowNodeTest.cpp
index cc67c4a1674..b8811813db4 100644
--- a/packages/react-native/ReactCommon/react/renderer/core/tests/ShadowNodeTest.cpp
+++ b/packages/react-native/ReactCommon/react/renderer/core/tests/ShadowNodeTest.cpp
@@ -281,3 +281,33 @@ TEST_F(ShadowNodeTest, handleState) {
       { secondNode->setStateData(TestState()); },
       "Attempt to mutate a sealed object.");
 }
+
+TEST_F(ShadowNodeTest, testCloneTree) {
+  auto& family = nodeABA_->getFamily();
+  auto newTraits = ShadowNodeTraits();
+  newTraits.set(ShadowNodeTraits::Trait::Reserved);
+  auto rootNode = nodeA_->cloneTree(
+      family,
+      [newTraits](ShadowNode const& oldShadowNode) {
+        return oldShadowNode.clone({.traits = newTraits});
+      },
+      newTraits);
+
+  EXPECT_TRUE(rootNode->getTraits().check(ShadowNodeTraits::Trait::Reserved));
+
+  EXPECT_FALSE(rootNode->getChildren()[0]->getTraits().check(
+      ShadowNodeTraits::Trait::Reserved));
+
+  auto const& firstLevelChild = *rootNode->getChildren()[1];
+
+  EXPECT_TRUE(
+      firstLevelChild.getTraits().check(ShadowNodeTraits::Trait::Reserved));
+
+  EXPECT_FALSE(firstLevelChild.getChildren()[1]->getTraits().check(
+      ShadowNodeTraits::Trait::Reserved));
+
+  auto const& secondLevelchild = *firstLevelChild.getChildren()[0];
+
+  EXPECT_TRUE(
+      secondLevelchild.getTraits().check(ShadowNodeTraits::Trait::Reserved));
+}