Add feature flag to use JSI NativeState instead of HostObject

React/Fabric/RCTSurfacePresenter.mm

@@ -277,6 +277,10 @@ - (RCTScheduler *)_createScheduler
     CoreFeatures::enablePropIteratorSetter = true;
   }
 
+  if (reactNativeConfig && reactNativeConfig->getBool("react_fabric:use_native_state")) {
+    CoreFeatures::useNativeState = true;
+  }
+
   auto componentRegistryFactory =
       [factory = wrapManagedObject(_mountingManager.componentViewRegistry.componentViewFactory)](
           EventDispatcher::Weak const &eventDispatcher, ContextContainer::Shared const &contextContainer) {

ReactAndroid/src/main/java/com/facebook/react/config/ReactFeatureFlags.java

@@ -118,4 +118,10 @@ public class ReactFeatureFlags {
    * SurfaceMountingManager.
    */
   public static boolean reduceDeleteCreateMutation = false;
+
+  /**
+   * Use JSI NativeState API to store references to native objects rather than the more expensive
+   * HostObject pattern
+   */
+  public static boolean useNativeState = false;
 }

ReactAndroid/src/main/jni/react/fabric/Binding.cpp

@@ -433,6 +433,9 @@ void Binding::installFabricUIManager(
   CoreFeatures::enablePropIteratorSetter =
       getFeatureFlagValue("enableCppPropsIteratorSetter");
 
+  // NativeState experiment
+  CoreFeatures::useNativeState = getFeatureFlagValue("useNativeState");
+
   // RemoveDelete mega-op
   ShadowViewMutation::PlatformSupportsRemoveDeleteTreeInstruction =
       getFeatureFlagValue("enableRemoveDeleteTreeInstruction");

ReactCommon/react/renderer/core/CoreFeatures.cpp

@@ -13,6 +13,7 @@ namespace react {
 bool CoreFeatures::enablePropIteratorSetter = false;
 bool CoreFeatures::enableMapBuffer = false;
 bool CoreFeatures::blockPaintForUseLayoutEffect = false;
+bool CoreFeatures::useNativeState = false;
 
 } // namespace react
 } // namespace facebook

ReactCommon/react/renderer/core/CoreFeatures.h

@@ -30,6 +30,10 @@ class CoreFeatures {
   // useLayoutEffect hooks to be processed. This changes affects scheduling of
   // when a transaction is mounted.
   static bool blockPaintForUseLayoutEffect;
+
+  // Whether to use Hermes' NativeState instead of HostObject
+  // in simple data passing scenarios with JS
+  static bool useNativeState;
 };
 
 } // namespace react

ReactCommon/react/renderer/core/ShadowNode.h

@@ -29,7 +29,9 @@ static constexpr const int kShadowNodeChildrenSmallVectorSize = 8;
 class ComponentDescriptor;
 struct ShadowNodeFragment;
 
-class ShadowNode : public Sealable, public DebugStringConvertible {
+class ShadowNode : public Sealable,
+                   public DebugStringConvertible,
+                   public jsi::NativeState {
  public:
   using Shared = std::shared_ptr<ShadowNode const>;
   using Weak = std::weak_ptr<ShadowNode const>;

ReactCommon/react/renderer/graphics/BUCK

@@ -8,7 +8,6 @@ load(
     "get_apple_compiler_flags",
     "get_apple_inspector_flags",
     "get_preprocessor_flags_for_build_mode",
-    "react_native_target",
     "react_native_xplat_target",
     "rn_xplat_cxx_library",
     "subdir_glob",
@@ -53,7 +52,7 @@ rn_xplat_cxx_library(
     fbandroid_allow_jni_merging = True,
     fbandroid_deps = [
         FBJNI_TARGET,
-        react_native_target("jni/react/jni:jni"),
+        "//xplat/folly:dynamic",
     ],
     fbandroid_exported_headers = subdir_glob(
         [

ReactCommon/react/renderer/runtimescheduler/BUCK

@@ -49,6 +49,7 @@ rn_xplat_cxx_library(
     deps = [
         "//xplat/folly:dynamic",
         react_native_xplat_target("runtimeexecutor:runtimeexecutor"),
+        react_native_xplat_target("react/renderer/core:core"),
         react_native_xplat_target("react/renderer/debug:debug"),
         react_native_xplat_target("butter:butter"),
         react_native_xplat_target("callinvoker:callinvoker"),

ReactCommon/react/renderer/runtimescheduler/Task.h

@@ -21,7 +21,7 @@ class TaskPriorityComparer;
 
 using RawCallback = std::function<void(jsi::Runtime &)>;
 
-struct Task final {
+struct Task final : public jsi::NativeState {
   Task(
       SchedulerPriority priority,
       jsi::Function callback,

ReactCommon/react/renderer/runtimescheduler/primitives.h

@@ -9,6 +9,7 @@
 
 #include <folly/dynamic.h>
 #include <jsi/jsi.h>
+#include <react/renderer/core/CoreFeatures.h>
 #include <react/renderer/runtimescheduler/Task.h>
 
 namespace facebook {
@@ -22,9 +23,15 @@ struct TaskWrapper : public jsi::HostObject {
 
 inline static jsi::Value valueFromTask(
     jsi::Runtime &runtime,
-    std::shared_ptr<Task> const &task) {
-  return jsi::Object::createFromHostObject(
-      runtime, std::make_shared<TaskWrapper>(task));
+    std::shared_ptr<Task> task) {
+  if (CoreFeatures::useNativeState) {
+    jsi::Object obj(runtime);
+    obj.setNativeState(runtime, std::move(task));
+    return obj;
+  } else {
+    return jsi::Object::createFromHostObject(
+        runtime, std::make_shared<TaskWrapper>(task));
+  }
 }
 
 inline static std::shared_ptr<Task> taskFromValue(
@@ -34,7 +41,11 @@ inline static std::shared_ptr<Task> taskFromValue(
     return nullptr;
   }
 
-  return value.getObject(runtime).getHostObject<TaskWrapper>(runtime)->task;
+  if (CoreFeatures::useNativeState) {
+    return value.getObject(runtime).getNativeState<Task>(runtime);
+  } else {
+    return value.getObject(runtime).getHostObject<TaskWrapper>(runtime)->task;
+  }
 }
 
 } // namespace react

ReactCommon/react/renderer/uimanager/primitives.h

@@ -11,6 +11,7 @@
 #include <jsi/JSIDynamic.h>
 #include <jsi/jsi.h>
 #include <react/debug/react_native_assert.h>
+#include <react/renderer/core/CoreFeatures.h>
 #include <react/renderer/core/EventHandler.h>
 #include <react/renderer/core/ShadowNode.h>
 
@@ -38,9 +39,9 @@ struct ShadowNodeWrapper : public jsi::HostObject {
   ShadowNode::Shared shadowNode;
 };
 
-struct ShadowNodeListWrapper : public jsi::HostObject {
+struct ShadowNodeListWrapper : public jsi::HostObject, public jsi::NativeState {
   ShadowNodeListWrapper(ShadowNode::UnsharedListOfShared shadowNodeList)
-      : shadowNodeList(shadowNodeList) {}
+      : shadowNodeList(std::move(shadowNodeList)) {}
 
   // The below method needs to be implemented out-of-line in order for the class
   // to have at least one "key function" (see
@@ -52,29 +53,63 @@ struct ShadowNodeListWrapper : public jsi::HostObject {
 
 inline static ShadowNode::Shared shadowNodeFromValue(
     jsi::Runtime &runtime,
-    jsi::Value const &value) {
+    const jsi::Value &value) {
   if (value.isNull()) {
     return nullptr;
   }
 
-  return value.getObject(runtime)
-      .getHostObject<ShadowNodeWrapper>(runtime)
-      ->shadowNode;
+  if (CoreFeatures::useNativeState) {
+    return value.getObject(runtime).getNativeState<ShadowNode>(runtime);
+  } else {
+    return value.getObject(runtime)
+        .getHostObject<ShadowNodeWrapper>(runtime)
+        ->shadowNode;
+  }
 }
 
 inline static jsi::Value valueFromShadowNode(
     jsi::Runtime &runtime,
     ShadowNode::Shared shadowNode) {
-  return jsi::Object::createFromHostObject(
-      runtime, std::make_shared<ShadowNodeWrapper>(std::move(shadowNode)));
+  if (CoreFeatures::useNativeState) {
+    jsi::Object obj(runtime);
+    // Need to const_cast since JSI only allows non-const pointees
+    obj.setNativeState(
+        runtime, std::const_pointer_cast<ShadowNode>(std::move(shadowNode)));
+    return obj;
+  } else {
+    return jsi::Object::createFromHostObject(
+        runtime, std::make_shared<ShadowNodeWrapper>(std::move(shadowNode)));
+  }
 }
 
 inline static ShadowNode::UnsharedListOfShared shadowNodeListFromValue(
     jsi::Runtime &runtime,
-    jsi::Value const &value) {
-  return value.getObject(runtime)
-      .getHostObject<ShadowNodeListWrapper>(runtime)
-      ->shadowNodeList;
+    const jsi::Value &value) {
+  if (CoreFeatures::useNativeState) {
+    return value.getObject(runtime)
+        .getNativeState<ShadowNodeListWrapper>(runtime)
+        ->shadowNodeList;
+  } else {
+    return value.getObject(runtime)
+        .getHostObject<ShadowNodeListWrapper>(runtime)
+        ->shadowNodeList;
+  }
+}
+
+inline static jsi::Value valueFromShadowNodeList(
+    jsi::Runtime &runtime,
+    ShadowNode::UnsharedListOfShared shadowNodeList) {
+  auto wrapper =
+      std::make_shared<ShadowNodeListWrapper>(std::move(shadowNodeList));
+  if (CoreFeatures::useNativeState) {
+    // Use the wrapper for NativeState too, otherwise we can't implement
+    // the marker interface. Could be simplifed to a simple struct wrapper.
+    jsi::Object obj(runtime);
+    obj.setNativeState(runtime, std::move(wrapper));
+    return obj;
+  } else {
+    return jsi::Object::createFromHostObject(runtime, std::move(wrapper));
+  }
 }
 
 inline static ShadowNode::UnsharedListOfShared shadowNodeListFromWeakList(
@@ -93,25 +128,14 @@ inline static ShadowNode::UnsharedListOfShared shadowNodeListFromWeakList(
 inline static ShadowNode::UnsharedListOfWeak weakShadowNodeListFromValue(
     jsi::Runtime &runtime,
     jsi::Value const &value) {
-  auto shadowNodeList = value.getObject(runtime)
-                            .getHostObject<ShadowNodeListWrapper>(runtime)
-                            ->shadowNodeList;
-
+  auto shadowNodeList = shadowNodeListFromValue(runtime, value);
   auto weakShadowNodeList = std::make_shared<ShadowNode::ListOfWeak>();
   for (auto const &shadowNode : *shadowNodeList) {
     weakShadowNodeList->push_back(shadowNode);
   }
-
   return weakShadowNodeList;
 }
 
-inline static jsi::Value valueFromShadowNodeList(
-    jsi::Runtime &runtime,
-    const ShadowNode::UnsharedListOfShared &shadowNodeList) {
-  return jsi::Object::createFromHostObject(
-      runtime, std::make_unique<ShadowNodeListWrapper>(shadowNodeList));
-}
-
 inline static Tag tagFromValue(jsi::Value const &value) {
   return (Tag)value.getNumber();
 }