Native Animated - Restore default values when removing props on Android

Summary:
Same as #11819 but for Android. I didn't notice the bug initially in my app because I was using different animations on Android which did not trigger this issue.

**Test plan**
Created a simple repro example and tested that this fixes it. https://gist.github.com/janicduplessis/0f3eb362dae63fedf99a0d3ee041796a
Closes #12842

Differential Revision: D5556439

Pulled By: hramos

fbshipit-source-id: d13f4ad258d03cca46c793751ebc49d942b99152

Author: janicduplessis

ReactAndroid/src/main/java/com/facebook/react/animated/NativeAnimatedModule.java

@@ -17,17 +17,20 @@
 import com.facebook.react.bridge.Arguments;
 import com.facebook.react.bridge.Callback;
 import com.facebook.react.bridge.LifecycleEventListener;
-import com.facebook.react.bridge.OnBatchCompleteListener;
 import com.facebook.react.bridge.ReactApplicationContext;
 import com.facebook.react.bridge.ReactContextBaseJavaModule;
 import com.facebook.react.bridge.ReactMethod;
 import com.facebook.react.bridge.ReadableMap;
 import com.facebook.react.bridge.WritableMap;
+import com.facebook.react.common.annotations.VisibleForTesting;
 import com.facebook.react.module.annotations.ReactModule;
 import com.facebook.react.modules.core.DeviceEventManagerModule;
 import com.facebook.react.modules.core.ReactChoreographer;
 import com.facebook.react.uimanager.GuardedFrameCallback;
+import com.facebook.react.uimanager.NativeViewHierarchyManager;
+import com.facebook.react.uimanager.UIBlock;
 import com.facebook.react.uimanager.UIManagerModule;
+import com.facebook.react.uimanager.UIManagerModuleListener;
 
 /**
  * Module that exposes interface for creating and managing animated nodes on the "native" side.
@@ -72,19 +75,18 @@
  */
 @ReactModule(name = NativeAnimatedModule.NAME)
 public class NativeAnimatedModule extends ReactContextBaseJavaModule implements
-    OnBatchCompleteListener, LifecycleEventListener {
+    LifecycleEventListener, UIManagerModuleListener {
 
   protected static final String NAME = "NativeAnimatedModule";
 
   private interface UIThreadOperation {
     void execute(NativeAnimatedNodesManager animatedNodesManager);
   }
 
-  private final Object mOperationsCopyLock = new Object();
   private final GuardedFrameCallback mAnimatedFrameCallback;
   private final ReactChoreographer mReactChoreographer;
   private ArrayList<UIThreadOperation> mOperations = new ArrayList<>();
-  private volatile @Nullable ArrayList<UIThreadOperation> mReadyOperations = null;
+  private ArrayList<UIThreadOperation> mPreOperations = new ArrayList<>();
 
   private @Nullable NativeAnimatedNodesManager mNodesManager;
 
@@ -95,26 +97,9 @@ public NativeAnimatedModule(ReactApplicationContext reactContext) {
     mAnimatedFrameCallback = new GuardedFrameCallback(reactContext) {
       @Override
       protected void doFrameGuarded(final long frameTimeNanos) {
-        if (mNodesManager == null) {
-          UIManagerModule uiManager = getReactApplicationContext()
-            .getNativeModule(UIManagerModule.class);
-          mNodesManager = new NativeAnimatedNodesManager(uiManager);
-        }
-
-        ArrayList<UIThreadOperation> operations;
-        synchronized (mOperationsCopyLock) {
-          operations = mReadyOperations;
-          mReadyOperations = null;
-        }
-
-        if (operations != null) {
-          for (int i = 0, size = operations.size(); i < size; i++) {
-            operations.get(i).execute(mNodesManager);
-          }
-        }
-
-        if (mNodesManager.hasActiveAnimations()) {
-          mNodesManager.runUpdates(frameTimeNanos);
+        NativeAnimatedNodesManager nodesManager = getNodesManager();
+        if (nodesManager.hasActiveAnimations()) {
+          nodesManager.runUpdates(frameTimeNanos);
         }
 
         // TODO: Would be great to avoid adding this callback in case there are no active animations
@@ -130,7 +115,10 @@ protected void doFrameGuarded(final long frameTimeNanos) {
 
   @Override
   public void initialize() {
-    getReactApplicationContext().addLifecycleEventListener(this);
+    ReactApplicationContext reactCtx = getReactApplicationContext();
+    UIManagerModule uiManager = reactCtx.getNativeModule(UIManagerModule.class);
+    reactCtx.addLifecycleEventListener(this);
+    uiManager.addUIManagerListener(this);
   }
 
   @Override
@@ -139,24 +127,32 @@ public void onHostResume() {
   }
 
   @Override
-  public void onBatchComplete() {
-    // Note: The order of executing onBatchComplete handler (especially in terms of onBatchComplete
-    // from the UIManagerModule) doesn't matter as we only enqueue operations for the UI thread to
-    // be executed from here. Thanks to ReactChoreographer all the operations from here are going
-    // to be executed *after* all the operations enqueued by UIManager as the callback type that we
-    // use for ReactChoreographer (CallbackType.NATIVE_ANIMATED_MODULE) is run after callbacks that
-    // UIManager uses.
-    ArrayList<UIThreadOperation> operations = mOperations.isEmpty() ? null : mOperations;
-    if (operations != null) {
-      mOperations = new ArrayList<>();
-      synchronized (mOperationsCopyLock) {
-        if (mReadyOperations == null) {
-          mReadyOperations = operations;
-        } else {
-          mReadyOperations.addAll(operations);
+  public void willDispatchViewUpdates(final UIManagerModule uiManager) {
+    if (mOperations.isEmpty() && mPreOperations.isEmpty()) {
+      return;
+    }
+    final ArrayList<UIThreadOperation> preOperations = mPreOperations;
+    final ArrayList<UIThreadOperation> operations = mOperations;
+    mPreOperations = new ArrayList<>();
+    mOperations = new ArrayList<>();
+    uiManager.prependUIBlock(new UIBlock() {
+      @Override
+      public void execute(NativeViewHierarchyManager nativeViewHierarchyManager) {
+        NativeAnimatedNodesManager nodesManager = getNodesManager();
+        for (UIThreadOperation operation : preOperations) {
+          operation.execute(nodesManager);
         }
       }
-    }
+    });
+    uiManager.addUIBlock(new UIBlock() {
+      @Override
+      public void execute(NativeViewHierarchyManager nativeViewHierarchyManager) {
+        NativeAnimatedNodesManager nodesManager = getNodesManager();
+        for (UIThreadOperation operation : operations) {
+          operation.execute(nodesManager);
+        }
+      }
+    });
   }
 
   @Override
@@ -174,6 +170,15 @@ public String getName() {
     return NAME;
   }
 
+  private NativeAnimatedNodesManager getNodesManager() {
+    if (mNodesManager == null) {
+      UIManagerModule uiManager = getReactApplicationContext().getNativeModule(UIManagerModule.class);
+      mNodesManager = new NativeAnimatedNodesManager(uiManager);
+    }
+
+    return mNodesManager;
+  }
+
   private void clearFrameCallback() {
     Assertions.assertNotNull(mReactChoreographer).removeFrameCallback(
       ReactChoreographer.CallbackType.NATIVE_ANIMATED_MODULE,
@@ -186,6 +191,11 @@ private void enqueueFrameCallback() {
       mAnimatedFrameCallback);
   }
 
+  @VisibleForTesting
+  public void setNodesManager(NativeAnimatedNodesManager nodesManager) {
+    mNodesManager = nodesManager;
+  }
+
   @ReactMethod
   public void createAnimatedNode(final int tag, final ReadableMap config) {
     mOperations.add(new UIThreadOperation() {
@@ -336,6 +346,12 @@ public void execute(NativeAnimatedNodesManager animatedNodesManager) {
 
   @ReactMethod
   public void disconnectAnimatedNodeFromView(final int animatedNodeTag, final int viewTag) {
+    mPreOperations.add(new UIThreadOperation() {
+      @Override
+      public void execute(NativeAnimatedNodesManager animatedNodesManager) {
+        animatedNodesManager.restoreDefaultValues(animatedNodeTag, viewTag);
+      }
+    });
     mOperations.add(new UIThreadOperation() {
       @Override
       public void execute(NativeAnimatedNodesManager animatedNodesManager) {

ReactAndroid/src/main/java/com/facebook/react/animated/NativeAnimatedNodesManager.java

@@ -92,7 +92,7 @@ public void createAnimatedNode(int tag, ReadableMap config) {
     } else if ("value".equals(type)) {
       node = new ValueAnimatedNode(config);
     } else if ("props".equals(type)) {
-      node = new PropsAnimatedNode(config, this);
+      node = new PropsAnimatedNode(config, this, mUIImplementation);
     } else if ("interpolation".equals(type)) {
       node = new InterpolationAnimatedNode(config);
     } else if ("addition".equals(type)) {
@@ -289,11 +289,7 @@ public void connectAnimatedNodeToView(int animatedNodeTag, int viewTag) {
         "of type " + PropsAnimatedNode.class.getName());
     }
     PropsAnimatedNode propsAnimatedNode = (PropsAnimatedNode) node;
-    if (propsAnimatedNode.mConnectedViewTag != -1) {
-      throw new JSApplicationIllegalArgumentException("Animated node " + animatedNodeTag + " is " +
-        "already attached to a view");
-    }
-    propsAnimatedNode.mConnectedViewTag = viewTag;
+    propsAnimatedNode.connectToView(viewTag);
     mUpdatedNodes.put(animatedNodeTag, node);
   }
 
@@ -308,11 +304,24 @@ public void disconnectAnimatedNodeFromView(int animatedNodeTag, int viewTag) {
         "of type " + PropsAnimatedNode.class.getName());
     }
     PropsAnimatedNode propsAnimatedNode = (PropsAnimatedNode) node;
-    if (propsAnimatedNode.mConnectedViewTag != viewTag) {
-      throw new JSApplicationIllegalArgumentException("Attempting to disconnect view that has " +
-        "not been connected with the given animated node");
+    propsAnimatedNode.disconnectFromView(viewTag);
+  }
+
+  public void restoreDefaultValues(int animatedNodeTag, int viewTag) {
+    AnimatedNode node = mAnimatedNodes.get(animatedNodeTag);
+    // Restoring default values needs to happen before UIManager operations so it is
+    // possible the node hasn't been created yet if it is being connected and
+    // disconnected in the same batch. In that case we don't need to restore
+    // default values since it will never actually update the view.
+    if (node == null) {
+      return;
     }
-    propsAnimatedNode.mConnectedViewTag = -1;
+    if (!(node instanceof PropsAnimatedNode)) {
+      throw new JSApplicationIllegalArgumentException("Animated node connected to view should be" +
+        "of type " + PropsAnimatedNode.class.getName());
+    }
+    PropsAnimatedNode propsAnimatedNode = (PropsAnimatedNode) node;
+    propsAnimatedNode.restoreDefaultValues();
   }
 
   public void addAnimatedEventToView(int viewTag, String eventName, ReadableMap eventMapping) {
@@ -513,7 +522,7 @@ private void updateNodes(List<AnimatedNode> nodes) {
       if (nextNode instanceof PropsAnimatedNode) {
         // Send property updates to native view manager
         try {
-          ((PropsAnimatedNode) nextNode).updateView(mUIImplementation);
+          ((PropsAnimatedNode) nextNode).updateView();
         } catch (IllegalViewOperationException e) {
             // An exception is thrown if the view hasn't been created yet. This can happen because views are
             // created in batches. If this particular view didn't make it into a batch yet, the view won't

ReactAndroid/src/main/java/com/facebook/react/animated/PropsAnimatedNode.java

@@ -9,6 +9,7 @@
 
 package com.facebook.react.animated;
 
+import com.facebook.react.bridge.JSApplicationIllegalArgumentException;
 import com.facebook.react.bridge.JavaOnlyMap;
 import com.facebook.react.bridge.ReadableMap;
 import com.facebook.react.bridge.ReadableMapKeySetIterator;
@@ -27,47 +28,78 @@
  */
 /*package*/ class PropsAnimatedNode extends AnimatedNode {
 
-  /*package*/ int mConnectedViewTag = -1;
-
+  private int mConnectedViewTag = -1;
   private final NativeAnimatedNodesManager mNativeAnimatedNodesManager;
-  private final Map<String, Integer> mPropMapping;
+  private final UIImplementation mUIImplementation;
+  private final Map<String, Integer> mPropNodeMapping;
+  // This is the backing map for `mDiffMap` we can mutate this to update it instead of having to
+  // create a new one for each update.
+  private final JavaOnlyMap mPropMap;
+  private final ReactStylesDiffMap mDiffMap;
 
-  PropsAnimatedNode(ReadableMap config, NativeAnimatedNodesManager nativeAnimatedNodesManager) {
+  PropsAnimatedNode(ReadableMap config, NativeAnimatedNodesManager nativeAnimatedNodesManager, UIImplementation uiImplementation) {
     ReadableMap props = config.getMap("props");
     ReadableMapKeySetIterator iter = props.keySetIterator();
-    mPropMapping = new HashMap<>();
+    mPropNodeMapping = new HashMap<>();
     while (iter.hasNextKey()) {
       String propKey = iter.nextKey();
       int nodeIndex = props.getInt(propKey);
-      mPropMapping.put(propKey, nodeIndex);
+      mPropNodeMapping.put(propKey, nodeIndex);
     }
+    mPropMap = new JavaOnlyMap();
+    mDiffMap = new ReactStylesDiffMap(mPropMap);
     mNativeAnimatedNodesManager = nativeAnimatedNodesManager;
+    mUIImplementation = uiImplementation;
+  }
+
+  public void connectToView(int viewTag) {
+    if (mConnectedViewTag != -1) {
+      throw new JSApplicationIllegalArgumentException("Animated node " + mTag + " is " +
+        "already attached to a view");
+    }
+    mConnectedViewTag = viewTag;
+  }
+
+  public void disconnectFromView(int viewTag) {
+    if (mConnectedViewTag != viewTag) {
+      throw new JSApplicationIllegalArgumentException("Attempting to disconnect view that has " +
+        "not been connected with the given animated node");
+    }
+
+    mConnectedViewTag = -1;
   }
 
-  public final void updateView(UIImplementation uiImplementation) {
+  public void restoreDefaultValues() {
+    ReadableMapKeySetIterator it = mPropMap.keySetIterator();
+    while(it.hasNextKey()) {
+      mPropMap.putNull(it.nextKey());
+    }
+
+    mUIImplementation.synchronouslyUpdateViewOnUIThread(
+      mConnectedViewTag,
+      mDiffMap);
+  }
+
+  public final void updateView() {
     if (mConnectedViewTag == -1) {
-      throw new IllegalStateException("Node has not been attached to a view");
+      return;
     }
-    JavaOnlyMap propsMap = new JavaOnlyMap();
-    for (Map.Entry<String, Integer> entry : mPropMapping.entrySet()) {
+    for (Map.Entry<String, Integer> entry : mPropNodeMapping.entrySet()) {
       @Nullable AnimatedNode node = mNativeAnimatedNodesManager.getNodeById(entry.getValue());
       if (node == null) {
         throw new IllegalArgumentException("Mapped property node does not exists");
       } else if (node instanceof StyleAnimatedNode) {
-        ((StyleAnimatedNode) node).collectViewUpdates(propsMap);
+        ((StyleAnimatedNode) node).collectViewUpdates(mPropMap);
       } else if (node instanceof ValueAnimatedNode) {
-        propsMap.putDouble(entry.getKey(), ((ValueAnimatedNode) node).getValue());
+        mPropMap.putDouble(entry.getKey(), ((ValueAnimatedNode) node).getValue());
       } else {
         throw new IllegalArgumentException("Unsupported type of node used in property node " +
             node.getClass());
       }
     }
-    // TODO: Reuse propsMap and stylesDiffMap objects - note that in subsequent animation steps
-    // for a given node most of the time we will be creating the same set of props (just with
-    // different values). We can take advantage on that and optimize the way we allocate property
-    // maps (we also know that updating view props doesn't retain a reference to the styles object).
-    uiImplementation.synchronouslyUpdateViewOnUIThread(
+
+    mUIImplementation.synchronouslyUpdateViewOnUIThread(
       mConnectedViewTag,
-      new ReactStylesDiffMap(propsMap));
+      mDiffMap);
   }
 }

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

@@ -881,6 +881,10 @@ public void addUIBlock(UIBlock block) {
     mOperationsQueue.enqueueUIBlock(block);
   }
 
+  public void prependUIBlock(UIBlock block) {
+    mOperationsQueue.prependUIBlock(block);
+  }
+
   public int resolveRootTagFromReactTag(int reactTag) {
     if (mShadowNodeRegistry.isRootNode(reactTag)) {
       return reactTag;