Provide the back callback instance when unregistering it from SysUI.

Bug: 228535747
Test: atest BackAnimationControllerTest
Test: Make sure back to home animation works on tablets and phones.
Change-Id: Iacbd581f382da50e96cfefbefe569800e9e1043b
diff --git a/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java b/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
index 921674a..4d854b6 100644
--- a/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
+++ b/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
@@ -46,6 +46,7 @@
 import com.android.systemui.shared.system.QuickStepContract;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
 import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat;
+
 /**
  * Controls the animation of swiping back and returning to launcher.
  *
@@ -88,6 +89,7 @@
     private boolean mAnimatorSetInProgress = false;
     private float mBackProgress = 0;
     private boolean mBackInProgress = false;
+    private IOnBackInvokedCallback mBackCallback;
 
     public LauncherBackAnimationController(
             BaseQuickstepLauncher launcher,
@@ -113,34 +115,35 @@
      * @param handler Handler to the thread to run the animations on.
      */
     public void registerBackCallbacks(Handler handler) {
-        SystemUiProxy.INSTANCE.get(mLauncher).setBackToLauncherCallback(
-                new IOnBackInvokedCallback.Stub() {
-                    @Override
-                    public void onBackCancelled() {
-                        handler.post(() -> resetPositionAnimated());
-                    }
+        mBackCallback = new IOnBackInvokedCallback.Stub() {
+            @Override
+            public void onBackCancelled() {
+                handler.post(() -> resetPositionAnimated());
+            }
 
-                    @Override
-                    public void onBackInvoked() {
-                        handler.post(() -> startTransition());
-                    }
+            @Override
+            public void onBackInvoked() {
+                handler.post(() -> startTransition());
+            }
 
-                    @Override
-                    public void onBackProgressed(BackEvent backEvent) {
-                        mBackProgress = backEvent.getProgress();
-                        // TODO: Update once the interpolation curve spec is finalized.
-                        mBackProgress =
-                                1 - (1 - mBackProgress) * (1 - mBackProgress) * (1
-                                        - mBackProgress);
-                        if (!mBackInProgress) {
-                            startBack(backEvent);
-                        } else {
-                            updateBackProgress(mBackProgress, backEvent);
-                        }
-                    }
+            @Override
+            public void onBackProgressed(BackEvent backEvent) {
+                mBackProgress = backEvent.getProgress();
+                // TODO: Update once the interpolation curve spec is finalized.
+                mBackProgress =
+                        1 - (1 - mBackProgress) * (1 - mBackProgress) * (1
+                                - mBackProgress);
+                if (!mBackInProgress) {
+                    startBack(backEvent);
+                } else {
+                    updateBackProgress(mBackProgress, backEvent);
+                }
+            }
 
-                    public void onBackStarted() { }
-                });
+            @Override
+            public void onBackStarted() { }
+        };
+        SystemUiProxy.INSTANCE.get(mLauncher).setBackToLauncherCallback(mBackCallback);
     }
 
     private void resetPositionAnimated() {
@@ -163,7 +166,10 @@
 
     /** Unregisters the back to launcher callback in shell. */
     public void unregisterBackCallbacks() {
-        SystemUiProxy.INSTANCE.get(mLauncher).clearBackToLauncherCallback();
+        if (mBackCallback != null) {
+            SystemUiProxy.INSTANCE.get(mLauncher).clearBackToLauncherCallback(mBackCallback);
+        }
+        mBackCallback = null;
     }
 
     private void startBack(BackEvent backEvent) {
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index 5ef89d3..fc4c23f 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -847,8 +847,14 @@
         }
     }
 
-    /** Clears the previously registered {@link IOnBackInvokedCallback}. */
-    public void clearBackToLauncherCallback() {
+    /** Clears the previously registered {@link IOnBackInvokedCallback}.
+     *
+     * @param callback The previously registered callback instance.
+     */
+    public void clearBackToLauncherCallback(IOnBackInvokedCallback callback) {
+        if (mBackToLauncherCallback != callback) {
+            return;
+        }
         mBackToLauncherCallback = null;
         if (mBackAnimation == null) {
             return;