Fix memory leak that occurs when removing root nodes

Reviewed By: sriramramani

Differential Revision: D4540427

fbshipit-source-id: 8bdc0d2318371315b95300cdc6ca10bd0dbb55dd

ReactAndroid/src/main/java/com/facebook/react/flat/FlatUIImplementation.java

@@ -37,7 +37,8 @@ public class FlatUIImplementation extends UIImplementation {
   public static FlatUIImplementation createInstance(
       ReactApplicationContext reactContext,
       List<ViewManager> viewManagers,
-      EventDispatcher eventDispatcher) {
+      EventDispatcher eventDispatcher,
+      boolean memoryImprovementEnabled) {
 
     RCTImageViewManager rctImageViewManager = findRCTImageManager(viewManagers);
     if (rctImageViewManager != null) {
@@ -61,7 +62,8 @@ public class FlatUIImplementation extends UIImplementation {
       rctImageViewManager,
       viewManagerRegistry,
       operationsQueue,
-      eventDispatcher
+      eventDispatcher,
+      memoryImprovementEnabled
     );
   }
 
@@ -73,17 +75,20 @@ public class FlatUIImplementation extends UIImplementation {
   private final ReactApplicationContext mReactContext;
   private @Nullable RCTImageViewManager mRCTImageViewManager;
   private final StateBuilder mStateBuilder;
+  private final boolean mMemoryImprovementEnabled;
 
   private FlatUIImplementation(
       ReactApplicationContext reactContext,
       @Nullable RCTImageViewManager rctImageViewManager,
       ViewManagerRegistry viewManagers,
       FlatUIViewOperationQueue operationsQueue,
-      EventDispatcher eventDispatcher) {
+      EventDispatcher eventDispatcher,
+      boolean memoryImprovementEnabled) {
     super(reactContext, viewManagers, operationsQueue, eventDispatcher);
     mReactContext = reactContext;
     mRCTImageViewManager = rctImageViewManager;
     mStateBuilder = new StateBuilder(operationsQueue);
+    mMemoryImprovementEnabled = memoryImprovementEnabled;
   }
 
   @Override
@@ -521,6 +526,9 @@ public class FlatUIImplementation extends UIImplementation {
 
   @Override
   public void removeRootView(int rootViewTag) {
+    if (mMemoryImprovementEnabled) {
+      removeRootShadowNode(rootViewTag);
+    }
     mStateBuilder.removeRootView(rootViewTag);
   }
 

ReactAndroid/src/main/java/com/facebook/react/flat/FlatUIImplementationProvider.java

@@ -20,11 +20,26 @@ import com.facebook.react.uimanager.events.EventDispatcher;
  * UIImplementationProvider that creates instances of {@link FlatUIImplementation}.
  */
 public final class FlatUIImplementationProvider extends UIImplementationProvider {
+
+  private final boolean mMemoryImprovementEnabled;
+
+  public FlatUIImplementationProvider() {
+    mMemoryImprovementEnabled = true;
+  }
+
+  public FlatUIImplementationProvider(boolean memoryImprovementEnabled) {
+    mMemoryImprovementEnabled = memoryImprovementEnabled;
+  }
+
   @Override
   public FlatUIImplementation createUIImplementation(
       ReactApplicationContext reactContext,
       List<ViewManager> viewManagers,
       EventDispatcher eventDispatcher) {
-    return FlatUIImplementation.createInstance(reactContext, viewManagers, eventDispatcher);
+    return FlatUIImplementation.createInstance(
+      reactContext,
+      viewManagers,
+      eventDispatcher,
+      mMemoryImprovementEnabled);
   }
 }

ReactAndroid/src/main/java/com/facebook/react/uimanager/UIImplementation.java

@@ -133,10 +133,17 @@ public class UIImplementation {
    * Unregisters a root node with a given tag.
    */
   public void removeRootView(int rootViewTag) {
-    mShadowNodeRegistry.removeRootNode(rootViewTag);
+    removeRootShadowNode(rootViewTag);
     mOperationsQueue.enqueueRemoveRootView(rootViewTag);
   }
 
+  /**
+   * Unregisters a root node with a given tag from the shadow node registry
+   */
+  public void removeRootShadowNode(int rootViewTag) {
+    mShadowNodeRegistry.removeRootNode(rootViewTag);
+  }
+
   /**
    * Invoked when native view that corresponds to a root node, or acts as a root view (ie. Modals)
    * has its size changed.