From 883e7e412b1e95d430dafdfb17c38172deeec4d8 Mon Sep 17 00:00:00 2001
From: Shan Huang <shanh@google.com>
Date: Fri, 8 Apr 2022 00:52:44 +0000
Subject: [PATCH] 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
---
 .../LauncherBackAnimationController.java      | 58 ++++++++++---------
 .../com/android/quickstep/SystemUiProxy.java  | 10 +++-
 2 files changed, 40 insertions(+), 28 deletions(-)

diff --git a/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java b/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
index 921674a2d4..4d854b6716 100644
--- a/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
+++ b/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
@@ -46,6 +46,7 @@ import com.android.quickstep.util.RectFSpringAnim;
 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 @@ public class LauncherBackAnimationController {
     private boolean mAnimatorSetInProgress = false;
     private float mBackProgress = 0;
     private boolean mBackInProgress = false;
+    private IOnBackInvokedCallback mBackCallback;
 
     public LauncherBackAnimationController(
             BaseQuickstepLauncher launcher,
@@ -113,34 +115,35 @@ public class LauncherBackAnimationController {
      * @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 @@ public class LauncherBackAnimationController {
 
     /** 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 5ef89d309a..fc4c23f814 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -847,8 +847,14 @@ public class SystemUiProxy implements ISystemUiProxy, DisplayController.DisplayI
         }
     }
 
-    /** 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;