Use MutableIntObjectMap for SurfaceMountingManager view registry (#56646)

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

@@ -15,6 +15,7 @@ import android.view.ViewGroup
 import android.view.ViewParent
 import androidx.annotation.AnyThread
 import androidx.annotation.UiThread
+import androidx.collection.MutableIntObjectMap
 import androidx.collection.SparseArrayCompat
 import androidx.core.graphics.drawable.toDrawable
 import com.facebook.common.logging.FLog
@@ -56,7 +57,10 @@ import java.util.ArrayDeque
 import java.util.LinkedList
 import java.util.Queue
 import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.locks.ReentrantReadWriteLock
 import kotlin.concurrent.Volatile
+import kotlin.concurrent.read
+import kotlin.concurrent.write
 
 /** Returns true if the collection contains [key]. */
 private operator fun <T> SparseArrayCompat<T>.contains(key: Int): Boolean = containsKey(key)
@@ -84,7 +88,20 @@ internal constructor(
   public var context: ThemedReactContext? = reactContext
     private set
 
-  private val tagToViewState: ConcurrentHashMap<Int, ViewState> = ConcurrentHashMap() // any thread
+  private val tagToViewState: ConcurrentHashMap<Int, ViewState>?
+  private val optimizedTagToViewState: MutableIntObjectMap<ViewState>?
+  private val registryLock = ReentrantReadWriteLock()
+
+  init {
+    if (ReactNativeFeatureFlags.useOptimizedViewRegistryOnAndroid()) {
+      tagToViewState = null
+      optimizedTagToViewState = MutableIntObjectMap()
+    } else {
+      tagToViewState = ConcurrentHashMap()
+      optimizedTagToViewState = null
+    }
+  }
+
   private val onViewAttachMountItems: Queue<MountItem> = ArrayDeque()
 
   // These are all non-null, until StopSurface is called
@@ -126,7 +143,7 @@ internal constructor(
       return
     }
 
-    tagToViewState[surfaceId] = ViewState(surfaceId, rootView, rootViewManager, true)
+    registryPut(surfaceId, ViewState(surfaceId, rootView, rootViewManager, true))
 
     val runnable: Runnable =
         object : GuardedRunnable(checkNotNull(context)) {
@@ -190,7 +207,7 @@ internal constructor(
     if (tagSetForStoppedSurface?.containsKey(tag) == true) {
       return true
     }
-    return tagToViewState.containsKey(tag)
+    return registryContains(tag)
   }
 
   @UiThread
@@ -236,7 +253,7 @@ internal constructor(
     // Reset all StateWrapper objects
     // Since this can happen on any thread, is it possible to race between StateWrapper destruction
     // and some accesses from View classes in the UI thread?
-    for (viewState in tagToViewState.values) {
+    registryForEachValue { viewState ->
       viewState.stateWrapper?.destroyState()
       viewState.stateWrapper = null
 
@@ -248,23 +265,37 @@ internal constructor(
       if (ReactNativeFeatureFlags.enableViewRecycling()) {
         viewManagerRegistry?.onSurfaceStopped(surfaceId)
       }
-      val tagSetForStoppedSurface =
-          SparseArrayCompat<Any>().also { this.tagSetForStoppedSurface = it }
-      for ((key, value) in tagToViewState) {
-        // Using this as a placeholder value in the map. We're using SparseArrayCompat
-        // since it can efficiently represent the list of pending tags
-        tagSetForStoppedSurface[key] = this
-
-        // We must call `onDropViewInstance` on all remaining Views
-        onViewStateDeleted(value)
+
+      if (optimizedTagToViewState != null) {
+        val viewStatesToDelete: ArrayList<ViewState>
+        registryLock.write {
+          val tagSetForStoppedSurface =
+              SparseArrayCompat<Any>().also { this.tagSetForStoppedSurface = it }
+          viewStatesToDelete = ArrayList(optimizedTagToViewState.size)
+          optimizedTagToViewState.forEach { key, value ->
+            tagSetForStoppedSurface[key] = this@SurfaceMountingManager
+            viewStatesToDelete.add(value)
+          }
+          optimizedTagToViewState.clear()
+        }
+        for (viewState in viewStatesToDelete) {
+          onViewStateDeleted(viewState)
+        }
+      } else {
+        val tagSetForStoppedSurface =
+            SparseArrayCompat<Any>().also { this.tagSetForStoppedSurface = it }
+        for ((key, value) in tagToViewState!!) {
+          tagSetForStoppedSurface[key] = this
+          onViewStateDeleted(value)
+        }
+        tagToViewState!!.clear()
       }
 
       // Evict all views from cache and memory
       jsResponderHandler = null
       rootViewManager = null
       mountItemExecutor = null
       context = null
-      tagToViewState.clear()
       onViewAttachMountItems.clear()
       tagToSynchronousMountProps.clear()
       FLog.e(TAG, "Surface [$surfaceId] was stopped on SurfaceMountingManager.")
@@ -572,7 +603,7 @@ internal constructor(
             this.stateWrapper = stateWrapper
             this.eventEmitter = eventEmitterWrapper
           }
-      tagToViewState[reactTag] = viewState
+      registryPut(reactTag, viewState)
 
       if (isLayoutable) {
         @Suppress("UNCHECKED_CAST")
@@ -921,13 +952,19 @@ internal constructor(
       return
     }
 
-    var viewState = tagToViewState[reactTag]
-    if (viewState == null) {
-      // TODO T62717437 - Use a flag to determine that these event emitters belong to virtual nodes
-      // only.
-      viewState = ViewState(reactTag)
-      tagToViewState[reactTag] = viewState
-    }
+    // TODO T62717437 - Use a flag to determine that these event emitters belong to virtual nodes
+    // only.
+    val viewState: ViewState =
+        if (optimizedTagToViewState != null) {
+          registryLock.write { optimizedTagToViewState.getOrPut(reactTag) { ViewState(reactTag) } }
+        } else {
+          var vs = tagToViewState!![reactTag]
+          if (vs == null) {
+            vs = ViewState(reactTag)
+            tagToViewState!![reactTag] = vs
+          }
+          vs
+        }
     val previousEventEmitterWrapper = viewState.eventEmitter
     viewState.eventEmitter = eventEmitter
 
@@ -1039,7 +1076,7 @@ internal constructor(
       // To delete we simply remove the tag from the registry.
       // We want to rely on the correct set of MountInstructions being sent to the platform,
       // or StopSurface being called, so we do not handle deleting descendants of the View.
-      tagToViewState.remove(reactTag)
+      registryRemove(reactTag)
 
       onViewStateDeleted(viewState)
     }
@@ -1084,12 +1121,52 @@ internal constructor(
   }
 
   private fun getViewState(reactTag: Int): ViewState =
-      tagToViewState[reactTag]
+      registryGet(reactTag)
           ?: throw RetryableMountingLayerException(
               "Unable to find viewState for tag $reactTag. Surface stopped: $isStopped"
           )
 
-  private fun getNullableViewState(reactTag: Int): ViewState? = tagToViewState[reactTag]
+  private fun getNullableViewState(reactTag: Int): ViewState? = registryGet(reactTag)
+
+  private fun registryGet(tag: Int): ViewState? {
+    return if (optimizedTagToViewState != null) {
+      registryLock.read { optimizedTagToViewState[tag] }
+    } else {
+      tagToViewState!![tag]
+    }
+  }
+
+  private fun registryPut(tag: Int, state: ViewState) {
+    if (optimizedTagToViewState != null) {
+      registryLock.write { optimizedTagToViewState[tag] = state }
+    } else {
+      tagToViewState!![tag] = state
+    }
+  }
+
+  private fun registryRemove(tag: Int) {
+    if (optimizedTagToViewState != null) {
+      registryLock.write { optimizedTagToViewState.remove(tag) }
+    } else {
+      tagToViewState!!.remove(tag)
+    }
+  }
+
+  private fun registryContains(tag: Int): Boolean {
+    return if (optimizedTagToViewState != null) {
+      registryLock.read { optimizedTagToViewState.containsKey(tag) }
+    } else {
+      tagToViewState!!.containsKey(tag)
+    }
+  }
+
+  private inline fun registryForEachValue(action: (ViewState) -> Unit) {
+    if (optimizedTagToViewState != null) {
+      registryLock.read { optimizedTagToViewState.forEachValue(action) }
+    } else {
+      tagToViewState!!.values.forEach(action)
+    }
+  }
 
   /** Applies a bitmap as the background of the view with the given tag, if it exists. */
   @UiThread
@@ -1100,7 +1177,7 @@ internal constructor(
 
   public fun printSurfaceState(): Unit {
     FLog.e(TAG, "Views created for surface $surfaceId:")
-    for (viewState in tagToViewState.values) {
+    registryForEachValue { viewState ->
       val viewManagerName = viewState.viewManager?.name
       val view = viewState.view
       val parent = if (view != null) view.parent as View? else null
@@ -1142,7 +1219,7 @@ internal constructor(
   ): Unit {
     // When the surface stopped we will reset the view state map. We are not going to enqueue
     // pending events as they are not expected to be dispatched anyways.
-    val viewState = tagToViewState[reactTag]
+    val viewState = registryGet(reactTag)
 
     if (viewState == null) {
       // Cannot queue event without view state. Do nothing here.

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/internal/featureflags/ReactNativeFeatureFlags.kt

@@ -4,7 +4,7 @@
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  *
- * @generated SignedSource<<cd8218c8b8588f3317bf63ce8d608548>>
+ * @generated SignedSource<<6ad566ffaa8330c696fa2088ff696a2b>>
  */
 
 /**
@@ -534,6 +534,12 @@ public object ReactNativeFeatureFlags {
   @JvmStatic
   public fun useNestedScrollViewAndroid(): Boolean = accessor.useNestedScrollViewAndroid()
 
+  /**
+   * Use MutableIntObjectMap with ReadWriteLock instead of ConcurrentHashMap for the view registry in SurfaceMountingManager to reduce memory overhead and GC pressure.
+   */
+  @JvmStatic
+  public fun useOptimizedViewRegistryOnAndroid(): Boolean = accessor.useOptimizedViewRegistryOnAndroid()
+
   /**
    * Use shared animation backend in C++ Animated
    */

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/internal/featureflags/ReactNativeFeatureFlagsCxxAccessor.kt

@@ -4,7 +4,7 @@
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  *
- * @generated SignedSource<<4f9f5c1c46217ed6802abd5f786aac19>>
+ * @generated SignedSource<<ff84a26e3306cc438fead82ad766ce53>>
  */
 
 /**
@@ -104,6 +104,7 @@ internal class ReactNativeFeatureFlagsCxxAccessor : ReactNativeFeatureFlagsAcces
   private var useLISAlgorithmInDifferentiatorCache: Boolean? = null
   private var useNativeViewConfigsInBridgelessModeCache: Boolean? = null
   private var useNestedScrollViewAndroidCache: Boolean? = null
+  private var useOptimizedViewRegistryOnAndroidCache: Boolean? = null
   private var useSharedAnimatedBackendCache: Boolean? = null
   private var useTraitHiddenOnAndroidCache: Boolean? = null
   private var useTurboModuleInteropCache: Boolean? = null
@@ -869,6 +870,15 @@ internal class ReactNativeFeatureFlagsCxxAccessor : ReactNativeFeatureFlagsAcces
     return cached
   }
 
+  override fun useOptimizedViewRegistryOnAndroid(): Boolean {
+    var cached = useOptimizedViewRegistryOnAndroidCache
+    if (cached == null) {
+      cached = ReactNativeFeatureFlagsCxxInterop.useOptimizedViewRegistryOnAndroid()
+      useOptimizedViewRegistryOnAndroidCache = cached
+    }
+    return cached
+  }
+
   override fun useSharedAnimatedBackend(): Boolean {
     var cached = useSharedAnimatedBackendCache
     if (cached == null) {

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/internal/featureflags/ReactNativeFeatureFlagsCxxInterop.kt

@@ -4,7 +4,7 @@
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  *
- * @generated SignedSource<<8916e9f4a938a69ff175c806db9835d4>>
+ * @generated SignedSource<<817790c7ceb9112376b4ab4ee338ff43>>
  */
 
 /**
@@ -196,6 +196,8 @@ public object ReactNativeFeatureFlagsCxxInterop {
 
   @DoNotStrip @JvmStatic public external fun useNestedScrollViewAndroid(): Boolean
 
+  @DoNotStrip @JvmStatic public external fun useOptimizedViewRegistryOnAndroid(): Boolean
+
   @DoNotStrip @JvmStatic public external fun useSharedAnimatedBackend(): Boolean
 
   @DoNotStrip @JvmStatic public external fun useTraitHiddenOnAndroid(): Boolean

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/internal/featureflags/ReactNativeFeatureFlagsDefaults.kt

@@ -4,7 +4,7 @@
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  *
- * @generated SignedSource<<a9a8ce443fa160a7494fc1c9e7baa02f>>
+ * @generated SignedSource<<1d764e41252e2709e574da5e0ac64bd7>>
  */
 
 /**
@@ -191,6 +191,8 @@ public open class ReactNativeFeatureFlagsDefaults : ReactNativeFeatureFlagsProvi
 
   override fun useNestedScrollViewAndroid(): Boolean = false
 
+  override fun useOptimizedViewRegistryOnAndroid(): Boolean = false
+
   override fun useSharedAnimatedBackend(): Boolean = false
 
   override fun useTraitHiddenOnAndroid(): Boolean = false

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/internal/featureflags/ReactNativeFeatureFlagsLocalAccessor.kt

@@ -4,7 +4,7 @@
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  *
- * @generated SignedSource<<cbe90c2bf8ba9d34804d97c31edfd31a>>
+ * @generated SignedSource<<299507458fe84339ddf816dd58671fd3>>
  */
 
 /**
@@ -108,6 +108,7 @@ internal class ReactNativeFeatureFlagsLocalAccessor : ReactNativeFeatureFlagsAcc
   private var useLISAlgorithmInDifferentiatorCache: Boolean? = null
   private var useNativeViewConfigsInBridgelessModeCache: Boolean? = null
   private var useNestedScrollViewAndroidCache: Boolean? = null
+  private var useOptimizedViewRegistryOnAndroidCache: Boolean? = null
   private var useSharedAnimatedBackendCache: Boolean? = null
   private var useTraitHiddenOnAndroidCache: Boolean? = null
   private var useTurboModuleInteropCache: Boolean? = null
@@ -957,6 +958,16 @@ internal class ReactNativeFeatureFlagsLocalAccessor : ReactNativeFeatureFlagsAcc
     return cached
   }
 
+  override fun useOptimizedViewRegistryOnAndroid(): Boolean {
+    var cached = useOptimizedViewRegistryOnAndroidCache
+    if (cached == null) {
+      cached = currentProvider.useOptimizedViewRegistryOnAndroid()
+      accessedFeatureFlags.add("useOptimizedViewRegistryOnAndroid")
+      useOptimizedViewRegistryOnAndroidCache = cached
+    }
+    return cached
+  }
+
   override fun useSharedAnimatedBackend(): Boolean {
     var cached = useSharedAnimatedBackendCache
     if (cached == null) {

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/internal/featureflags/ReactNativeFeatureFlagsProvider.kt

@@ -4,7 +4,7 @@
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  *
- * @generated SignedSource<<536a5156deea17740dd24782bf79feb4>>
+ * @generated SignedSource<<8f887fb839df553b23886a61537e6991>>
  */
 
 /**
@@ -191,6 +191,8 @@ public interface ReactNativeFeatureFlagsProvider {
 
   @DoNotStrip public fun useNestedScrollViewAndroid(): Boolean
 
+  @DoNotStrip public fun useOptimizedViewRegistryOnAndroid(): Boolean
+
   @DoNotStrip public fun useSharedAnimatedBackend(): Boolean
 
   @DoNotStrip public fun useTraitHiddenOnAndroid(): Boolean