diff --git a/core/java/android/app/TaskInfo.java b/core/java/android/app/TaskInfo.java
index 938ce0d56933d325e5cabfba166932534cc791db..9019ddf941d94b97a2165284a2abea980dca8085 100644
--- a/core/java/android/app/TaskInfo.java
+++ b/core/java/android/app/TaskInfo.java
@@ -343,7 +343,9 @@ public class TaskInfo {
                 // TopActivityToken and bounds are important if top activity is in size compat
                 && (!topActivityInSizeCompat || topActivityToken.equals(that.topActivityToken))
                 && (!topActivityInSizeCompat || configuration.windowConfiguration.getBounds()
-                    .equals(that.configuration.windowConfiguration.getBounds()));
+                    .equals(that.configuration.windowConfiguration.getBounds()))
+                && (!topActivityInSizeCompat || configuration.getLayoutDirection()
+                    == that.configuration.getLayoutDirection());
     }
 
     /**
diff --git a/libs/WindowManager/Shell/res/layout/size_compat_ui.xml b/libs/WindowManager/Shell/res/layout/size_compat_ui.xml
new file mode 100644
index 0000000000000000000000000000000000000000..cd3153145be341f53baf593faf03fb847eb5e32d
--- /dev/null
+++ b/libs/WindowManager/Shell/res/layout/size_compat_ui.xml
@@ -0,0 +1,30 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2021 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.
+  -->
+<com.android.wm.shell.sizecompatui.SizeCompatRestartButton
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="wrap_content"
+    android:layout_height="wrap_content">
+
+    <ImageButton
+        android:id="@+id/size_compat_restart_button"
+        android:layout_width="@dimen/size_compat_button_size"
+        android:layout_height="@dimen/size_compat_button_size"
+        android:layout_gravity="center"
+        android:src="@drawable/size_compat_restart_button"
+        android:contentDescription="@string/restart_button_description"/>
+
+</com.android.wm.shell.sizecompatui.SizeCompatRestartButton>
diff --git a/libs/WindowManager/Shell/res/values/dimen.xml b/libs/WindowManager/Shell/res/values/dimen.xml
index cb2aada2a9cffc7e6a62f8e377cd386767689fab..583964b2f4a4d8e02b9e491110881403f4f65d61 100644
--- a/libs/WindowManager/Shell/res/values/dimen.xml
+++ b/libs/WindowManager/Shell/res/values/dimen.xml
@@ -174,6 +174,9 @@
     <!-- The width/height of the icon view on staring surface. -->
     <dimen name="starting_surface_icon_size">108dp</dimen>
 
+    <!-- The width/height of the size compat restart button. -->
+    <dimen name="size_compat_button_size">48dp</dimen>
+
     <!-- The width of the brand image on staring surface. -->
     <dimen name="starting_surface_brand_image_width">200dp</dimen>
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/FullscreenTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/FullscreenTaskListener.java
index afe523af7cb0268315b007684a4d89f9f8d70b3d..59fa9948bc2d9c196ebf4622a127de8e0eb8264f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/FullscreenTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/FullscreenTaskListener.java
@@ -94,10 +94,22 @@ public class FullscreenTaskListener implements ShellTaskOrganizer.TaskListener {
                 taskInfo.taskId);
     }
 
+    @Override
+    public boolean supportSizeCompatUI() {
+        return true;
+    }
+
+    @Override
+    public void attachChildSurfaceToTask(int taskId, SurfaceControl.Builder b) {
+        if (!mLeashByTaskId.contains(taskId)) {
+            throw new IllegalArgumentException("There is no surface for taskId=" + taskId);
+        }
+        b.setParent(mLeashByTaskId.get(taskId));
+    }
+
     @Override
     public void dump(@NonNull PrintWriter pw, String prefix) {
         final String innerPrefix = prefix + "  ";
-        final String childPrefix = innerPrefix + "  ";
         pw.println(prefix + this);
         pw.println(innerPrefix + mLeashByTaskId.size() + " Tasks");
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
index 03724171da2170ad7cbcfa1686c0be7c48fec298..34a7157dbb3ce99f1a63ef0629a9f6141f88d046 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
@@ -82,6 +82,15 @@ public class ShellTaskOrganizer extends TaskOrganizer {
         default void onTaskInfoChanged(RunningTaskInfo taskInfo) {}
         default void onTaskVanished(RunningTaskInfo taskInfo) {}
         default void onBackPressedOnTaskRoot(RunningTaskInfo taskInfo) {}
+        /** Whether this task listener supports size compat UI. */
+        default boolean supportSizeCompatUI() {
+            return false;
+        }
+        /** Attaches the a child window surface to the task surface. */
+        default void attachChildSurfaceToTask(int taskId, SurfaceControl.Builder b) {
+            throw new IllegalStateException(
+                    "This task listener doesn't support child surface attachment.");
+        }
         default void dump(@NonNull PrintWriter pw, String prefix) {};
     }
 
@@ -358,8 +367,10 @@ public class ShellTaskOrganizer extends TaskOrganizer {
             return;
         }
 
-        // The task is vanished, notify to remove size compat UI on this Task if there is any.
-        if (taskListener == null) {
+        // The task is vanished or doesn't support size compat UI, notify to remove size compat UI
+        // on this Task if there is any.
+        if (taskListener == null || !taskListener.supportSizeCompatUI()
+                || !taskInfo.topActivityInSizeCompat) {
             mSizeCompatUI.onSizeCompatInfoChanged(taskInfo.displayId, taskInfo.taskId,
                     null /* taskConfig */, null /* sizeCompatActivity*/,
                     null /* taskListener */);
@@ -367,10 +378,7 @@ public class ShellTaskOrganizer extends TaskOrganizer {
         }
 
         mSizeCompatUI.onSizeCompatInfoChanged(taskInfo.displayId, taskInfo.taskId,
-                taskInfo.configuration.windowConfiguration.getBounds(),
-                // null if the top activity not in size compat.
-                taskInfo.topActivityInSizeCompat ? taskInfo.topActivityToken : null,
-                taskListener);
+                taskInfo.configuration, taskInfo.topActivityToken, taskListener);
     }
 
     private TaskListener getTaskListener(RunningTaskInfo runningTaskInfo) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatRestartButton.java b/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatRestartButton.java
index e47e1ac71c73332798eb5e1149d14f303e5223a6..9094d7de8d632ce366fe70c0d6d9729e0774ea3f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatRestartButton.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatRestartButton.java
@@ -16,101 +16,80 @@
 
 package com.android.wm.shell.sizecompatui;
 
-import android.app.ActivityClient;
 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.os.IBinder;
-import android.util.Log;
-import android.view.Gravity;
+import android.util.AttributeSet;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.WindowManager;
 import android.widget.Button;
+import android.widget.FrameLayout;
 import android.widget.ImageButton;
 import android.widget.LinearLayout;
 import android.widget.PopupWindow;
 
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.wm.shell.R;
 
 /** Button to restart the size compat activity. */
-class SizeCompatRestartButton extends ImageButton implements View.OnClickListener,
+public class SizeCompatRestartButton extends FrameLayout implements View.OnClickListener,
         View.OnLongClickListener {
-    private static final String TAG = "SizeCompatRestartButton";
-
-    final WindowManager.LayoutParams mWinParams;
-    final boolean mShouldShowHint;
-    final int mDisplayId;
-    final int mPopupOffsetX;
-    final int mPopupOffsetY;
 
-    private IBinder mLastActivityToken;
-    private PopupWindow mShowingHint;
+    private SizeCompatUILayout mLayout;
+    private ImageButton mRestartButton;
+    @VisibleForTesting
+    PopupWindow mShowingHint;
+    private WindowManager.LayoutParams mWinParams;
 
-    SizeCompatRestartButton(Context context, int displayId, boolean hasShownHint) {
+    public SizeCompatRestartButton(@NonNull Context context) {
         super(context);
-        mDisplayId = displayId;
-        mShouldShowHint = !hasShownHint;
-        final Drawable drawable = context.getDrawable(R.drawable.size_compat_restart_button);
-        setImageDrawable(drawable);
-        setContentDescription(context.getString(R.string.restart_button_description));
+    }
 
-        final int drawableW = drawable.getIntrinsicWidth();
-        final int drawableH = drawable.getIntrinsicHeight();
-        mPopupOffsetX = drawableW / 2;
-        mPopupOffsetY = drawableH * 2;
+    public SizeCompatRestartButton(@NonNull Context context, @Nullable AttributeSet attrs) {
+        super(context, attrs);
+    }
 
-        final ColorStateList color = ColorStateList.valueOf(Color.LTGRAY);
-        final 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.SYSTEM_FLAG_SHOW_FOR_ALL_USERS;
-        mWinParams.setTitle(SizeCompatRestartButton.class.getSimpleName()
-                + context.getDisplayId());
+    public SizeCompatRestartButton(@NonNull Context context, @Nullable AttributeSet attrs,
+            int defStyleAttr) {
+        super(context, attrs, defStyleAttr);
     }
 
-    void updateLastTargetActivity(IBinder activityToken) {
-        mLastActivityToken = activityToken;
+    public SizeCompatRestartButton(@NonNull Context context, @Nullable AttributeSet attrs,
+            int defStyleAttr, int defStyleRes) {
+        super(context, attrs, defStyleAttr, defStyleRes);
     }
 
-    /** @return {@code false} if the target display is invalid. */
-    boolean show() {
-        try {
-            getContext().getSystemService(WindowManager.class).addView(this, mWinParams);
-        } catch (WindowManager.InvalidDisplayException e) {
-            // The target display may have been removed when the callback has just arrived.
-            Log.w(TAG, "Cannot show on display " + getContext().getDisplayId(), e);
-            return false;
-        }
-        return true;
+    void inject(SizeCompatUILayout layout) {
+        mLayout = layout;
+        mWinParams = layout.getWindowLayoutParams();
     }
 
     void remove() {
-        if (mShowingHint != null) {
-            mShowingHint.dismiss();
-        }
-        getContext().getSystemService(WindowManager.class).removeViewImmediate(this);
+        dismissHint();
+    }
+
+    @Override
+    protected void onFinishInflate() {
+        super.onFinishInflate();
+        mRestartButton = findViewById(R.id.size_compat_restart_button);
+        final ColorStateList color = ColorStateList.valueOf(Color.LTGRAY);
+        final GradientDrawable mask = new GradientDrawable();
+        mask.setShape(GradientDrawable.OVAL);
+        mask.setColor(color);
+        mRestartButton.setBackground(new RippleDrawable(color, null /* content */, mask));
+        mRestartButton.setOnClickListener(this);
+        mRestartButton.setOnLongClickListener(this);
     }
 
     @Override
     public void onClick(View v) {
-        ActivityClient.getInstance().restartActivityProcessIfVisible(mLastActivityToken);
+        mLayout.onRestartButtonClicked();
     }
 
     @Override
@@ -122,20 +101,26 @@ class SizeCompatRestartButton extends ImageButton implements View.OnClickListene
     @Override
     protected void onAttachedToWindow() {
         super.onAttachedToWindow();
-        if (mShouldShowHint) {
+        if (mLayout.mShouldShowHint) {
+            mLayout.mShouldShowHint = false;
             showHint();
         }
     }
 
+    @Override
+    public void setVisibility(@Visibility int visibility) {
+        if (visibility == View.GONE && mShowingHint != null) {
+            // Also dismiss the popup.
+            dismissHint();
+        }
+        super.setVisibility(visibility);
+    }
+
     @Override
     public void setLayoutDirection(int layoutDirection) {
-        final int gravity = getGravity(layoutDirection);
+        final int gravity = SizeCompatUILayout.getGravity(layoutDirection);
         if (mWinParams.gravity != gravity) {
             mWinParams.gravity = gravity;
-            if (mShowingHint != null) {
-                mShowingHint.dismiss();
-                showHint();
-            }
             getContext().getSystemService(WindowManager.class).updateViewLayout(this,
                     mWinParams);
         }
@@ -147,8 +132,10 @@ class SizeCompatRestartButton extends ImageButton implements View.OnClickListene
             return;
         }
 
+        // TODO: popup is not attached to the button surface. Need to handle this differently for
+        // non-fullscreen task.
         final View popupView = LayoutInflater.from(getContext()).inflate(
-                R.layout.size_compat_mode_hint, null /* root */);
+                R.layout.size_compat_mode_hint, null);
         final PopupWindow popupWindow = new PopupWindow(popupView,
                 LinearLayout.LayoutParams.WRAP_CONTENT, LinearLayout.LayoutParams.WRAP_CONTENT);
         popupWindow.setWindowLayoutType(mWinParams.type);
@@ -161,12 +148,15 @@ class SizeCompatRestartButton extends ImageButton implements View.OnClickListene
         final 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);
+        gotItButton.setOnClickListener(view -> dismissHint());
+        popupWindow.showAtLocation(mRestartButton, mWinParams.gravity, mLayout.mPopupOffsetX,
+                mLayout.mPopupOffsetY);
     }
 
-    private static int getGravity(int layoutDirection) {
-        return Gravity.BOTTOM
-                | (layoutDirection == View.LAYOUT_DIRECTION_RTL ? Gravity.START : Gravity.END);
+    void dismissHint() {
+        if (mShowingHint != null) {
+            mShowingHint.dismiss();
+            mShowingHint = null;
+        }
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUIController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUIController.java
index 48ee86c4954fd25165c577e40e7a439001f4916b..a3880f497ff337e87f3fb1a06071c52b4aa50854 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUIController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUIController.java
@@ -18,134 +18,166 @@ package com.android.wm.shell.sizecompatui;
 
 import android.annotation.Nullable;
 import android.content.Context;
-import android.graphics.Rect;
+import android.content.res.Configuration;
 import android.hardware.display.DisplayManager;
 import android.os.IBinder;
+import android.util.ArraySet;
 import android.util.Log;
 import android.util.SparseArray;
 import android.view.Display;
-import android.view.View;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.DisplayImeController;
-import com.android.wm.shell.common.ShellExecutor;
+import com.android.wm.shell.common.DisplayLayout;
+import com.android.wm.shell.common.SyncTransactionQueue;
 
 import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.function.Consumer;
 
 /**
- * Shows a restart-activity button on Task when the foreground activity is in size compatibility
- * mode.
+ * Controls to show/update restart-activity buttons on Tasks based on whether the foreground
+ * activities are in size compatibility mode.
  */
 public class SizeCompatUIController implements DisplayController.OnDisplaysChangedListener,
         DisplayImeController.ImePositionProcessor {
-    private static final String TAG = "SizeCompatUI";
+    private static final String TAG = "SizeCompatUIController";
+
+    /** Whether the IME is shown on display id. */
+    private final Set<Integer> mDisplaysWithIme = new ArraySet<>(1);
 
     /** The showing buttons by task id. */
-    private final SparseArray<SizeCompatRestartButton> mActiveButtons = new SparseArray<>(1);
+    private final SparseArray<SizeCompatUILayout> mActiveLayouts = new SparseArray<>(0);
+
     /** Avoid creating display context frequently for non-default display. */
     private final SparseArray<WeakReference<Context>> mDisplayContextCache = new SparseArray<>(0);
 
-    @VisibleForTesting
     private final Context mContext;
-    private final ShellExecutor mMainExecutor;
     private final DisplayController mDisplayController;
     private final DisplayImeController mImeController;
+    private final SyncTransactionQueue mSyncQueue;
 
     /** Only show once automatically in the process life. */
     private boolean mHasShownHint;
 
-    @VisibleForTesting
     public SizeCompatUIController(Context context,
             DisplayController displayController,
             DisplayImeController imeController,
-            ShellExecutor mainExecutor) {
+            SyncTransactionQueue syncQueue) {
         mContext = context;
-        mMainExecutor = mainExecutor;
         mDisplayController = displayController;
         mImeController = imeController;
+        mSyncQueue = syncQueue;
         mDisplayController.addDisplayWindowListener(this);
         mImeController.addPositionProcessor(this);
     }
 
-    public void onSizeCompatInfoChanged(int displayId, int taskId, @Nullable Rect taskBounds,
-            @Nullable IBinder sizeCompatActivity,
+    /**
+     * Called when the Task info changed. Creates and updates the restart button if there is an
+     * activity in size compat, or removes the restart button if there is no size compat activity.
+     *
+     * @param displayId display the task and activity are in.
+     * @param taskId task the activity is in.
+     * @param taskConfig task config to place the restart button with.
+     * @param sizeCompatActivity the size compat activity in the task. Can be {@code null} if the
+     *                           top activity in this Task is not in size compat.
+     * @param taskListener listener to handle the Task Surface placement.
+     */
+    public void onSizeCompatInfoChanged(int displayId, int taskId,
+            @Nullable Configuration taskConfig, @Nullable IBinder sizeCompatActivity,
             @Nullable ShellTaskOrganizer.TaskListener taskListener) {
-        // TODO Draw button on Task surface
-        if (taskBounds == null || sizeCompatActivity == null || taskListener == null) {
+        if (taskConfig == null || sizeCompatActivity == null || taskListener == null) {
             // Null token means the current foreground activity is not in size compatibility mode.
-            removeRestartButton(taskId);
+            removeLayout(taskId);
+        } else if (mActiveLayouts.contains(taskId)) {
+            // Button already exists, update the button layout.
+            updateLayout(taskId, taskConfig, sizeCompatActivity, taskListener);
         } else {
-            updateRestartButton(displayId, taskId, sizeCompatActivity);
+            // Create a new restart button.
+            createLayout(displayId, taskId, taskConfig, sizeCompatActivity, taskListener);
         }
     }
 
     @Override
     public void onDisplayRemoved(int displayId) {
         mDisplayContextCache.remove(displayId);
-        for (int i = 0; i < mActiveButtons.size(); i++) {
-            final int taskId = mActiveButtons.keyAt(i);
-            final SizeCompatRestartButton button = mActiveButtons.get(taskId);
-            if (button != null && button.mDisplayId == displayId) {
-                removeRestartButton(taskId);
-            }
+
+        // Remove all buttons on the removed display.
+        final List<Integer> toRemoveTaskIds = new ArrayList<>();
+        forAllLayoutsOnDisplay(displayId, layout -> toRemoveTaskIds.add(layout.getTaskId()));
+        for (int i = toRemoveTaskIds.size() - 1; i >= 0; i--) {
+            removeLayout(toRemoveTaskIds.get(i));
         }
     }
 
     @Override
-    public void onImeVisibilityChanged(int displayId, boolean isShowing) {
-        final int newVisibility = isShowing ? View.GONE : View.VISIBLE;
-        for (int i = 0; i < mActiveButtons.size(); i++) {
-            final int taskId = mActiveButtons.keyAt(i);
-            final SizeCompatRestartButton button = mActiveButtons.get(taskId);
-            if (button == null || button.mDisplayId != displayId) {
-                continue;
-            }
+    public void onDisplayConfigurationChanged(int displayId, Configuration newConfig) {
+        final DisplayLayout displayLayout = mDisplayController.getDisplayLayout(displayId);
+        forAllLayoutsOnDisplay(displayId, layout -> layout.updateDisplayLayout(displayLayout));
+    }
 
-            // Hide the button when input method is showing.
-            if (button.getVisibility() != newVisibility) {
-                button.setVisibility(newVisibility);
-            }
+    @Override
+    public void onImeVisibilityChanged(int displayId, boolean isShowing) {
+        if (isShowing) {
+            mDisplaysWithIme.add(displayId);
+        } else {
+            mDisplaysWithIme.remove(displayId);
         }
+
+        // Hide the button when input method is showing.
+        forAllLayoutsOnDisplay(displayId, layout -> layout.updateImeVisibility(isShowing));
     }
 
-    private void updateRestartButton(int displayId, int taskId, IBinder activityToken) {
-        SizeCompatRestartButton restartButton = mActiveButtons.get(taskId);
-        if (restartButton != null) {
-            restartButton.updateLastTargetActivity(activityToken);
-            return;
-        }
+    private boolean isImeShowingOnDisplay(int displayId) {
+        return mDisplaysWithIme.contains(displayId);
+    }
 
+    private void createLayout(int displayId, int taskId, Configuration taskConfig,
+            IBinder activityToken, ShellTaskOrganizer.TaskListener taskListener) {
         final Context context = getOrCreateDisplayContext(displayId);
         if (context == null) {
-            Log.i(TAG, "Cannot get context for display " + displayId);
+            Log.e(TAG, "Cannot get context for display " + displayId);
             return;
         }
 
-        restartButton = createRestartButton(context, displayId);
-        restartButton.updateLastTargetActivity(activityToken);
-        if (restartButton.show()) {
-            mActiveButtons.append(taskId, restartButton);
-        } else {
-            onDisplayRemoved(displayId);
-        }
+        final SizeCompatUILayout layout = createLayout(context, displayId, taskId, taskConfig,
+                activityToken, taskListener);
+        mActiveLayouts.put(taskId, layout);
+        layout.createSizeCompatButton(isImeShowingOnDisplay(displayId));
     }
 
     @VisibleForTesting
-    SizeCompatRestartButton createRestartButton(Context context, int displayId) {
-        final SizeCompatRestartButton button = new SizeCompatRestartButton(context, displayId,
+    SizeCompatUILayout createLayout(Context context, int displayId, int taskId,
+            Configuration taskConfig, IBinder activityToken,
+            ShellTaskOrganizer.TaskListener taskListener) {
+        final SizeCompatUILayout layout = new SizeCompatUILayout(mSyncQueue, context, taskConfig,
+                taskId, activityToken, taskListener, mDisplayController.getDisplayLayout(displayId),
                 mHasShownHint);
         // Only show hint for the first time.
         mHasShownHint = true;
-        return button;
+        return layout;
     }
 
-    private void removeRestartButton(int taskId) {
-        final SizeCompatRestartButton button = mActiveButtons.get(taskId);
-        if (button != null) {
-            button.remove();
-            mActiveButtons.remove(taskId);
+    private void updateLayout(int taskId, Configuration taskConfig,
+            IBinder sizeCompatActivity,
+            ShellTaskOrganizer.TaskListener taskListener) {
+        final SizeCompatUILayout layout = mActiveLayouts.get(taskId);
+        if (layout == null) {
+            return;
+        }
+        layout.updateSizeCompatInfo(taskConfig, sizeCompatActivity, taskListener,
+                isImeShowingOnDisplay(layout.getDisplayId()));
+    }
+
+    private void removeLayout(int taskId) {
+        final SizeCompatUILayout layout = mActiveLayouts.get(taskId);
+        if (layout != null) {
+            layout.release();
+            mActiveLayouts.remove(taskId);
         }
     }
 
@@ -167,4 +199,14 @@ public class SizeCompatUIController implements DisplayController.OnDisplaysChang
         }
         return context;
     }
+
+    private void forAllLayoutsOnDisplay(int displayId, Consumer<SizeCompatUILayout> callback) {
+        for (int i = 0; i < mActiveLayouts.size(); i++) {
+            final int taskId = mActiveLayouts.keyAt(i);
+            final SizeCompatUILayout layout = mActiveLayouts.get(taskId);
+            if (layout != null && layout.getDisplayId() == displayId) {
+                callback.accept(layout);
+            }
+        }
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUILayout.java b/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUILayout.java
new file mode 100644
index 0000000000000000000000000000000000000000..5924b53f822cbd5edc142588d30713f20c25a26d
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUILayout.java
@@ -0,0 +1,235 @@
+/*
+ * Copyright (C) 2021 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.wm.shell.sizecompatui;
+
+import static android.view.WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE;
+import static android.view.WindowManager.LayoutParams.FLAG_NOT_TOUCH_MODAL;
+import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_NO_MOVE_ANIMATION;
+import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_TRUSTED_OVERLAY;
+import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
+
+import android.annotation.Nullable;
+import android.app.ActivityClient;
+import android.content.Context;
+import android.content.res.Configuration;
+import android.graphics.PixelFormat;
+import android.graphics.Rect;
+import android.os.Binder;
+import android.os.IBinder;
+import android.view.Gravity;
+import android.view.SurfaceControl;
+import android.view.View;
+import android.view.WindowManager;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.wm.shell.R;
+import com.android.wm.shell.ShellTaskOrganizer;
+import com.android.wm.shell.common.DisplayLayout;
+import com.android.wm.shell.common.SyncTransactionQueue;
+
+/**
+ * Records and handles layout of size compat UI on a task with size compat activity. Helps to
+ * calculate proper bounds when configuration or button position changes.
+ */
+class SizeCompatUILayout {
+    private static final String TAG = "SizeCompatUILayout";
+
+    private final SyncTransactionQueue mSyncQueue;
+    private Context mContext;
+    private Configuration mTaskConfig;
+    private final int mDisplayId;
+    private final int mTaskId;
+    private IBinder mActivityToken;
+    private ShellTaskOrganizer.TaskListener mTaskListener;
+    private DisplayLayout mDisplayLayout;
+    @VisibleForTesting
+    final SizeCompatUIWindowManager mWindowManager;
+
+    @VisibleForTesting
+    @Nullable
+    SizeCompatRestartButton mButton;
+    final int mButtonSize;
+    final int mPopupOffsetX;
+    final int mPopupOffsetY;
+    boolean mShouldShowHint;
+
+    SizeCompatUILayout(SyncTransactionQueue syncQueue, Context context, Configuration taskConfig,
+            int taskId, IBinder activityToken, ShellTaskOrganizer.TaskListener taskListener,
+            DisplayLayout displayLayout, boolean hasShownHint) {
+        mSyncQueue = syncQueue;
+        mContext = context.createConfigurationContext(taskConfig);
+        mTaskConfig = taskConfig;
+        mDisplayId = mContext.getDisplayId();
+        mTaskId = taskId;
+        mActivityToken = activityToken;
+        mTaskListener = taskListener;
+        mDisplayLayout = displayLayout;
+        mShouldShowHint = !hasShownHint;
+        mWindowManager = new SizeCompatUIWindowManager(mContext, taskConfig, this);
+
+        mButtonSize =
+                mContext.getResources().getDimensionPixelSize(R.dimen.size_compat_button_size);
+        mPopupOffsetX = mButtonSize / 4;
+        mPopupOffsetY = mButtonSize;
+    }
+
+    /** Creates the button window. */
+    void createSizeCompatButton(boolean isImeShowing) {
+        if (isImeShowing || mButton != null) {
+            // When ime is showing, wait until ime is dismiss to create UI.
+            return;
+        }
+        mButton = mWindowManager.createSizeCompatUI();
+        updateSurfacePosition();
+    }
+
+    /** Releases the button window. */
+    void release() {
+        mButton.remove();
+        mButton = null;
+        mWindowManager.release();
+    }
+
+    /** Called when size compat info changed. */
+    void updateSizeCompatInfo(Configuration taskConfig, IBinder activityToken,
+            ShellTaskOrganizer.TaskListener taskListener, boolean isImeShowing) {
+        final Configuration prevTaskConfig = mTaskConfig;
+        final ShellTaskOrganizer.TaskListener prevTaskListener = mTaskListener;
+        mTaskConfig = taskConfig;
+        mActivityToken = activityToken;
+        mTaskListener = taskListener;
+
+        // Update configuration.
+        mContext = mContext.createConfigurationContext(taskConfig);
+        mWindowManager.setConfiguration(taskConfig);
+
+        if (mButton == null || prevTaskListener != taskListener) {
+            // TaskListener changed, recreate the button for new surface parent.
+            release();
+            createSizeCompatButton(isImeShowing);
+            return;
+        }
+
+        if (!taskConfig.windowConfiguration.getBounds()
+                .equals(prevTaskConfig.windowConfiguration.getBounds())) {
+            // Reposition the button surface.
+            updateSurfacePosition();
+        }
+
+        if (taskConfig.getLayoutDirection() != prevTaskConfig.getLayoutDirection()) {
+            // Update layout for RTL.
+            mButton.setLayoutDirection(taskConfig.getLayoutDirection());
+            updateSurfacePosition();
+        }
+    }
+
+    /** Called when display layout changed. */
+    void updateDisplayLayout(DisplayLayout displayLayout) {
+        if (displayLayout == mDisplayLayout) {
+            return;
+        }
+
+        final Rect prevStableBounds = new Rect();
+        final Rect curStableBounds = new Rect();
+        mDisplayLayout.getStableBounds(prevStableBounds);
+        displayLayout.getStableBounds(curStableBounds);
+        mDisplayLayout = displayLayout;
+        if (!prevStableBounds.equals(curStableBounds)) {
+            // Stable bounds changed, update button surface position.
+            updateSurfacePosition();
+        }
+    }
+
+    /** Called when IME visibility changed. */
+    void updateImeVisibility(boolean isImeShowing) {
+        if (mButton == null) {
+            // Button may not be created because ime is previous showing.
+            createSizeCompatButton(isImeShowing);
+            return;
+        }
+
+        final int newVisibility = isImeShowing ? View.GONE : View.VISIBLE;
+        if (mButton.getVisibility() != newVisibility) {
+            mButton.setVisibility(newVisibility);
+        }
+    }
+
+    /** Gets the layout params for restart button. */
+    WindowManager.LayoutParams getWindowLayoutParams() {
+        final WindowManager.LayoutParams winParams = new WindowManager.LayoutParams(
+                mButtonSize, mButtonSize,
+                TYPE_APPLICATION_OVERLAY,
+                FLAG_NOT_FOCUSABLE | FLAG_NOT_TOUCH_MODAL,
+                PixelFormat.TRANSLUCENT);
+        winParams.gravity = getGravity(getLayoutDirection());
+        winParams.token = new Binder();
+        winParams.setTitle(SizeCompatRestartButton.class.getSimpleName() + mContext.getDisplayId());
+        winParams.privateFlags |= PRIVATE_FLAG_NO_MOVE_ANIMATION | PRIVATE_FLAG_TRUSTED_OVERLAY;
+        return winParams;
+    }
+
+    /** Called when it is ready to be placed button surface button. */
+    void attachToParentSurface(SurfaceControl.Builder b) {
+        mTaskListener.attachChildSurfaceToTask(mTaskId, b);
+    }
+
+    /** Called when the restart button is clicked. */
+    void onRestartButtonClicked() {
+        ActivityClient.getInstance().restartActivityProcessIfVisible(mActivityToken);
+    }
+
+    @VisibleForTesting
+    void updateSurfacePosition() {
+        if (mButton == null || mWindowManager.getSurfaceControl() == null) {
+            return;
+        }
+        // The hint popup won't be at the correct position.
+        mButton.dismissHint();
+
+        // Use stable bounds to prevent the button from overlapping with system bars.
+        final Rect taskBounds = mTaskConfig.windowConfiguration.getBounds();
+        final Rect stableBounds = new Rect();
+        mDisplayLayout.getStableBounds(stableBounds);
+        stableBounds.intersect(taskBounds);
+
+        // Position of the button in the container coordinate.
+        final int positionX = getLayoutDirection() == View.LAYOUT_DIRECTION_RTL
+                ? stableBounds.left - taskBounds.left
+                : stableBounds.right - taskBounds.left - mButtonSize;
+        final int positionY = stableBounds.bottom - taskBounds.top - mButtonSize;
+
+        mSyncQueue.runInSync(t ->
+                t.setPosition(mWindowManager.getSurfaceControl(), positionX, positionY));
+    }
+
+    int getDisplayId() {
+        return mDisplayId;
+    }
+
+    int getTaskId() {
+        return mTaskId;
+    }
+
+    private int getLayoutDirection() {
+        return mContext.getResources().getConfiguration().getLayoutDirection();
+    }
+
+    static int getGravity(int layoutDirection) {
+        return Gravity.BOTTOM
+                | (layoutDirection == View.LAYOUT_DIRECTION_RTL ? Gravity.START : Gravity.END);
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUIWindowManager.java b/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUIWindowManager.java
new file mode 100644
index 0000000000000000000000000000000000000000..a7ad982a47367505bfa41d3d6b69ae67427976f9
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUIWindowManager.java
@@ -0,0 +1,103 @@
+/*
+ * Copyright (C) 2021 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.wm.shell.sizecompatui;
+
+import android.annotation.Nullable;
+import android.content.Context;
+import android.content.res.Configuration;
+import android.view.IWindow;
+import android.view.LayoutInflater;
+import android.view.SurfaceControl;
+import android.view.SurfaceControlViewHost;
+import android.view.SurfaceSession;
+import android.view.WindowlessWindowManager;
+
+import com.android.wm.shell.R;
+
+/**
+ * Holds view hierarchy of a root surface and helps to inflate {@link SizeCompatRestartButton}.
+ */
+class SizeCompatUIWindowManager extends WindowlessWindowManager {
+
+    private Context mContext;
+    private final SizeCompatUILayout mLayout;
+
+    @Nullable
+    private SurfaceControlViewHost mViewHost;
+    @Nullable
+    private SurfaceControl mLeash;
+
+    SizeCompatUIWindowManager(Context context, Configuration config, SizeCompatUILayout layout) {
+        super(config, null /* rootSurface */, null /* hostInputToken */);
+        mContext = context;
+        mLayout = layout;
+    }
+
+    @Override
+    public void setConfiguration(Configuration configuration) {
+        super.setConfiguration(configuration);
+        mContext = mContext.createConfigurationContext(configuration);
+    }
+
+    @Override
+    protected void attachToParentSurface(IWindow window, SurfaceControl.Builder b) {
+        // Can't set position for the ViewRootImpl SC directly. Create a leash to manipulate later.
+        final SurfaceControl.Builder builder = new SurfaceControl.Builder(new SurfaceSession())
+                .setContainerLayer()
+                .setName("SizeCompatUILeash")
+                .setHidden(false)
+                .setCallsite("SizeCompatUIWindowManager#attachToParentSurface");
+        mLayout.attachToParentSurface(builder);
+        mLeash = builder.build();
+        b.setParent(mLeash);
+    }
+
+    /** Inflates {@link SizeCompatRestartButton} on to the root surface. */
+    SizeCompatRestartButton createSizeCompatUI() {
+        if (mViewHost == null) {
+            mViewHost = new SurfaceControlViewHost(mContext, mContext.getDisplay(), this);
+        }
+
+        final SizeCompatRestartButton button = (SizeCompatRestartButton)
+                LayoutInflater.from(mContext).inflate(R.layout.size_compat_ui, null);
+        button.inject(mLayout);
+        mViewHost.setView(button, mLayout.getWindowLayoutParams());
+        return button;
+    }
+
+    /** Releases the surface control and tears down the view hierarchy. */
+    void release() {
+        if (mViewHost != null) {
+            mViewHost.release();
+            mViewHost = null;
+        }
+
+        if (mLeash != null) {
+            new SurfaceControl.Transaction().remove(mLeash).apply();
+            mLeash = null;
+        }
+    }
+
+    /**
+     * Gets {@link SurfaceControl} of the surface holding size compat UI view. @return {@code null}
+     * if not feasible.
+     */
+    @Nullable
+    SurfaceControl getSurfaceControl() {
+        return mLeash;
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
index c815e65189e61a23ed82d6b1ed2f9527ed2dcb40..733e4845dd6246a97da44408f78890404c7e9623 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
@@ -31,6 +31,7 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 
@@ -105,6 +106,11 @@ public class ShellTaskOrganizerTests {
         public void onTaskVanished(RunningTaskInfo taskInfo) {
             vanished.add(taskInfo);
         }
+
+        @Override
+        public boolean supportSizeCompatUI() {
+            return true;
+        }
     }
 
     @Before
@@ -282,10 +288,10 @@ public class ShellTaskOrganizerTests {
 
         // sizeCompatActivity is null if top activity is not in size compat.
         verify(mSizeCompatUI).onSizeCompatInfoChanged(taskInfo1.displayId, taskInfo1.taskId,
-                taskInfo1.configuration.windowConfiguration.getBounds(),
-                null /* sizeCompatActivity*/ , taskListener);
+                null /* taskConfig */, null /* sizeCompatActivity*/, null /* taskListener */);
 
         // sizeCompatActivity is non-null if top activity is in size compat.
+        clearInvocations(mSizeCompatUI);
         final RunningTaskInfo taskInfo2 =
                 createTaskInfo(taskInfo1.taskId, taskInfo1.getWindowingMode());
         taskInfo2.displayId = taskInfo1.displayId;
@@ -293,14 +299,12 @@ public class ShellTaskOrganizerTests {
         taskInfo2.topActivityInSizeCompat = true;
         mOrganizer.onTaskInfoChanged(taskInfo2);
         verify(mSizeCompatUI).onSizeCompatInfoChanged(taskInfo1.displayId, taskInfo1.taskId,
-                taskInfo1.configuration.windowConfiguration.getBounds(),
-                taskInfo1.topActivityToken,
-                taskListener);
+                taskInfo1.configuration, taskInfo1.topActivityToken, taskListener);
 
+        clearInvocations(mSizeCompatUI);
         mOrganizer.onTaskVanished(taskInfo1);
         verify(mSizeCompatUI).onSizeCompatInfoChanged(taskInfo1.displayId, taskInfo1.taskId,
-                null /* taskConfig */, null /* sizeCompatActivity*/,
-                null /* taskListener */);
+                null /* taskConfig */, null /* sizeCompatActivity*/, null /* taskListener */);
     }
 
     private static RunningTaskInfo createTaskInfo(int taskId, int windowingMode) {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatRestartButtonTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatRestartButtonTest.java
new file mode 100644
index 0000000000000000000000000000000000000000..d9086a6ccdc1cf6631cd4ed2ee34d4955b24c826
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatRestartButtonTest.java
@@ -0,0 +1,115 @@
+/*
+ * Copyright (C) 2021 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.wm.shell.sizecompatui;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+import android.content.res.Configuration;
+import android.os.IBinder;
+import android.testing.AndroidTestingRunner;
+import android.view.LayoutInflater;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.wm.shell.R;
+import com.android.wm.shell.ShellTaskOrganizer;
+import com.android.wm.shell.ShellTestCase;
+import com.android.wm.shell.common.DisplayLayout;
+import com.android.wm.shell.common.SyncTransactionQueue;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Tests for {@link SizeCompatRestartButton}.
+ *
+ * Build/Install/Run:
+ *  atest WMShellUnitTests:SizeCompatRestartButtonTest
+ */
+@RunWith(AndroidTestingRunner.class)
+@SmallTest
+public class SizeCompatRestartButtonTest extends ShellTestCase {
+
+    @Mock private SyncTransactionQueue mSyncTransactionQueue;
+    @Mock private IBinder mActivityToken;
+    @Mock private ShellTaskOrganizer.TaskListener mTaskListener;
+    @Mock private DisplayLayout mDisplayLayout;
+
+    private SizeCompatUILayout mLayout;
+    private SizeCompatRestartButton mButton;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+
+        final int taskId = 1;
+        mLayout = new SizeCompatUILayout(mSyncTransactionQueue, mContext, new Configuration(),
+                taskId, mActivityToken, mTaskListener, mDisplayLayout, false /* hasShownHint*/);
+        mButton = (SizeCompatRestartButton)
+                LayoutInflater.from(mContext).inflate(R.layout.size_compat_ui, null);
+        mButton.inject(mLayout);
+
+        spyOn(mLayout);
+        spyOn(mButton);
+        doNothing().when(mButton).showHint();
+    }
+
+    @Test
+    public void testOnClick() {
+        doNothing().when(mLayout).onRestartButtonClicked();
+
+        mButton.onClick(mButton);
+
+        verify(mLayout).onRestartButtonClicked();
+    }
+
+    @Test
+    public void testOnLongClick() {
+        verify(mButton, never()).showHint();
+
+        mButton.onLongClick(mButton);
+
+        verify(mButton).showHint();
+    }
+
+    @Test
+    public void testOnAttachedToWindow_showHint() {
+        mLayout.mShouldShowHint = false;
+        mButton.onAttachedToWindow();
+
+        verify(mButton, never()).showHint();
+
+        mLayout.mShouldShowHint = true;
+        mButton.onAttachedToWindow();
+
+        verify(mButton).showHint();
+    }
+
+    @Test
+    public void testRemove() {
+        mButton.remove();
+
+        verify(mButton).dismissHint();
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUIControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUIControllerTest.java
index 0eb64e5963d17bbd7d447182126e83f712fc060a..806a90b7832afe56cff16af73c52e6b257004122 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUIControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUIControllerTest.java
@@ -16,23 +16,28 @@
 
 package com.android.wm.shell.sizecompatui;
 
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
 import android.content.Context;
-import android.graphics.Rect;
+import android.content.res.Configuration;
 import android.os.IBinder;
 import android.testing.AndroidTestingRunner;
-import android.view.View;
 
 import androidx.test.filters.SmallTest;
 
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
-import com.android.wm.shell.TestShellExecutor;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.DisplayImeController;
+import com.android.wm.shell.common.DisplayLayout;
+import com.android.wm.shell.common.SyncTransactionQueue;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -50,28 +55,34 @@ import org.mockito.MockitoAnnotations;
 @SmallTest
 public class SizeCompatUIControllerTest extends ShellTestCase {
     private static final int DISPLAY_ID = 0;
-
-    private final TestShellExecutor mShellMainExecutor = new TestShellExecutor();
+    private static final int TASK_ID = 12;
 
     private SizeCompatUIController mController;
     private @Mock DisplayController mMockDisplayController;
+    private @Mock DisplayLayout mMockDisplayLayout;
     private @Mock DisplayImeController mMockImeController;
-    private @Mock SizeCompatRestartButton mMockButton;
     private @Mock IBinder mMockActivityToken;
     private @Mock ShellTaskOrganizer.TaskListener mMockTaskListener;
+    private @Mock SyncTransactionQueue mMockSyncQueue;
+    private @Mock SizeCompatUILayout mMockLayout;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        doReturn(true).when(mMockButton).show();
 
+        doReturn(mMockDisplayLayout).when(mMockDisplayController).getDisplayLayout(anyInt());
+        doReturn(DISPLAY_ID).when(mMockLayout).getDisplayId();
+        doReturn(TASK_ID).when(mMockLayout).getTaskId();
         mController = new SizeCompatUIController(mContext, mMockDisplayController,
-                mMockImeController, mShellMainExecutor) {
+                mMockImeController, mMockSyncQueue) {
             @Override
-            SizeCompatRestartButton createRestartButton(Context context, int displayId) {
-                return mMockButton;
+            SizeCompatUILayout createLayout(Context context, int displayId, int taskId,
+                    Configuration taskConfig, IBinder activityToken,
+                    ShellTaskOrganizer.TaskListener taskListener) {
+                return mMockLayout;
             }
         };
+        spyOn(mController);
     }
 
     @Test
@@ -82,42 +93,72 @@ public class SizeCompatUIControllerTest extends ShellTestCase {
 
     @Test
     public void testOnSizeCompatInfoChanged() {
-        final int taskId = 12;
-        final Rect taskBounds = new Rect(0, 0, 1000, 2000);
+        final Configuration taskConfig = new Configuration();
+
+        // Verify that the restart button is added with non-null size compat info.
+        mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, taskConfig,
+                mMockActivityToken, mMockTaskListener);
+
+        verify(mController).createLayout(any(), eq(DISPLAY_ID), eq(TASK_ID), eq(taskConfig),
+                eq(mMockActivityToken), eq(mMockTaskListener));
 
-        // Verify that the restart button is added with non-null size compat activity.
-        mController.onSizeCompatInfoChanged(DISPLAY_ID, taskId, taskBounds,
+        // Verify that the restart button is updated with non-null new size compat info.
+        final Configuration newTaskConfig = new Configuration();
+        mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, newTaskConfig,
                 mMockActivityToken, mMockTaskListener);
-        mShellMainExecutor.flushAll();
 
-        verify(mMockButton).show();
-        verify(mMockButton).updateLastTargetActivity(eq(mMockActivityToken));
+        verify(mMockLayout).updateSizeCompatInfo(taskConfig, mMockActivityToken, mMockTaskListener,
+                false /* isImeShowing */);
+
+        // Verify that the restart button is removed with null size compat info.
+        mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, null, null, mMockTaskListener);
+
+        verify(mMockLayout).release();
+    }
+
+    @Test
+    public void testOnDisplayRemoved() {
+        final Configuration taskConfig = new Configuration();
+        mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, taskConfig,
+                mMockActivityToken, mMockTaskListener);
+
+        mController.onDisplayRemoved(DISPLAY_ID + 1);
+
+        verify(mMockLayout, never()).release();
+
+        mController.onDisplayRemoved(DISPLAY_ID);
+
+        verify(mMockLayout).release();
+    }
+
+    @Test
+    public void testOnDisplayConfigurationChanged() {
+        final Configuration taskConfig = new Configuration();
+        mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, taskConfig,
+                mMockActivityToken, mMockTaskListener);
+
+        final Configuration newTaskConfig = new Configuration();
+        mController.onDisplayConfigurationChanged(DISPLAY_ID + 1, newTaskConfig);
+
+        verify(mMockLayout, never()).updateDisplayLayout(any());
 
-        // Verify that the restart button is removed with null size compat activity.
-        mController.onSizeCompatInfoChanged(DISPLAY_ID, taskId, null, null, null);
+        mController.onDisplayConfigurationChanged(DISPLAY_ID, newTaskConfig);
 
-        mShellMainExecutor.flushAll();
-        verify(mMockButton).remove();
+        verify(mMockLayout).updateDisplayLayout(mMockDisplayLayout);
     }
 
     @Test
     public void testChangeButtonVisibilityOnImeShowHide() {
-        final int taskId = 12;
-        final Rect taskBounds = new Rect(0, 0, 1000, 2000);
-        mController.onSizeCompatInfoChanged(DISPLAY_ID, taskId, taskBounds,
+        final Configuration taskConfig = new Configuration();
+        mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, taskConfig,
                 mMockActivityToken, mMockTaskListener);
-        mShellMainExecutor.flushAll();
 
-        // Verify that the restart button is hidden when IME is visible.
-        doReturn(View.VISIBLE).when(mMockButton).getVisibility();
         mController.onImeVisibilityChanged(DISPLAY_ID, true /* isShowing */);
 
-        verify(mMockButton).setVisibility(eq(View.GONE));
+        verify(mMockLayout).updateImeVisibility(true);
 
-        // Verify that the restart button is visible when IME is hidden.
-        doReturn(View.GONE).when(mMockButton).getVisibility();
         mController.onImeVisibilityChanged(DISPLAY_ID, false /* isShowing */);
 
-        verify(mMockButton).setVisibility(eq(View.VISIBLE));
+        verify(mMockLayout).updateImeVisibility(false);
     }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUILayoutTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUILayoutTest.java
new file mode 100644
index 0000000000000000000000000000000000000000..236db44bdce0f3eb0310fe0e723b60eea4814949
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUILayoutTest.java
@@ -0,0 +1,206 @@
+/*
+ * Copyright (C) 2021 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.wm.shell.sizecompatui;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+import android.app.ActivityClient;
+import android.content.res.Configuration;
+import android.graphics.Rect;
+import android.os.IBinder;
+import android.testing.AndroidTestingRunner;
+import android.view.DisplayInfo;
+import android.view.SurfaceControl;
+import android.view.View;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.wm.shell.ShellTaskOrganizer;
+import com.android.wm.shell.ShellTestCase;
+import com.android.wm.shell.common.DisplayLayout;
+import com.android.wm.shell.common.SyncTransactionQueue;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Tests for {@link SizeCompatUILayout}.
+ *
+ * Build/Install/Run:
+ *  atest WMShellUnitTests:SizeCompatUILayoutTest
+ */
+@RunWith(AndroidTestingRunner.class)
+@SmallTest
+public class SizeCompatUILayoutTest extends ShellTestCase {
+
+    private static final int TASK_ID = 1;
+
+    @Mock private SyncTransactionQueue mSyncTransactionQueue;
+    @Mock private IBinder mActivityToken;
+    @Mock private ShellTaskOrganizer.TaskListener mTaskListener;
+    @Mock private DisplayLayout mDisplayLayout;
+    @Mock private SizeCompatRestartButton mButton;
+    private Configuration mTaskConfig;
+
+    private SizeCompatUILayout mLayout;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mTaskConfig = new Configuration();
+
+        mLayout = new SizeCompatUILayout(mSyncTransactionQueue, mContext, new Configuration(),
+                TASK_ID, mActivityToken, mTaskListener, mDisplayLayout, false /* hasShownHint*/);
+
+        spyOn(mLayout);
+        spyOn(mLayout.mWindowManager);
+        doReturn(mButton).when(mLayout.mWindowManager).createSizeCompatUI();
+    }
+
+    @Test
+    public void testCreateSizeCompatButton() {
+        // Not create button if IME is showing.
+        mLayout.createSizeCompatButton(true /* isImeShowing */);
+
+        verify(mLayout.mWindowManager, never()).createSizeCompatUI();
+        assertNull(mLayout.mButton);
+
+        mLayout.createSizeCompatButton(false /* isImeShowing */);
+
+        verify(mLayout.mWindowManager).createSizeCompatUI();
+        assertNotNull(mLayout.mButton);
+    }
+
+    @Test
+    public void testRelease() {
+        mLayout.createSizeCompatButton(false /* isImeShowing */);
+
+        mLayout.release();
+
+        assertNull(mLayout.mButton);
+        verify(mButton).remove();
+        verify(mLayout.mWindowManager).release();
+    }
+
+    @Test
+    public void testUpdateSizeCompatInfo() {
+        mLayout.createSizeCompatButton(false /* isImeShowing */);
+
+        // No diff
+        clearInvocations(mLayout);
+        mLayout.updateSizeCompatInfo(mTaskConfig, mActivityToken, mTaskListener,
+                false /* isImeShowing */);
+
+        verify(mLayout, never()).updateSurfacePosition();
+        verify(mLayout, never()).release();
+        verify(mLayout, never()).createSizeCompatButton(anyBoolean());
+
+        // Change task listener, recreate button.
+        clearInvocations(mLayout);
+        final ShellTaskOrganizer.TaskListener newTaskListener = mock(
+                ShellTaskOrganizer.TaskListener.class);
+        mLayout.updateSizeCompatInfo(mTaskConfig, mActivityToken, newTaskListener,
+                false /* isImeShowing */);
+
+        verify(mLayout).release();
+        verify(mLayout).createSizeCompatButton(anyBoolean());
+
+        // Change task bounds, update position.
+        clearInvocations(mLayout);
+        final Configuration newTaskConfiguration = new Configuration();
+        newTaskConfiguration.windowConfiguration.setBounds(new Rect(0, 1000, 0, 2000));
+        mLayout.updateSizeCompatInfo(newTaskConfiguration, mActivityToken, newTaskListener,
+                false /* isImeShowing */);
+
+        verify(mLayout).updateSurfacePosition();
+    }
+
+    @Test
+    public void testUpdateDisplayLayout() {
+        final DisplayInfo displayInfo = new DisplayInfo();
+        displayInfo.logicalWidth = 1000;
+        displayInfo.logicalHeight = 2000;
+        final DisplayLayout displayLayout1 = new DisplayLayout(displayInfo,
+                mContext.getResources(), false, false);
+
+        mLayout.updateDisplayLayout(displayLayout1);
+        verify(mLayout).updateSurfacePosition();
+
+        // No update if the display bounds is the same.
+        clearInvocations(mLayout);
+        final DisplayLayout displayLayout2 = new DisplayLayout(displayInfo,
+                mContext.getResources(), false, false);
+        mLayout.updateDisplayLayout(displayLayout2);
+        verify(mLayout, never()).updateSurfacePosition();
+    }
+
+    @Test
+    public void testUpdateImeVisibility() {
+        // Create button if it is not created.
+        mLayout.mButton = null;
+        mLayout.updateImeVisibility(false /* isImeShowing */);
+
+        verify(mLayout).createSizeCompatButton(false /* isImeShowing */);
+
+        // Hide button if ime is shown.
+        clearInvocations(mLayout);
+        doReturn(View.VISIBLE).when(mButton).getVisibility();
+        mLayout.updateImeVisibility(true /* isImeShowing */);
+
+        verify(mLayout, never()).createSizeCompatButton(anyBoolean());
+        verify(mButton).setVisibility(View.GONE);
+
+        // Show button if ime is not shown.
+        doReturn(View.GONE).when(mButton).getVisibility();
+        mLayout.updateImeVisibility(false /* isImeShowing */);
+
+        verify(mLayout, never()).createSizeCompatButton(anyBoolean());
+        verify(mButton).setVisibility(View.VISIBLE);
+    }
+
+    @Test
+    public void testAttachToParentSurface() {
+        final SurfaceControl.Builder b = new SurfaceControl.Builder();
+        mLayout.attachToParentSurface(b);
+
+        verify(mTaskListener).attachChildSurfaceToTask(TASK_ID, b);
+    }
+
+    @Test
+    public void testOnRestartButtonClicked() {
+        spyOn(ActivityClient.getInstance());
+        doNothing().when(ActivityClient.getInstance()).restartActivityProcessIfVisible(any());
+
+        mLayout.onRestartButtonClicked();
+
+        verify(ActivityClient.getInstance()).restartActivityProcessIfVisible(mActivityToken);
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/wmshell/WMShellBaseModule.java b/packages/SystemUI/src/com/android/systemui/wmshell/WMShellBaseModule.java
index 449db61a0fbba2aafa36395bc5eccd577e06a3dc..fba0b0079012ae61e852a1dac3cf984538f090fd 100644
--- a/packages/SystemUI/src/com/android/systemui/wmshell/WMShellBaseModule.java
+++ b/packages/SystemUI/src/com/android/systemui/wmshell/WMShellBaseModule.java
@@ -20,7 +20,6 @@ import static android.os.Process.THREAD_PRIORITY_DISPLAY;
 
 import android.animation.AnimationHandler;
 import android.app.ActivityTaskManager;
-import android.app.IActivityManager;
 import android.content.Context;
 import android.content.pm.LauncherApps;
 import android.content.pm.PackageManager;
@@ -73,7 +72,6 @@ import com.android.wm.shell.pip.PipMediaController;
 import com.android.wm.shell.pip.PipSurfaceTransactionHelper;
 import com.android.wm.shell.pip.PipUiEventLogger;
 import com.android.wm.shell.pip.phone.PipAppOpsListener;
-import com.android.wm.shell.pip.phone.PipController;
 import com.android.wm.shell.pip.phone.PipTouchHandler;
 import com.android.wm.shell.sizecompatui.SizeCompatUIController;
 import com.android.wm.shell.splitscreen.SplitScreen;
@@ -211,8 +209,8 @@ public abstract class WMShellBaseModule {
     @Provides
     static SizeCompatUIController provideSizeCompatUIController(Context context,
             DisplayController displayController, DisplayImeController imeController,
-            @ShellMainThread ShellExecutor mainExecutor) {
-        return new SizeCompatUIController(context, displayController, imeController, mainExecutor);
+            SyncTransactionQueue syncQueue) {
+        return new SizeCompatUIController(context, displayController, imeController, syncQueue);
     }
 
     @WMSingleton