From cf33f1c5c2e12d6257d9579b9a9e33c9d954d183 Mon Sep 17 00:00:00 2001
From: Riddle Hsu <riddlehsu@google.com>
Date: Mon, 18 Feb 2019 21:20:51 +0800
Subject: [PATCH] Add restart button for size compatibility mode activity

- The floating restart button will show when an size compatibility
  mode activity shown with non-native screen configuration. e.g.
  display size changed, move to another display.
- Consolidate onDisplayRemoved into CommandQueue.Callback so the
  components which implement CommandQueue.Callbacks don't need to
  register display listener individually. The leakage of
  AutoHideController when removing display is also fixed by the way.

Bug: 112288258
Test: runtest systemui -c \
      com.android.systemui.SizeCompatModeActivityControllerTest

Change-Id: Ib04efe983ae0d8d21b33fb9fd9c60e7f6f0dc92e
---
 .../SystemUI/res/drawable/btn_restart.xml     |  28 ++
 .../res/layout/size_compat_mode_hint.xml      |  48 +++
 packages/SystemUI/res/values/config.xml       |   1 +
 packages/SystemUI/res/values/strings.xml      |   3 +
 .../system/TaskStackChangeListener.java       |   2 +
 .../system/TaskStackChangeListeners.java      |  15 +
 .../SizeCompatModeActivityController.java     | 285 ++++++++++++++++++
 .../systemui/statusbar/CommandQueue.java      |  13 +-
 .../statusbar/NavigationBarController.java    |  10 +-
 .../statusbar/phone/AutoHideController.java   |   7 +
 .../SizeCompatModeActivityControllerTest.java | 102 +++++++
 11 files changed, 503 insertions(+), 11 deletions(-)
 create mode 100644 packages/SystemUI/res/drawable/btn_restart.xml
 create mode 100644 packages/SystemUI/res/layout/size_compat_mode_hint.xml
 create mode 100644 packages/SystemUI/src/com/android/systemui/SizeCompatModeActivityController.java
 create mode 100644 packages/SystemUI/tests/src/com/android/systemui/SizeCompatModeActivityControllerTest.java

diff --git a/packages/SystemUI/res/drawable/btn_restart.xml b/packages/SystemUI/res/drawable/btn_restart.xml
new file mode 100644
index 0000000000000..73a48d31a8149
--- /dev/null
+++ b/packages/SystemUI/res/drawable/btn_restart.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2019 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24">
+    <path
+        android:fillColor="#aa000000"
+        android:pathData="M0,12 a12,12 0 1,0 24,0 a12,12 0 1,0 -24,0" />
+    <path
+        android:fillColor="@android:color/white"
+        android:pathData="M17.65,6.35c-1.63,-1.63 -3.94,-2.57 -6.48,-2.31c-3.67,0.37 -6.69,3.35 -7.1,7.02C3.52,15.91 7.27,20 12,20c3.19,0 5.93,-1.87 7.21,-4.57c0.31,-0.66 -0.16,-1.43 -0.89,-1.43h-0.01c-0.37,0 -0.72,0.2 -0.88,0.53c-1.13,2.43 -3.84,3.97 -6.81,3.32c-2.22,-0.49 -4.01,-2.3 -4.49,-4.52C5.31,9.44 8.26,6 12,6c1.66,0 3.14,0.69 4.22,1.78l-2.37,2.37C13.54,10.46 13.76,11 14.21,11H19c0.55,0 1,-0.45 1,-1V5.21c0,-0.45 -0.54,-0.67 -0.85,-0.35L17.65,6.35z"/>
+</vector>
diff --git a/packages/SystemUI/res/layout/size_compat_mode_hint.xml b/packages/SystemUI/res/layout/size_compat_mode_hint.xml
new file mode 100644
index 0000000000000..347c2b47767e1
--- /dev/null
+++ b/packages/SystemUI/res/layout/size_compat_mode_hint.xml
@@ -0,0 +1,48 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2019 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="wrap_content"
+    android:layout_height="wrap_content"
+    android:background="@android:color/background_light"
+    android:orientation="vertical">
+
+    <TextView
+        android:layout_width="180dp"
+        android:layout_height="wrap_content"
+        android:paddingLeft="10dp"
+        android:paddingRight="10dp"
+        android:paddingTop="10dp"
+        android:text="@string/restart_button_description"
+        android:textAlignment="viewStart"
+        android:textColor="@android:color/primary_text_light"
+        android:textSize="16sp" />
+
+    <Button
+        android:id="@+id/got_it"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:includeFontPadding="false"
+        android:layout_gravity="end"
+        android:minHeight="36dp"
+        android:background="?android:attr/selectableItemBackground"
+        android:text="@string/got_it"
+        android:textAllCaps="true"
+        android:textColor="#3c78d8"
+        android:textSize="16sp"
+        android:textStyle="bold" />
+
+</LinearLayout>
diff --git a/packages/SystemUI/res/values/config.xml b/packages/SystemUI/res/values/config.xml
index b4dc0eff0761d..1d566809b7315 100644
--- a/packages/SystemUI/res/values/config.xml
+++ b/packages/SystemUI/res/values/config.xml
@@ -311,6 +311,7 @@
         <item>com.android.systemui.ScreenDecorations</item>
         <item>com.android.systemui.biometrics.BiometricDialogImpl</item>
         <item>com.android.systemui.SliceBroadcastRelayHandler</item>
+        <item>com.android.systemui.SizeCompatModeActivityController</item>
     </string-array>
 
     <!-- SystemUI vender service, used in config_systemUIServiceComponents. -->
diff --git a/packages/SystemUI/res/values/strings.xml b/packages/SystemUI/res/values/strings.xml
index 5338bb4792d9e..dc46049eb6fcd 100644
--- a/packages/SystemUI/res/values/strings.xml
+++ b/packages/SystemUI/res/values/strings.xml
@@ -2373,6 +2373,9 @@
     <!-- What to show on the ambient display player when song doesn't have a title. [CHAR LIMIT=20] -->
     <string name="music_controls_no_title">No title</string>
 
+    <!-- Description of the restart button in the hint of size compatibility mode. [CHAR LIMIT=NONE] -->
+    <string name="restart_button_description">Tap to restart this app and go full screen.</string>
+
     <!-- Text used for content description of deep link button in the header of expanded bubble
          view. [CHAR_LIMIT=NONE] -->
     <string name="bubbles_deep_link_button_description">Open <xliff:g id="app_name" example="YouTube">%1$s</xliff:g></string>
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/system/TaskStackChangeListener.java b/packages/SystemUI/shared/src/com/android/systemui/shared/system/TaskStackChangeListener.java
index c54a4699964d6..f2a961d0b681e 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/system/TaskStackChangeListener.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/system/TaskStackChangeListener.java
@@ -18,6 +18,7 @@ package com.android.systemui.shared.system;
 
 import android.app.ActivityManager.RunningTaskInfo;
 import android.content.ComponentName;
+import android.os.IBinder;
 import android.os.UserHandle;
 import android.util.Log;
 
@@ -73,6 +74,7 @@ public abstract class TaskStackChangeListener {
     }
 
     public void onActivityRequestedOrientationChanged(int taskId, int requestedOrientation) { }
+    public void onSizeCompatModeActivityChanged(int displayId, IBinder activityToken) { }
 
     /**
      * Checks that the current user matches the process. Since
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/system/TaskStackChangeListeners.java b/packages/SystemUI/shared/src/com/android/systemui/shared/system/TaskStackChangeListeners.java
index 7e4ab854adcfc..6e3eb7c397bcf 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/system/TaskStackChangeListeners.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/system/TaskStackChangeListeners.java
@@ -23,6 +23,7 @@ import android.app.IActivityManager;
 import android.app.TaskStackListener;
 import android.content.ComponentName;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.Looper;
 import android.os.Message;
 import android.os.RemoteException;
@@ -180,6 +181,12 @@ public class TaskStackChangeListeners extends TaskStackListener {
                 requestedOrientation).sendToTarget();
     }
 
+    @Override
+    public void onSizeCompatModeActivityChanged(int displayId, IBinder activityToken) {
+        mHandler.obtainMessage(H.ON_SIZE_COMPAT_MODE_ACTIVITY_CHANGED, displayId, 0 /* unused */,
+                activityToken).sendToTarget();
+    }
+
     private final class H extends Handler {
         private static final int ON_TASK_STACK_CHANGED = 1;
         private static final int ON_TASK_SNAPSHOT_CHANGED = 2;
@@ -197,6 +204,7 @@ public class TaskStackChangeListeners extends TaskStackListener {
         private static final int ON_TASK_MOVED_TO_FRONT = 14;
         private static final int ON_ACTIVITY_REQUESTED_ORIENTATION_CHANGE = 15;
         private static final int ON_ACTIVITY_LAUNCH_ON_SECONDARY_DISPLAY_REROUTED = 16;
+        private static final int ON_SIZE_COMPAT_MODE_ACTIVITY_CHANGED = 17;
 
 
         public H(Looper looper) {
@@ -319,6 +327,13 @@ public class TaskStackChangeListeners extends TaskStackListener {
                         }
                         break;
                     }
+                    case ON_SIZE_COMPAT_MODE_ACTIVITY_CHANGED: {
+                        for (int i = mTaskStackListeners.size() - 1; i >= 0; i--) {
+                            mTaskStackListeners.get(i).onSizeCompatModeActivityChanged(
+                                    msg.arg1, (IBinder) msg.obj);
+                        }
+                        break;
+                    }
                 }
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/SizeCompatModeActivityController.java b/packages/SystemUI/src/com/android/systemui/SizeCompatModeActivityController.java
new file mode 100644
index 0000000000000..7a543442dd832
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/SizeCompatModeActivityController.java
@@ -0,0 +1,285 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui;
+
+import android.app.ActivityTaskManager;
+import android.content.Context;
+import android.content.res.ColorStateList;
+import android.graphics.Color;
+import android.graphics.PixelFormat;
+import android.graphics.drawable.Drawable;
+import android.graphics.drawable.GradientDrawable;
+import android.graphics.drawable.RippleDrawable;
+import android.hardware.display.DisplayManager;
+import android.inputmethodservice.InputMethodService;
+import android.os.IBinder;
+import android.os.RemoteException;
+import android.util.Log;
+import android.util.SparseArray;
+import android.view.Display;
+import android.view.Gravity;
+import android.view.LayoutInflater;
+import android.view.View;
+import android.view.WindowManager;
+import android.widget.Button;
+import android.widget.ImageButton;
+import android.widget.LinearLayout;
+import android.widget.PopupWindow;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.systemui.shared.system.ActivityManagerWrapper;
+import com.android.systemui.shared.system.TaskStackChangeListener;
+import com.android.systemui.statusbar.CommandQueue;
+
+import java.lang.ref.WeakReference;
+
+/** Shows a restart-activity button when the foreground activity is in size compatibility mode. */
+public class SizeCompatModeActivityController extends SystemUI implements CommandQueue.Callbacks {
+    private static final String TAG = "SizeCompatMode";
+
+    /** The showing buttons by display id. */
+    private final SparseArray<RestartActivityButton> mActiveButtons = new SparseArray<>(1);
+    /** Avoid creating display context frequently for non-default display. */
+    private final SparseArray<WeakReference<Context>> mDisplayContextCache = new SparseArray<>(0);
+
+    /** Only show once automatically in the process life. */
+    private boolean mHasShownHint;
+
+    public SizeCompatModeActivityController() {
+        this(ActivityManagerWrapper.getInstance());
+    }
+
+    @VisibleForTesting
+    SizeCompatModeActivityController(ActivityManagerWrapper am) {
+        am.registerTaskStackListener(new TaskStackChangeListener() {
+            @Override
+            public void onSizeCompatModeActivityChanged(int displayId, IBinder activityToken) {
+                // Note the callback already runs on main thread.
+                updateRestartButton(displayId, activityToken);
+            }
+        });
+    }
+
+    @Override
+    public void start() {
+        SysUiServiceProvider.getComponent(mContext, CommandQueue.class).addCallback(this);
+    }
+
+    @Override
+    public void setImeWindowStatus(int displayId, IBinder token, int vis, int backDisposition,
+            boolean showImeSwitcher) {
+        RestartActivityButton button = mActiveButtons.get(displayId);
+        if (button == null) {
+            return;
+        }
+        boolean imeShown = (vis & InputMethodService.IME_VISIBLE) != 0;
+        int newVisibility = imeShown ? View.GONE : View.VISIBLE;
+        // Hide the button when input method is showing.
+        if (button.getVisibility() != newVisibility) {
+            button.setVisibility(newVisibility);
+        }
+    }
+
+    @Override
+    public void onDisplayRemoved(int displayId) {
+        mDisplayContextCache.remove(displayId);
+        removeRestartButton(displayId);
+    }
+
+    private void removeRestartButton(int displayId) {
+        RestartActivityButton button = mActiveButtons.get(displayId);
+        if (button != null) {
+            button.remove();
+            mActiveButtons.remove(displayId);
+        }
+    }
+
+    private void updateRestartButton(int displayId, IBinder activityToken) {
+        if (activityToken == null) {
+            // Null token means the current foreground activity is not in size compatibility mode.
+            removeRestartButton(displayId);
+            return;
+        }
+
+        RestartActivityButton restartButton = mActiveButtons.get(displayId);
+        if (restartButton != null) {
+            restartButton.updateLastTargetActivity(activityToken);
+            return;
+        }
+
+        Context context = getOrCreateDisplayContext(displayId);
+        if (context == null) {
+            Log.i(TAG, "Cannot get context for display " + displayId);
+            return;
+        }
+
+        restartButton = createRestartButton(context);
+        restartButton.updateLastTargetActivity(activityToken);
+        restartButton.show();
+        mActiveButtons.append(displayId, restartButton);
+    }
+
+    @VisibleForTesting
+    RestartActivityButton createRestartButton(Context context) {
+        RestartActivityButton button = new RestartActivityButton(context, mHasShownHint);
+        mHasShownHint = true;
+        return button;
+    }
+
+    private Context getOrCreateDisplayContext(int displayId) {
+        if (displayId == Display.DEFAULT_DISPLAY) {
+            return mContext;
+        }
+        Context context = null;
+        WeakReference<Context> ref = mDisplayContextCache.get(displayId);
+        if (ref != null) {
+            context = ref.get();
+        }
+        if (context == null) {
+            Display display = mContext.getSystemService(DisplayManager.class).getDisplay(displayId);
+            if (display != null) {
+                context = mContext.createDisplayContext(display);
+                mDisplayContextCache.put(displayId, new WeakReference<Context>(context));
+            }
+        }
+        return context;
+    }
+
+    @VisibleForTesting
+    static class RestartActivityButton extends ImageButton implements View.OnClickListener,
+            View.OnLongClickListener {
+
+        final WindowManager.LayoutParams mWinParams;
+        final boolean mShouldShowHint;
+        IBinder mLastActivityToken;
+
+        final int mPopupOffsetX;
+        final int mPopupOffsetY;
+        PopupWindow mShowingHint;
+
+        RestartActivityButton(Context context, boolean hasShownHint) {
+            super(context);
+            mShouldShowHint = !hasShownHint;
+            Drawable drawable = context.getDrawable(R.drawable.btn_restart);
+            setImageDrawable(drawable);
+            setContentDescription(context.getString(R.string.restart_button_description));
+
+            int drawableW = drawable.getIntrinsicWidth();
+            int drawableH = drawable.getIntrinsicHeight();
+            mPopupOffsetX = drawableW / 2;
+            mPopupOffsetY = drawableH * 2;
+
+            ColorStateList color = ColorStateList.valueOf(Color.LTGRAY);
+            GradientDrawable mask = new GradientDrawable();
+            mask.setShape(GradientDrawable.OVAL);
+            mask.setColor(color);
+            setBackground(new RippleDrawable(color, null /* content */, mask));
+            setOnClickListener(this);
+            setOnLongClickListener(this);
+
+            mWinParams = new WindowManager.LayoutParams();
+            mWinParams.gravity = getGravity(getResources().getConfiguration().getLayoutDirection());
+            mWinParams.width = drawableW * 2;
+            mWinParams.height = drawableH * 2;
+            mWinParams.type = WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
+            mWinParams.flags = WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE
+                    | WindowManager.LayoutParams.FLAG_NOT_TOUCH_MODAL;
+            mWinParams.format = PixelFormat.TRANSLUCENT;
+            mWinParams.privateFlags |= WindowManager.LayoutParams.PRIVATE_FLAG_SHOW_FOR_ALL_USERS;
+            mWinParams.setTitle(SizeCompatModeActivityController.class.getSimpleName()
+                    + context.getDisplayId());
+        }
+
+        void updateLastTargetActivity(IBinder activityToken) {
+            mLastActivityToken = activityToken;
+        }
+
+        void show() {
+            getContext().getSystemService(WindowManager.class).addView(this, mWinParams);
+        }
+
+        void remove() {
+            getContext().getSystemService(WindowManager.class).removeViewImmediate(this);
+        }
+
+        @Override
+        public void onClick(View v) {
+            try {
+                ActivityTaskManager.getService().restartActivityProcessIfVisible(
+                        mLastActivityToken);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Unable to restart activity", e);
+            }
+        }
+
+        @Override
+        public boolean onLongClick(View v) {
+            showHint();
+            return true;
+        }
+
+        @Override
+        protected void onAttachedToWindow() {
+            super.onAttachedToWindow();
+            if (mShouldShowHint) {
+                showHint();
+            }
+        }
+
+        @Override
+        public void setLayoutDirection(int layoutDirection) {
+            int gravity = getGravity(layoutDirection);
+            if (mWinParams.gravity != gravity) {
+                mWinParams.gravity = gravity;
+                if (mShowingHint != null) {
+                    mShowingHint.dismiss();
+                    showHint();
+                }
+                getContext().getSystemService(WindowManager.class).updateViewLayout(this,
+                        mWinParams);
+            }
+            super.setLayoutDirection(layoutDirection);
+        }
+
+        void showHint() {
+            if (mShowingHint != null) {
+                return;
+            }
+
+            View popupView = LayoutInflater.from(getContext()).inflate(
+                    R.layout.size_compat_mode_hint, null /* root */);
+            PopupWindow popupWindow = new PopupWindow(popupView,
+                    LinearLayout.LayoutParams.WRAP_CONTENT, LinearLayout.LayoutParams.WRAP_CONTENT);
+            popupWindow.setElevation(getResources().getDimension(R.dimen.bubble_elevation));
+            popupWindow.setAnimationStyle(android.R.style.Animation_InputMethod);
+            popupWindow.setClippingEnabled(false);
+            popupWindow.setOnDismissListener(() -> mShowingHint = null);
+            mShowingHint = popupWindow;
+
+            Button gotItButton = popupView.findViewById(R.id.got_it);
+            gotItButton.setBackground(new RippleDrawable(ColorStateList.valueOf(Color.LTGRAY),
+                    null /* content */, null /* mask */));
+            gotItButton.setOnClickListener(view -> popupWindow.dismiss());
+            popupWindow.showAtLocation(this, mWinParams.gravity, mPopupOffsetX, mPopupOffsetY);
+        }
+
+        private static int getGravity(int layoutDirection) {
+            return Gravity.BOTTOM
+                    | (layoutDirection == View.LAYOUT_DIRECTION_RTL ? Gravity.START : Gravity.END);
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/CommandQueue.java b/packages/SystemUI/src/com/android/systemui/statusbar/CommandQueue.java
index 015085205460b..295abcbcfd176 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/CommandQueue.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/CommandQueue.java
@@ -277,9 +277,13 @@ public class CommandQueue extends IStatusBar.Stub implements CallbackController<
         default void hideBiometricDialog() { }
 
         /**
-        * @see IStatusBar#onDisplayReady(int)
-        */
+         * @see IStatusBar#onDisplayReady(int)
+         */
         default void onDisplayReady(int displayId) { }
+        /**
+         * @see DisplayManager.DisplayListener#onDisplayRemoved(int)
+         */
+        default void onDisplayRemoved(int displayId) { }
     }
 
     @VisibleForTesting
@@ -297,6 +301,11 @@ public class CommandQueue extends IStatusBar.Stub implements CallbackController<
         synchronized (mLock) {
             mDisplayDisabled.remove(displayId);
         }
+        // This callback is registered with {@link #mHandler} that already posts to run on main
+        // thread, so it is safe to dispatch directly.
+        for (int i = mCallbacks.size() - 1; i >= 0; i--) {
+            mCallbacks.get(i).onDisplayRemoved(displayId);
+        }
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NavigationBarController.java b/packages/SystemUI/src/com/android/systemui/statusbar/NavigationBarController.java
index 9f382eabc1804..14dc272d1c5e8 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NavigationBarController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NavigationBarController.java
@@ -23,7 +23,6 @@ import static com.android.systemui.SysUiServiceProvider.getComponent;
 
 import android.content.Context;
 import android.hardware.display.DisplayManager;
-import android.hardware.display.DisplayManager.DisplayListener;
 import android.os.Handler;
 import android.os.RemoteException;
 import android.util.Log;
@@ -50,7 +49,7 @@ import javax.inject.Singleton;
 
 /** A controller to handle navigation bars. */
 @Singleton
-public class NavigationBarController implements DisplayListener, Callbacks {
+public class NavigationBarController implements Callbacks {
 
     private static final String TAG = NavigationBarController.class.getName();
 
@@ -66,21 +65,14 @@ public class NavigationBarController implements DisplayListener, Callbacks {
         mContext = context;
         mHandler = handler;
         mDisplayManager = (DisplayManager) mContext.getSystemService(Context.DISPLAY_SERVICE);
-        mDisplayManager.registerDisplayListener(this, mHandler);
         getComponent(mContext, CommandQueue.class).addCallback(this);
     }
 
-    @Override
-    public void onDisplayAdded(int displayId) { }
-
     @Override
     public void onDisplayRemoved(int displayId) {
         removeNavigationBar(displayId);
     }
 
-    @Override
-    public void onDisplayChanged(int displayId) { }
-
     @Override
     public void onDisplayReady(int displayId) {
         Display display = mDisplayManager.getDisplay(displayId);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/AutoHideController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/AutoHideController.java
index b9425d4745d22..2bbc53c05cbe3 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/AutoHideController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/AutoHideController.java
@@ -75,6 +75,13 @@ public class AutoHideController implements CommandQueue.Callbacks {
         mDisplayId = context.getDisplayId();
     }
 
+    @Override
+    public void onDisplayRemoved(int displayId) {
+        if (displayId == mDisplayId) {
+            mCommandQueue.removeCallback(this);
+        }
+    }
+
     void setStatusBar(StatusBar statusBar) {
         mStatusBar = statusBar;
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/SizeCompatModeActivityControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/SizeCompatModeActivityControllerTest.java
new file mode 100644
index 0000000000000..587005ea6e6a4
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/SizeCompatModeActivityControllerTest.java
@@ -0,0 +1,102 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui;
+
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.verify;
+
+import android.content.Context;
+import android.inputmethodservice.InputMethodService;
+import android.os.IBinder;
+import android.testing.AndroidTestingRunner;
+import android.view.View;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.systemui.SizeCompatModeActivityController.RestartActivityButton;
+import com.android.systemui.shared.system.ActivityManagerWrapper;
+import com.android.systemui.shared.system.TaskStackChangeListener;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * runtest systemui -c com.android.systemui.SizeCompatModeActivityControllerTest
+ */
+@RunWith(AndroidTestingRunner.class)
+@SmallTest
+public class SizeCompatModeActivityControllerTest extends SysuiTestCase {
+    private static final int DISPLAY_ID = 0;
+
+    private SizeCompatModeActivityController mController;
+    private TaskStackChangeListener mTaskStackListener;
+    private @Mock ActivityManagerWrapper mMockAm;
+    private @Mock RestartActivityButton mMockButton;
+    private @Mock IBinder mMockActivityToken;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+
+        mController = new SizeCompatModeActivityController(mMockAm) {
+            @Override
+            RestartActivityButton createRestartButton(Context context) {
+                return mMockButton;
+            };
+        };
+        mController.mContext = mContext;
+
+        ArgumentCaptor<TaskStackChangeListener> listenerCaptor =
+                ArgumentCaptor.forClass(TaskStackChangeListener.class);
+        verify(mMockAm).registerTaskStackListener(listenerCaptor.capture());
+        mTaskStackListener = listenerCaptor.getValue();
+    }
+
+    @Test
+    public void testOnSizeCompatModeActivityChanged() {
+        // Verifies that the restart button is added with non-null component name.
+        mTaskStackListener.onSizeCompatModeActivityChanged(DISPLAY_ID, mMockActivityToken);
+        verify(mMockButton).show();
+        verify(mMockButton).updateLastTargetActivity(eq(mMockActivityToken));
+
+        // Verifies that the restart button is removed with null component name.
+        mTaskStackListener.onSizeCompatModeActivityChanged(DISPLAY_ID, null /* activityToken */);
+        verify(mMockButton).remove();
+    }
+
+    @Test
+    public void testChangeButtonVisibilityOnImeShowHide() {
+        mTaskStackListener.onSizeCompatModeActivityChanged(DISPLAY_ID, mMockActivityToken);
+
+        // Verifies that the restart button is hidden when IME is visible.
+        doReturn(View.VISIBLE).when(mMockButton).getVisibility();
+        mController.setImeWindowStatus(DISPLAY_ID, null /* token */, InputMethodService.IME_VISIBLE,
+                0 /* backDisposition */, false /* showImeSwitcher */);
+        verify(mMockButton).setVisibility(eq(View.GONE));
+
+        // Verifies that the restart button is visible when IME is hidden.
+        doReturn(View.GONE).when(mMockButton).getVisibility();
+        mController.setImeWindowStatus(DISPLAY_ID, null /* token */, 0 /* vis */,
+                0 /* backDisposition */, false /* showImeSwitcher */);
+        verify(mMockButton).setVisibility(eq(View.VISIBLE));
+    }
+}