Add ability to clean up initial states after removing animation state…

… for undeclared transitions

Summary:
# Context
We are trying to resolve the crash associated to this task (T209041862) which happens due to an inconsistency between `initialStatesToRestore` and `animatedStates` in Litho's `TransitionManager`. Concretely, there is a missing `animatedState` for a given entry on `initialStatesToRestore`.

# This diff
We have identified two sections of the code where we remove "unused" `AnimationStates` right after the `initialStatesToRestore` are populated, and before the  actual transitions are ran.

This means that there is space for that cleanup to remove an animation that somehow is associated to an `initialStateToRestore`. Therefore, I'm adding a gated mechanism that will remove from `initialStatesToRestore` any property handles associated to those cleaned up animation states.

{F1973764319}

For now I only add this integration for IG, since the crash is happening mostly there. I just want to validate that this would indeed mitigate it.

Reviewed By: jettbow

Differential Revision: D67079132

fbshipit-source-id: e3ef250cde5be0c4c1e01448cdbfa4e0b807cce7

litho-rendercore-transitions/src/main/java/com/facebook/litho/TransitionManager.java

@@ -99,6 +99,21 @@ public class TransitionManager {
 
   @Nullable private final AnimationInconsistencyDebugger mAnimationInconsistencyDebugger;
 
+  /**
+   * Whether we should remove the {@link PropertyState} in the {@link
+   * TransitionManager#mInitialStatesToRestore} for transitions that are cleaned up by {@link
+   * TransitionManager#finishUndeclaredTransitions()} or {@link
+   * TransitionManager#cleanupNonAnimatingAnimationStates()}.
+   *
+   * <p>There is a suspicion that this is leading to crashes where we try to restore an animation
+   * state, but there is no longer an animation state for it.
+   */
+  private static boolean sRemoveStateToRestoreIfAnimationIsCleanedUp = false;
+
+  public static void enableInitialStateCleanupAfterAnimationCleanup(boolean enabled) {
+    sRemoveStateToRestoreIfAnimationIsCleanedUp = enabled;
+  }
+
   /**
    * Whether a piece of content identified by a transition key is appearing, disappearing, or just
    * possibly changing some properties.
@@ -306,23 +321,59 @@ void setupTransitions(
   //       which change without a change transition declared. Also the flag should probably belong
   //       to the properties and not to the AnimationState.
   void finishUndeclaredTransitions() {
-    for (AnimationState animationState : new ArrayList<>(mAnimationStates.values())) {
-      if (animationState.shouldFinishUndeclaredAnimation) {
-        animationState.shouldFinishUndeclaredAnimation = false;
-
-        for (PropertyState propertyState :
-            new ArrayList<>(animationState.propertyStates.values())) {
-          final AnimationBinding animationBinding = propertyState.animation;
-          if (animationBinding != null) {
-            animationBinding.stop();
-            mAnimationBindingListener.finishAnimation(
-                animationBinding, AnimationCleanupTrigger.UNDECLARED_TRANSITIONS);
+    if (sRemoveStateToRestoreIfAnimationIsCleanedUp) {
+      Set<TransitionId> toRemove = new HashSet<>();
+      for (TransitionId transitionId : new ArrayList<>(mAnimationStates.ids())) {
+        AnimationState animationState = mAnimationStates.get(transitionId);
+        if (animationState.shouldFinishUndeclaredAnimation) {
+          animationState.shouldFinishUndeclaredAnimation = false;
+
+          for (PropertyState propertyState :
+              new ArrayList<>(animationState.propertyStates.values())) {
+            final AnimationBinding animationBinding = propertyState.animation;
+            if (animationBinding != null) {
+              animationBinding.stop();
+              mAnimationBindingListener.finishAnimation(
+                  animationBinding, AnimationCleanupTrigger.UNDECLARED_TRANSITIONS);
+            }
+          }
+
+          toRemove.add(transitionId);
+        }
+      }
+
+      removeStatesToRestoreForTransitionIds(toRemove);
+    } else {
+      for (AnimationState animationState : new ArrayList<>(mAnimationStates.values())) {
+        if (animationState.shouldFinishUndeclaredAnimation) {
+          animationState.shouldFinishUndeclaredAnimation = false;
+
+          for (PropertyState propertyState :
+              new ArrayList<>(animationState.propertyStates.values())) {
+            final AnimationBinding animationBinding = propertyState.animation;
+            if (animationBinding != null) {
+              animationBinding.stop();
+              mAnimationBindingListener.finishAnimation(
+                  animationBinding, AnimationCleanupTrigger.UNDECLARED_TRANSITIONS);
+            }
           }
         }
       }
     }
   }
 
+  private void removeStatesToRestoreForTransitionIds(Set<TransitionId> toRemove) {
+    if (!toRemove.isEmpty()) {
+      for (PropertyHandle propertyHandle : new ArrayList<>(mInitialStatesToRestore.keySet())) {
+        TransitionId transitionId =
+            propertyHandle != null ? propertyHandle.getTransitionId() : null;
+        if (transitionId != null && toRemove.contains(transitionId)) {
+          mInitialStatesToRestore.remove(propertyHandle);
+        }
+      }
+    }
+  }
+
   /**
    * Called after {@link #setupTransitions} has been called and the new layout has been mounted.
    * This restores the state of the previous layout for content that will animate and then starts
@@ -725,7 +776,6 @@ private void createAnimationsForTransitionUnit(
     final List<PropertyHandle> animatedPropertyHandles = new ArrayList<>();
     animatedPropertyHandles.add(propertyHandle);
     mAnimationsToPropertyHandles.put(animation, animatedPropertyHandles);
-
     mInitialStatesToRestore.put(propertyHandle, startValue);
 
     if (!TextUtils.isEmpty(transition.getTraceName())) {
@@ -906,6 +956,10 @@ private void cleanupNonAnimatingAnimationStates() {
     for (TransitionId transitionId : toRemove) {
       removeAnimationState(transitionId, AnimationCleanupTrigger.NON_ANIMATING_CLEANUP);
     }
+
+    if (sRemoveStateToRestoreIfAnimationIsCleanedUp) {
+      removeStatesToRestoreForTransitionIds(toRemove);
+    }
   }
 
   private void trackAnimationStateRemoval(