Add unstable ReferenceStateWrapper API

packages/react-native/ReactAndroid/api/ReactAndroid.api

@@ -2384,15 +2384,6 @@ public final class com/facebook/react/fabric/FabricUIManagerProviderImpl : com/f
 	public fun createUIManager (Lcom/facebook/react/bridge/ReactApplicationContext;)Lcom/facebook/react/bridge/UIManager;
 }
 
-public final class com/facebook/react/fabric/StateWrapperImpl : com/facebook/jni/HybridClassBase, com/facebook/react/uimanager/StateWrapper {
-	public fun destroyState ()V
-	public fun getStateData ()Lcom/facebook/react/bridge/ReadableNativeMap;
-	public fun getStateDataMapBuffer ()Lcom/facebook/react/common/mapbuffer/ReadableMapBuffer;
-	public fun toString ()Ljava/lang/String;
-	public fun updateState (Lcom/facebook/react/bridge/WritableMap;)V
-	public final fun updateStateImpl (Lcom/facebook/react/bridge/NativeMap;)V
-}
-
 public final class com/facebook/react/fabric/events/EventBeatManager : com/facebook/jni/HybridClassBase, com/facebook/react/uimanager/events/BatchEventDispatchedListener {
 	public fun <init> ()V
 	public fun onBatchEventDispatched ()V

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/fabric/StateWrapperImpl.kt

@@ -15,22 +15,24 @@ import com.facebook.react.bridge.NativeMap
 import com.facebook.react.bridge.ReadableNativeMap
 import com.facebook.react.bridge.WritableMap
 import com.facebook.react.common.mapbuffer.ReadableMapBuffer
-import com.facebook.react.uimanager.StateWrapper
+import com.facebook.react.uimanager.ReferenceStateWrapper
 
 /**
  * This class holds reference to the C++ EventEmitter object. Instances of this class are created on
  * the Bindings.cpp, where the pointer to the C++ event emitter is set.
  */
 @SuppressLint("MissingNativeLoadLibrary")
 @DoNotStripAny
-public class StateWrapperImpl private constructor() : HybridClassBase(), StateWrapper {
+internal class StateWrapperImpl private constructor() : HybridClassBase(), ReferenceStateWrapper {
 
   private external fun initHybrid()
 
   private external fun getStateDataImpl(): ReadableNativeMap?
 
   private external fun getStateMapBufferDataImpl(): ReadableMapBuffer?
 
+  private external fun getStateDataReferenceImpl(): Any?
+
   public external fun updateStateImpl(map: NativeMap)
 
   public override val stateDataMapBuffer: ReadableMapBuffer?
@@ -51,6 +53,15 @@ public class StateWrapperImpl private constructor() : HybridClassBase(), StateWr
       return getStateDataImpl()
     }
 
+  public override val stateDataReference: Any?
+    get() {
+      if (!isValid) {
+        FLog.e(TAG, "Race between StateWrapperImpl destruction and getState")
+        return null
+      }
+      return getStateDataReferenceImpl()
+    }
+
   init {
     initHybrid()
   }

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/uimanager/ReferenceStateWrapper.kt

@@ -0,0 +1,13 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+package com.facebook.react.uimanager
+
+internal interface ReferenceStateWrapper : StateWrapper {
+  /** Returns state data backed by JNI reference. The underlying object should not be modified. */
+  public val stateDataReference: Any?
+}

packages/react-native/ReactAndroid/src/main/jni/react/fabric/StateWrapperImpl.cpp

@@ -43,6 +43,13 @@ StateWrapperImpl::getStateMapBufferDataImpl() {
   }
 }
 
+jni::local_ref<jobject> StateWrapperImpl::getStateDataReferenceImpl() {
+  if (state_) {
+    return state_->getJNIReference();
+  }
+  return nullptr;
+}
+
 void StateWrapperImpl::updateStateImpl(NativeMap* map) {
   if (state_) {
     // Get folly::dynamic from map
@@ -68,6 +75,9 @@ void StateWrapperImpl::registerNatives() {
       makeNativeMethod(
           "getStateMapBufferDataImpl",
           StateWrapperImpl::getStateMapBufferDataImpl),
+      makeNativeMethod(
+          "getStateDataReferenceImpl",
+          StateWrapperImpl::getStateDataReferenceImpl),
   });
 }
 

packages/react-native/ReactAndroid/src/main/jni/react/fabric/StateWrapperImpl.h

@@ -27,6 +27,7 @@ class StateWrapperImpl : public jni::HybridClass<StateWrapperImpl> {
 
   jni::local_ref<JReadableMapBuffer::jhybridobject> getStateMapBufferDataImpl();
   jni::local_ref<ReadableNativeMap::jhybridobject> getStateDataImpl();
+  jni::local_ref<jobject> getStateDataReferenceImpl();
   void updateStateImpl(NativeMap* map);
   void setState(std::shared_ptr<const State> state);
   std::shared_ptr<const State> getState() const;

packages/react-native/ReactCommon/react/renderer/core/ConcreteState.h

@@ -14,6 +14,7 @@
 #include <react/renderer/core/State.h>
 
 #ifdef ANDROID
+#include <fbjni/fbjni.h>
 #include <react/renderer/mapbuffer/MapBuffer.h>
 #include <react/renderer/mapbuffer/MapBufferBuilder.h>
 #endif
@@ -25,6 +26,11 @@ template <typename StateDataT>
 concept StateDataWithMapBuffer = requires(StateDataT stateData) {
   { stateData.getMapBuffer() } -> std::same_as<MapBuffer>;
 };
+
+template <typename StateDataT>
+concept StateDataWithJNIReference = requires(StateDataT stateData) {
+  { stateData.getJNIReference() } -> std::same_as<jni::local_ref<jobject>>;
+};
 #endif
 
 /*
@@ -119,6 +125,14 @@ class ConcreteState : public State {
       return MapBufferBuilder::EMPTY();
     }
   }
+
+  jni::local_ref<jobject> getJNIReference() const override {
+    if constexpr (StateDataWithJNIReference<DataT>) {
+      return getData().getJNIReference();
+    } else {
+      return nullptr;
+    }
+  }
 #endif
 };
 

packages/react-native/ReactCommon/react/renderer/core/State.h

@@ -8,6 +8,7 @@
 #pragma once
 
 #ifdef ANDROID
+#include <fbjni/fbjni.h>
 #include <folly/dynamic.h>
 #include <react/renderer/mapbuffer/MapBuffer.h>
 #endif
@@ -66,6 +67,7 @@ class State {
 #ifdef ANDROID
   virtual folly::dynamic getDynamic() const = 0;
   virtual MapBuffer getMapBuffer() const = 0;
+  virtual jni::local_ref<jobject> getJNIReference() const = 0;
   virtual void updateState(folly::dynamic&& data) const = 0;
 #endif