From 9c805062e780b323a08747568a1efda367cbec13 Mon Sep 17 00:00:00 2001
From: David Vacca <dvacca@fb.com>
Date: Tue, 27 Mar 2018 18:27:13 -0700
Subject: [PATCH] Refactor RootShadowNodeRegistry to make it ThreadSafe

Reviewed By: achen1

Differential Revision: D7417965

fbshipit-source-id: 90fa007242d2f00a315a4db25d2b6a0949e4b0d3
---
 .../react/fabric/FabricUIManager.java         |  6 ++--
 .../react/fabric/RootShadowNodeRegistry.java  | 28 +++++++++++++------
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/ReactAndroid/src/main/java/com/facebook/react/fabric/FabricUIManager.java b/ReactAndroid/src/main/java/com/facebook/react/fabric/FabricUIManager.java
index 47359b44f..6959c275f 100644
--- a/ReactAndroid/src/main/java/com/facebook/react/fabric/FabricUIManager.java
+++ b/ReactAndroid/src/main/java/com/facebook/react/fabric/FabricUIManager.java
@@ -270,7 +270,7 @@ public class FabricUIManager implements UIManager {
       mUIViewOperationQueue.dispatchViewUpdates(
         mCurrentBatch++, System.currentTimeMillis(), System.currentTimeMillis());
 
-      mRootShadowNodeRegistry.addNode(currentRootShadowNode);
+      mRootShadowNodeRegistry.replaceNode(currentRootShadowNode);
     } catch (Exception e) {
       handleException(getRootNode(rootTag), e);
     }
@@ -320,7 +320,7 @@ public class FabricUIManager implements UIManager {
   }
 
   @Override
-  public synchronized <T extends SizeMonitoringFrameLayout & MeasureSpecProvider> int addRootView(
+  public <T extends SizeMonitoringFrameLayout & MeasureSpecProvider> int addRootView(
       final T rootView) {
     int rootTag = ReactRootViewTagGenerator.getNextRootViewTag();
     ThemedReactContext themedRootContext =
@@ -332,7 +332,7 @@ public class FabricUIManager implements UIManager {
     int heightMeasureSpec = rootView.getHeightMeasureSpec();
     updateRootView(rootShadowNode, widthMeasureSpec, heightMeasureSpec);
 
-    mRootShadowNodeRegistry.addNode(rootShadowNode);
+    mRootShadowNodeRegistry.registerNode(rootShadowNode);
     mUIViewOperationQueue.addRootView(rootTag, rootView, themedRootContext);
     return rootTag;
   }
diff --git a/ReactAndroid/src/main/java/com/facebook/react/fabric/RootShadowNodeRegistry.java b/ReactAndroid/src/main/java/com/facebook/react/fabric/RootShadowNodeRegistry.java
index 49884cfc1..1d66d4b50 100644
--- a/ReactAndroid/src/main/java/com/facebook/react/fabric/RootShadowNodeRegistry.java
+++ b/ReactAndroid/src/main/java/com/facebook/react/fabric/RootShadowNodeRegistry.java
@@ -7,30 +7,40 @@
 
 package com.facebook.react.fabric;
 
-import android.util.SparseArray;
 import com.facebook.react.uimanager.ReactShadowNode;
+import java.util.concurrent.ConcurrentHashMap;
+import javax.annotation.concurrent.ThreadSafe;
 
 /**
  * Simple container class to keep track of {@link ReactShadowNode}s that represents the Root
- * Shadow Nodes of a {@link FabricUIManager}
+ * Shadow Nodes of a {@link FabricUIManager}.
  */
+@ThreadSafe
 public class RootShadowNodeRegistry {
 
-  private final SparseArray<ReactShadowNode> mTagsToRootNodes;
+  private final ConcurrentHashMap<Integer, ReactShadowNode> mTagsToRootNodes = new ConcurrentHashMap<>();
 
-  public RootShadowNodeRegistry() {
-    mTagsToRootNodes = new SparseArray<>();
-  }
-
-  public void addNode(ReactShadowNode node) {
+  /**
+   * Registers the {@link ReactShadowNode} received as a parameter as a RootShadowNode.
+   */
+  public synchronized void registerNode(ReactShadowNode node) {
     mTagsToRootNodes.put(node.getReactTag(), node);
   }
 
-  public void removeNode(int tag) {
+  /**
+   * Register the {@link ReactShadowNode} received as a parameter as a RootShadowNode, replacing
+   * the previous RootShadowNode associated for the {@link ReactShadowNode#getReactTag()}
+   */
+  public void replaceNode(ReactShadowNode node) {
+    mTagsToRootNodes.replace(node.getReactTag(), node);
+  }
+
+  public void removeNode(Integer tag) {
     mTagsToRootNodes.remove(tag);
   }
 
   public ReactShadowNode getNode(int tag) {
     return mTagsToRootNodes.get(tag);
   }
+
 }