From 6db19b036e622bb06ffa048f40b2ad3a9e8fea44 Mon Sep 17 00:00:00 2001
From: Samuel Susla <samuelsusla@fb.com>
Date: Mon, 10 Aug 2020 12:49:36 -0700
Subject: [PATCH] Introduce Element<>.stateData API

Summary:
Changelog: [Internal]

Previous interface `Element<>.state` wasn't usable because creating ConcreteState  requires ownership of component descriptor and family. Family isn't created until later and it isn't accessible to the caller.

To work around this shortcoming, we create `stateData` rather than state.

Reviewed By: JoshuaGross

Differential Revision: D23028296

fbshipit-source-id: fba35ea8e6986b77379b1dddaa37012f4234f86e
---
 .../react/renderer/element/ComponentBuilder.cpp | 16 ++++++++++++----
 ReactCommon/react/renderer/element/Element.h    | 17 +++++++----------
 .../react/renderer/element/ElementFragment.h    |  3 ++-
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/ReactCommon/react/renderer/element/ComponentBuilder.cpp b/ReactCommon/react/renderer/element/ComponentBuilder.cpp
index 08817a5ccf7..242b8d854e6 100644
--- a/ReactCommon/react/renderer/element/ComponentBuilder.cpp
+++ b/ReactCommon/react/renderer/element/ComponentBuilder.cpp
@@ -30,10 +30,8 @@ ShadowNode::Unshared ComponentBuilder::build(
           elementFragment.tag, elementFragment.surfaceId, nullptr},
       nullptr);
 
-  auto state = elementFragment.state
-      ? elementFragment.state
-      : componentDescriptor.createInitialState(
-            ShadowNodeFragment{elementFragment.props}, family);
+  auto state = componentDescriptor.createInitialState(
+      ShadowNodeFragment{elementFragment.props}, family);
 
   auto constShadowNode = componentDescriptor.createShadowNode(
       ShadowNodeFragment{
@@ -42,6 +40,16 @@ ShadowNode::Unshared ComponentBuilder::build(
           state},
       family);
 
+  if (elementFragment.stateCallback) {
+    auto newState = componentDescriptor.createState(
+        *family, elementFragment.stateCallback());
+    constShadowNode = componentDescriptor.cloneShadowNode(
+        *constShadowNode,
+        ShadowNodeFragment{ShadowNodeFragment::propsPlaceholder(),
+                           ShadowNodeFragment::childrenPlaceholder(),
+                           newState});
+  }
+
   auto shadowNode = std::const_pointer_cast<ShadowNode>(constShadowNode);
 
   if (elementFragment.referenceCallback) {
diff --git a/ReactCommon/react/renderer/element/Element.h b/ReactCommon/react/renderer/element/Element.h
index babe33cd6e4..e0bf5d7a20b 100644
--- a/ReactCommon/react/renderer/element/Element.h
+++ b/ReactCommon/react/renderer/element/Element.h
@@ -33,6 +33,7 @@ class Element final {
   using ConcreteProps = typename ShadowNodeT::ConcreteProps;
   using SharedConcreteProps = std::shared_ptr<ConcreteProps const>;
   using ConcreteState = typename ShadowNodeT::ConcreteState;
+  using ConcreteStateData = typename ShadowNodeT::ConcreteStateData;
   using SharedConcreteState = std::shared_ptr<ConcreteState const>;
   using ConcreteShadowNode = ShadowNodeT;
   using ConcreteUnsharedShadowNode = std::shared_ptr<ConcreteShadowNode>;
@@ -88,19 +89,15 @@ class Element final {
     return *this;
   }
 
-  /*
-   * Sets `state`.
-   */
-  Element &state(SharedConcreteState state) {
-    fragment_.state = state;
-    return *this;
-  }
-
   /*
    * Sets `state` using callback.
    */
-  Element &state(std::function<SharedConcreteState()> callback) {
-    fragment_.state = state();
+  Element &stateData(std::function<void(ConcreteStateData&)> callback) {
+    fragment_.stateCallback = [&]() -> StateData::Shared {
+      auto stateData = ConcreteStateData();
+      callback(stateData);
+      return std::make_shared<ConcreteStateData>(stateData);
+    };
     return *this;
   }
 
diff --git a/ReactCommon/react/renderer/element/ElementFragment.h b/ReactCommon/react/renderer/element/ElementFragment.h
index 7fee1e01be6..fd0807a2cec 100644
--- a/ReactCommon/react/renderer/element/ElementFragment.h
+++ b/ReactCommon/react/renderer/element/ElementFragment.h
@@ -30,6 +30,7 @@ class ElementFragment final {
   using ReferenceCallback =
       std::function<void(ShadowNode::Unshared const &shadowNode)>;
   using FinalizeCallback = std::function<void(ShadowNode &shadowNode)>;
+  using StateCallback = std::function<StateData::Shared()>;
 
   /*
    * ComponentDescriptor part (describes the type)
@@ -47,7 +48,6 @@ class ElementFragment final {
    * ShadowNode part (describes the instance)
    */
   Props::Shared props;
-  State::Shared state;
   List children;
 
   /*
@@ -55,6 +55,7 @@ class ElementFragment final {
    */
   ReferenceCallback referenceCallback;
   FinalizeCallback finalizeCallback;
+  StateCallback stateCallback;
 };
 
 } // namespace react