diff --git a/quickstep/res/drawable/ic_sysbar_rotate_button_ccw_start_0.xml b/quickstep/res/drawable/ic_sysbar_rotate_button_ccw_start_0.xml
new file mode 100644
index 0000000..ff5cb9e
--- /dev/null
+++ b/quickstep/res/drawable/ic_sysbar_rotate_button_ccw_start_0.xml
@@ -0,0 +1,187 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2020 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.
+-->
+<animated-vector xmlns:android="http://schemas.android.com/apk/res/android" xmlns:aapt="http://schemas.android.com/aapt">
+    <aapt:attr name="android:drawable">
+        <vector android:name="root"
+                android:width="28dp"
+                android:height="28dp"
+                android:viewportWidth="28.0"
+                android:viewportHeight="28.0">
+            <!-- Use scaleX to flip icon so arrows always point in the direction of motion -->
+            <group android:name="icon" android:pivotX="14" android:pivotY="14"
+                   android:scaleX="1">
+                <!-- Tint color to be set directly -->
+                <path android:fillColor="#FFFFFFFF"
+                      android:pathData="M12.02,10.83L9.25,8.06l2.77,-2.77l1.12,1.12l-0.85,0.86h5.16c0.72,0 1.31,0.56 1.31,1.26v9.16l-1.58,-1.58V8.85h-4.89l0.86,0.86L12.02,10.83zM15.98,17.17l-1.12,1.12l0.85,0.86h-4.88v-7.26L9.25,10.3v9.17c0,0.7 0.59,1.26 1.31,1.26h5.16v0.01l-0.85,0.85l1.12,1.12l2.77,-2.77L15.98,17.17z"/>
+            </group>
+        </vector>
+    </aapt:attr>
+
+    <!-- Repeat all animations 5 times but don't fade out at the end -->
+    <target android:name="root">
+        <aapt:attr name="android:animation">
+            <set android:ordering="sequentially">
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+            </set>
+        </aapt:attr>
+    </target>
+    <target android:name="icon">
+        <aapt:attr name="android:animation">
+            <set android:ordering="sequentially">
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="100"
+                                android:duration="600"
+                                android:valueFrom="0"
+                                android:valueTo="-90">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="0"
+                                android:valueTo="0"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="0"
+                                android:valueTo="-90">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="0"
+                                android:valueTo="0"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="0"
+                                android:valueTo="-90">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="0"
+                                android:valueTo="0"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="0"
+                                android:valueTo="-90">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="0"
+                                android:valueTo="0"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="0"
+                                android:valueTo="-90">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+            </set>
+        </aapt:attr>
+    </target>
+</animated-vector>
\ No newline at end of file
diff --git a/quickstep/res/drawable/ic_sysbar_rotate_button_ccw_start_90.xml b/quickstep/res/drawable/ic_sysbar_rotate_button_ccw_start_90.xml
new file mode 100644
index 0000000..90fedb1
--- /dev/null
+++ b/quickstep/res/drawable/ic_sysbar_rotate_button_ccw_start_90.xml
@@ -0,0 +1,187 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2020 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.
+-->
+<animated-vector xmlns:android="http://schemas.android.com/apk/res/android" xmlns:aapt="http://schemas.android.com/aapt">
+    <aapt:attr name="android:drawable">
+        <vector android:name="root"
+                android:width="28dp"
+                android:height="28dp"
+                android:viewportWidth="28.0"
+                android:viewportHeight="28.0">
+            <!-- Use scaleX to flip icon so arrows always point in the direction of motion -->
+            <group android:name="icon" android:pivotX="14" android:pivotY="14"
+                   android:scaleX="1">
+                <!-- Tint color to be set directly -->
+                <path android:fillColor="#FFFFFFFF"
+                      android:pathData="M12.02,10.83L9.25,8.06l2.77,-2.77l1.12,1.12l-0.85,0.86h5.16c0.72,0 1.31,0.56 1.31,1.26v9.16l-1.58,-1.58V8.85h-4.89l0.86,0.86L12.02,10.83zM15.98,17.17l-1.12,1.12l0.85,0.86h-4.88v-7.26L9.25,10.3v9.17c0,0.7 0.59,1.26 1.31,1.26h5.16v0.01l-0.85,0.85l1.12,1.12l2.77,-2.77L15.98,17.17z"/>
+            </group>
+        </vector>
+    </aapt:attr>
+
+    <!-- Repeat all animations 5 times but don't fade out at the end -->
+    <target android:name="root">
+        <aapt:attr name="android:animation">
+            <set android:ordering="sequentially">
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+            </set>
+        </aapt:attr>
+    </target>
+    <target android:name="icon">
+        <aapt:attr name="android:animation">
+            <set android:ordering="sequentially">
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="100"
+                                android:duration="600"
+                                android:valueFrom="90"
+                                android:valueTo="0">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="90"
+                                android:valueTo="90"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="90"
+                                android:valueTo="0">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="90"
+                                android:valueTo="90"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="90"
+                                android:valueTo="0">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="90"
+                                android:valueTo="90"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="90"
+                                android:valueTo="0">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="90"
+                                android:valueTo="90"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="90"
+                                android:valueTo="0">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+            </set>
+        </aapt:attr>
+    </target>
+</animated-vector>
\ No newline at end of file
diff --git a/quickstep/res/drawable/ic_sysbar_rotate_button_cw_start_0.xml b/quickstep/res/drawable/ic_sysbar_rotate_button_cw_start_0.xml
new file mode 100644
index 0000000..a89e7a3
--- /dev/null
+++ b/quickstep/res/drawable/ic_sysbar_rotate_button_cw_start_0.xml
@@ -0,0 +1,187 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2020 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.
+-->
+<animated-vector xmlns:android="http://schemas.android.com/apk/res/android" xmlns:aapt="http://schemas.android.com/aapt">
+    <aapt:attr name="android:drawable">
+        <vector android:name="root"
+                android:width="28dp"
+                android:height="28dp"
+                android:viewportWidth="28.0"
+                android:viewportHeight="28.0">
+            <!-- Use scaleX to flip icon so arrows always point in the direction of motion -->
+            <group android:name="icon" android:pivotX="14" android:pivotY="14"
+                   android:scaleX="-1">
+                <!-- Tint color to be set directly -->
+                <path android:fillColor="#FFFFFFFF"
+                      android:pathData="M12.02,10.83L9.25,8.06l2.77,-2.77l1.12,1.12l-0.85,0.86h5.16c0.72,0 1.31,0.56 1.31,1.26v9.16l-1.58,-1.58V8.85h-4.89l0.86,0.86L12.02,10.83zM15.98,17.17l-1.12,1.12l0.85,0.86h-4.88v-7.26L9.25,10.3v9.17c0,0.7 0.59,1.26 1.31,1.26h5.16v0.01l-0.85,0.85l1.12,1.12l2.77,-2.77L15.98,17.17z"/>
+            </group>
+        </vector>
+    </aapt:attr>
+
+    <!-- Repeat all animations 5 times but don't fade out at the end -->
+    <target android:name="root">
+        <aapt:attr name="android:animation">
+            <set android:ordering="sequentially">
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+            </set>
+        </aapt:attr>
+    </target>
+    <target android:name="icon">
+        <aapt:attr name="android:animation">
+            <set android:ordering="sequentially">
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="100"
+                                android:duration="600"
+                                android:valueFrom="0"
+                                android:valueTo="90">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="0"
+                                android:valueTo="0"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="0"
+                                android:valueTo="90">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="0"
+                                android:valueTo="0"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="0"
+                                android:valueTo="90">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="0"
+                                android:valueTo="0"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="0"
+                                android:valueTo="90">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="0"
+                                android:valueTo="0"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="0"
+                                android:valueTo="90">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+            </set>
+        </aapt:attr>
+    </target>
+</animated-vector>
\ No newline at end of file
diff --git a/quickstep/res/drawable/ic_sysbar_rotate_button_cw_start_90.xml b/quickstep/res/drawable/ic_sysbar_rotate_button_cw_start_90.xml
new file mode 100644
index 0000000..0dc67b0
--- /dev/null
+++ b/quickstep/res/drawable/ic_sysbar_rotate_button_cw_start_90.xml
@@ -0,0 +1,187 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2020 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.
+-->
+<animated-vector xmlns:android="http://schemas.android.com/apk/res/android" xmlns:aapt="http://schemas.android.com/aapt">
+    <aapt:attr name="android:drawable">
+        <vector android:name="root"
+                android:width="28dp"
+                android:height="28dp"
+                android:viewportWidth="28.0"
+                android:viewportHeight="28.0">
+            <!-- Use scaleX to flip icon so arrows always point in the direction of motion -->
+            <group android:name="icon" android:pivotX="14" android:pivotY="14"
+                   android:scaleX="-1">
+                <!-- Tint color to be set directly -->
+                <path android:fillColor="#FFFFFFFF"
+                      android:pathData="M12.02,10.83L9.25,8.06l2.77,-2.77l1.12,1.12l-0.85,0.86h5.16c0.72,0 1.31,0.56 1.31,1.26v9.16l-1.58,-1.58V8.85h-4.89l0.86,0.86L12.02,10.83zM15.98,17.17l-1.12,1.12l0.85,0.86h-4.88v-7.26L9.25,10.3v9.17c0,0.7 0.59,1.26 1.31,1.26h5.16v0.01l-0.85,0.85l1.12,1.12l2.77,-2.77L15.98,17.17z"/>
+            </group>
+        </vector>
+    </aapt:attr>
+
+    <!-- Repeat all animations 5 times but don't fade out at the end -->
+    <target android:name="root">
+        <aapt:attr name="android:animation">
+            <set android:ordering="sequentially">
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+                <!-- Linear fade out -->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="1700"
+                                android:valueFrom="1"
+                                android:valueTo="0"
+                                android:interpolator="@android:anim/linear_interpolator"/>
+                <!-- Linear fade in-->
+                <objectAnimator android:propertyName="alpha"
+                                android:duration="100"
+                                android:startOffset="100"
+                                android:valueFrom="0"
+                                android:valueTo="1"
+                                android:interpolator="@android:anim/linear_interpolator" />
+            </set>
+        </aapt:attr>
+    </target>
+    <target android:name="icon">
+        <aapt:attr name="android:animation">
+            <set android:ordering="sequentially">
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="100"
+                                android:duration="600"
+                                android:valueFrom="90"
+                                android:valueTo="180">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="90"
+                                android:valueTo="90"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="90"
+                                android:valueTo="180">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="90"
+                                android:valueTo="90"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="90"
+                                android:valueTo="180">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="90"
+                                android:valueTo="90"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="90"
+                                android:valueTo="180">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+
+                <!-- Reset rotation position for fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:startOffset="1300"
+                                android:duration="100"
+                                android:valueFrom="90"
+                                android:valueTo="90"/>
+
+                <!-- Icon rotation with start timing offset after fade in -->
+                <objectAnimator android:propertyName="rotation"
+                                android:duration="600"
+                                android:valueFrom="90"
+                                android:valueTo="180">
+                    <aapt:attr name="android:interpolator">
+                        <pathInterpolator android:pathData="M 0.0,0.0 c0.408,1.181 0.674,1.08 1.0,1.0"/>
+                    </aapt:attr>
+                </objectAnimator>
+            </set>
+        </aapt:attr>
+    </target>
+</animated-vector>
\ No newline at end of file
diff --git a/quickstep/res/layout/taskbar.xml b/quickstep/res/layout/taskbar.xml
index e680233..d32c115 100644
--- a/quickstep/res/layout/taskbar.xml
+++ b/quickstep/res/layout/taskbar.xml
@@ -29,7 +29,7 @@
         android:layout_gravity="bottom" >
 
         <LinearLayout
-            android:id="@+id/system_button_layout"
+            android:id="@+id/nav_button_layout"
             android:layout_width="wrap_content"
             android:layout_height="wrap_content"
             android:paddingLeft="@dimen/taskbar_nav_buttons_spacing"
@@ -44,6 +44,15 @@
             android:forceHasOverlappingRendering="false"
             android:gravity="center" />
 
+        <LinearLayout
+            android:id="@+id/contextual_button_layout"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:paddingLeft="@dimen/taskbar_nav_buttons_spacing"
+            android:paddingRight="@dimen/taskbar_nav_buttons_spacing"
+            android:forceHasOverlappingRendering="false"
+            android:gravity="center" />
+
     </com.android.launcher3.taskbar.TaskbarView>
 
     <com.android.launcher3.taskbar.ImeBarView
diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index c85d3dc..dd63cfa 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -16,7 +16,7 @@
 
 <resources>
     <dimen name="task_thumbnail_icon_size">48dp</dimen>
-    <dimen name="task_thumbnail_icon_size_grid">32dp</dimen>
+    <dimen name="task_thumbnail_icon_size_grid">40dp</dimen>
     <!-- For screens without rounded corners -->
     <dimen name="task_corner_radius_small">2dp</dimen>
     <!-- For Launchers that want to override the default dialog corner radius -->
@@ -37,10 +37,10 @@
     <dimen name="overview_actions_horizontal_margin">16dp</dimen>
 
     <dimen name="overview_grid_top_margin">77dp</dimen>
-    <dimen name="overview_grid_bottom_margin">90dp</dimen>
+    <dimen name="overview_grid_bottom_margin">70dp</dimen>
     <dimen name="overview_grid_side_margin">54dp</dimen>
     <dimen name="overview_grid_row_spacing">42dp</dimen>
-    <dimen name="overview_grid_focus_vertical_margin">90dp</dimen>
+    <dimen name="overview_grid_focus_vertical_margin">40dp</dimen>
     <dimen name="split_placeholder_size">110dp</dimen>
 
     <!-- These speeds are in dp/s -->
diff --git a/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index 7c97b93..5471e49 100644
--- a/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -48,6 +48,7 @@
 import com.android.launcher3.shadows.ShadowDeviceFlag;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.ItemInfoMatcher;
 import com.android.launcher3.util.LauncherModelHelper;
 import com.android.launcher3.util.ViewOnDrawExecutor;
@@ -239,8 +240,8 @@
         }
 
         @Override
-        public int getPageToBindSynchronously() {
-            return 0;
+        public IntSet getPagesToBindSynchronously() {
+            return IntSet.wrap(0);
         }
 
         @Override
@@ -259,7 +260,7 @@
         public void finishFirstPageBind(ViewOnDrawExecutor executor) { }
 
         @Override
-        public void finishBindingItems(int pageBoundFirst) { }
+        public void finishBindingItems(IntSet pagesBoundFirst) { }
 
         @Override
         public void preAddApps() { }
@@ -287,7 +288,7 @@
         public void bindAllWidgets(List<WidgetsListBaseEntry> widgets) { }
 
         @Override
-        public void onPageBoundSynchronously(int page) { }
+        public void onPagesBoundSynchronously(IntSet pages) { }
 
         @Override
         public void executeOnNextDraw(ViewOnDrawExecutor executor) { }
diff --git a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
index c13225a..2a86e81 100644
--- a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
@@ -55,6 +55,7 @@
 import com.android.launcher3.taskbar.TaskbarStateHandler;
 import com.android.launcher3.uioverrides.RecentsViewStateController;
 import com.android.launcher3.util.ActivityOptionsWrapper;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.ObjectWrapper;
 import com.android.launcher3.util.UiThreadHelper;
 import com.android.quickstep.RecentsModel;
@@ -411,8 +412,8 @@
     }
 
     @Override
-    public void finishBindingItems(int pageBoundFirst) {
-        super.finishBindingItems(pageBoundFirst);
+    public void finishBindingItems(IntSet pagesBoundFirst) {
+        super.finishBindingItems(pagesBoundFirst);
         // Instantiate and initialize WellbeingModel now that its loading won't interfere with
         // populating workspace.
         // TODO: Find a better place for this
diff --git a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
index 36322ce..43d03b3 100644
--- a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
+++ b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
@@ -38,7 +38,6 @@
 import static com.android.launcher3.config.FeatureFlags.SEPARATE_RECENTS_ACTIVITY;
 import static com.android.launcher3.dragndrop.DragLayer.ALPHA_INDEX_TRANSITIONS;
 import static com.android.launcher3.statehandlers.DepthController.DEPTH;
-import static com.android.quickstep.TaskUtils.taskIsATargetWithMode;
 import static com.android.quickstep.TaskViewUtils.findTaskViewToLaunch;
 import static com.android.systemui.shared.system.QuickStepContract.getWindowCornerRadius;
 import static com.android.systemui.shared.system.QuickStepContract.supportsRoundedCornersOnWindows;
@@ -1049,7 +1048,15 @@
     }
 
     private boolean launcherIsATargetWithMode(RemoteAnimationTargetCompat[] targets, int mode) {
-        return taskIsATargetWithMode(targets, mLauncher.getTaskId(), mode);
+        for (RemoteAnimationTargetCompat target : targets) {
+            if (target.mode == mode && target.taskInfo != null
+                    // Compare component name instead of task-id because transitions will promote
+                    // the target up to the root task while getTaskId returns the leaf.
+                    && target.taskInfo.topActivity.equals(mLauncher.getComponentName())) {
+                return true;
+            }
+        }
+        return false;
     }
 
     /**
diff --git a/quickstep/src/com/android/launcher3/taskbar/ButtonProvider.java b/quickstep/src/com/android/launcher3/taskbar/ButtonProvider.java
index 540f748..86ac39f 100644
--- a/quickstep/src/com/android/launcher3/taskbar/ButtonProvider.java
+++ b/quickstep/src/com/android/launcher3/taskbar/ButtonProvider.java
@@ -27,6 +27,7 @@
 
 import com.android.launcher3.R;
 import com.android.launcher3.taskbar.TaskbarNavButtonController.TaskbarButton;
+import com.android.launcher3.taskbar.contextual.RotationContextButton;
 
 /**
  * Creates Buttons for Taskbar for 3 button nav.
@@ -68,6 +69,11 @@
         return getButtonForDrawable(R.drawable.ic_ime_switcher, BUTTON_IME_SWITCH);
     }
 
+    public RotationContextButton getContextualRotation() {
+        // Rotation suggestion button
+        return new RotationContextButton(mContext);
+    }
+
     private View getButtonForDrawable(@DrawableRes int drawableId, @TaskbarButton int buttonType) {
         ImageView buttonView = new ImageView(mContext);
         buttonView.setImageResource(drawableId);
diff --git a/quickstep/src/com/android/launcher3/taskbar/ImeBarView.java b/quickstep/src/com/android/launcher3/taskbar/ImeBarView.java
index 287caab..d581302 100644
--- a/quickstep/src/com/android/launcher3/taskbar/ImeBarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/ImeBarView.java
@@ -25,7 +25,6 @@
 
 public class ImeBarView extends RelativeLayout {
 
-    private ButtonProvider mButtonProvider;
     private View mImeView;
 
     public ImeBarView(Context context) {
@@ -41,8 +40,8 @@
     }
 
     public void init(ButtonProvider buttonProvider) {
-        mButtonProvider = buttonProvider;
-
+        // TODO (b/187966005), maybe need to replace ime switcher button with
+        //  RotationContextButton when device rotates
         ActivityContext context = getActivityContext();
         RelativeLayout.LayoutParams imeParams = new RelativeLayout.LayoutParams(
                 context.getDeviceProfile().iconSizePx,
@@ -56,13 +55,13 @@
         downParams.addRule(ALIGN_PARENT_START);
 
         // Down Arrow
-        View downView = mButtonProvider.getDown();
+        View downView = buttonProvider.getDown();
         downView.setLayoutParams(downParams);
         downView.setRotation(-90);
         addView(downView);
 
         // IME switcher button
-        mImeView = mButtonProvider.getImeSwitcher();
+        mImeView = buttonProvider.getImeSwitcher();
         mImeView.setLayoutParams(imeParams);
         addView(mImeView);
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
index c2d107c..f124de7 100644
--- a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
@@ -32,7 +32,6 @@
 import com.android.launcher3.anim.PendingAnimation;
 import com.android.launcher3.states.StateAnimationConfig;
 
-
 /**
  * A data source which integrates with a Launcher instance
  * TODO: Rename to have Launcher prefix
@@ -51,6 +50,7 @@
 
     private @Nullable Animator mAnimator;
     private boolean mIsAnimatingToLauncher;
+    private ContextualRotationNotifier mContextualRotationNotifier;
 
     public LauncherTaskbarUIController(
             BaseQuickstepLauncher launcher, TaskbarActivityContext context) {
@@ -67,7 +67,8 @@
     }
 
     @Override
-    protected void onCreate() {
+    protected void onCreate(ContextualRotationNotifier notifier) {
+        mContextualRotationNotifier = notifier;
         mTaskbarStateHandler.setAnimationController(mTaskbarAnimationController);
         mTaskbarAnimationController.init();
         mHotseatController.init();
@@ -82,6 +83,7 @@
             // End this first, in case it relies on properties that are about to be cleaned up.
             mAnimator.end();
         }
+        mContextualRotationNotifier = null;
         mTaskbarStateHandler.setAnimationController(null);
         mTaskbarAnimationController.cleanup();
         mHotseatController.cleanup();
@@ -105,6 +107,9 @@
             @Override
             public void updateTaskbarVisibilityAlpha(float alpha) {
                 mTaskbarView.setAlpha(alpha);
+                if (mContextualRotationNotifier != null) {
+                    mContextualRotationNotifier.onTaskbarVisibilityChanged(alpha == 1);
+                }
             }
 
             @Override
@@ -222,7 +227,7 @@
     }
 
     public boolean isDraggingItem() {
-        return mTaskbarView.isDraggingItem();
+        return mContext.getDragController().isDragging();
     }
 
     /**
@@ -272,4 +277,8 @@
         void updateTaskbarScale(float scale);
         void updateTaskbarTranslationY(float translationY);
     }
+
+    public interface ContextualRotationNotifier {
+        void onTaskbarVisibilityChanged(boolean showing);
+    }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 4ba0ee0..d51506c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -28,9 +28,7 @@
 import android.content.Intent;
 import android.content.pm.LauncherApps;
 import android.graphics.PixelFormat;
-import android.graphics.Point;
 import android.graphics.Rect;
-import android.graphics.drawable.Drawable;
 import android.os.Process;
 import android.os.SystemProperties;
 import android.util.Log;
@@ -43,24 +41,17 @@
 import android.widget.Toast;
 
 import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
 
 import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.DeviceProfile;
-import com.android.launcher3.DragSource;
-import com.android.launcher3.DropTarget;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.R;
-import com.android.launcher3.dragndrop.DragController;
-import com.android.launcher3.dragndrop.DragOptions;
-import com.android.launcher3.dragndrop.DragView;
-import com.android.launcher3.dragndrop.DraggableView;
 import com.android.launcher3.folder.Folder;
 import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.model.data.FolderInfo;
-import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.taskbar.TaskbarNavButtonController.TaskbarButton;
+import com.android.launcher3.taskbar.contextual.RotationButtonController;
 import com.android.launcher3.touch.ItemClickHandler;
 import com.android.launcher3.util.PackageManagerHelper;
 import com.android.launcher3.util.Themes;
@@ -89,13 +80,18 @@
     private final LayoutInflater mLayoutInflater;
     private final TaskbarDragLayer mDragLayer;
     private final TaskbarIconController mIconController;
-    private final MyDragController mDragController;
+    private final TaskbarDragController mDragController;
 
     private final WindowManager mWindowManager;
     private WindowManager.LayoutParams mWindowLayoutParams;
+    private boolean mIsFullscreen;
+    // The size we should return to when we call setTaskbarWindowFullscreen(false)
+    private int mLastRequestedNonFullscreenHeight;
 
     private final SysUINavigationMode.Mode mNavMode;
+    private final SystemTaskbarNotificationManager mSystemTaskbarNotificationManager;
     private final TaskbarNavButtonController mNavButtonController;
+    private final RotationButtonController mRotationButtonController;
 
     private final boolean mIsSafeModeEnabled;
 
@@ -105,17 +101,48 @@
     private final View.OnClickListener mOnTaskbarIconClickListener;
     private final View.OnLongClickListener mOnTaskbarIconLongClickListener;
 
+    private final TaskbarManager.SystemTaskbarNotifier mSystemTaskbarNotifier =
+            new TaskbarManager.SystemTaskbarNotifier() {
+                @Override
+                public void updateImeStatus(int displayId, int vis, int backDisposition,
+                        boolean showImeSwitcher) {
+                    /*
+                     * When in 3 button nav, sysui flags don't get called since we prevent
+                     *  sysui nav bar from instantiating at all, which is what's responsible for
+                     * sending sysui state flags over.
+                     */
+                    mIconController.updateImeStatus(displayId, vis, showImeSwitcher);
+                }
+
+                @Override
+                public void onRotationProposal(int rotation, boolean isValid) {
+                    mRotationButtonController.onRotationProposal(rotation, isValid);
+                }
+
+                @Override
+                public void disable(int displayId, int state1, int state2, boolean animate) {
+                    mRotationButtonController.onDisable2FlagChanged(state2);
+                }
+
+                @Override
+                public void onSystemBarAttributesChanged(int displayId, int behavior) {
+                    mRotationButtonController.onBehaviorChanged(displayId, behavior);
+                }
+            };
+
     public TaskbarActivityContext(Context windowContext, DeviceProfile dp,
-            TaskbarNavButtonController buttonController) {
+            TaskbarNavButtonController buttonController,
+            SystemTaskbarNotificationManager systemTaskbarNotificationManager) {
         super(windowContext, Themes.getActivityThemeRes(windowContext));
         mDeviceProfile = dp;
         mNavButtonController = buttonController;
         mNavMode = SysUINavigationMode.getMode(windowContext);
+        mSystemTaskbarNotificationManager = systemTaskbarNotificationManager;
         mIsSafeModeEnabled = TraceHelper.allowIpcs("isSafeMode",
                 () -> getPackageManager().isSafeMode());
 
-        mOnTaskbarIconLongClickListener =
-                new TaskbarDragController(this)::startSystemDragOnLongClick;
+        mDragController = new TaskbarDragController(this);
+        mOnTaskbarIconLongClickListener = mDragController::startDragOnLongClick;
         mOnTaskbarIconClickListener = this::onTaskbarIconClicked;
 
         float taskbarIconSize = getResources().getDimension(R.dimen.taskbar_icon_size);
@@ -125,8 +152,11 @@
         mLayoutInflater = LayoutInflater.from(this).cloneInContext(this);
         mDragLayer = (TaskbarDragLayer) mLayoutInflater
                 .inflate(R.layout.taskbar, null, false);
-        mIconController = new TaskbarIconController(this, mDragLayer);
-        mDragController = new MyDragController(this);
+
+        mRotationButtonController = new RotationButtonController(this,
+                R.color.popup_color_primary_light, R.color.popup_color_primary_light);
+        mIconController = new TaskbarIconController(this, mDragLayer,
+                mRotationButtonController);
 
         Display display = windowContext.getDisplay();
         Context c = display.getDisplayId() == Display.DEFAULT_DISPLAY
@@ -136,9 +166,10 @@
     }
 
     public void init() {
+        mLastRequestedNonFullscreenHeight = mDeviceProfile.taskbarSize;
         mWindowLayoutParams = new WindowManager.LayoutParams(
                 MATCH_PARENT,
-                mDeviceProfile.taskbarSize,
+                mLastRequestedNonFullscreenHeight,
                 TYPE_APPLICATION_OVERLAY,
                 WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE,
                 PixelFormat.TRANSLUCENT);
@@ -156,19 +187,13 @@
                 new int[] { ITYPE_EXTRA_NAVIGATION_BAR, ITYPE_BOTTOM_TAPPABLE_ELEMENT }
         );
 
-        mIconController.init(mOnTaskbarIconClickListener, mOnTaskbarIconLongClickListener);
+        mIconController.init(mOnTaskbarIconClickListener, mOnTaskbarIconLongClickListener,
+                mNavMode);
         mWindowManager.addView(mDragLayer, mWindowLayoutParams);
-    }
-
-    /**
-     * Updates the TaskbarContainer height (pass deviceProfile.taskbarSize to reset).
-     */
-    public void setTaskbarWindowHeight(int height) {
-        if (mWindowLayoutParams.height == height) {
-            return;
+        if (mNavMode == Mode.THREE_BUTTONS) {
+            mSystemTaskbarNotificationManager
+                    .registerSystemTaskbarNotifications(mSystemTaskbarNotifier);
         }
-        mWindowLayoutParams.height = height;
-        mWindowManager.updateViewLayout(mDragLayer, mWindowLayoutParams);
     }
 
     public boolean canShowNavButtons() {
@@ -196,7 +221,7 @@
     }
 
     @Override
-    public DragController getDragController() {
+    public TaskbarDragController getDragController() {
         return mDragController;
     }
 
@@ -207,7 +232,10 @@
         mUIController.onDestroy();
         mUIController = uiController;
         mIconController.setUIController(mUIController);
-        mUIController.onCreate();
+        mUIController.onCreate(mRotationButtonController::onTaskBarVisibilityChange);
+        if (mNavMode == Mode.THREE_BUTTONS) {
+            mRotationButtonController.init();
+        }
     }
 
     /**
@@ -217,6 +245,11 @@
         setUIController(TaskbarUIController.DEFAULT);
         mIconController.onDestroy();
         mWindowManager.removeViewImmediate(mDragLayer);
+        if (mNavMode == Mode.THREE_BUTTONS) {
+            mSystemTaskbarNotificationManager.removeSystemTaskbarNotifications(
+                    mSystemTaskbarNotifier);
+            mRotationButtonController.cleanup();
+        }
     }
 
     void onNavigationButtonClick(@TaskbarButton int buttonType) {
@@ -231,20 +264,32 @@
     }
 
     /**
-     * When in 3 button nav, the above doesn't get called since we prevent sysui nav bar from
-     * instantiating at all, which is what's responsible for sending sysui state flags over.
-     *
-     * @param vis IME visibility flag
+     * Updates the TaskbarContainer to MATCH_PARENT vs original Taskbar size.
      */
-    public void updateImeStatus(int displayId, int vis, boolean showImeSwitcher) {
-        mIconController.updateImeStatus(displayId, vis, showImeSwitcher);
+    public void setTaskbarWindowFullscreen(boolean fullscreen) {
+        mIsFullscreen = fullscreen;
+        setTaskbarWindowHeight(fullscreen ? MATCH_PARENT : mLastRequestedNonFullscreenHeight);
     }
 
     /**
-     * Updates the TaskbarContainer to MATCH_PARENT vs original Taskbar size.
+     * Updates the TaskbarContainer height (pass deviceProfile.taskbarSize to reset).
      */
-    protected void setTaskbarWindowFullscreen(boolean fullscreen) {
-        setTaskbarWindowHeight(fullscreen ? MATCH_PARENT : getDeviceProfile().taskbarSize);
+    public void setTaskbarWindowHeight(int height) {
+        if (mWindowLayoutParams.height == height) {
+            return;
+        }
+        if (height != MATCH_PARENT) {
+            mLastRequestedNonFullscreenHeight = height;
+            if (mIsFullscreen) {
+                // We still need to be fullscreen, so defer any change to our height until we call
+                // setTaskbarWindowFullscreen(false). For example, this could happen when dragging
+                // from the gesture region, as the drag will cancel the gesture and reset launcher's
+                // state, which in turn normally would reset the taskbar window height as well.
+                return;
+            }
+        }
+        mWindowLayoutParams.height = height;
+        mWindowManager.updateViewLayout(mDragLayer, mWindowLayoutParams);
     }
 
     protected void onTaskbarIconClicked(View view) {
@@ -307,27 +352,4 @@
 
         AbstractFloatingView.closeAllOpenViews(this);
     }
-
-    private static class MyDragController extends DragController<TaskbarActivityContext> {
-        MyDragController(TaskbarActivityContext activity) {
-            super(activity);
-        }
-
-        @Override
-        protected DragView startDrag(@Nullable Drawable drawable, @Nullable View view,
-                DraggableView originalView, int dragLayerX, int dragLayerY, DragSource source,
-                ItemInfo dragInfo, Point dragOffset, Rect dragRegion, float initialDragViewScale,
-                float dragViewScaleOnDrop, DragOptions options) {
-            return null;
-        }
-
-        @Override
-        protected void exitDrag() {
-        }
-
-        @Override
-        protected DropTarget getDefaultDropTarget(int[] dropCoordinates) {
-            return null;
-        }
-    }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
index ee44927..855c507 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
@@ -20,19 +20,35 @@
 
 import android.content.ClipData;
 import android.content.ClipDescription;
-import android.content.Context;
 import android.content.Intent;
 import android.content.pm.LauncherApps;
 import android.content.res.Resources;
 import android.graphics.Canvas;
 import android.graphics.Point;
+import android.graphics.Rect;
+import android.graphics.drawable.Drawable;
 import android.os.UserHandle;
 import android.view.DragEvent;
+import android.view.MotionEvent;
 import android.view.View;
 
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.BubbleTextView;
+import com.android.launcher3.DragSource;
+import com.android.launcher3.DropTarget;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.R;
+import com.android.launcher3.accessibility.DragViewStateAnnouncer;
+import com.android.launcher3.dragndrop.DragController;
+import com.android.launcher3.dragndrop.DragDriver;
+import com.android.launcher3.dragndrop.DragOptions;
+import com.android.launcher3.dragndrop.DragView;
+import com.android.launcher3.dragndrop.DraggableView;
+import com.android.launcher3.graphics.DragPreviewProvider;
+import com.android.launcher3.icons.FastBitmapDrawable;
+import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.ClipDescriptionCompat;
@@ -41,14 +57,20 @@
 /**
  * Handles long click on Taskbar items to start a system drag and drop operation.
  */
-public class TaskbarDragController {
+public class TaskbarDragController extends DragController<TaskbarActivityContext>  {
 
-    private final Context mContext;
     private final int mDragIconSize;
+    private final int[] mTempXY = new int[2];
 
-    public TaskbarDragController(Context context) {
-        mContext = context;
-        Resources resources = mContext.getResources();
+    // Where the initial touch was relative to the dragged icon.
+    private int mRegistrationX;
+    private int mRegistrationY;
+
+    private boolean mIsSystemDragInProgress;
+
+    public TaskbarDragController(TaskbarActivityContext activity) {
+        super(activity);
+        Resources resources = mActivity.getResources();
         mDragIconSize = resources.getDimensionPixelSize(R.dimen.taskbar_icon_drag_icon_size);
     }
 
@@ -57,18 +79,146 @@
      * generate the ClipDescription and Intent.
      * @return Whether {@link View#startDragAndDrop} started successfully.
      */
-    protected boolean startSystemDragOnLongClick(View view) {
+    protected boolean startDragOnLongClick(View view) {
         if (!(view instanceof BubbleTextView)) {
             return false;
         }
 
         BubbleTextView btv = (BubbleTextView) view;
-        View.DragShadowBuilder shadowBuilder = new View.DragShadowBuilder(view) {
+
+        mActivity.setTaskbarWindowFullscreen(true);
+        view.post(() -> {
+            startInternalDrag(btv);
+            btv.setVisibility(INVISIBLE);
+        });
+        return true;
+    }
+
+    private void startInternalDrag(BubbleTextView btv) {
+        float iconScale = 1f;
+        Drawable icon = btv.getIcon();
+        if (icon instanceof FastBitmapDrawable) {
+            iconScale = ((FastBitmapDrawable) icon).getAnimatedScale();
+        }
+
+        // Clear the pressed state if necessary
+        btv.clearFocus();
+        btv.setPressed(false);
+        btv.clearPressedBackground();
+
+        final DragPreviewProvider previewProvider = new DragPreviewProvider(btv);
+        final Drawable drawable = previewProvider.createDrawable();
+        final float scale = previewProvider.getScaleAndPosition(drawable, mTempXY);
+        int dragLayerX = mTempXY[0];
+        int dragLayerY = mTempXY[1];
+
+        Rect dragRect = new Rect();
+        btv.getSourceVisualDragBounds(dragRect);
+        dragLayerY += dragRect.top;
+
+        DragOptions dragOptions = new DragOptions();
+        // TODO: open popup/pre-drag
+        // PopupContainerWithArrow popupContainer = PopupContainerWithArrow.showForIcon(view);
+        // if (popupContainer != null) {
+        //     dragOptions.preDragCondition = popupContainer.createPreDragCondition();
+        // }
+
+        startDrag(
+                drawable,
+                /* view = */ null,
+                /* originalView = */ btv,
+                dragLayerX,
+                dragLayerY,
+                (View target, DropTarget.DragObject d, boolean success) -> {} /* DragSource */,
+                (WorkspaceItemInfo) btv.getTag(),
+                /* dragVisualizeOffset = */ null,
+                dragRect,
+                scale * iconScale,
+                scale,
+                dragOptions);
+    }
+
+    @Override
+    protected DragView startDrag(@Nullable Drawable drawable, @Nullable View view,
+            DraggableView originalView, int dragLayerX, int dragLayerY, DragSource source,
+            ItemInfo dragInfo, Point dragOffset, Rect dragRegion, float initialDragViewScale,
+            float dragViewScaleOnDrop, DragOptions options) {
+        mOptions = options;
+
+        mRegistrationX = mMotionDown.x - dragLayerX;
+        mRegistrationY = mMotionDown.y - dragLayerY;
+
+        final int dragRegionLeft = dragRegion == null ? 0 : dragRegion.left;
+        final int dragRegionTop = dragRegion == null ? 0 : dragRegion.top;
+
+        mLastDropTarget = null;
+
+        mDragObject = new DropTarget.DragObject(mActivity.getApplicationContext());
+        mDragObject.originalView = originalView;
+
+        mIsInPreDrag = mOptions.preDragCondition != null
+                && !mOptions.preDragCondition.shouldStartDrag(0);
+
+        float scalePx = mDragIconSize - dragRegion.width();
+        final DragView dragView = mDragObject.dragView = new TaskbarDragView(
+                mActivity,
+                drawable,
+                mRegistrationX,
+                mRegistrationY,
+                initialDragViewScale,
+                dragViewScaleOnDrop,
+                scalePx);
+        dragView.setItemInfo(dragInfo);
+        mDragObject.dragComplete = false;
+
+        mDragObject.xOffset = mMotionDown.x - (dragLayerX + dragRegionLeft);
+        mDragObject.yOffset = mMotionDown.y - (dragLayerY + dragRegionTop);
+
+        mDragDriver = DragDriver.create(this, mOptions, /* secondaryEventConsumer = */ ev -> {});
+        if (!mOptions.isAccessibleDrag) {
+            mDragObject.stateAnnouncer = DragViewStateAnnouncer.createFor(dragView);
+        }
+
+        mDragObject.dragSource = source;
+        mDragObject.dragInfo = dragInfo;
+        mDragObject.originalDragInfo = mDragObject.dragInfo.makeShallowCopy();
+
+        if (dragRegion != null) {
+            dragView.setDragRegion(new Rect(dragRegion));
+        }
+
+        dragView.show(mLastTouch.x, mLastTouch.y);
+        mDistanceSinceScroll = 0;
+
+        if (!mIsInPreDrag) {
+            callOnDragStart();
+        } else if (mOptions.preDragCondition != null) {
+            mOptions.preDragCondition.onPreDragStart(mDragObject);
+        }
+
+        handleMoveEvent(mLastTouch.x, mLastTouch.y);
+
+        return dragView;
+    }
+
+    @Override
+    protected void callOnDragStart() {
+        super.callOnDragStart();
+        // Pre-drag has ended, start the global system drag.
+        AbstractFloatingView.closeAllOpenViews(mActivity);
+        startSystemDrag((BubbleTextView) mDragObject.originalView);
+    }
+
+    private void startSystemDrag(BubbleTextView btv) {
+        View.DragShadowBuilder shadowBuilder = new View.DragShadowBuilder(btv) {
+
             @Override
             public void onProvideShadowMetrics(Point shadowSize, Point shadowTouchPoint) {
                 shadowSize.set(mDragIconSize, mDragIconSize);
-                // TODO: should be based on last touch point on the icon.
-                shadowTouchPoint.set(shadowSize.x / 2, shadowSize.y / 2);
+                // The registration point was taken before the icon scaled to mDragIconSize, so
+                // offset the registration to where the touch is on the new size.
+                int offset = (mDragIconSize - btv.getIconSize()) / 2;
+                shadowTouchPoint.set(mRegistrationX + offset, mRegistrationY + offset);
             }
 
             @Override
@@ -81,12 +231,12 @@
             }
         };
 
-        Object tag = view.getTag();
+        Object tag = btv.getTag();
         ClipDescription clipDescription = null;
         Intent intent = null;
         if (tag instanceof WorkspaceItemInfo) {
             WorkspaceItemInfo item = (WorkspaceItemInfo) tag;
-            LauncherApps launcherApps = mContext.getSystemService(LauncherApps.class);
+            LauncherApps launcherApps = mActivity.getSystemService(LauncherApps.class);
             clipDescription = new ClipDescription(item.title,
                     new String[] {
                             item.itemType == LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT
@@ -116,28 +266,89 @@
 
         if (clipDescription != null && intent != null) {
             ClipData clipData = new ClipData(clipDescription, new ClipData.Item(intent));
-            view.setOnDragListener(getDraggedViewDragListener());
-            return view.startDragAndDrop(clipData, shadowBuilder, null /* localState */,
-                    View.DRAG_FLAG_GLOBAL | View.DRAG_FLAG_OPAQUE);
+            if (btv.startDragAndDrop(clipData, shadowBuilder, null /* localState */,
+                    View.DRAG_FLAG_GLOBAL | View.DRAG_FLAG_OPAQUE)) {
+                onSystemDragStarted();
+            }
         }
-        return false;
     }
 
-    /**
-     * Hide the original Taskbar item while it is being dragged.
-     */
-    private View.OnDragListener getDraggedViewDragListener() {
-        return (view, dragEvent) -> {
+    private void onSystemDragStarted() {
+        mIsSystemDragInProgress = true;
+        mActivity.getDragLayer().setOnDragListener((view, dragEvent) -> {
             switch (dragEvent.getAction()) {
                 case DragEvent.ACTION_DRAG_STARTED:
-                    view.setVisibility(INVISIBLE);
+                    // Return true to tell system we are interested in events, so we get DRAG_ENDED.
                     return true;
                 case DragEvent.ACTION_DRAG_ENDED:
-                    view.setVisibility(VISIBLE);
-                    view.setOnDragListener(null);
+                    mIsSystemDragInProgress = false;
+                    maybeOnDragEnd();
                     return true;
             }
             return false;
-        };
+        });
+    }
+
+    @Override
+    public boolean isDragging() {
+        return super.isDragging() || mIsSystemDragInProgress;
+    }
+
+    /**
+     * Whether we started dragging the given view and the drag is still in progress.
+     */
+    public boolean isDraggingView(View child) {
+        return isDragging() && mDragObject != null && mDragObject.originalView == child;
+    }
+
+    private void maybeOnDragEnd() {
+        if (!isDragging()) {
+            ((View) mDragObject.originalView).setVisibility(VISIBLE);
+        }
+    }
+
+    @Override
+    protected void callOnDragEnd() {
+        super.callOnDragEnd();
+        maybeOnDragEnd();
+    }
+
+    @Override
+    protected float getX(MotionEvent ev) {
+        // We will resize to fill the screen while dragging, so use screen coordinates. This ensures
+        // we start at the correct position even though touch down is on the smaller DragLayer size.
+        return ev.getRawX();
+    }
+
+    @Override
+    protected float getY(MotionEvent ev) {
+        // We will resize to fill the screen while dragging, so use screen coordinates. This ensures
+        // we start at the correct position even though touch down is on the smaller DragLayer size.
+        return ev.getRawY();
+    }
+
+    @Override
+    protected Point getClampedDragLayerPos(float x, float y) {
+        // No need to clamp, as we will take up the entire screen.
+        mTmpPoint.set(Math.round(x), Math.round(y));
+        return mTmpPoint;
+    }
+
+    @Override
+    protected void exitDrag() {
+        if (mDragObject != null) {
+            mActivity.getDragLayer().removeView(mDragObject.dragView);
+        }
+    }
+
+    @Override
+    public void addDropTarget(DropTarget target) {
+        // No-op as Taskbar currently doesn't support any drop targets internally.
+        // Note: if we do add internal DropTargets, we'll still need to ignore Folder.
+    }
+
+    @Override
+    protected DropTarget getDefaultDropTarget(int[] dropCoordinates) {
+        return null;
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragLayer.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragLayer.java
index 45ec911..2469f95 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragLayer.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragLayer.java
@@ -40,7 +40,7 @@
     private final int mFolderMargin;
     private final Paint mTaskbarBackgroundPaint;
 
-    private TaskbarIconController.Callbacks mControllerCallbacks;
+    private TaskbarIconController.TaskbarDragLayerCallbacks mControllerCallbacks;
     private TaskbarView mTaskbarView;
 
     private final OnComputeInsetsListener mTaskbarInsetsComputer = this::onComputeTaskbarInsets;
@@ -69,10 +69,11 @@
 
     @Override
     public void recreateControllers() {
-        mControllers = new TouchController[0];
+        mControllers = new TouchController[] {mActivity.getDragController()};
     }
 
-    public void init(TaskbarIconController.Callbacks callbacks, TaskbarView taskbarView) {
+    public void init(TaskbarIconController.TaskbarDragLayerCallbacks callbacks,
+            TaskbarView taskbarView) {
         mControllerCallbacks = callbacks;
         mTaskbarView = taskbarView;
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragView.java
new file mode 100644
index 0000000..cf28eff
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragView.java
@@ -0,0 +1,56 @@
+/*
+ * 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.launcher3.taskbar;
+
+import android.graphics.drawable.Drawable;
+
+import com.android.launcher3.R;
+import com.android.launcher3.dragndrop.DragView;
+
+/**
+ * A DragView drawn/used by the Taskbar. Note that this is only for the internal drag-and-drop,
+ * while the pre-drag is still in progress (i.e. when the long press popup is still open). After
+ * that ends, we switch to a system drag and drop view instead.
+ */
+public class TaskbarDragView extends DragView<TaskbarActivityContext> {
+    public TaskbarDragView(TaskbarActivityContext launcher, Drawable drawable, int registrationX,
+            int registrationY, float initialScale, float scaleOnDrop, float finalScaleDps) {
+        super(launcher, drawable, registrationX, registrationY, initialScale, scaleOnDrop,
+                finalScaleDps);
+    }
+
+    @Override
+    public void animateTo(int toTouchX, int toTouchY, Runnable onCompleteRunnable, int duration) {
+        Runnable onAnimationEnd = () -> {
+            if (onCompleteRunnable != null) {
+                onCompleteRunnable.run();
+            }
+            mActivity.getDragLayer().removeView(this);
+        };
+
+        duration = Math.max(duration,
+                getResources().getInteger(R.integer.config_dropAnimMinDuration));
+
+        animate()
+                .translationX(toTouchX - mRegistrationX)
+                .translationY(toTouchY - mRegistrationY)
+                .scaleX(mScaleOnDrop)
+                .scaleY(mScaleOnDrop)
+                .withEndAction(onAnimationEnd)
+                .setDuration(duration)
+                .start();
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarIconController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarIconController.java
index 683a5b9..5d4b8b7 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarIconController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarIconController.java
@@ -31,6 +31,8 @@
 
 import com.android.launcher3.R;
 import com.android.launcher3.anim.AlphaUpdateListener;
+import com.android.launcher3.taskbar.contextual.RotationButtonController;
+import com.android.quickstep.SysUINavigationMode;
 import com.android.systemui.shared.system.ViewTreeObserverWrapper.InsetsInfo;
 
 /**
@@ -45,27 +47,35 @@
 
     private final TaskbarView mTaskbarView;
     private final ImeBarView mImeBarView;
+    private final RotationButtonController mRotationButtonController;
 
     @NonNull
     private TaskbarUIController mUIController = TaskbarUIController.DEFAULT;
 
-    TaskbarIconController(TaskbarActivityContext activity, TaskbarDragLayer dragLayer) {
+    TaskbarIconController(TaskbarActivityContext activity, TaskbarDragLayer dragLayer,
+            RotationButtonController rotationButtonController) {
         mActivity = activity;
         mDragLayer = dragLayer;
         mTaskbarView = mDragLayer.findViewById(R.id.taskbar_view);
         mImeBarView = mDragLayer.findViewById(R.id.ime_bar_view);
+        mRotationButtonController = rotationButtonController;
     }
 
-    public void init(OnClickListener clickListener, OnLongClickListener longClickListener) {
+    public void init(OnClickListener clickListener, OnLongClickListener longClickListener,
+            SysUINavigationMode.Mode navMode) {
         mDragLayer.addOnLayoutChangeListener((v, a, b, c, d, e, f, g, h) ->
                 mUIController.alignRealHotseatWithTaskbar());
 
         ButtonProvider buttonProvider = new ButtonProvider(mActivity);
         mImeBarView.init(buttonProvider);
-        mTaskbarView.construct(clickListener, longClickListener, buttonProvider);
+        mTaskbarView.init(new TaskbarViewCallbacks(), clickListener, longClickListener,
+                buttonProvider);
         mTaskbarView.getLayoutParams().height = mActivity.getDeviceProfile().taskbarSize;
 
-        mDragLayer.init(new Callbacks(), mTaskbarView);
+        mDragLayer.init(new TaskbarDragLayerCallbacks(), mTaskbarView);
+        if (navMode == SysUINavigationMode.Mode.THREE_BUTTONS) {
+            mRotationButtonController.setRotationButton(mTaskbarView.getContextualRotationButton());
+        }
     }
 
     public void onDestroy() {
@@ -102,7 +112,7 @@
     /**
      * Callbacks for {@link TaskbarDragLayer} to interact with the icon controller
      */
-    public class Callbacks {
+    public class TaskbarDragLayerCallbacks {
 
         /**
          * Called to update the touchable insets
@@ -126,6 +136,12 @@
                             mTaskbarView.mSystemButtonContainer, mTempRect);
                     insetsInfo.touchableRegion.set(mTempRect);
                 }
+                if (mTaskbarView.mContextualButtonContainer.getVisibility() == VISIBLE) {
+                    mDragLayer.getDescendantRectRelativeToSelf(
+                            mTaskbarView.mContextualButtonContainer, mTempRect);
+                    insetsInfo.touchableRegion.union(mTempRect);
+                }
+
                 insetsInfo.setTouchableInsets(TOUCHABLE_INSETS_REGION);
             }
 
@@ -160,4 +176,16 @@
             mImeBarView.setVisibility(alpha == 0 ? GONE : VISIBLE);
         }
     }
+
+    /**
+     * Callbacks for {@link TaskbarView} to interact with the icon controller
+     */
+    public class TaskbarViewCallbacks {
+        /**
+         * Returns whether no other controller is currently handling the given View's visibility.
+         */
+        public boolean canUpdateViewVisibility(View child) {
+            return !mActivity.getDragController().isDraggingView(child);
+        }
+    }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index d026bfb..b9acee8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -40,11 +40,14 @@
 import com.android.quickstep.SysUINavigationMode.Mode;
 import com.android.quickstep.TouchInteractionService;
 
+import java.util.ArrayList;
+import java.util.List;
+
 /**
- * Class to manager taskbar lifecycle
+ * Class to manage taskbar lifecycle
  */
 public class TaskbarManager implements DisplayController.DisplayInfoChangeListener,
-        SysUINavigationMode.NavigationModeChangeListener {
+        SysUINavigationMode.NavigationModeChangeListener, SystemTaskbarNotificationManager {
 
     private final Context mContext;
     private final DisplayController mDisplayController;
@@ -59,6 +62,8 @@
 
     private boolean mUserUnlocked = false;
 
+    private List<SystemTaskbarNotifier> mSystemTaskbarNotifiers = new ArrayList<>();
+
     public TaskbarManager(TouchInteractionService service) {
         mDisplayController = DisplayController.INSTANCE.get(service);
         mSysUINavigationMode = SysUINavigationMode.INSTANCE.get(service);
@@ -124,7 +129,7 @@
             return;
         }
         mTaskbarActivityContext = new TaskbarActivityContext(
-                mContext, dp.copy(mContext), mNavButtonController);
+                mContext, dp.copy(mContext), mNavButtonController, this);
         mTaskbarActivityContext.init();
         if (mLauncher != null) {
             mTaskbarActivityContext.setUIController(
@@ -132,6 +137,9 @@
         }
     }
 
+    // TODO - I don't think this is the best place for these pass through methods,
+    //  maybe directly in TaskbarIconController?
+
     /**
      * See {@link com.android.systemui.shared.system.QuickStepContract.SystemUiStateFlags}
      * @param systemUiStateFlags The latest SystemUiStateFlags
@@ -143,6 +151,16 @@
         }
     }
 
+    public void registerSystemTaskbarNotifications(SystemTaskbarNotifier notifier) {
+        if (!mSystemTaskbarNotifiers.contains(notifier)) {
+            mSystemTaskbarNotifiers.add(notifier);
+        }
+    }
+
+    public void removeSystemTaskbarNotifications(SystemTaskbarNotifier notifier) {
+        mSystemTaskbarNotifiers.remove(notifier);
+    }
+
     /**
      * When in 3 button nav, the above doesn't get called since we prevent sysui nav bar from
      * instantiating at all, which is what's responsible for sending sysui state flags over.
@@ -153,8 +171,26 @@
      */
     public void updateImeStatus(int displayId, int vis, int backDisposition,
             boolean showImeSwitcher) {
-        if (mTaskbarActivityContext != null) {
-            mTaskbarActivityContext.updateImeStatus(displayId, vis, showImeSwitcher);
+        for (SystemTaskbarNotifier notifier : mSystemTaskbarNotifiers) {
+            notifier.updateImeStatus(displayId, vis, backDisposition, showImeSwitcher);
+        }
+    }
+
+    public void onRotationProposal(int rotation, boolean isValid) {
+        for (SystemTaskbarNotifier notifier : mSystemTaskbarNotifiers) {
+            notifier.onRotationProposal(rotation, isValid);
+        }
+    }
+
+    public void disable(int displayId, int state1, int state2, boolean animate) {
+        for (SystemTaskbarNotifier notifier : mSystemTaskbarNotifiers) {
+            notifier.disable(displayId, state1, state2, animate);
+        }
+    }
+
+    public void onSystemBarAttributesChanged(int displayId, int behavior) {
+        for (SystemTaskbarNotifier notifier : mSystemTaskbarNotifiers) {
+            notifier.onSystemBarAttributesChanged(displayId, behavior);
         }
     }
 
@@ -166,4 +202,18 @@
         mDisplayController.removeChangeListener(this);
         mSysUINavigationMode.removeModeChangeListener(this);
     }
+
+    public interface SystemTaskbarNotifier {
+        void updateImeStatus(int displayId, int vis, int backDisposition,
+                boolean showImeSwitcher);
+        void onRotationProposal(int rotation, boolean isValid);
+        void disable(int displayId, int state1, int state2, boolean animate);
+        void onSystemBarAttributesChanged(int displayId, int behavior);
+
+    }
 }
+
+interface SystemTaskbarNotificationManager {
+    void registerSystemTaskbarNotifications(TaskbarManager.SystemTaskbarNotifier notifier);
+    void removeSystemTaskbarNotifications(TaskbarManager.SystemTaskbarNotifier notifier);
+}
\ No newline at end of file
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
index 50adead..f7a5618 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
@@ -27,7 +27,7 @@
      */
     public void alignRealHotseatWithTaskbar() { }
 
-    protected void onCreate() { }
+    protected void onCreate(LauncherTaskbarUIController.ContextualRotationNotifier notifier) { }
 
     protected void onDestroy() { }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
index c6573a6..5c89f8c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
@@ -24,7 +24,6 @@
 import android.graphics.Rect;
 import android.graphics.RectF;
 import android.util.AttributeSet;
-import android.view.DragEvent;
 import android.view.Gravity;
 import android.view.MotionEvent;
 import android.view.View;
@@ -35,7 +34,6 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
-import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.BubbleTextView;
 import com.android.launcher3.Insettable;
 import com.android.launcher3.R;
@@ -44,6 +42,7 @@
 import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.taskbar.contextual.RotationContextButton;
 import com.android.launcher3.views.ActivityContext;
 
 /**
@@ -62,12 +61,14 @@
 
     private final TaskbarActivityContext mActivityContext;
 
-    // Initialized in TaskbarController constructor.
+    // Initialized in init.
+    private TaskbarIconController.TaskbarViewCallbacks mControllerCallbacks;
     private View.OnClickListener mIconClickListener;
     private View.OnLongClickListener mIconLongClickListener;
 
     LinearLayout mSystemButtonContainer;
     LinearLayout mHotseatIconsContainer;
+    LinearLayout mContextualButtonContainer;
 
     // Delegate touches to the closest view if within mIconTouchSize.
     private boolean mDelegateTargeted;
@@ -75,12 +76,12 @@
     // Prevents dispatching touches to children if true
     private boolean mTouchEnabled = true;
 
-    private boolean mIsDraggingItem;
     // Only non-null when the corresponding Folder is open.
     private @Nullable FolderIcon mLeaveBehindFolderIcon;
 
     /** Provider of buttons added to taskbar in 3 button nav */
     private ButtonProvider mButtonProvider;
+    private RotationContextButton mContextualRotationButton;
 
     private boolean mDisableRelayout;
     private boolean mAreHolesAllowed;
@@ -114,12 +115,15 @@
     @Override
     protected void onFinishInflate() {
         super.onFinishInflate();
-        mSystemButtonContainer = findViewById(R.id.system_button_layout);
+        mSystemButtonContainer = findViewById(R.id.nav_button_layout);
         mHotseatIconsContainer = findViewById(R.id.hotseat_icons_layout);
+        mContextualButtonContainer = findViewById(R.id.contextual_button_layout);
     }
 
-    protected void construct(OnClickListener clickListener, OnLongClickListener longClickListener,
-                ButtonProvider buttonProvider) {
+    protected void init(TaskbarIconController.TaskbarViewCallbacks callbacks,
+            OnClickListener clickListener, OnLongClickListener longClickListener,
+            ButtonProvider buttonProvider) {
+        mControllerCallbacks = callbacks;
         mIconClickListener = clickListener;
         mIconLongClickListener = longClickListener;
         mButtonProvider = buttonProvider;
@@ -132,6 +136,10 @@
 
         int numHotseatIcons = mActivityContext.getDeviceProfile().numShownHotseatIcons;
         updateHotseatItems(new ItemInfo[numHotseatIcons]);
+
+        if (mActivityContext.canShowNavButtons()) {
+            createContextualRegion();
+        }
     }
 
     /**
@@ -225,6 +233,9 @@
     }
 
     private void updateHotseatItemVisibility(View hotseatView) {
+        if (!mControllerCallbacks.canUpdateViewVisibility(hotseatView)) {
+            return;
+        }
         hotseatView.setVisibility(
                 hotseatView.getTag() != null ? VISIBLE : (mAreHolesAllowed ? INVISIBLE : GONE));
     }
@@ -349,24 +360,6 @@
         mSystemButtonContainer.addView(mButtonProvider.getRecents(), buttonParams);
     }
 
-    @Override
-    public boolean onDragEvent(DragEvent event) {
-        switch (event.getAction()) {
-            case DragEvent.ACTION_DRAG_STARTED:
-                mIsDraggingItem = true;
-                AbstractFloatingView.closeAllOpenViews(mActivityContext);
-                return true;
-            case DragEvent.ACTION_DRAG_ENDED:
-                mIsDraggingItem = false;
-                break;
-        }
-        return super.onDragEvent(event);
-    }
-
-    public boolean isDraggingItem() {
-        return mIsDraggingItem;
-    }
-
     /**
      * @return The bounding box of where the hotseat elements are relative to this TaskbarView.
      */
@@ -393,6 +386,16 @@
         }
     }
 
+    private void createContextualRegion() {
+        mContextualRotationButton = mButtonProvider.getContextualRotation();
+        mContextualRotationButton.setVisibility(GONE);
+        mContextualButtonContainer.addView(mContextualRotationButton);
+    }
+
+    @Nullable
+    public RotationContextButton getContextualRotationButton() {
+        return mContextualRotationButton;
+    }
     // FolderIconParent implemented methods.
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/taskbar/contextual/RotationButton.java b/quickstep/src/com/android/launcher3/taskbar/contextual/RotationButton.java
new file mode 100644
index 0000000..d421077
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/taskbar/contextual/RotationButton.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright 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.launcher3.taskbar.contextual;
+
+import android.graphics.drawable.AnimatedVectorDrawable;
+import android.view.View;
+
+/**
+ * Interface of a rotation button that interacts {@link RotationButtonController}.
+ * This interface exists because of the two different styles of rotation button in Sysui,
+ * one in contextual for 3 button nav and a floating rotation button for gestural.
+ * Keeping the interface for eventual migration of floating button, so some methods are
+ * pass through to "super" while others are trivially implemented.
+ *
+ * Changes:
+ *  * Directly use AnimatedVectorDrawable instead of KeyButtonDrawable
+ */
+public interface RotationButton {
+    void setRotationButtonController(RotationButtonController rotationButtonController);
+    View getCurrentView();
+    boolean show();
+    boolean hide();
+    boolean isVisible();
+    void updateIcon(int lightIconColor, int darkIconColor);
+    void setOnClickListener(View.OnClickListener onClickListener);
+    void setOnHoverListener(View.OnHoverListener onHoverListener);
+    AnimatedVectorDrawable getImageDrawable();
+    void setDarkIntensity(float darkIntensity);
+    default boolean acceptRotationProposal() {
+        return getCurrentView() != null;
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/taskbar/contextual/RotationButtonController.java b/quickstep/src/com/android/launcher3/taskbar/contextual/RotationButtonController.java
new file mode 100644
index 0000000..6f6abc2
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/taskbar/contextual/RotationButtonController.java
@@ -0,0 +1,512 @@
+/*
+ * Copyright 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.launcher3.taskbar.contextual;
+
+import static android.view.Display.DEFAULT_DISPLAY;
+
+import static com.android.internal.view.RotationPolicy.NATURAL_ROTATION;
+import static com.android.launcher3.anim.Interpolators.LINEAR;
+
+import android.animation.Animator;
+import android.animation.AnimatorListenerAdapter;
+import android.animation.ObjectAnimator;
+import android.annotation.ColorInt;
+import android.annotation.DrawableRes;
+import android.annotation.SuppressLint;
+import android.app.StatusBarManager;
+import android.content.ContentResolver;
+import android.content.Context;
+import android.graphics.drawable.AnimatedVectorDrawable;
+import android.os.Handler;
+import android.os.Looper;
+import android.os.RemoteException;
+import android.provider.Settings;
+import android.util.Log;
+import android.view.IRotationWatcher;
+import android.view.MotionEvent;
+import android.view.Surface;
+import android.view.View;
+import android.view.WindowInsetsController;
+import android.view.WindowManagerGlobal;
+import android.view.accessibility.AccessibilityManager;
+
+import com.android.internal.logging.UiEvent;
+import com.android.internal.logging.UiEventLogger;
+import com.android.internal.logging.UiEventLoggerImpl;
+import com.android.internal.view.RotationPolicy;
+import com.android.launcher3.R;
+import com.android.launcher3.util.DisplayController;
+import com.android.systemui.shared.recents.utilities.Utilities;
+import com.android.systemui.shared.recents.utilities.ViewRippler;
+import com.android.systemui.shared.system.ActivityManagerWrapper;
+import com.android.systemui.shared.system.TaskStackChangeListener;
+import com.android.systemui.shared.system.TaskStackChangeListeners;
+
+import java.util.Optional;
+
+/**
+ * Copied over from the SysUI equivalent class. Known issues/things not ported over
+ *  * When rotation button visible and in auto-hide mode, we ask auto-hide controller to
+ *    keep the navbar around longer. Will need to implement if we use auto-hide on taskbar
+ *
+ * Contains logic that deals with showing a rotate suggestion button with animation.
+ */
+public class RotationButtonController {
+
+    private static final String TAG = "StatusBar/RotationButtonController";
+    private static final int BUTTON_FADE_IN_OUT_DURATION_MS = 100;
+    private static final int NAVBAR_HIDDEN_PENDING_ICON_TIMEOUT_MS = 20000;
+
+    private static final int NUM_ACCEPTED_ROTATION_SUGGESTIONS_FOR_INTRODUCTION = 3;
+
+    private final Context mContext;
+    private final Handler mMainThreadHandler = new Handler(Looper.getMainLooper());
+    private final UiEventLogger mUiEventLogger = new UiEventLoggerImpl();
+    private final ViewRippler mViewRippler = new ViewRippler();
+    private final DisplayController mDisplayController;
+    private RotationButton mRotationButton;
+
+    private int mLastRotationSuggestion;
+    private boolean mPendingRotationSuggestion;
+    private boolean mHoveringRotationSuggestion;
+    private final AccessibilityManager mAccessibilityManager;
+    private final TaskStackListenerImpl mTaskStackListener;
+    private boolean mListenersRegistered = false;
+    private boolean mIsTaskbarShowing;
+    @SuppressLint("InlinedApi")
+    private @WindowInsetsController.Behavior
+    int mBehavior = WindowInsetsController.BEHAVIOR_DEFAULT;
+    private boolean mSkipOverrideUserLockPrefsOnce;
+    private final int mLightIconColor;
+    private final int mDarkIconColor;
+    private int mIconResId = R.drawable.ic_sysbar_rotate_button_ccw_start_90;
+
+    private final Runnable mRemoveRotationProposal =
+            () -> setRotateSuggestionButtonState(false /* visible */);
+    private final Runnable mCancelPendingRotationProposal =
+            () -> mPendingRotationSuggestion = false;
+    private Animator mRotateHideAnimator;
+
+
+    private final IRotationWatcher.Stub mRotationWatcher = new IRotationWatcher.Stub() {
+        @Override
+        public void onRotationChanged(final int rotation) throws RemoteException {
+            // We need this to be scheduled as early as possible to beat the redrawing of
+            // window in response to the orientation change.
+            mMainThreadHandler.postAtFrontOfQueue(() -> {
+                // If the screen rotation changes while locked, potentially update lock to flow with
+                // new screen rotation and hide any showing suggestions.
+                if (isRotationLocked()) {
+                    if (shouldOverrideUserLockPrefs(rotation)) {
+                        setRotationLockedAtAngle(rotation);
+                    }
+                    setRotateSuggestionButtonState(false /* visible */, true /* forced */);
+                }
+            });
+        }
+    };
+
+    /**
+     * Determines if rotation suggestions disabled2 flag exists in flag
+     * @param disable2Flags see if rotation suggestion flag exists in this flag
+     * @return whether flag exists
+     */
+    static boolean hasDisable2RotateSuggestionFlag(int disable2Flags) {
+        return (disable2Flags & StatusBarManager.DISABLE2_ROTATE_SUGGESTIONS) != 0;
+    }
+
+    public RotationButtonController(Context context, @ColorInt int lightIconColor,
+            @ColorInt int darkIconColor) {
+        mContext = context;
+        mLightIconColor = lightIconColor;
+        mDarkIconColor = darkIconColor;
+
+        mAccessibilityManager = AccessibilityManager.getInstance(context);
+        mTaskStackListener = new TaskStackListenerImpl();
+        mDisplayController = DisplayController.INSTANCE.getNoCreate();
+    }
+
+    public void setRotationButton(RotationButton rotationButton) {
+        mRotationButton = rotationButton;
+        mRotationButton.setRotationButtonController(this);
+        mRotationButton.setOnClickListener(this::onRotateSuggestionClick);
+        mRotationButton.setOnHoverListener(this::onRotateSuggestionHover);
+    }
+
+    public void init() {
+        registerListeners();
+        if (mDisplayController.getInfo().id != DEFAULT_DISPLAY) {
+            // Currently there is no accelerometer sensor on non-default display, disable fixed
+            // rotation for non-default display
+            onDisable2FlagChanged(StatusBarManager.DISABLE2_ROTATE_SUGGESTIONS);
+        }
+    }
+
+    public void cleanup() {
+        unregisterListeners();
+    }
+
+    private void registerListeners() {
+        if (mListenersRegistered) {
+            return;
+        }
+
+        mListenersRegistered = true;
+        try {
+            WindowManagerGlobal.getWindowManagerService()
+                    .watchRotation(mRotationWatcher, mDisplayController.getInfo().id);
+        } catch (IllegalArgumentException e) {
+            mListenersRegistered = false;
+            Log.w(TAG, "RegisterListeners for the display failed");
+        } catch (RemoteException e) {
+            Log.e(TAG, "RegisterListeners caught a RemoteException", e);
+            return;
+        }
+
+        TaskStackChangeListeners.getInstance().registerTaskStackListener(mTaskStackListener);
+    }
+
+    void unregisterListeners() {
+        if (!mListenersRegistered) {
+            return;
+        }
+
+        mListenersRegistered = false;
+        try {
+            WindowManagerGlobal.getWindowManagerService().removeRotationWatcher(mRotationWatcher);
+        } catch (RemoteException e) {
+            Log.e(TAG, "UnregisterListeners caught a RemoteException", e);
+            return;
+        }
+
+        TaskStackChangeListeners.getInstance().unregisterTaskStackListener(mTaskStackListener);
+    }
+
+    void setRotationLockedAtAngle(int rotationSuggestion) {
+        RotationPolicy.setRotationLockAtAngle(mContext, true, rotationSuggestion);
+    }
+
+    public boolean isRotationLocked() {
+        return RotationPolicy.isRotationLocked(mContext);
+    }
+
+    public void setRotateSuggestionButtonState(boolean visible) {
+        setRotateSuggestionButtonState(visible, false /* force */);
+    }
+
+    void setRotateSuggestionButtonState(final boolean visible, final boolean force) {
+        // At any point the button can become invisible because an a11y service became active.
+        // Similarly, a call to make the button visible may be rejected because an a11y service is
+        // active. Must account for this.
+        // Rerun a show animation to indicate change but don't rerun a hide animation
+        if (!visible && !mRotationButton.isVisible()) return;
+
+        final View view = mRotationButton.getCurrentView();
+        if (view == null) return;
+
+        final AnimatedVectorDrawable currentDrawable = mRotationButton.getImageDrawable();
+        if (currentDrawable == null) return;
+
+        // Clear any pending suggestion flag as it has either been nullified or is being shown
+        mPendingRotationSuggestion = false;
+        mMainThreadHandler.removeCallbacks(mCancelPendingRotationProposal);
+
+        // Handle the visibility change and animation
+        if (visible) { // Appear and change (cannot force)
+            // Stop and clear any currently running hide animations
+            if (mRotateHideAnimator != null && mRotateHideAnimator.isRunning()) {
+                mRotateHideAnimator.cancel();
+            }
+            mRotateHideAnimator = null;
+
+            // Reset the alpha if any has changed due to hide animation
+            view.setAlpha(1f);
+
+            // Run the rotate icon's animation if it has one
+            currentDrawable.reset();
+            currentDrawable.start();
+
+            // TODO(b/187754252): No idea why this doesn't work. If we remove the "false"
+            //  we see the animation show the pressed state... but it only shows the first time.
+            if (!isRotateSuggestionIntroduced()) mViewRippler.start(view);
+
+            // Set visibility unless a11y service is active.
+            mRotationButton.show();
+        } else { // Hide
+            mViewRippler.stop(); // Prevent any pending ripples, force hide or not
+
+            if (force) {
+                // If a hide animator is running stop it and make invisible
+                if (mRotateHideAnimator != null && mRotateHideAnimator.isRunning()) {
+                    mRotateHideAnimator.pause();
+                }
+                mRotationButton.hide();
+                return;
+            }
+
+            // Don't start any new hide animations if one is running
+            if (mRotateHideAnimator != null && mRotateHideAnimator.isRunning()) return;
+
+            ObjectAnimator fadeOut = ObjectAnimator.ofFloat(view, "alpha", 0f);
+            fadeOut.setDuration(BUTTON_FADE_IN_OUT_DURATION_MS);
+            fadeOut.setInterpolator(LINEAR);
+            fadeOut.addListener(new AnimatorListenerAdapter() {
+                @Override
+                public void onAnimationEnd(Animator animation) {
+                    mRotationButton.hide();
+                }
+            });
+
+            mRotateHideAnimator = fadeOut;
+            fadeOut.start();
+        }
+    }
+
+    void setDarkIntensity(float darkIntensity) {
+        mRotationButton.setDarkIntensity(darkIntensity);
+    }
+
+    public void onRotationProposal(int rotation, boolean isValid) {
+        int windowRotation = mDisplayController.getInfo().rotation;
+
+        if (!mRotationButton.acceptRotationProposal()) {
+            return;
+        }
+
+        // This method will be called on rotation suggestion changes even if the proposed rotation
+        // is not valid for the top app. Use invalid rotation choices as a signal to remove the
+        // rotate button if shown.
+        if (!isValid) {
+            setRotateSuggestionButtonState(false /* visible */);
+            return;
+        }
+
+        // If window rotation matches suggested rotation, remove any current suggestions
+        if (rotation == windowRotation) {
+            mMainThreadHandler.removeCallbacks(mRemoveRotationProposal);
+            setRotateSuggestionButtonState(false /* visible */);
+            return;
+        }
+
+        // Prepare to show the navbar icon by updating the icon style to change anim params
+        mLastRotationSuggestion = rotation; // Remember rotation for click
+        final boolean rotationCCW = Utilities.isRotationAnimationCCW(windowRotation, rotation);
+        if (windowRotation == Surface.ROTATION_0 || windowRotation == Surface.ROTATION_180) {
+            mIconResId = rotationCCW
+                    ? R.drawable.ic_sysbar_rotate_button_ccw_start_90
+                    : R.drawable.ic_sysbar_rotate_button_cw_start_90;
+        } else { // 90 or 270
+            mIconResId = rotationCCW
+                    ? R.drawable.ic_sysbar_rotate_button_ccw_start_0
+                    : R.drawable.ic_sysbar_rotate_button_ccw_start_0;
+        }
+        mRotationButton.updateIcon(mLightIconColor, mDarkIconColor);
+
+        if (canShowRotationButton()) {
+            // The navbar is visible / it's in visual immersive mode, so show the icon right away
+            showAndLogRotationSuggestion();
+        } else {
+            // If the navbar isn't shown, flag the rotate icon to be shown should the navbar become
+            // visible given some time limit.
+            mPendingRotationSuggestion = true;
+            mMainThreadHandler.removeCallbacks(mCancelPendingRotationProposal);
+            mMainThreadHandler.postDelayed(mCancelPendingRotationProposal,
+                    NAVBAR_HIDDEN_PENDING_ICON_TIMEOUT_MS);
+        }
+    }
+
+    public void onDisable2FlagChanged(int state2) {
+        final boolean rotateSuggestionsDisabled = hasDisable2RotateSuggestionFlag(state2);
+        if (rotateSuggestionsDisabled) onRotationSuggestionsDisabled();
+    }
+
+    public void onBehaviorChanged(int displayId, @WindowInsetsController.Behavior int behavior) {
+        if (mDisplayController.getInfo().id != displayId) {
+            return;
+        }
+
+        if (mBehavior != behavior) {
+            mBehavior = behavior;
+            showPendingRotationButtonIfNeeded();
+        }
+    }
+
+    public void onTaskBarVisibilityChange(boolean showing) {
+        if (mIsTaskbarShowing != showing) {
+            mIsTaskbarShowing = showing;
+            showPendingRotationButtonIfNeeded();
+        }
+    }
+
+    private void showPendingRotationButtonIfNeeded() {
+        if (canShowRotationButton() && mPendingRotationSuggestion) {
+            showAndLogRotationSuggestion();
+        }
+    }
+
+    /** Return true when either the task bar is visible or it's in visual immersive mode. */
+    @SuppressLint("InlinedApi")
+    private boolean canShowRotationButton() {
+        return mIsTaskbarShowing || mBehavior == WindowInsetsController.BEHAVIOR_DEFAULT;
+    }
+
+    public @DrawableRes
+    int getIconResId() {
+        return mIconResId;
+    }
+
+    public @ColorInt int getLightIconColor() {
+        return mLightIconColor;
+    }
+
+    public @ColorInt int getDarkIconColor() {
+        return mDarkIconColor;
+    }
+
+    private void onRotateSuggestionClick(View v) {
+        mUiEventLogger.log(RotationButtonEvent.ROTATION_SUGGESTION_ACCEPTED);
+        incrementNumAcceptedRotationSuggestionsIfNeeded();
+        setRotationLockedAtAngle(mLastRotationSuggestion);
+    }
+
+    private boolean onRotateSuggestionHover(View v, MotionEvent event) {
+        final int action = event.getActionMasked();
+        mHoveringRotationSuggestion = (action == MotionEvent.ACTION_HOVER_ENTER)
+                || (action == MotionEvent.ACTION_HOVER_MOVE);
+        rescheduleRotationTimeout(true /* reasonHover */);
+        return false; // Must return false so a11y hover events are dispatched correctly.
+    }
+
+    private void onRotationSuggestionsDisabled() {
+        // Immediately hide the rotate button and clear any planned removal
+        setRotateSuggestionButtonState(false /* visible */, true /* force */);
+        mMainThreadHandler.removeCallbacks(mRemoveRotationProposal);
+    }
+
+    private void showAndLogRotationSuggestion() {
+        setRotateSuggestionButtonState(true /* visible */);
+        rescheduleRotationTimeout(false /* reasonHover */);
+        mUiEventLogger.log(RotationButtonEvent.ROTATION_SUGGESTION_SHOWN);
+    }
+
+    /**
+     * Makes {@link #shouldOverrideUserLockPrefs} always return {@code false} once. It is used to
+     * avoid losing original user rotation when display rotation is changed by entering the fixed
+     * orientation overview.
+     */
+    void setSkipOverrideUserLockPrefsOnce() {
+        mSkipOverrideUserLockPrefsOnce = true;
+    }
+
+    private boolean shouldOverrideUserLockPrefs(final int rotation) {
+        if (mSkipOverrideUserLockPrefsOnce) {
+            mSkipOverrideUserLockPrefsOnce = false;
+            return false;
+        }
+        // Only override user prefs when returning to the natural rotation (normally portrait).
+        // Don't let apps that force landscape or 180 alter user lock.
+        return rotation == NATURAL_ROTATION;
+    }
+
+    private void rescheduleRotationTimeout(final boolean reasonHover) {
+        // May be called due to a new rotation proposal or a change in hover state
+        if (reasonHover) {
+            // Don't reschedule if a hide animator is running
+            if (mRotateHideAnimator != null && mRotateHideAnimator.isRunning()) return;
+            // Don't reschedule if not visible
+            if (!mRotationButton.isVisible()) return;
+        }
+
+        // Stop any pending removal
+        mMainThreadHandler.removeCallbacks(mRemoveRotationProposal);
+        // Schedule timeout
+        mMainThreadHandler.postDelayed(mRemoveRotationProposal,
+                computeRotationProposalTimeout());
+    }
+
+    private int computeRotationProposalTimeout() {
+        return mAccessibilityManager.getRecommendedTimeoutMillis(
+                mHoveringRotationSuggestion ? 16000 : 5000,
+                AccessibilityManager.FLAG_CONTENT_CONTROLS);
+    }
+
+    private boolean isRotateSuggestionIntroduced() {
+        ContentResolver cr = mContext.getContentResolver();
+        return Settings.Secure.getInt(cr, Settings.Secure.NUM_ROTATION_SUGGESTIONS_ACCEPTED, 0)
+                >= NUM_ACCEPTED_ROTATION_SUGGESTIONS_FOR_INTRODUCTION;
+    }
+
+    private void incrementNumAcceptedRotationSuggestionsIfNeeded() {
+        // Get the number of accepted suggestions
+        ContentResolver cr = mContext.getContentResolver();
+        final int numSuggestions = Settings.Secure.getInt(cr,
+                Settings.Secure.NUM_ROTATION_SUGGESTIONS_ACCEPTED, 0);
+
+        // Increment the number of accepted suggestions only if it would change intro mode
+        if (numSuggestions < NUM_ACCEPTED_ROTATION_SUGGESTIONS_FOR_INTRODUCTION) {
+            Settings.Secure.putInt(cr, Settings.Secure.NUM_ROTATION_SUGGESTIONS_ACCEPTED,
+                    numSuggestions + 1);
+        }
+    }
+
+    private class TaskStackListenerImpl extends TaskStackChangeListener {
+        // Invalidate any rotation suggestion on task change or activity orientation change
+        // Note: all callbacks happen on main thread
+
+        @Override
+        public void onTaskStackChanged() {
+            setRotateSuggestionButtonState(false /* visible */);
+        }
+
+        @Override
+        public void onTaskRemoved(int taskId) {
+            setRotateSuggestionButtonState(false /* visible */);
+        }
+
+        @Override
+        public void onTaskMovedToFront(int taskId) {
+            setRotateSuggestionButtonState(false /* visible */);
+        }
+
+        @Override
+        public void onActivityRequestedOrientationChanged(int taskId, int requestedOrientation) {
+            // Only hide the icon if the top task changes its requestedOrientation
+            // Launcher can alter its requestedOrientation while it's not on top, don't hide on this
+            Optional.ofNullable(ActivityManagerWrapper.getInstance())
+                    .map(ActivityManagerWrapper::getRunningTask)
+                    .ifPresent(a -> {
+                        if (a.id == taskId) setRotateSuggestionButtonState(false /* visible */);
+                    });
+        }
+    }
+
+    enum RotationButtonEvent implements UiEventLogger.UiEventEnum {
+        @UiEvent(doc = "The rotation button was shown")
+        ROTATION_SUGGESTION_SHOWN(206),
+        @UiEvent(doc = "The rotation button was clicked")
+        ROTATION_SUGGESTION_ACCEPTED(207);
+
+        private final int mId;
+        RotationButtonEvent(int id) {
+            mId = id;
+        }
+        @Override public int getId() {
+            return mId;
+        }
+    }
+}
+
diff --git a/quickstep/src/com/android/launcher3/taskbar/contextual/RotationContextButton.java b/quickstep/src/com/android/launcher3/taskbar/contextual/RotationContextButton.java
new file mode 100644
index 0000000..7ad3191
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/taskbar/contextual/RotationContextButton.java
@@ -0,0 +1,112 @@
+/*
+ * Copyright 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.launcher3.taskbar.contextual;
+
+import android.content.Context;
+import android.graphics.drawable.AnimatedVectorDrawable;
+import android.view.View;
+import android.widget.ImageView;
+
+import com.android.launcher3.R;
+
+/** Containing logic for the rotation button in nav bar. */
+public class RotationContextButton extends ImageView implements RotationButton {
+
+    private AnimatedVectorDrawable mImageDrawable;
+
+    public RotationContextButton(Context context) {
+        super(context);
+        setBackgroundResource(R.drawable.taskbar_icon_click_feedback_roundrect);
+    }
+
+    @Override
+    public void setRotationButtonController(RotationButtonController rotationButtonController) {
+        // TODO(b/187754252) UI polish, different icons based on light/dark context, etc
+        mImageDrawable = (AnimatedVectorDrawable) getContext()
+                .getDrawable(rotationButtonController.getIconResId());
+        setImageDrawable(mImageDrawable);
+        mImageDrawable.setCallback(this);
+    }
+
+    @Override
+    public View getCurrentView() {
+        return this;
+    }
+
+    @Override
+    public boolean show() {
+        setVisibility(VISIBLE);
+        return true;
+    }
+
+    @Override
+    public boolean hide() {
+        setVisibility(GONE);
+        return true;
+    }
+
+    @Override
+    public boolean isVisible() {
+        return getVisibility() == VISIBLE;
+    }
+
+    @Override
+    public void updateIcon(int lightIconColor, int darkIconColor) {
+        // TODO(b/187754252): UI Polish
+    }
+
+    @Override
+    public void setOnClickListener(View.OnClickListener onClickListener) {
+        super.setOnClickListener(onClickListener);
+    }
+
+    @Override
+    public void setOnHoverListener(View.OnHoverListener onHoverListener) {
+        super.setOnHoverListener(onHoverListener);
+    }
+
+    @Override
+    public AnimatedVectorDrawable getImageDrawable() {
+        return mImageDrawable;
+    }
+
+    @Override
+    public void setDarkIntensity(float darkIntensity) {
+        // TODO(b/187754252) UI polish
+    }
+
+    @Override
+    public void setVisibility(int visibility) {
+        super.setVisibility(visibility);
+
+        if (visibility != View.VISIBLE && mImageDrawable != null) {
+            mImageDrawable.clearAnimationCallbacks();
+            mImageDrawable.reset();
+        }
+
+        // Start the rotation animation once it becomes visible
+        if (visibility == View.VISIBLE && mImageDrawable != null) {
+            mImageDrawable.reset();
+            mImageDrawable.start();
+        }
+    }
+
+    @Override
+    public boolean acceptRotationProposal() {
+        return isAttachedToWindow();
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/uioverrides/RecentsViewStateController.java b/quickstep/src/com/android/launcher3/uioverrides/RecentsViewStateController.java
index 996d36a..6cad3dd 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/RecentsViewStateController.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/RecentsViewStateController.java
@@ -106,8 +106,7 @@
         float clearAllButtonAlpha = state.areElementsVisible(mLauncher, CLEAR_ALL_BUTTON) ? 1 : 0;
         propertySetter.setFloat(mRecentsView.getClearAllButton(), ClearAllButton.VISIBILITY_ALPHA,
                 clearAllButtonAlpha, LINEAR);
-        float overviewButtonAlpha = state.areElementsVisible(mLauncher, OVERVIEW_ACTIONS)
-                && mRecentsView.shouldShowOverviewActionsForState(state) ? 1 : 0;
+        float overviewButtonAlpha = state.areElementsVisible(mLauncher, OVERVIEW_ACTIONS) ? 1 : 0;
         propertySetter.setFloat(mLauncher.getActionsView().getVisibilityAlpha(),
                 MultiValueAlpha.VALUE, overviewButtonAlpha, config.getInterpolator(
                         ANIM_OVERVIEW_ACTIONS_FADE, LINEAR));
diff --git a/quickstep/src/com/android/quickstep/BaseActivityInterface.java b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
index 4ae6fa8..803f036 100644
--- a/quickstep/src/com/android/quickstep/BaseActivityInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
@@ -208,17 +208,18 @@
             Rect gridRect = new Rect();
             calculateGridSize(context, dp, gridRect);
 
-            int verticalMargin = res.getDimensionPixelSize(
-                    R.dimen.overview_grid_focus_vertical_margin);
-            float taskHeight = gridRect.height() - verticalMargin * 2;
+            int verticalMargin = Math.max(
+                    res.getDimensionPixelSize(R.dimen.overview_grid_focus_vertical_margin),
+                    res.getDimensionPixelSize(R.dimen.overview_actions_height));
+            float taskHeight =
+                    gridRect.height() - verticalMargin * 2 - dp.overviewTaskThumbnailTopMarginPx;
 
             PointF taskDimension = getTaskDimension(context, dp);
-            float scale = taskHeight / Math.max(taskDimension.x, taskDimension.y);
+            float scale = taskHeight / taskDimension.y;
             int outWidth = Math.round(scale * taskDimension.x);
             int outHeight = Math.round(scale * taskDimension.y);
 
-            int gravity = Gravity.CENTER_VERTICAL;
-            gravity |= orientedState.getRecentsRtlSetting(res) ? Gravity.RIGHT : Gravity.LEFT;
+            int gravity = Gravity.CENTER;
             Gravity.apply(gravity, outWidth, outHeight, gridRect, outRect);
         } else {
             int taskMargin = dp.overviewTaskMarginPx;
@@ -307,8 +308,7 @@
         float rowHeight = (gridRect.height() - rowSpacing) / 2f;
 
         PointF taskDimension = getTaskDimension(context, dp);
-        float scale = (rowHeight - dp.overviewTaskThumbnailTopMarginPx) / Math.max(
-                taskDimension.x, taskDimension.y);
+        float scale = (rowHeight - dp.overviewTaskThumbnailTopMarginPx) / taskDimension.y;
         int outWidth = Math.round(scale * taskDimension.x);
         int outHeight = Math.round(scale * taskDimension.y);
 
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index e52405b..afafce7 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -261,7 +261,7 @@
         }
 
         @Override
-        public void onSplitScreenSecondaryBoundsChanged(Rect bounds, Rect insets)  {
+        public void onSplitScreenSecondaryBoundsChanged(Rect bounds, Rect insets) {
             WindowBounds wb = new WindowBounds(bounds, insets);
             MAIN_EXECUTOR.execute(() -> SplitScreenBounds.INSTANCE.setSecondaryWindowBounds(wb));
         }
@@ -269,8 +269,34 @@
         @Override
         public void onImeWindowStatusChanged(int displayId, IBinder token, int vis,
                 int backDisposition, boolean showImeSwitcher) {
-            MAIN_EXECUTOR.execute(() -> mTaskbarManager.updateImeStatus(
-                    displayId, vis, backDisposition, showImeSwitcher));
+            executeForTaskbarManager(() -> mTaskbarManager
+                    .updateImeStatus(displayId, vis, backDisposition, showImeSwitcher));
+        }
+
+        @Override
+        public void onRotationProposal(int rotation, boolean isValid) {
+            executeForTaskbarManager(() -> mTaskbarManager.onRotationProposal(rotation, isValid));
+        }
+
+        @Override
+        public void disable(int displayId, int state1, int state2, boolean animate) {
+            executeForTaskbarManager(() -> mTaskbarManager
+                    .disable(displayId, state1, state2, animate));
+        }
+
+        @Override
+        public void onSystemBarAttributesChanged(int displayId, int behavior) {
+            executeForTaskbarManager(() -> mTaskbarManager
+                    .onSystemBarAttributesChanged(displayId, behavior));
+        }
+
+        private void executeForTaskbarManager(final Runnable r) {
+            MAIN_EXECUTOR.execute(() -> {
+                if (mTaskbarManager == null) {
+                    return;
+                }
+                r.run();
+            });
         }
 
         public TaskbarManager getTaskbarManager() {
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsStateController.java b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsStateController.java
index f0364eb..52083bb 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsStateController.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsStateController.java
@@ -77,9 +77,7 @@
         float clearAllButtonAlpha = state.hasClearAllButton() ? 1 : 0;
         setter.setFloat(mRecentsView.getClearAllButton(), ClearAllButton.VISIBILITY_ALPHA,
                 clearAllButtonAlpha, LINEAR);
-        float overviewButtonAlpha =
-                state.hasOverviewActions() && mRecentsView.shouldShowOverviewActionsForState(state)
-                        ? 1 : 0;
+        float overviewButtonAlpha = state.hasOverviewActions() ? 1 : 0;
         setter.setFloat(mActivity.getActionsView().getVisibilityAlpha(),
                 MultiValueAlpha.VALUE, overviewButtonAlpha, LINEAR);
 
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
index ac3fb27..efce650 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
@@ -26,14 +26,12 @@
 import android.content.Context;
 import android.os.Build;
 import android.util.AttributeSet;
-import android.util.Log;
 
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.anim.AnimatorPlaybackController;
 import com.android.launcher3.anim.PendingAnimation;
 import com.android.launcher3.statemanager.StateManager.StateListener;
-import com.android.launcher3.testing.TestProtocol;
 import com.android.quickstep.FallbackActivityInterface;
 import com.android.quickstep.GestureState;
 import com.android.quickstep.RecentsActivity;
@@ -122,6 +120,12 @@
         }
     }
 
+    @Nullable
+    @Override
+    protected TaskView getHomeTaskView() {
+        return mHomeTaskInfo != null ? getTaskView(mHomeTaskInfo.taskId) : null;
+    }
+
     @Override
     protected boolean shouldAddStubTaskView(RunningTaskInfo runningTaskInfo) {
         if (mHomeTaskInfo != null && runningTaskInfo != null &&
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index a0bba86..7cd8d60 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -32,7 +32,6 @@
 import static com.android.launcher3.Utilities.squaredHypot;
 import static com.android.launcher3.Utilities.squaredTouchSlop;
 import static com.android.launcher3.anim.Interpolators.ACCEL;
-import static com.android.launcher3.anim.Interpolators.ACCEL_0_5;
 import static com.android.launcher3.anim.Interpolators.ACCEL_0_75;
 import static com.android.launcher3.anim.Interpolators.ACCEL_DEACCEL;
 import static com.android.launcher3.anim.Interpolators.FAST_OUT_SLOW_IN;
@@ -341,6 +340,7 @@
     private static final int ADDITION_TASK_DURATION = 200;
     private static final float INITIAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET = 0.55f;
     private static final float ADDITIONAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET = 0.05f;
+    private static final float ANIMATION_DISMISS_PROGRESS_MIDPOINT = 0.5f;
 
     protected final RecentsOrientedState mOrientationState;
     protected final BaseActivityInterface<STATE_TYPE, ACTIVITY_TYPE> mSizeStrategy;
@@ -348,6 +348,13 @@
     protected SurfaceTransactionApplier mSyncTransactionApplier;
     protected int mTaskWidth;
     protected int mTaskHeight;
+    // Used to position the top of a task in the top row of the grid
+    private float mTaskGridVerticalDiff;
+    // The vertical space one grid task takes + space between top and bottom row.
+    private float mTopBottomRowHeightDiff;
+    // mTaskGridVerticalDiff and mTopBottomRowHeightDiff summed together provides the top
+    // position for bottom row of grid tasks.
+
     protected final TransformParams mLiveTileParams = new TransformParams();
     protected final TaskViewSimulator mLiveTileTaskViewSimulator;
     protected final Rect mLastComputedTaskSize = new Rect();
@@ -487,7 +494,6 @@
     protected boolean mRunningTaskTileHidden;
     private Task mTmpRunningTask;
     protected int mFocusedTaskId = -1;
-    private float mFocusedTaskRatio;
 
     private boolean mRunningTaskIconScaledDown = false;
 
@@ -812,7 +818,7 @@
         // taskview for entering split screen, we only pretend to dismiss the task
         if (child instanceof TaskView && child != mSplitHiddenTaskView) {
             TaskView taskView = (TaskView) child;
-            mHasVisibleTaskData.delete(taskView.getTask().key.id);
+            mHasVisibleTaskData.delete(taskView.getTaskId());
             mTaskViewPool.recycle(taskView);
             mActionsView.updateHiddenFlags(HIDDEN_NO_TASKS, getTaskViewCount() == 0);
         }
@@ -919,7 +925,7 @@
     public TaskView getTaskView(int taskId) {
         for (int i = 0; i < getTaskViewCount(); i++) {
             TaskView taskView = getTaskViewAt(i);
-            if (taskView.hasTaskId(taskId)) {
+            if (taskView.getTaskId() == taskId) {
                 return taskView;
             }
         }
@@ -1111,6 +1117,9 @@
             final TaskView taskView = (TaskView) getChildAt(pageIndex);
             taskView.bind(task, mOrientationState);
         }
+        if (mFocusedTaskId == -1 && getTaskViewCount() > 0) {
+            mFocusedTaskId = getTaskViewAt(0).getTaskId();
+        }
         updateTaskSize();
 
         if (mNextPage == INVALID_PAGE) {
@@ -1167,7 +1176,7 @@
     public void resetTaskVisuals() {
         for (int i = getTaskViewCount() - 1; i >= 0; i--) {
             TaskView taskView = getTaskViewAt(i);
-            if (mIgnoreResetTaskId != taskView.getTask().key.id) {
+            if (mIgnoreResetTaskId != taskView.getTaskId()) {
                 taskView.resetViewTransforms();
                 taskView.setStableAlpha(mContentAlpha);
                 taskView.setFullscreenProgress(mFullscreenProgress);
@@ -1286,6 +1295,11 @@
         mSizeStrategy.calculateGridTaskSize(mActivity, mActivity.getDeviceProfile(),
                 mLastComputedGridTaskSize, mOrientationHandler);
 
+        mTaskGridVerticalDiff = mLastComputedGridTaskSize.top - mLastComputedTaskSize.top;
+        mTopBottomRowHeightDiff =
+                mLastComputedGridTaskSize.height() + dp.overviewTaskThumbnailTopMarginPx
+                        + mRowSpacing;
+
         // Force TaskView to update size from thumbnail
         updateTaskSize();
 
@@ -1316,6 +1330,15 @@
      * Updates TaskView scaling and translation required to support variable width.
      */
     private void updateTaskSize() {
+        updateTaskSize(false);
+    }
+
+    /**
+     * Updates TaskView scaling and translation required to support variable width.
+     *
+     * @param isTaskDismissal indicates if update was called due to task dismissal
+     */
+    private void updateTaskSize(boolean isTaskDismissal) {
         final int taskCount = getTaskViewCount();
         if (taskCount == 0) {
             return;
@@ -1336,7 +1359,7 @@
 
         mClearAllButton.setFullscreenTranslationPrimary(accumulatedTranslationX);
 
-        updateGridProperties();
+        updateGridProperties(isTaskDismissal);
     }
 
     public void getTaskSize(Rect outRect) {
@@ -1351,19 +1374,7 @@
     public Point getSelectedTaskSize() {
         mSizeStrategy.calculateTaskSize(mActivity, mActivity.getDeviceProfile(), mTempRect,
                 mOrientationHandler);
-        int taskWidth = mTempRect.width();
-        int taskHeight = mTempRect.height();
-        if (mRunningTaskId != -1) {
-            int boxLength = Math.max(taskWidth, taskHeight);
-            if (mFocusedTaskRatio > 1) {
-                taskWidth = boxLength;
-                taskHeight = (int) (boxLength / mFocusedTaskRatio);
-            } else {
-                taskWidth = (int) (boxLength * mFocusedTaskRatio);
-                taskHeight = boxLength;
-            }
-        }
-        return new Point(taskWidth, taskHeight);
+        return new Point(mTempRect.width(), mTempRect.height());
     }
 
     /** Gets the last computed task size */
@@ -1395,13 +1406,7 @@
             loadVisibleTaskData(TaskView.FLAG_UPDATE_ALL);
 
             // After scrolling, update ActionsView's visibility.
-            TaskView focusedTaskView = getFocusedTaskView();
-            if (focusedTaskView != null) {
-                float scrollDiff = Math.abs(getScrollForPage(indexOfChild(focusedTaskView))
-                        - mOrientationHandler.getPrimaryScroll(this));
-                float delta = (mGridSideMargin - scrollDiff) / (float) mGridSideMargin;
-                mActionsView.getScrollAlpha().setValue(Utilities.boundToRange(delta, 0, 1));
-            }
+            updateActionsViewScrollAlpha();
         }
 
         // Update the high res thumbnail loader state
@@ -1409,6 +1414,20 @@
         return scrolling;
     }
 
+    private void updateActionsViewScrollAlpha() {
+        float scrollAlpha = 1f;
+        if (showAsGrid()) {
+            TaskView focusedTaskView = getFocusedTaskView();
+            if (focusedTaskView != null) {
+                float scrollDiff = Math.abs(getScrollForPage(indexOfChild(focusedTaskView))
+                        - mOrientationHandler.getPrimaryScroll(this));
+                float delta = (mGridSideMargin - scrollDiff) / (float) mGridSideMargin;
+                scrollAlpha = Utilities.boundToRange(delta, 0, 1);
+            }
+        }
+        mActionsView.getScrollAlpha().setValue(scrollAlpha);
+    }
+
     /**
      * Scales and adjusts translation of adjacent pages as if on a curved carousel.
      */
@@ -1551,7 +1570,7 @@
         setCurrentTask(-1);
         mIgnoreResetTaskId = -1;
         mTaskListChangeId = -1;
-        mFocusedTaskId = -1;
+        mFocusedTaskId = getTaskViewCount() > 0 ? getTaskViewAt(0).getTaskId() : -1;
 
         if (mRecentsAnimationController != null) {
             if (LIVE_TILE.get() && mEnableDrawingLiveTile) {
@@ -1593,11 +1612,8 @@
         return getTaskView(mFocusedTaskId);
     }
 
-    /**
-     * Returns the width to height ratio of the focused {@link TaskView}.
-     */
-    public float getFocusedTaskRatio() {
-        return mFocusedTaskRatio;
+    protected @Nullable TaskView getHomeTaskView() {
+        return null;
     }
 
     /**
@@ -1732,11 +1748,7 @@
         }
         setRunningTaskHidden(false);
         animateUpRunningTaskIconScale();
-
-        if (mCurrentGestureEndTarget == GestureState.GestureEndTarget.RECENTS
-                && (!showAsGrid() || getFocusedTaskView() != null)) {
-            animateActionsViewIn();
-        }
+        animateActionsViewIn();
 
         mCurrentGestureEndTarget = null;
     }
@@ -1778,9 +1790,7 @@
         boolean runningTaskTileHidden = mRunningTaskTileHidden;
         int runningTaskId = runningTaskInfo == null ? -1 : runningTaskInfo.taskId;
         setCurrentTask(runningTaskId);
-        if (mActivity.getDeviceProfile().isTablet && FeatureFlags.ENABLE_OVERVIEW_GRID.get()) {
-            setFocusedTask(runningTaskId);
-        }
+        mFocusedTaskId = runningTaskId;
         setCurrentPage(getRunningTaskIndex());
         setRunningTaskViewShowScreenshot(false);
         setRunningTaskHidden(runningTaskTileHidden);
@@ -1809,15 +1819,6 @@
     }
 
     /**
-     * Sets the focused task id and store the width to height ratio of the focused task.
-     */
-    protected void setFocusedTask(int focusedTaskId) {
-        mFocusedTaskId = focusedTaskId;
-        mFocusedTaskRatio =
-                mLastComputedTaskSize.width() / (float) mLastComputedTaskSize.height();
-    }
-
-    /**
      * Hides the tile associated with {@link #mRunningTaskId}
      */
     public void setRunningTaskHidden(boolean isHidden) {
@@ -1866,13 +1867,6 @@
         anim.start();
     }
 
-    private void animateActionsViewOut() {
-        ObjectAnimator anim = ObjectAnimator.ofFloat(
-                mActionsView.getVisibilityAlpha(), MultiValueAlpha.VALUE, 1, 0);
-        anim.setDuration(TaskView.SCALE_ICON_DURATION);
-        anim.start();
-    }
-
     public void animateUpRunningTaskIconScale() {
         mRunningTaskIconScaledDown = false;
         TaskView firstTask = getRunningTaskView();
@@ -1906,20 +1900,8 @@
             return;
         }
 
-        final int boxLength = Math.max(mLastComputedGridTaskSize.width(),
-                mLastComputedGridTaskSize.height());
         int taskTopMargin = mActivity.getDeviceProfile().overviewTaskThumbnailTopMarginPx;
 
-        /*
-         * taskGridVerticalDiff is used to position the top of a task in the top row of the grid
-         * heightOffset is the vertical space one grid task takes + space between top and
-         *   bottom row
-         * Summed together they provide the top position for bottom row of grid tasks
-         */
-        final float taskGridVerticalDiff =
-                mLastComputedGridTaskSize.top - mLastComputedTaskSize.top;
-        final float heightOffset = (boxLength + taskTopMargin) + mRowSpacing;
-
         int topRowWidth = 0;
         int bottomRowWidth = 0;
         float topAccumulatedTranslationX = 0;
@@ -1939,6 +1921,8 @@
         int snappedTaskRowWidth = 0;
         int snappedPage = getNextPage();
         TaskView snappedTaskView = getTaskViewAtByAbsoluteIndex(snappedPage);
+        TaskView homeTaskView = getHomeTaskView();
+        TaskView nextFocusedTaskView = null;
 
         if (!isTaskDismissal) {
             mTopRowIdSet.clear();
@@ -1976,15 +1960,20 @@
                     // calculate the distance focused task need to shift.
                     focusedTaskShift += mIsRtl ? taskWidthAndSpacing : -taskWidthAndSpacing;
                 }
-                int taskId = taskView.getTask().key.id;
+                int taskId = taskView.getTaskId();
                 boolean isTopRow = isTaskDismissal ? mTopRowIdSet.contains(taskId)
                         : topRowWidth <= bottomRowWidth;
                 if (isTopRow) {
-                    topRowWidth += taskWidthAndSpacing;
+                    if (homeTaskView != null && nextFocusedTaskView == null) {
+                        // TaskView will be focused when swipe up, don't count towards row width.
+                        nextFocusedTaskView = taskView;
+                    } else {
+                        topRowWidth += taskWidthAndSpacing;
+                    }
                     topSet.add(i);
                     mTopRowIdSet.add(taskId);
 
-                    taskView.setGridTranslationY(taskGridVerticalDiff);
+                    taskView.setGridTranslationY(mTaskGridVerticalDiff);
 
                     // Move horizontally into empty space.
                     float widthOffset = 0;
@@ -2003,7 +1992,7 @@
                     bottomSet.add(i);
 
                     // Move into bottom row.
-                    taskView.setGridTranslationY(heightOffset + taskGridVerticalDiff);
+                    taskView.setGridTranslationY(mTopBottomRowHeightDiff + mTaskGridVerticalDiff);
 
                     // Move horizontally into empty space.
                     float widthOffset = 0;
@@ -2086,8 +2075,10 @@
         // of swiping up after quick switch.
         if (snappedTaskView != null) {
             int distanceFromClearAll = longRowWidth - snappedTaskRowWidth;
+            // ClearAllButton should be off screen when snapped task is in its snapped position.
             int minimumDistance =
-                    mLastComputedGridSize.width() - snappedTaskView.getLayoutParams().width;
+                    mTaskWidth - snappedTaskView.getLayoutParams().width
+                            + (mLastComputedGridSize.width() - mTaskWidth) / 2;
             if (distanceFromClearAll < minimumDistance) {
                 int distanceDifference = minimumDistance - distanceFromClearAll;
                 clearAllTotalTranslationX += mIsRtl ? -distanceDifference : distanceDifference;
@@ -2107,8 +2098,8 @@
         if (taskView1 == null || taskView2 == null) {
             return false;
         }
-        int taskId1 = taskView1.getTask().key.id;
-        int taskId2 = taskView2.getTask().key.id;
+        int taskId1 = taskView1.getTaskId();
+        int taskId2 = taskView2.getTaskId();
         if (taskId1 == mFocusedTaskId || taskId2 == mFocusedTaskId) {
             return false;
         }
@@ -2237,8 +2228,16 @@
         }
     }
 
-    public PendingAnimation createTaskDismissAnimation(TaskView taskView, boolean animateTaskView,
-            boolean shouldRemoveTask, long duration) {
+    /**
+     * Creates a {@link PendingAnimation} for dismissing the specified {@link TaskView}.
+     * @param dismissedTaskView the {@link TaskView} to be dismissed
+     * @param animateTaskView whether the {@link TaskView} to be dismissed should be animated
+     * @param shouldRemoveTask whether the associated {@link Task} should be removed from
+     *                         ActivityManager after dismissal
+     * @param duration duration of the animation
+     */
+    public PendingAnimation createTaskDismissAnimation(TaskView dismissedTaskView,
+            boolean animateTaskView, boolean shouldRemoveTask, long duration) {
         if (mPendingAnimation != null) {
             mPendingAnimation.createPlaybackController().dispatchOnCancel();
         }
@@ -2249,30 +2248,65 @@
             return anim;
         }
 
+        boolean showAsGrid = showAsGrid();
+        int taskCount = getTaskViewCount();
+        int dismissedIndex = indexOfChild(dismissedTaskView);
+        int dismissedTaskId = dismissedTaskView.getTaskId();
+
+        // Grid specific properties.
+        boolean isFocusedTaskDismissed = false;
+        TaskView nextFocusedTaskView = null;
+        boolean nextFocusedTaskFromTop = false;
+        float dismissedTaskWidth = 0;
+        float nextFocusedTaskWidth = 0;
+
+        // Non-grid specific properties.
         int[] oldScroll = new int[count];
         int[] newScroll = new int[count];
-        getPageScrolls(oldScroll, false, SIMPLE_SCROLL_LOGIC);
-        getPageScrolls(newScroll, false, (v) -> v.getVisibility() != GONE && v != taskView);
-        int taskCount = getTaskViewCount();
         int scrollDiffPerPage = 0;
-        if (count > 1) {
-            scrollDiffPerPage = Math.abs(oldScroll[1] - oldScroll[0]);
-        }
-        int draggedIndex = indexOfChild(taskView);
-
-        boolean isFocusedTaskDismissed = taskView.getTask().key.id == mFocusedTaskId;
-        if (isFocusedTaskDismissed && showAsGrid()) {
-            anim.setFloat(mActionsView, VIEW_ALPHA, 0, clampToProgress(ACCEL_0_5, 0, 0.5f));
-        }
-        float dismissedTaskWidth = taskView.getLayoutParams().width + mPageSpacing;
         boolean needsCurveUpdates = false;
+
+        if (showAsGrid) {
+            dismissedTaskWidth = dismissedTaskView.getLayoutParams().width + mPageSpacing;
+            isFocusedTaskDismissed = dismissedTaskId == mFocusedTaskId;
+            if (isFocusedTaskDismissed) {
+                nextFocusedTaskFromTop =
+                        mTopRowIdSet.size() > 0 && mTopRowIdSet.size() >= (taskCount - 1) / 2f;
+                // Pick the next focused task from the preferred row.
+                for (int i = 0; i < taskCount; i++) {
+                    TaskView taskView = getTaskViewAt(i);
+                    if (taskView == dismissedTaskView) {
+                        continue;
+                    }
+                    boolean isTopRow = mTopRowIdSet.contains(taskView.getTaskId());
+                    if ((nextFocusedTaskFromTop && isTopRow
+                            || (!nextFocusedTaskFromTop && !isTopRow))) {
+                        nextFocusedTaskView = taskView;
+                        break;
+                    }
+                }
+                if (nextFocusedTaskView != null) {
+                    nextFocusedTaskWidth =
+                            nextFocusedTaskView.getLayoutParams().width + mPageSpacing;
+                }
+            }
+        } else {
+            getPageScrolls(oldScroll, false, SIMPLE_SCROLL_LOGIC);
+            getPageScrolls(newScroll, false,
+                    v -> v.getVisibility() != GONE && v != dismissedTaskView);
+            if (count > 1) {
+                scrollDiffPerPage = Math.abs(oldScroll[1] - oldScroll[0]);
+            }
+        }
+
+        int distanceFromDismissedTask = 0;
         for (int i = 0; i < count; i++) {
             View child = getChildAt(i);
-            if (child == taskView) {
+            if (child == dismissedTaskView) {
                 if (animateTaskView) {
-                    addDismissedTaskAnimations(taskView, duration, anim);
+                    addDismissedTaskAnimations(dismissedTaskView, duration, anim);
                 }
-            } else if (!showAsGrid()) {
+            } else if (!showAsGrid) {
                 // Compute scroll offsets from task dismissal for animation.
                 // If we just take newScroll - oldScroll, everything to the right of dragged task
                 // translates to the left. We need to offset this in some cases:
@@ -2281,15 +2315,15 @@
                 // - Current page is rightmost page (leftmost for RTL)
                 // - Dragging an adjacent page on the left side (right side for RTL)
                 int offset = mIsRtl ? scrollDiffPerPage : 0;
-                if (mCurrentPage == draggedIndex) {
+                if (mCurrentPage == dismissedIndex) {
                     int lastPage = taskCount - 1;
                     if (mCurrentPage == lastPage) {
                         offset += mIsRtl ? -scrollDiffPerPage : scrollDiffPerPage;
                     }
                 } else {
-                    // Dragging an adjacent page.
+                    // Dismissing an adjacent page.
                     int negativeAdjacent = mCurrentPage - 1; // (Right in RTL, left in LTR)
-                    if (draggedIndex == negativeAdjacent) {
+                    if (dismissedIndex == negativeAdjacent) {
                         offset += mIsRtl ? -scrollDiffPerPage : scrollDiffPerPage;
                     }
                 }
@@ -2302,28 +2336,53 @@
 
                     float additionalDismissDuration =
                             ADDITIONAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET * Math.abs(
-                                    i - draggedIndex);
+                                    i - dismissedIndex);
                     anim.setFloat(child, translationProperty, scrollDiff, clampToProgress(LINEAR,
                             Utilities.boundToRange(INITIAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET
                                     + additionalDismissDuration, 0f, 1f), 1));
                     needsCurveUpdates = true;
                 }
             } else if (child instanceof TaskView) {
+                TaskView taskView = (TaskView) child;
+                if (isFocusedTaskDismissed) {
+                    if (!isSameGridRow(taskView, nextFocusedTaskView)) {
+                        continue;
+                    }
+                } else {
+                    if (i < dismissedIndex || !isSameGridRow(taskView, dismissedTaskView)) {
+                        continue;
+                    }
+                }
                 // Animate task with index >= dismissed index and in the same row as the
-                // dismissed index, or if the dismissed task was the focused task. Offset
-                // successive task dismissal durations for a staggered effect.
-                if (isFocusedTaskDismissed || (i >= draggedIndex && isSameGridRow((TaskView) child,
-                        taskView))) {
-                    FloatProperty translationProperty =
-                            ((TaskView) child).getPrimaryDismissTranslationProperty();
-                    float additionalDismissDuration =
-                            ADDITIONAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET * Math.abs(
-                                    i - draggedIndex);
-                    anim.setFloat(child, translationProperty,
-                            !mIsRtl ? -dismissedTaskWidth : dismissedTaskWidth,
-                            clampToProgress(LINEAR, Utilities.boundToRange(
-                                    INITIAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET
-                                            + additionalDismissDuration, 0f, 1f), 1));
+                // dismissed index or next focused index. Offset successive task dismissal
+                // durations for a staggered effect.
+                float animationStartProgress = Utilities.boundToRange(
+                        INITIAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET
+                                + ADDITIONAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET
+                                * ++distanceFromDismissedTask, 0f, 1f);
+                if (taskView == nextFocusedTaskView) {
+                    // Enlarge the task to be focused next, and translate into focus position.
+                    float scale = mTaskWidth / (float) mLastComputedGridTaskSize.width();
+                    anim.setFloat(taskView, TaskView.SNAPSHOT_SCALE, scale,
+                            clampToProgress(LINEAR, animationStartProgress, 1f));
+                    anim.setFloat(taskView, taskView.getPrimaryDismissTranslationProperty(),
+                            mIsRtl ? dismissedTaskWidth : -dismissedTaskWidth,
+                            clampToProgress(LINEAR, animationStartProgress, 1f));
+                    float secondaryTranslation = -mTaskGridVerticalDiff;
+                    if (!nextFocusedTaskFromTop) {
+                        secondaryTranslation -= mTopBottomRowHeightDiff;
+                    }
+                    anim.setFloat(taskView, taskView.getSecondaryDissmissTranslationProperty(),
+                            secondaryTranslation,
+                            clampToProgress(LINEAR, animationStartProgress, 1f));
+                    anim.setFloat(taskView, TaskView.FOCUS_TRANSITION, 0f,
+                            clampToProgress(LINEAR, 0f, ANIMATION_DISMISS_PROGRESS_MIDPOINT));
+                } else {
+                    float primaryTranslation =
+                            isFocusedTaskDismissed ? nextFocusedTaskWidth : dismissedTaskWidth;
+                    anim.setFloat(taskView, taskView.getPrimaryDismissTranslationProperty(),
+                            mIsRtl ? primaryTranslation : -primaryTranslation,
+                            clampToProgress(LINEAR, animationStartProgress, 1f));
                 }
             }
         }
@@ -2334,14 +2393,15 @@
 
         // Add a tiny bit of translation Z, so that it draws on top of other views
         if (animateTaskView) {
-            taskView.setTranslationZ(0.1f);
+            dismissedTaskView.setTranslationZ(0.1f);
         }
 
         mPendingAnimation = anim;
+        final TaskView finalNextFocusedTaskView = nextFocusedTaskView;
         mPendingAnimation.addEndListener(new Consumer<Boolean>() {
             @Override
             public void accept(Boolean success) {
-                if (LIVE_TILE.get() && mEnableDrawingLiveTile && taskView.isRunningTask()
+                if (LIVE_TILE.get() && mEnableDrawingLiveTile && dismissedTaskView.isRunningTask()
                         && success) {
                     finishRecentsAnimation(true /* toHome */, () -> onEnd(success));
                 } else {
@@ -2353,11 +2413,11 @@
             private void onEnd(boolean success) {
                 if (success) {
                     if (shouldRemoveTask) {
-                        if (taskView.getTask() != null) {
+                        if (dismissedTaskView.getTask() != null) {
                             UI_HELPER_EXECUTOR.execute(() -> ActivityManagerWrapper.getInstance()
-                                    .removeTask(taskView.getTask().key.id));
+                                    .removeTask(dismissedTaskId));
                             mActivity.getStatsLogManager().logger()
-                                    .withItemInfo(taskView.getItemInfo())
+                                    .withItemInfo(dismissedTaskView.getItemInfo())
                                     .log(LAUNCHER_TASK_DISMISS_SWIPE_UP);
                         }
                     }
@@ -2367,31 +2427,31 @@
                     resetTaskVisuals();
 
                     int pageToSnapTo = mCurrentPage;
-                    // Snap to start if focused task was dismissed, as after quick switch it could
-                    // be at any page but the focused task always displays at the start.
-                    if (taskView.getTask().key.id == mFocusedTaskId) {
-                        pageToSnapTo = mTaskViewStartIndex;
-                    } else if (draggedIndex < pageToSnapTo || pageToSnapTo == (getTaskViewCount()
-                            - 1)) {
+                    if (finalNextFocusedTaskView != null) {
+                        pageToSnapTo = indexOfChild(finalNextFocusedTaskView);
+                    }
+                    if (dismissedIndex < pageToSnapTo || pageToSnapTo == (taskCount - 1)) {
                         pageToSnapTo -= 1;
                     }
-                    removeViewInLayout(taskView);
+                    removeViewInLayout(dismissedTaskView);
+                    mTopRowIdSet.remove(dismissedTaskId);
 
-                    if (getTaskViewCount() == 0) {
+                    if (taskCount == 1) {
                         removeViewInLayout(mClearAllButton);
                         startHome();
                     } else {
-                        snapToPageImmediately(pageToSnapTo);
-                        // Grid got messed up, reapply.
-                        updateGridProperties(true);
-                        if (showAsGrid() && getFocusedTaskView() == null
-                                && mActionsView.getVisibilityAlpha().getValue() == 1) {
-                            animateActionsViewOut();
+                        // Update focus task and its size.
+                        if (finalNextFocusedTaskView != null) {
+                            mFocusedTaskId = finalNextFocusedTaskView.getTaskId();
+                            mTopRowIdSet.remove(mFocusedTaskId);
+                            finalNextFocusedTaskView.animateIconScaleAndDimIntoView();
+                            setCurrentPage(pageToSnapTo);
                         }
+                        updateTaskSize(true);
+                        // Update scroll and snap to page.
+                        updateScrollSynchronously();
+                        snapToPageImmediately(pageToSnapTo);
                     }
-                    // Update the layout synchronously so that the position of next view is
-                    // immediately available.
-                    onLayout(false /*  changed */, getLeft(), getTop(), getRight(), getBottom());
                 }
                 onDismissAnimationEnds();
                 mPendingAnimation = null;
@@ -2427,7 +2487,8 @@
         int taskCount = getTaskViewCount();
         for (int i = 0; i < taskCount; i++) {
             TaskView taskView = getTaskViewAt(i);
-            if (taskView == mSplitHiddenTaskView && taskView != getFocusedTaskView()) {
+            if (taskView == mSplitHiddenTaskView
+                    && !(showAsGrid() && taskView == getFocusedTaskView())) {
                 // Case where the hidden task view would have overlapped w/ placeholder,
                 // but because it's going to hide we don't care
                 // TODO (b/187312247) edge case for thumbnails that are off screen but scroll on
@@ -2572,7 +2633,7 @@
         mContentAlpha = alpha;
         for (int i = getTaskViewCount() - 1; i >= 0; i--) {
             TaskView child = getTaskViewAt(i);
-            if (!mRunningTaskTileHidden || child.getTask().key.id != mRunningTaskId) {
+            if (!mRunningTaskTileHidden || child.getTaskId() != mRunningTaskId) {
                 child.setStableAlpha(alpha);
             }
         }
@@ -3398,9 +3459,13 @@
     }
 
     /**
-     * Updates page scroll synchronously and layout child views.
+     * Updates page scroll synchronously after measure and layout child views.
      */
     public void updateScrollSynchronously() {
+        // onMeasure is needed to update child's measured width which is used in scroll calculation,
+        // in case TaskView sizes has changed when being focused/unfocused.
+        onMeasure(makeMeasureSpec(getMeasuredWidth(), EXACTLY),
+                makeMeasureSpec(getMeasuredHeight(), EXACTLY));
         onLayout(false /*  changed */, getLeft(), getTop(), getRight(), getBottom());
         updateMinAndMaxScrollX();
     }
@@ -3598,6 +3663,7 @@
     public void setOverviewGridEnabled(boolean overviewGridEnabled) {
         if (mOverviewGridEnabled != overviewGridEnabled) {
             mOverviewGridEnabled = overviewGridEnabled;
+            updateActionsViewScrollAlpha();
             // Request layout to ensure scroll position is recalculated with updated mGridProgress.
             requestLayout();
         }
@@ -3735,11 +3801,6 @@
                 && mCurrentGestureEndTarget != GestureState.GestureEndTarget.RECENTS;
     }
 
-    public boolean shouldShowOverviewActionsForState(STATE_TYPE state) {
-        return !state.displayOverviewTasksAsGrid(mActivity.getDeviceProfile())
-                || getFocusedTaskView() != null;
-    }
-
     /**
      * Used to register callbacks for when our empty message state changes.
      *
diff --git a/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java b/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java
index f730d2d..df195d7 100644
--- a/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java
@@ -421,7 +421,9 @@
      */
     public static class PreviewPositionHelper {
 
-        // Contains the portion of the thumbnail that is clipped when fullscreen progress = 0.
+        private static final RectF EMPTY_RECT_F = new RectF();
+
+        // Contains the portion of the thumbnail that is unclipped when fullscreen progress = 1.
         private final RectF mClippedInsets = new RectF();
         private final Matrix mMatrix = new Matrix();
         private boolean mIsOrientationChanged;
@@ -626,15 +628,17 @@
                     break;
             }
             mClippedInsets.offsetTo(newLeftInset * scale, newTopInset * scale);
-            mMatrix.postTranslate(translateX - mClippedInsets.left,
-                    translateY - mClippedInsets.top);
+            mMatrix.postTranslate(translateX, translateY);
+            if (TaskView.FULL_THUMBNAIL) {
+                mMatrix.postTranslate(-mClippedInsets.left, -mClippedInsets.top);
+            }
         }
 
         /**
          * Insets to used for clipping the thumbnail (in case it is drawing outside its own space)
          */
         public RectF getInsetsToDrawInFullscreen() {
-            return mClippedInsets;
+            return TaskView.FULL_THUMBNAIL ? mClippedInsets : EMPTY_RECT_F;
         }
     }
 }
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index b2a7b1b..ea37d70 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -147,6 +147,11 @@
      */
     public static final boolean CLIP_STATUS_AND_NAV_BARS = false;
 
+    /**
+     * Should the TaskView scale down to fit whole thumbnail in fullscreen.
+     */
+    public static final boolean FULL_THUMBNAIL = false;
+
     private static final float EDGE_SCALE_DOWN_FACTOR_CAROUSEL = 0.03f;
     private static final float EDGE_SCALE_DOWN_FACTOR_GRID = 0.00f;
 
@@ -167,7 +172,7 @@
     private static final List<Rect> SYSTEM_GESTURE_EXCLUSION_RECT =
             Collections.singletonList(new Rect());
 
-    private static final FloatProperty<TaskView> FOCUS_TRANSITION =
+    public static final FloatProperty<TaskView> FOCUS_TRANSITION =
             new FloatProperty<TaskView>("focusTransition") {
                 @Override
                 public void setValue(TaskView taskView, float v) {
@@ -336,6 +341,19 @@
                 }
             };
 
+    public static final FloatProperty<TaskView> SNAPSHOT_SCALE =
+            new FloatProperty<TaskView>("snapshotScale") {
+                @Override
+                public void setValue(TaskView taskView, float v) {
+                    taskView.setSnapshotScale(v);
+                }
+
+                @Override
+                public Float get(TaskView taskView) {
+                    return taskView.mSnapshotView.getScaleX();
+                }
+            };
+
     private final TaskOutlineProvider mOutlineProvider;
 
     private Task mTask;
@@ -524,8 +542,8 @@
         return mTask;
     }
 
-    public boolean hasTaskId(int taskId) {
-        return mTask != null && mTask.key != null && mTask.key.id == taskId;
+    public int getTaskId() {
+        return mTask != null && mTask.key != null ? mTask.key.id : -1;
     }
 
     public TaskThumbnailView getThumbnail() {
@@ -847,6 +865,7 @@
                 mSplitSelectTranslationX = 0f;
         mDismissTranslationY = mTaskOffsetTranslationY = mTaskResistanceTranslationY =
                 mSplitSelectTranslationY = 0f;
+        setSnapshotScale(1f);
         applyTranslationX();
         applyTranslationY();
         setTranslationZ(0);
@@ -926,6 +945,9 @@
         if (mActivity.getDeviceProfile().isTablet && FeatureFlags.ENABLE_OVERVIEW_GRID.get()) {
             setPivotX(getLayoutDirection() == LAYOUT_DIRECTION_RTL ? 0 : right - left);
             setPivotY(mSnapshotView.getTop());
+            mSnapshotView.setPivotX(
+                    getLayoutDirection() == LAYOUT_DIRECTION_RTL ? 0 : right - left);
+            mSnapshotView.setPivotY(0);
         } else {
             setPivotX((right - left) * 0.5f);
             setPivotY(mSnapshotView.getTop() + mSnapshotView.getHeight() * 0.5f);
@@ -956,6 +978,11 @@
         return mFullscreenScale;
     }
 
+    private void setSnapshotScale(float dismissScale) {
+        mSnapshotView.setScaleX(dismissScale);
+        mSnapshotView.setScaleY(dismissScale);
+    }
+
     /**
      * Moves TaskView between carousel and 2 row grid.
      *
@@ -1349,46 +1376,26 @@
 
             int boxWidth;
             int boxHeight;
-            float thumbnailRatio;
             boolean isFocusedTask = isFocusedTask();
             if (isFocusedTask) {
                 // Task will be focused and should use focused task size. Use focusTaskRatio
                 // that is associated with the original orientation of the focused task.
                 boxWidth = taskWidth;
                 boxHeight = taskHeight;
-                thumbnailRatio = getRecentsView().getFocusedTaskRatio();
             } else {
                 // Otherwise task is in grid, and should use lastComputedGridTaskSize.
                 Rect lastComputedGridTaskSize = getRecentsView().getLastComputedGridTaskSize();
                 boxWidth = lastComputedGridTaskSize.width();
                 boxHeight = lastComputedGridTaskSize.height();
-                thumbnailRatio = mTask != null ? mTask.getVisibleThumbnailRatio(
-                        TaskView.CLIP_STATUS_AND_NAV_BARS) : 0f;
             }
-            int boxLength = Math.max(boxWidth, boxHeight);
 
             // Bound width/height to the box size.
-            if (thumbnailRatio == 0f) {
-                expectedWidth = boxWidth;
-                expectedHeight = boxHeight + thumbnailPadding;
-            } else if (thumbnailRatio > 1) {
-                expectedWidth = boxLength;
-                expectedHeight = (int) (boxLength / thumbnailRatio) + thumbnailPadding;
-            } else {
-                expectedWidth = (int) (boxLength * thumbnailRatio);
-                expectedHeight = boxLength + thumbnailPadding;
-            }
+            expectedWidth = boxWidth;
+            expectedHeight = boxHeight + thumbnailPadding;
 
             // Scale to to fit task Rect.
             fullscreenScale = taskWidth / (float) boxWidth;
 
-            // In full screen, scale back TaskView to original size.
-            if (expectedWidth > boxWidth) {
-                fullscreenScale *= boxWidth / (float) expectedWidth;
-            } else if (expectedHeight - thumbnailPadding > boxHeight) {
-                fullscreenScale *= boxHeight / (float) (expectedHeight - thumbnailPadding);
-            }
-
             // Align to top of task Rect.
             boxTranslationY = (expectedHeight - thumbnailPadding - taskHeight) / 2.0f;
         } else {
diff --git a/quickstep/tests/src/com/android/quickstep/NavigationBarRotationContextTest.java b/quickstep/tests/src/com/android/quickstep/NavigationBarRotationContextTest.java
new file mode 100644
index 0000000..af5819a
--- /dev/null
+++ b/quickstep/tests/src/com/android/quickstep/NavigationBarRotationContextTest.java
@@ -0,0 +1,123 @@
+/*
+ * Copyright 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.quickstep;
+
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.content.Context;
+import android.graphics.drawable.AnimatedVectorDrawable;
+import android.view.View;
+import android.view.WindowInsetsController;
+
+import androidx.test.InstrumentationRegistry;
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.launcher3.taskbar.contextual.RotationButton;
+import com.android.launcher3.taskbar.contextual.RotationButtonController;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockitoAnnotations;
+
+/** SysUI equivalent */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class NavigationBarRotationContextTest {
+    private static final int DEFAULT_ROTATE = 0;
+    private static final int DEFAULT_DISPLAY = 0;
+
+
+    private RotationButtonController mRotationButtonController;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+        Context mTargetContext = InstrumentationRegistry.getTargetContext();
+        final View view = new View(mTargetContext);
+        RotationButton rotationButton = mock(RotationButton.class);
+        mRotationButtonController = new RotationButtonController(mTargetContext, 0, 0);
+        mRotationButtonController.setRotationButton(rotationButton);
+        // Due to a mockito issue, only spy the object after setting the initial state
+        mRotationButtonController = spy(mRotationButtonController);
+        final AnimatedVectorDrawable kbd = mock(AnimatedVectorDrawable.class);
+        doReturn(view).when(rotationButton).getCurrentView();
+        doReturn(true).when(rotationButton).acceptRotationProposal();
+    }
+
+    @Test
+    public void testOnInvalidRotationProposal() {
+        mRotationButtonController.onRotationProposal(DEFAULT_ROTATE + 1,
+                false /* isValid */);
+        verify(mRotationButtonController, times(1))
+                .setRotateSuggestionButtonState(false /* visible */);
+    }
+
+    @Test
+    public void testOnSameRotationProposal() {
+        mRotationButtonController.onRotationProposal(DEFAULT_ROTATE,
+                true /* isValid */);
+        verify(mRotationButtonController, times(1))
+                .setRotateSuggestionButtonState(false /* visible */);
+    }
+
+    @Test
+    public void testOnRotationProposalShowButtonShowNav() {
+        // No navigation bar should not call to set visibility state
+        mRotationButtonController.onBehaviorChanged(DEFAULT_DISPLAY,
+                WindowInsetsController.BEHAVIOR_SHOW_TRANSIENT_BARS_BY_SWIPE);
+        mRotationButtonController.onTaskBarVisibilityChange(false /* showing */);
+        verify(mRotationButtonController, times(0)).setRotateSuggestionButtonState(
+                false /* visible */);
+        verify(mRotationButtonController, times(0)).setRotateSuggestionButtonState(
+                true /* visible */);
+
+        // No navigation bar with rotation change should not call to set visibility state
+        mRotationButtonController.onRotationProposal(DEFAULT_ROTATE + 1,
+                true /* isValid */);
+        verify(mRotationButtonController, times(0)).setRotateSuggestionButtonState(
+                false /* visible */);
+        verify(mRotationButtonController, times(0)).setRotateSuggestionButtonState(
+                true /* visible */);
+
+        // Since rotation has changed rotation should be pending, show mButton when showing nav bar
+        mRotationButtonController.onTaskBarVisibilityChange(true /* showing */);
+        verify(mRotationButtonController, times(1)).setRotateSuggestionButtonState(
+                true /* visible */);
+    }
+
+    @Test
+    public void testOnRotationProposalShowButton() {
+        // Navigation bar being visible should not call to set visibility state
+        mRotationButtonController.onTaskBarVisibilityChange(true /* showing */);
+        verify(mRotationButtonController, times(0))
+                .setRotateSuggestionButtonState(false /* visible */);
+        verify(mRotationButtonController, times(0))
+                .setRotateSuggestionButtonState(true /* visible */);
+
+        // Navigation bar is visible and rotation requested
+        mRotationButtonController.onRotationProposal(DEFAULT_ROTATE + 1,
+                true /* isValid */);
+        verify(mRotationButtonController, times(1))
+                .setRotateSuggestionButtonState(true /* visible */);
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java b/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
index a2abfd5..275cf81 100644
--- a/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
@@ -26,12 +26,12 @@
 
 import android.os.Process;
 
-import com.android.launcher3.PagedView;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.shadows.ShadowLooperExecutor;
 import com.android.launcher3.util.Executors;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.LauncherLayoutBuilder;
 import com.android.launcher3.util.LauncherModelHelper;
 import com.android.launcher3.util.LooperExecutor;
@@ -92,7 +92,7 @@
         // Add a new callback
         cb1.reset();
         MyCallbacks cb2 = spy(MyCallbacks.class);
-        cb2.mPageToBindSync = 2;
+        cb2.mPageToBindSync = IntSet.wrap(2);
         mModelHelper.getModel().addCallbacksAndLoad(cb2);
 
         waitForLoaderAndTempMainThread();
@@ -178,16 +178,16 @@
     private abstract static class MyCallbacks implements Callbacks {
 
         final List<ItemInfo> mItems = new ArrayList<>();
-        int mPageToBindSync = 0;
-        int mPageBoundSync = PagedView.INVALID_PAGE;
+        IntSet mPageToBindSync = IntSet.wrap(0);
+        IntSet mPageBoundSync = new IntSet();
         ViewOnDrawExecutor mDeferredExecutor;
         AppInfo[] mAppInfos;
 
         MyCallbacks() { }
 
         @Override
-        public void onPageBoundSynchronously(int page) {
-            mPageBoundSync = page;
+        public void onPagesBoundSynchronously(IntSet pages) {
+            mPageBoundSync = pages;
         }
 
         @Override
@@ -206,13 +206,13 @@
         }
 
         @Override
-        public int getPageToBindSynchronously() {
+        public IntSet getPagesToBindSynchronously() {
             return mPageToBindSync;
         }
 
         public void reset() {
             mItems.clear();
-            mPageBoundSync = PagedView.INVALID_PAGE;
+            mPageBoundSync = new IntSet();
             mDeferredExecutor = null;
             mAppInfos = null;
         }
diff --git a/robolectric_tests/src/com/android/launcher3/util/LauncherPageRestoreHelperTest.java b/robolectric_tests/src/com/android/launcher3/util/LauncherPageRestoreHelperTest.java
new file mode 100644
index 0000000..51f5851
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/util/LauncherPageRestoreHelperTest.java
@@ -0,0 +1,224 @@
+/**
+ * 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.launcher3.util;
+
+import android.os.Bundle;
+
+import com.android.launcher3.LauncherPageRestoreHelper;
+import com.android.launcher3.Workspace;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.when;
+
+@RunWith(RobolectricTestRunner.class)
+public class LauncherPageRestoreHelperTest {
+
+    // Type: int
+    private static final String RUNTIME_STATE_CURRENT_SCREEN = "launcher.current_screen";
+    // Type: int
+    private static final String RUNTIME_STATE_CURRENT_SCREEN_COUNT =
+            "launcher.current_screen_count";
+
+    private LauncherPageRestoreHelper mPageRestoreHelper;
+    private Bundle mState;
+
+    @Mock
+    private Workspace mWorkspace;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mPageRestoreHelper = new LauncherPageRestoreHelper(mWorkspace);
+        mState = new Bundle();
+    }
+
+    @Test
+    public void givenNoChildrenInWorkspace_whenSavePages_thenNothingSaved() {
+        when(mWorkspace.getChildCount()).thenReturn(0);
+
+        mPageRestoreHelper.savePagesToRestore(mState);
+
+        assertFalse(mState.containsKey(RUNTIME_STATE_CURRENT_SCREEN_COUNT));
+        assertFalse(mState.containsKey(RUNTIME_STATE_CURRENT_SCREEN));
+    }
+
+    @Test
+    public void givenMultipleCurrentPages_whenSavePages_thenSavedCorrectly() {
+        when(mWorkspace.getChildCount()).thenReturn(5);
+        when(mWorkspace.getCurrentPage()).thenReturn(2);
+        givenPanelCount(2);
+
+        mPageRestoreHelper.savePagesToRestore(mState);
+
+        assertEquals(5, mState.getInt(RUNTIME_STATE_CURRENT_SCREEN_COUNT));
+        assertEquals(2, mState.getInt(RUNTIME_STATE_CURRENT_SCREEN));
+    }
+
+    @Test
+    public void givenNullSavedState_whenRestorePages_thenReturnEmptyIntSet() {
+        IntSet result = mPageRestoreHelper.getPagesToRestore(null);
+
+        assertTrue(result.isEmpty());
+    }
+
+    @Test
+    public void givenTotalPageCountMissing_whenRestorePages_thenReturnEmptyIntSet() {
+        givenSavedCurrentPage(1);
+        givenPanelCount(1);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertTrue(result.isEmpty());
+    }
+
+    @Test
+    public void givenCurrentPageMissing_whenRestorePages_thenReturnEmptyIntSet() {
+        givenSavedPageCount(3);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertTrue(result.isEmpty());
+    }
+
+    @Test
+    public void givenOnePanel_whenRestorePages_thenReturnThatPage() {
+        givenSavedCurrentPage(2);
+        givenSavedPageCount(5);
+        givenPanelCount(1);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(1, result.size());
+        assertEquals(2, result.getArray().get(0));
+    }
+
+    @Test
+    public void givenTwoPanelOnFirstPages_whenRestorePages_thenReturnThosePages() {
+        givenSavedCurrentPage(0, 1);
+        givenSavedPageCount(2);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(0, 1), result);
+    }
+
+    @Test
+    public void givenTwoPanelOnMiddlePages_whenRestorePages_thenReturnThosePages() {
+        givenSavedCurrentPage(2, 3);
+        givenSavedPageCount(5);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(2, 3), result);
+    }
+
+    @Test
+    public void givenTwoPanelOnLastPage_whenRestorePages_thenReturnOnlyLastPage() {
+        // The device has two panel home but the current page is the last page, so we don't have
+        // a right panel, only the left one.
+        givenSavedCurrentPage(2);
+        givenSavedPageCount(3);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(2), result);
+    }
+
+    @Test
+    public void givenOnlyOnePageAndPhoneFolding_whenRestorePages_thenReturnOnlyOnePage() {
+        givenSavedCurrentPage(0);
+        givenSavedPageCount(1);
+        givenPanelCount(1);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(0), result);
+    }
+
+    @Test
+    public void givenPhoneFolding_whenRestorePages_thenReturnOnlyTheFirstCurrentPage() {
+        givenSavedCurrentPage(2, 3);
+        givenSavedPageCount(4);
+        givenPanelCount(1);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(2), result);
+    }
+
+    @Test
+    public void givenPhoneUnfolding_whenRestorePages_thenReturnCurrentPagePlusTheNextOne() {
+        givenSavedCurrentPage(2);
+        givenSavedPageCount(4);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(2, 3), result);
+    }
+
+    @Test
+    public void givenPhoneUnfoldingOnLastPage_whenRestorePages_thenReturnOnlyLastPage() {
+        givenSavedCurrentPage(4);
+        givenSavedPageCount(5);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(4), result);
+    }
+
+    @Test
+    public void givenOnlyOnePageAndPhoneUnfolding_whenRestorePages_thenReturnOnlyOnePage() {
+        givenSavedCurrentPage(0);
+        givenSavedPageCount(1);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(0), result);
+    }
+
+    private void givenPanelCount(int panelCount) {
+        when(mWorkspace.getPanelCount()).thenReturn(panelCount);
+        when(mWorkspace.getLeftmostVisiblePageForIndex(anyInt())).thenAnswer(invocation -> {
+            int pageIndex = invocation.getArgument(0);
+            return pageIndex * panelCount / panelCount;
+        });
+    }
+
+    private void givenSavedPageCount(int pageCount) {
+        mState.putInt(RUNTIME_STATE_CURRENT_SCREEN_COUNT, pageCount);
+    }
+
+    private void givenSavedCurrentPage(int... pages) {
+        mState.putInt(RUNTIME_STATE_CURRENT_SCREEN, pages[0]);
+    }
+}
diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index 4f0ef12..d49dd73 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -403,7 +403,7 @@
         }
     }
 
-    void clearPressedBackground() {
+    public void clearPressedBackground() {
         setPressed(false);
         setStayPressed(false);
     }
@@ -849,8 +849,9 @@
         switch (display) {
             case DISPLAY_ALL_APPS:
                 return grid.allAppsIconSizePx;
-            case DISPLAY_WORKSPACE:
             case DISPLAY_FOLDER:
+                return grid.folderChildIconSizePx;
+            case DISPLAY_WORKSPACE:
             default:
                 return grid.iconSizePx;
         }
diff --git a/src/com/android/launcher3/DeleteDropTarget.java b/src/com/android/launcher3/DeleteDropTarget.java
index 80ec192..ba55834 100644
--- a/src/com/android/launcher3/DeleteDropTarget.java
+++ b/src/com/android/launcher3/DeleteDropTarget.java
@@ -33,6 +33,7 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.views.Snackbar;
 
 public class DeleteDropTarget extends ButtonDropTarget {
@@ -131,7 +132,7 @@
             onAccessibilityDrop(null, item);
             ModelWriter modelWriter = mLauncher.getModelWriter();
             Runnable onUndoClicked = () -> {
-                mLauncher.setPageToBindSynchronously(itemPage);
+                mLauncher.setPagesToBindSynchronously(IntSet.wrap(itemPage));
                 modelWriter.abortDelete();
                 mLauncher.getStatsLogManager().logger().log(LAUNCHER_UNDO);
             };
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 8889e60..892fb6d 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -105,6 +105,7 @@
 import android.widget.Toast;
 
 import androidx.annotation.CallSuper;
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.StringRes;
 import androidx.annotation.VisibleForTesting;
@@ -164,6 +165,7 @@
 import com.android.launcher3.util.ActivityTracker;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.ItemInfoMatcher;
 import com.android.launcher3.util.MultiValueAlpha;
 import com.android.launcher3.util.MultiValueAlpha.AlphaProperty;
@@ -245,8 +247,6 @@
     protected static final int REQUEST_LAST = 100;
 
     // Type: int
-    private static final String RUNTIME_STATE_CURRENT_SCREEN = "launcher.current_screen";
-    // Type: int
     private static final String RUNTIME_STATE = "launcher.state";
     // Type: PendingRequestArgs
     private static final String RUNTIME_STATE_PENDING_REQUEST_ARGS = "launcher.request_args";
@@ -284,6 +284,8 @@
     private WidgetManagerHelper mAppWidgetManager;
     private LauncherAppWidgetHost mAppWidgetHost;
 
+    private LauncherPageRestoreHelper mPageRestoreHelper;
+
     private final int[] mTmpAddItemCellCoordinates = new int[2];
 
     @Thunk
@@ -319,8 +321,8 @@
 
     private PopupDataProvider mPopupDataProvider;
 
-    private int mSynchronouslyBoundPage = PagedView.INVALID_PAGE;
-    private int mPageToBindSynchronously = PagedView.INVALID_PAGE;
+    private IntSet mSynchronouslyBoundPages = new IntSet();
+    private IntSet mPagesToBindSynchronously = new IntSet();
 
     // We only want to get the SharedPreferences once since it does an FS stat each time we get
     // it from the context.
@@ -455,13 +457,10 @@
         restoreState(savedInstanceState);
         mStateManager.reapplyState();
 
-        // We only load the page synchronously if the user rotates (or triggers a
-        // configuration change) while launcher is in the foreground
-        int currentScreen = PagedView.INVALID_PAGE;
+        mPageRestoreHelper = new LauncherPageRestoreHelper(mWorkspace);
         if (savedInstanceState != null) {
-            currentScreen = savedInstanceState.getInt(RUNTIME_STATE_CURRENT_SCREEN, currentScreen);
+            mPagesToBindSynchronously = mPageRestoreHelper.getPagesToRestore(savedInstanceState);
         }
-        mPageToBindSynchronously = currentScreen;
 
         if (!mModel.addCallbacksAndLoad(this)) {
             if (!internalStateHandled) {
@@ -1290,7 +1289,8 @@
         }
     }
 
-    public FolderIcon findFolderIcon(final int folderIconId) {
+    @Override
+    public @Nullable FolderIcon findFolderIcon(final int folderIconId) {
         return (FolderIcon) mWorkspace.getHomescreenIconByItemId(folderIconId);
     }
 
@@ -1525,18 +1525,17 @@
     @Override
     public void onRestoreInstanceState(Bundle state) {
         super.onRestoreInstanceState(state);
-        mWorkspace.restoreInstanceStateForChild(mSynchronouslyBoundPage);
+        if (mSynchronouslyBoundPages != null) {
+            mSynchronouslyBoundPages.forEach(page -> mWorkspace.restoreInstanceStateForChild(page));
+        }
     }
 
     @Override
     protected void onSaveInstanceState(Bundle outState) {
-        if (mWorkspace.getChildCount() > 0) {
-            outState.putInt(RUNTIME_STATE_CURRENT_SCREEN, mWorkspace.getNextPage());
+        mPageRestoreHelper.savePagesToRestore(outState);
 
-        }
         outState.putInt(RUNTIME_STATE, mStateManager.getState().ordinal);
 
-
         AbstractFloatingView widgets = AbstractFloatingView
                 .getOpenView(this, AbstractFloatingView.TYPE_WIDGETS_FULL_SHEET);
         if (widgets != null) {
@@ -2015,24 +2014,24 @@
     }
 
     /**
-     * Sets the next page to bind synchronously on next bind.
-     * @param page
+     * Sets the next pages to bind synchronously on next bind.
+     * @param pages should not be null.
      */
-    public void setPageToBindSynchronously(int page) {
-        mPageToBindSynchronously = page;
+    public void setPagesToBindSynchronously(@NonNull IntSet pages) {
+        mPagesToBindSynchronously = pages;
     }
 
     /**
      * Implementation of the method from LauncherModel.Callbacks.
      */
     @Override
-    public int getPageToBindSynchronously() {
-        if (mPageToBindSynchronously != PagedView.INVALID_PAGE) {
-            return mPageToBindSynchronously;
-        } else  if (mWorkspace != null) {
-            return mWorkspace.getCurrentPage();
+    public IntSet getPagesToBindSynchronously() {
+        if (mPagesToBindSynchronously != null && !mPagesToBindSynchronously.isEmpty()) {
+            return mPagesToBindSynchronously;
+        } else if (mWorkspace != null) {
+            return mWorkspace.getVisiblePageIndices();
         } else {
-            return 0;
+            return new IntSet();
         }
     }
 
@@ -2448,10 +2447,10 @@
         return info;
     }
 
-    public void onPageBoundSynchronously(int page) {
-        mSynchronouslyBoundPage = page;
-        mWorkspace.setCurrentPage(page);
-        mPageToBindSynchronously = PagedView.INVALID_PAGE;
+    public void onPagesBoundSynchronously(IntSet pages) {
+        mSynchronouslyBoundPages = pages;
+        mWorkspace.setCurrentPage(pages.getArray().get(0));
+        mPagesToBindSynchronously = new IntSet();
     }
 
     @Override
@@ -2497,7 +2496,7 @@
      *
      * Implementation of the method from LauncherModel.Callbacks.
      */
-    public void finishBindingItems(int pageBoundFirst) {
+    public void finishBindingItems(IntSet pagesBoundFirst) {
         Object traceToken = TraceHelper.INSTANCE.beginSection("finishBindingItems");
         mWorkspace.restoreInstanceStateForRemainingPages();
 
@@ -2512,11 +2511,13 @@
         ItemInstallQueue.INSTANCE.get(this)
                 .resumeModelPush(FLAG_LOADER_RUNNING);
 
+        int currentPage = pagesBoundFirst != null && !pagesBoundFirst.isEmpty()
+                ? pagesBoundFirst.getArray().get(0) : PagedView.INVALID_PAGE;
         // When undoing the removal of the last item on a page, return to that page.
         // Since we are just resetting the current page without user interaction,
         // override the previous page so we don't log the page switch.
-        mWorkspace.setCurrentPage(pageBoundFirst, pageBoundFirst /* overridePrevPage */);
-        mPageToBindSynchronously = PagedView.INVALID_PAGE;
+        mWorkspace.setCurrentPage(currentPage, currentPage /* overridePrevPage */);
+        mPagesToBindSynchronously = new IntSet();
 
         // Cache one page worth of icons
         getViewCache().setCacheSize(R.layout.folder_application,
diff --git a/src/com/android/launcher3/LauncherPageRestoreHelper.java b/src/com/android/launcher3/LauncherPageRestoreHelper.java
new file mode 100644
index 0000000..e679a12
--- /dev/null
+++ b/src/com/android/launcher3/LauncherPageRestoreHelper.java
@@ -0,0 +1,92 @@
+/**
+ * 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.launcher3;
+
+import android.os.Bundle;
+import android.util.Log;
+
+import androidx.annotation.VisibleForTesting;
+
+import com.android.launcher3.util.IntSet;
+
+import static androidx.annotation.VisibleForTesting.PACKAGE_PRIVATE;
+
+/**
+ * There's a logic which prioritizes the binding for the current page and defers the other pages'
+ * binding. If two panel home is enabled, we want to bind both pages together.
+ * LauncherPageRestoreHelper's purpose is to contain the logic for persisting, restoring and
+ * calculating which pages to load immediately.
+ */
+public class LauncherPageRestoreHelper {
+
+    public static final String TAG = "LauncherPageRestoreHelper";
+
+    // Type: int
+    private static final String RUNTIME_STATE_CURRENT_SCREEN = "launcher.current_screen";
+    // Type: int
+    private static final String RUNTIME_STATE_CURRENT_SCREEN_COUNT =
+            "launcher.current_screen_count";
+
+    private Workspace mWorkspace;
+
+    public LauncherPageRestoreHelper(Workspace workspace) {
+        this.mWorkspace = workspace;
+    }
+
+    /**
+     * Some configuration changes trigger Launcher to recreate itself, and we want to give more
+     * priority to the currently active pages in the restoration process.
+     */
+    @VisibleForTesting(otherwise = PACKAGE_PRIVATE)
+    public IntSet getPagesToRestore(Bundle savedInstanceState) {
+        IntSet pagesToRestore = new IntSet();
+
+        if (savedInstanceState == null) {
+            return pagesToRestore;
+        }
+
+        int currentPage = savedInstanceState.getInt(RUNTIME_STATE_CURRENT_SCREEN, -1);
+        int totalPageCount = savedInstanceState.getInt(RUNTIME_STATE_CURRENT_SCREEN_COUNT, -1);
+        int panelCount = mWorkspace.getPanelCount();
+
+        if (totalPageCount <= 0 || currentPage < 0) {
+            Log.e(TAG, "getPagesToRestore: Invalid input: " + totalPageCount + ", " + currentPage);
+            return pagesToRestore;
+        }
+
+        int newCurrentPage = mWorkspace.getLeftmostVisiblePageForIndex(currentPage);
+        for (int page = newCurrentPage; page < newCurrentPage + panelCount
+                && page < totalPageCount; page++) {
+            pagesToRestore.add(page);
+        }
+
+        return pagesToRestore;
+    }
+
+    /**
+     * This should be called from Launcher's onSaveInstanceState method to persist everything that
+     * is necessary to calculate later which pages need to be initialized first after a
+     * configuration change.
+     */
+    @VisibleForTesting(otherwise = PACKAGE_PRIVATE)
+    public void savePagesToRestore(Bundle outState) {
+        int pageCount = mWorkspace.getChildCount();
+        if (pageCount > 0) {
+            outState.putInt(RUNTIME_STATE_CURRENT_SCREEN, mWorkspace.getCurrentPage());
+            outState.putInt(RUNTIME_STATE_CURRENT_SCREEN_COUNT, pageCount);
+        }
+    }
+}
diff --git a/src/com/android/launcher3/PagedView.java b/src/com/android/launcher3/PagedView.java
index b423871..97a44c1 100644
--- a/src/com/android/launcher3/PagedView.java
+++ b/src/com/android/launcher3/PagedView.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3;
 
+import static androidx.annotation.VisibleForTesting.PACKAGE_PRIVATE;
 import static com.android.launcher3.anim.Interpolators.SCROLL;
 import static com.android.launcher3.compat.AccessibilityManagerCompat.isAccessibilityEnabled;
 import static com.android.launcher3.compat.AccessibilityManagerCompat.isObservedEventType;
@@ -48,6 +49,7 @@
 import android.widget.ScrollView;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.compat.AccessibilityManagerCompat;
 import com.android.launcher3.config.FeatureFlags;
@@ -55,6 +57,7 @@
 import com.android.launcher3.touch.PagedOrientationHandler;
 import com.android.launcher3.touch.PagedOrientationHandler.ChildBounds;
 import com.android.launcher3.util.EdgeEffectCompat;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.Thunk;
 import com.android.launcher3.views.ActivityContext;
 
@@ -291,9 +294,15 @@
         return newPage;
     }
 
-    private int getLeftmostVisiblePageForIndex(int pageIndex) {
+    /**
+     * In most cases where panelCount is 1, this method will just return the page index that was
+     * passed in.
+     * But for example when two panel home is enabled we might need the leftmost visible page index
+     * because that page is the current page.
+     */
+    public int getLeftmostVisiblePageForIndex(int pageIndex) {
         int panelCount = getPanelCount();
-        return (pageIndex / panelCount) * panelCount;
+        return pageIndex - pageIndex % panelCount;
     }
 
     /**
@@ -304,16 +313,34 @@
     }
 
     /**
+     * Returns an IntSet with the indices of the currently visible pages
+     */
+    @VisibleForTesting(otherwise = PACKAGE_PRIVATE)
+    public IntSet getVisiblePageIndices() {
+        IntSet visiblePageIndices = new IntSet();
+        int panelCount = getPanelCount();
+        int pageCount = getPageCount();
+
+        // If a device goes from one panel to two panel (i.e. unfolding a foldable device) while
+        // an odd indexed page is the current page, then the new leftmost visible page will be
+        // different from the old mCurrentPage.
+        int currentPage = getLeftmostVisiblePageForIndex(mCurrentPage);
+        for (int page = currentPage; page < currentPage + panelCount && page < pageCount; page++) {
+            visiblePageIndices.add(page);
+        }
+        return visiblePageIndices;
+    }
+
+    /**
      * Executes the callback against each visible page
      */
     public void forEachVisiblePage(Consumer<View> callback) {
-        int panelCount = getPanelCount();
-        for (int i = mCurrentPage; i < mCurrentPage + panelCount; i++) {
-            View page = getPageAt(i);
+        getVisiblePageIndices().forEach(pageIndex -> {
+            View page = getPageAt(pageIndex);
             if (page != null) {
                 callback.accept(page);
             }
-        }
+        });
     }
 
     /**
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index cb9e1f3..becbb27 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -85,6 +85,7 @@
 import com.android.launcher3.shortcuts.ShortcutRequest;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.PackageManagerHelper;
+import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
 import com.android.launcher3.widget.PendingAddShortcutInfo;
 
@@ -651,25 +652,26 @@
      * @param outObj this is set to the internal data associated with {@param info},
      *               eg {@link LauncherActivityInfo} or {@link ShortcutInfo}.
      */
-    public static Drawable getFullDrawable(Launcher launcher, ItemInfo info, int width, int height,
+    public static Drawable getFullDrawable(Context context, ItemInfo info, int width, int height,
             Object[] outObj) {
-        Drawable icon = loadFullDrawableWithoutTheme(launcher, info, width, height, outObj);
+        Drawable icon = loadFullDrawableWithoutTheme(context, info, width, height, outObj);
         if (icon instanceof BitmapInfo.Extender) {
-            icon = ((BitmapInfo.Extender) icon).getThemedDrawable(launcher);
+            icon = ((BitmapInfo.Extender) icon).getThemedDrawable(context);
         }
         return icon;
     }
 
-    private static Drawable loadFullDrawableWithoutTheme(Launcher launcher, ItemInfo info,
+    private static Drawable loadFullDrawableWithoutTheme(Context context, ItemInfo info,
             int width, int height, Object[] outObj) {
-        LauncherAppState appState = LauncherAppState.getInstance(launcher);
+        ActivityContext activity = ActivityContext.lookupContext(context);
+        LauncherAppState appState = LauncherAppState.getInstance(context);
         if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_APPLICATION) {
-            LauncherActivityInfo activityInfo = launcher.getSystemService(LauncherApps.class)
+            LauncherActivityInfo activityInfo = context.getSystemService(LauncherApps.class)
                     .resolveActivity(info.getIntent(), info.user);
             outObj[0] = activityInfo;
-            return activityInfo == null ? null : LauncherAppState.getInstance(launcher)
+            return activityInfo == null ? null : LauncherAppState.getInstance(context)
                     .getIconProvider().getIcon(
-                            activityInfo, launcher.getDeviceProfile().inv.fillResIconDpi);
+                            activityInfo, activity.getDeviceProfile().inv.fillResIconDpi);
         } else if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT) {
             if (info instanceof PendingAddShortcutInfo) {
                 ShortcutConfigActivityInfo activityInfo =
@@ -678,18 +680,18 @@
                 return activityInfo.getFullResIcon(appState.getIconCache());
             }
             List<ShortcutInfo> si = ShortcutKey.fromItemInfo(info)
-                    .buildRequest(launcher)
+                    .buildRequest(context)
                     .query(ShortcutRequest.ALL);
             if (si.isEmpty()) {
                 return null;
             } else {
                 outObj[0] = si.get(0);
-                return ShortcutCachingLogic.getIcon(launcher, si.get(0),
+                return ShortcutCachingLogic.getIcon(context, si.get(0),
                         appState.getInvariantDeviceProfile().fillResIconDpi);
             }
         } else if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_FOLDER) {
             FolderAdaptiveIcon icon = FolderAdaptiveIcon.createFolderAdaptiveIcon(
-                    launcher, info.id, new Point(width, height));
+                    activity, info.id, new Point(width, height));
             if (icon == null) {
                 return null;
             }
@@ -707,8 +709,8 @@
      * badge. When dragged from workspace or folder, it may contain app AND/OR work profile badge
      **/
     @TargetApi(Build.VERSION_CODES.O)
-    public static Drawable getBadge(Launcher launcher, ItemInfo info, Object obj) {
-        LauncherAppState appState = LauncherAppState.getInstance(launcher);
+    public static Drawable getBadge(Context context, ItemInfo info, Object obj) {
+        LauncherAppState appState = LauncherAppState.getInstance(context);
         int iconSize = appState.getInvariantDeviceProfile().iconBitmapSize;
         if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT) {
             boolean iconBadged = (info instanceof ItemInfoWithIcon)
@@ -728,7 +730,7 @@
         } else if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_FOLDER) {
             return ((FolderAdaptiveIcon) obj).getBadge();
         } else {
-            return launcher.getPackageManager()
+            return context.getPackageManager()
                     .getUserBadgedIcon(new FixedSizeEmptyDrawable(iconSize), info.user);
         }
     }
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index d136cda..9a8f3dd 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3;
 
+import static androidx.annotation.VisibleForTesting.PROTECTED;
 import static com.android.launcher3.LauncherAnimUtils.SPRING_LOADED_EXIT_DELAY;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
 import static com.android.launcher3.LauncherState.ALL_APPS;
@@ -63,6 +64,8 @@
 import android.view.accessibility.AccessibilityNodeInfo;
 import android.widget.Toast;
 
+import androidx.annotation.VisibleForTesting;
+
 import com.android.launcher3.accessibility.AccessibleDragListenerAdapter;
 import com.android.launcher3.accessibility.WorkspaceAccessibilityHelper;
 import com.android.launcher3.anim.Interpolators;
@@ -461,7 +464,8 @@
     }
 
     @Override
-    protected int getPanelCount() {
+    @VisibleForTesting(otherwise = PROTECTED)
+    public int getPanelCount() {
         return isTwoPanelEnabled() ? 2 : super.getPanelCount();
     }
 
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 40dcb1e..cebdc1f 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -214,14 +214,14 @@
             + "predictions to be updated while they are visible to the user.");
 
     public static final BooleanFlag ENABLE_TASKBAR = getDebugFlag(
-            "ENABLE_TASKBAR", false, "Allows a system Taskbar to be shown on larger devices.");
+            "ENABLE_TASKBAR", true, "Allows a system Taskbar to be shown on larger devices.");
 
     public static final BooleanFlag ENABLE_OVERVIEW_GRID = getDebugFlag(
-            "ENABLE_OVERVIEW_GRID", false, "Uses grid overview layout. "
+            "ENABLE_OVERVIEW_GRID", true, "Uses grid overview layout. "
             + "Only applicable on large screen devices.");
 
     public static final BooleanFlag ENABLE_TWO_PANEL_HOME = getDebugFlag(
-            "ENABLE_TWO_PANEL_HOME", false,
+            "ENABLE_TWO_PANEL_HOME", true,
             "Uses two panel on home screen. Only applicable on large screen devices.");
 
     public static final BooleanFlag ENABLE_SPLIT_SELECT = getDebugFlag(
diff --git a/src/com/android/launcher3/dragndrop/DragController.java b/src/com/android/launcher3/dragndrop/DragController.java
index 5731db4..1e0edac 100644
--- a/src/com/android/launcher3/dragndrop/DragController.java
+++ b/src/com/android/launcher3/dragndrop/DragController.java
@@ -74,7 +74,7 @@
     /** Coordinate for last touch event **/
     protected final Point mLastTouch = new Point();
 
-    private final Point mTmpPoint = new Point();
+    protected final Point mTmpPoint = new Point();
 
     protected DropTarget.DragObject mDragObject;
 
@@ -317,7 +317,7 @@
         mDragObject.dragView.animateTo(mMotionDown.x, mMotionDown.y, onCompleteRunnable, duration);
     }
 
-    private void callOnDragEnd() {
+    protected void callOnDragEnd() {
         if (mIsInPreDrag && mOptions.preDragCondition != null) {
             mOptions.preDragCondition.onPreDragEnd(mDragObject, false /* dragStarted*/);
         }
@@ -343,7 +343,7 @@
     /**
      * Clamps the position to the drag layer bounds.
      */
-    private Point getClampedDragLayerPos(float x, float y) {
+    protected Point getClampedDragLayerPos(float x, float y) {
         mActivity.getDragLayer().getLocalVisibleRect(mRectTemp);
         mTmpPoint.x = (int) Math.max(mRectTemp.left, Math.min(x, mRectTemp.right - 1));
         mTmpPoint.y = (int) Math.max(mRectTemp.top, Math.min(y, mRectTemp.bottom - 1));
@@ -390,7 +390,7 @@
             return false;
         }
 
-        Point dragLayerPos = getClampedDragLayerPos(ev.getX(), ev.getY());
+        Point dragLayerPos = getClampedDragLayerPos(getX(ev), getY(ev));
         mLastTouch.set(dragLayerPos.x,  dragLayerPos.y);
         if (ev.getAction() == MotionEvent.ACTION_DOWN) {
             // Remember location of down touch
@@ -403,6 +403,14 @@
         return mDragDriver != null && mDragDriver.onInterceptTouchEvent(ev);
     }
 
+    protected float getX(MotionEvent ev) {
+        return ev.getX();
+    }
+
+    protected float getY(MotionEvent ev) {
+        return ev.getY();
+    }
+
     /**
      * Call this from a drag source view.
      */
diff --git a/src/com/android/launcher3/dragndrop/DragDriver.java b/src/com/android/launcher3/dragndrop/DragDriver.java
index d4ce308..72e47e5 100644
--- a/src/com/android/launcher3/dragndrop/DragDriver.java
+++ b/src/com/android/launcher3/dragndrop/DragDriver.java
@@ -165,8 +165,11 @@
      * Class for driving an internal (i.e. not using framework) drag/drop operation.
      */
     static class InternalDragDriver extends DragDriver {
+        private final DragController mDragController;
+
         InternalDragDriver(DragController dragController, Consumer<MotionEvent> sec) {
             super(dragController, sec);
+            mDragController = dragController;
         }
 
         @Override
@@ -176,11 +179,14 @@
 
             switch (action) {
                 case MotionEvent.ACTION_MOVE:
-                    mEventListener.onDriverDragMove(ev.getX(), ev.getY());
+                    mEventListener.onDriverDragMove(mDragController.getX(ev),
+                            mDragController.getY(ev));
                     break;
                 case MotionEvent.ACTION_UP:
-                    mEventListener.onDriverDragMove(ev.getX(), ev.getY());
-                    mEventListener.onDriverDragEnd(ev.getX(), ev.getY());
+                    mEventListener.onDriverDragMove(mDragController.getX(ev),
+                            mDragController.getY(ev));
+                    mEventListener.onDriverDragEnd(mDragController.getX(ev),
+                            mDragController.getY(ev));
                     break;
                 case MotionEvent.ACTION_CANCEL:
                     mEventListener.onDriverDragCancel();
@@ -197,7 +203,8 @@
 
             switch (action) {
                 case MotionEvent.ACTION_UP:
-                    mEventListener.onDriverDragEnd(ev.getX(), ev.getY());
+                    mEventListener.onDriverDragEnd(mDragController.getX(ev),
+                            mDragController.getY(ev));
                     break;
                 case MotionEvent.ACTION_CANCEL:
                     mEventListener.onDriverDragCancel();
diff --git a/src/com/android/launcher3/dragndrop/DragView.java b/src/com/android/launcher3/dragndrop/DragView.java
index 1664980..fb1a6be 100644
--- a/src/com/android/launcher3/dragndrop/DragView.java
+++ b/src/com/android/launcher3/dragndrop/DragView.java
@@ -53,22 +53,19 @@
 import androidx.dynamicanimation.animation.SpringAnimation;
 import androidx.dynamicanimation.animation.SpringForce;
 
-import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherSettings;
-import com.android.launcher3.LauncherState;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.Interpolators;
 import com.android.launcher3.icons.FastBitmapDrawable;
 import com.android.launcher3.icons.LauncherIcons;
 import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.util.RunnableList;
-import com.android.launcher3.util.Thunk;
+import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
 
 /** A custom view for rendering an icon, folder, shortcut or widget during drag-n-drop. */
-public class DragView extends FrameLayout implements StateListener<LauncherState> {
+public abstract class DragView<T extends Context & ActivityContext> extends FrameLayout {
 
     public static final int VIEW_ZOOM_DURATION = 150;
 
@@ -81,19 +78,18 @@
     private final int mHeight;
 
     private final int mBlurSizeOutline;
-    private final int mRegistrationX;
-    private final int mRegistrationY;
+    protected final int mRegistrationX;
+    protected final int mRegistrationY;
     private final float mInitialScale;
-    private final float mScaleOnDrop;
-    private final int[] mTempLoc = new int[2];
+    protected final float mScaleOnDrop;
+    protected final int[] mTempLoc = new int[2];
 
     private final RunnableList mOnDragStartCallback = new RunnableList();
 
     private Point mDragVisualizeOffset = null;
     private Rect mDragRegion = null;
-    private final Launcher mLauncher;
-    private final DragLayer mDragLayer;
-    @Thunk final DragController mDragController;
+    protected final T mActivity;
+    private final BaseDragLayer<T> mDragLayer;
     private boolean mHasDrawn = false;
 
     final ValueAnimator mAnim;
@@ -109,7 +105,7 @@
     private Path mScaledMaskPath;
     private Drawable mBadge;
 
-    public DragView(Launcher launcher, Drawable drawable, int registrationX,
+    public DragView(T launcher, Drawable drawable, int registrationX,
             int registrationY, final float initialScale, final float scaleOnDrop,
             final float finalScaleDps) {
         this(launcher, getViewFromDrawable(launcher, drawable),
@@ -122,7 +118,7 @@
      * <p>
      * The registration point is the point inside our view that the touch events should
      * be centered upon.
-     * @param launcher The Launcher instance
+     * @param activity The Launcher instance/ActivityContext this DragView is in.
      * @param content the view content that is attached to the drag view.
      * @param width the width of the dragView
      * @param height the height of the dragView
@@ -132,13 +128,12 @@
      * @param scaleOnDrop the scale used in the drop animation.
      * @param finalScaleDps the scale used in the zoom out animation when the drag view is shown.
      */
-    public DragView(Launcher launcher, View content, int width, int height, int registrationX,
+    public DragView(T activity, View content, int width, int height, int registrationX,
             int registrationY, final float initialScale, final float scaleOnDrop,
             final float finalScaleDps) {
-        super(launcher);
-        mLauncher = launcher;
-        mDragLayer = launcher.getDragLayer();
-        mDragController = launcher.getDragController();
+        super(activity);
+        mActivity = activity;
+        mDragLayer = activity.getDragLayer();
 
         mContent = content;
         mWidth = width;
@@ -187,24 +182,6 @@
         setWillNotDraw(false);
     }
 
-    @Override
-    protected void onAttachedToWindow() {
-        super.onAttachedToWindow();
-        mLauncher.getStateManager().addStateListener(this);
-    }
-
-    @Override
-    protected void onDetachedFromWindow() {
-        super.onDetachedFromWindow();
-        mLauncher.getStateManager().removeStateListener(this);
-    }
-
-    @Override
-    public void onStateTransitionComplete(LauncherState finalState) {
-        setVisibility((finalState == LauncherState.NORMAL
-                || finalState == LauncherState.SPRING_LOADED) ? VISIBLE : INVISIBLE);
-    }
-
     /**
      * Initialize {@code #mIconDrawable} if the item can be represented using
      * an {@link AdaptiveIconDrawable} or {@link FolderAdaptiveIcon}.
@@ -221,10 +198,10 @@
             Object[] outObj = new Object[1];
             int w = mWidth;
             int h = mHeight;
-            Drawable dr = Utilities.getFullDrawable(mLauncher, info, w, h, outObj);
+            Drawable dr = Utilities.getFullDrawable(mActivity, info, w, h, outObj);
 
             if (dr instanceof AdaptiveIconDrawable) {
-                int blurMargin = (int) mLauncher.getResources()
+                int blurMargin = (int) mActivity.getResources()
                         .getDimension(R.dimen.blur_size_medium_outline) / 2;
 
                 Rect bounds = new Rect(0, 0, w, h);
@@ -232,13 +209,13 @@
                 // Badge is applied after icon normalization so the bounds for badge should not
                 // be scaled down due to icon normalization.
                 Rect badgeBounds = new Rect(bounds);
-                mBadge = getBadge(mLauncher, info, outObj[0]);
+                mBadge = getBadge(mActivity, info, outObj[0]);
                 mBadge.setBounds(badgeBounds);
 
                 // Do not draw the background in case of folder as its translucent
                 final boolean shouldDrawBackground = !(dr instanceof FolderAdaptiveIcon);
 
-                try (LauncherIcons li = LauncherIcons.obtain(mLauncher)) {
+                try (LauncherIcons li = LauncherIcons.obtain(mActivity)) {
                     Drawable nDr; // drawable to be normalized
                     if (shouldDrawBackground) {
                         nDr = dr;
@@ -429,12 +406,11 @@
         applyTranslation();
     }
 
-    public void animateTo(int toTouchX, int toTouchY, Runnable onCompleteRunnable, int duration) {
-        mTempLoc[0] = toTouchX - mRegistrationX;
-        mTempLoc[1] = toTouchY - mRegistrationY;
-        mDragLayer.animateViewIntoPosition(this, mTempLoc, 1f, mScaleOnDrop, mScaleOnDrop,
-                DragLayer.ANIMATION_END_DISAPPEAR, onCompleteRunnable, duration);
-    }
+    /**
+     * Animate this DragView to the given DragLayer coordinates and then remove it.
+     */
+    public abstract void animateTo(int toTouchX, int toTouchY, Runnable onCompleteRunnable,
+            int duration);
 
     public void animateShift(final int shiftX, final int shiftY) {
         if (mAnim.isStarted()) {
@@ -470,7 +446,7 @@
             Picture picture = new Picture();
             mContent.draw(picture.beginRecording(mWidth, mHeight));
             picture.endRecording();
-            View view = new View(mLauncher);
+            View view = new View(mActivity);
             view.setClipToOutline(mContent.getClipToOutline());
             view.setOutlineProvider(mContent.getOutlineProvider());
             view.setBackground(new PictureDrawable(picture));
diff --git a/src/com/android/launcher3/dragndrop/FolderAdaptiveIcon.java b/src/com/android/launcher3/dragndrop/FolderAdaptiveIcon.java
index ea1fbdb..6a6603c 100644
--- a/src/com/android/launcher3/dragndrop/FolderAdaptiveIcon.java
+++ b/src/com/android/launcher3/dragndrop/FolderAdaptiveIcon.java
@@ -32,12 +32,12 @@
 
 import androidx.annotation.Nullable;
 
-import com.android.launcher3.Launcher;
 import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.folder.PreviewBackground;
 import com.android.launcher3.graphics.ShiftedBitmapDrawable;
 import com.android.launcher3.icons.BitmapRenderer;
 import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.views.ActivityContext;
 
 /**
  * {@link AdaptiveIconDrawable} representation of a {@link FolderIcon}
@@ -70,14 +70,14 @@
     }
 
     public static @Nullable FolderAdaptiveIcon createFolderAdaptiveIcon(
-            Launcher launcher, int folderId, Point dragViewSize) {
+            ActivityContext activity, int folderId, Point dragViewSize) {
         Preconditions.assertNonUiThread();
 
         // Create the actual drawable on the UI thread to avoid race conditions with
         // FolderIcon draw pass
         try {
             return MAIN_EXECUTOR.submit(() -> {
-                FolderIcon icon = launcher.findFolderIcon(folderId);
+                FolderIcon icon = activity.findFolderIcon(folderId);
                 return icon == null ? null : createDrawableOnUiThread(icon, dragViewSize);
 
             }).get();
diff --git a/src/com/android/launcher3/dragndrop/LauncherDragController.java b/src/com/android/launcher3/dragndrop/LauncherDragController.java
index a98d70c..0e8b0a5 100644
--- a/src/com/android/launcher3/dragndrop/LauncherDragController.java
+++ b/src/com/android/launcher3/dragndrop/LauncherDragController.java
@@ -96,7 +96,7 @@
         final float scaleDps = mIsInPreDrag
                 ? res.getDimensionPixelSize(R.dimen.pre_drag_view_scale) : 0f;
         final DragView dragView = mDragObject.dragView = drawable != null
-                ? new DragView(
+                ? new LauncherDragView(
                 mActivity,
                 drawable,
                 registrationX,
@@ -104,7 +104,7 @@
                 initialDragViewScale,
                 dragViewScaleOnDrop,
                 scaleDps)
-                : new DragView(
+                : new LauncherDragView(
                         mActivity,
                         view,
                         view.getMeasuredWidth(),
diff --git a/src/com/android/launcher3/dragndrop/LauncherDragView.java b/src/com/android/launcher3/dragndrop/LauncherDragView.java
new file mode 100644
index 0000000..cc68e2e
--- /dev/null
+++ b/src/com/android/launcher3/dragndrop/LauncherDragView.java
@@ -0,0 +1,70 @@
+/*
+ * 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.launcher3.dragndrop;
+
+import android.graphics.drawable.Drawable;
+import android.view.View;
+
+import com.android.launcher3.Launcher;
+import com.android.launcher3.LauncherState;
+import com.android.launcher3.statemanager.StateManager;
+
+/**
+ * A DragView drawn/used by the Launcher activity.
+ */
+public class LauncherDragView extends DragView<Launcher>
+        implements StateManager.StateListener<LauncherState> {
+
+
+    public LauncherDragView(Launcher launcher, Drawable drawable, int registrationX,
+            int registrationY, float initialScale, float scaleOnDrop, float finalScaleDps) {
+        super(launcher, drawable, registrationX, registrationY, initialScale, scaleOnDrop,
+                finalScaleDps);
+    }
+
+    public LauncherDragView(Launcher launcher, View content, int width, int height,
+            int registrationX, int registrationY, float initialScale, float scaleOnDrop,
+            float finalScaleDps) {
+        super(launcher, content, width, height, registrationX, registrationY, initialScale,
+                scaleOnDrop, finalScaleDps);
+    }
+
+    @Override
+    protected void onAttachedToWindow() {
+        super.onAttachedToWindow();
+        mActivity.getStateManager().addStateListener(this);
+    }
+
+    @Override
+    protected void onDetachedFromWindow() {
+        super.onDetachedFromWindow();
+        mActivity.getStateManager().removeStateListener(this);
+    }
+
+    @Override
+    public void onStateTransitionComplete(LauncherState finalState) {
+        setVisibility((finalState == LauncherState.NORMAL
+                || finalState == LauncherState.SPRING_LOADED) ? VISIBLE : INVISIBLE);
+    }
+
+    @Override
+    public void animateTo(int toTouchX, int toTouchY, Runnable onCompleteRunnable, int duration) {
+        mTempLoc[0] = toTouchX - mRegistrationX;
+        mTempLoc[1] = toTouchY - mRegistrationY;
+        mActivity.getDragLayer().animateViewIntoPosition(this, mTempLoc, 1f, mScaleOnDrop,
+                mScaleOnDrop, DragLayer.ANIMATION_END_DISAPPEAR, onCompleteRunnable, duration);
+    }
+}
diff --git a/src/com/android/launcher3/graphics/DragPreviewProvider.java b/src/com/android/launcher3/graphics/DragPreviewProvider.java
index a549750..f027b33 100644
--- a/src/com/android/launcher3/graphics/DragPreviewProvider.java
+++ b/src/com/android/launcher3/graphics/DragPreviewProvider.java
@@ -32,13 +32,13 @@
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.BubbleTextView;
-import com.android.launcher3.Launcher;
 import com.android.launcher3.R;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.dragndrop.DraggableView;
 import com.android.launcher3.icons.BitmapRenderer;
 import com.android.launcher3.icons.FastBitmapDrawable;
 import com.android.launcher3.util.SafeCloseable;
+import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
 
 import java.nio.ByteBuffer;
@@ -150,7 +150,7 @@
     }
 
     public float getScaleAndPosition(Drawable preview, int[] outPos) {
-        float scale = Launcher.getLauncher(mView.getContext())
+        float scale = ActivityContext.lookupContext(mView.getContext())
                 .getDragLayer().getLocationInDragLayer(mView, outPos);
         if (mView instanceof LauncherAppWidgetHostView) {
             // App widgets are technically scaled, but are drawn at their expected size -- so the
@@ -167,7 +167,7 @@
 
     /** Returns the scale and position of a given view for drag-n-drop. */
     public float getScaleAndPosition(View view, int[] outPos) {
-        float scale = Launcher.getLauncher(mView.getContext())
+        float scale = ActivityContext.lookupContext(mView.getContext())
                 .getDragLayer().getLocationInDragLayer(mView, outPos);
         if (mView instanceof LauncherAppWidgetHostView) {
             // App widgets are technically scaled, but are drawn at their expected size -- so the
@@ -201,7 +201,7 @@
         public void run() {
             Bitmap preview = convertPreviewToAlphaBitmap(mPreviewSnapshot);
             if (mIsIcon) {
-                int size = Launcher.getLauncher(mContext).getDeviceProfile().iconSizePx;
+                int size = ActivityContext.lookupContext(mContext).getDeviceProfile().iconSizePx;
                 preview = Bitmap.createScaledBitmap(preview, size, size, false);
             }
             //else case covers AppWidgetHost (doesn't drag/drop across different device profiles)
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 2a1aec8..952b850 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -79,6 +79,7 @@
 import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
@@ -391,11 +392,14 @@
         ArrayList<ItemInfo> otherWorkspaceItems = new ArrayList<>();
         ArrayList<LauncherAppWidgetInfo> currentAppWidgets = new ArrayList<>();
         ArrayList<LauncherAppWidgetInfo> otherAppWidgets = new ArrayList<>();
-        filterCurrentWorkspaceItems(0 /* currentScreenId */,
-                dataModel.workspaceItems, currentWorkspaceItems,
-                otherWorkspaceItems);
-        filterCurrentWorkspaceItems(0 /* currentScreenId */, dataModel.appWidgets,
-                currentAppWidgets, otherAppWidgets);
+
+        IntSet currentScreenIds = IntSet.wrap(0);
+        // TODO(b/185508060): support two panel preview.
+        filterCurrentWorkspaceItems(currentScreenIds, dataModel.workspaceItems,
+                currentWorkspaceItems, otherWorkspaceItems);
+        filterCurrentWorkspaceItems(currentScreenIds, dataModel.appWidgets, currentAppWidgets,
+                otherAppWidgets);
+
         sortWorkspaceItemsSpatially(mIdp, currentWorkspaceItems);
         for (ItemInfo itemInfo : currentWorkspaceItems) {
             switch (itemInfo.itemType) {
diff --git a/src/com/android/launcher3/model/BaseLoaderResults.java b/src/com/android/launcher3/model/BaseLoaderResults.java
index 5c85bab..12ee676 100644
--- a/src/com/android/launcher3/model/BaseLoaderResults.java
+++ b/src/com/android/launcher3/model/BaseLoaderResults.java
@@ -24,13 +24,13 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel.CallbackTask;
-import com.android.launcher3.PagedView;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.LooperExecutor;
 import com.android.launcher3.util.LooperIdleLock;
 import com.android.launcher3.util.ViewOnDrawExecutor;
@@ -160,20 +160,26 @@
         }
 
         private void bind() {
-            final int currentScreen;
+            IntSet currentScreenIndices;
             {
                 // Create an anonymous scope to calculate currentScreen as it has to be a
                 // final variable.
-                int currScreen = mCallbacks.getPageToBindSynchronously();
-                if (currScreen >= mOrderedScreenIds.size()) {
-                    // There may be no workspace screens (just hotseat items and an empty page).
-                    currScreen = PagedView.INVALID_PAGE;
+                IntSet screenIndices = mCallbacks.getPagesToBindSynchronously();
+                if (screenIndices == null || screenIndices.isEmpty()
+                        || screenIndices.getArray().get(screenIndices.size() - 1)
+                        >= mOrderedScreenIds.size()) {
+                    // There maybe no workspace screens (just hotseat items and an empty page).
+                    // Also we want to prevent IndexOutOfBoundsExceptions.
+                    screenIndices = new IntSet();
                 }
-                currentScreen = currScreen;
+                currentScreenIndices = screenIndices;
             }
-            final boolean validFirstPage = currentScreen >= 0;
-            final int currentScreenId =
-                    validFirstPage ? mOrderedScreenIds.get(currentScreen) : INVALID_SCREEN_ID;
+
+            final boolean validFirstPage = !currentScreenIndices.isEmpty();
+
+            IntSet currentScreenIds  = new IntSet();
+            currentScreenIndices.forEach(
+                    index -> currentScreenIds.add(mOrderedScreenIds.get(index)));
 
             // Separate the items that are on the current screen, and all the other remaining items
             ArrayList<ItemInfo> currentWorkspaceItems = new ArrayList<>();
@@ -181,9 +187,9 @@
             ArrayList<LauncherAppWidgetInfo> currentAppWidgets = new ArrayList<>();
             ArrayList<LauncherAppWidgetInfo> otherAppWidgets = new ArrayList<>();
 
-            filterCurrentWorkspaceItems(currentScreenId, mWorkspaceItems, currentWorkspaceItems,
+            filterCurrentWorkspaceItems(currentScreenIds, mWorkspaceItems, currentWorkspaceItems,
                     otherWorkspaceItems);
-            filterCurrentWorkspaceItems(currentScreenId, mAppWidgets, currentAppWidgets,
+            filterCurrentWorkspaceItems(currentScreenIds, mAppWidgets, currentAppWidgets,
                     otherAppWidgets);
             final InvariantDeviceProfile idp = mApp.getInvariantDeviceProfile();
             sortWorkspaceItemsSpatially(idp, currentWorkspaceItems);
@@ -220,14 +226,14 @@
             bindWorkspaceItems(otherWorkspaceItems, deferredExecutor);
             bindAppWidgets(otherAppWidgets, deferredExecutor);
             // Tell the workspace that we're done binding items
-            executeCallbacksTask(c -> c.finishBindingItems(currentScreen), deferredExecutor);
+            executeCallbacksTask(c -> c.finishBindingItems(currentScreenIndices), deferredExecutor);
 
             if (validFirstPage) {
                 executeCallbacksTask(c -> {
                     // We are loading synchronously, which means, some of the pages will be
                     // bound after first draw. Inform the mCallbacks that page binding is
                     // not complete, and schedule the remaining pages.
-                    c.onPageBoundSynchronously(currentScreen);
+                    c.onPagesBoundSynchronously(currentScreenIndices);
                     c.executeOnNextDraw((ViewOnDrawExecutor) deferredExecutor);
 
                 }, mUiExecutor);
diff --git a/src/com/android/launcher3/model/BgDataModel.java b/src/com/android/launcher3/model/BgDataModel.java
index 1d7d1a2..037f408 100644
--- a/src/com/android/launcher3/model/BgDataModel.java
+++ b/src/com/android/launcher3/model/BgDataModel.java
@@ -446,15 +446,16 @@
         int FLAG_QUIET_MODE_CHANGE_PERMISSION = 1 << 2;
 
         /**
-         * Returns the page number to bind first, synchronously if possible or -1
+         * Returns an IntSet of page numbers to bind first, synchronously if possible
+         * or an empty IntSet
          */
-        int getPageToBindSynchronously();
+        IntSet getPagesToBindSynchronously();
         void clearPendingBinds();
         void startBinding();
         void bindItems(List<ItemInfo> shortcuts, boolean forceAnimateIcons);
         void bindScreens(IntArray orderedScreenIds);
         void finishFirstPageBind(ViewOnDrawExecutor executor);
-        void finishBindingItems(int pageBoundFirst);
+        void finishBindingItems(IntSet pagesBoundFirst);
         void preAddApps();
         void bindAppsAdded(IntArray newScreens,
                 ArrayList<ItemInfo> addNotAnimated, ArrayList<ItemInfo> addAnimated);
@@ -468,7 +469,7 @@
         void bindRestoreItemsChange(HashSet<ItemInfo> updates);
         void bindWorkspaceComponentsRemoved(ItemInfoMatcher matcher);
         void bindAllWidgets(List<WidgetsListBaseEntry> widgets);
-        void onPageBoundSynchronously(int page);
+        void onPagesBoundSynchronously(IntSet pages);
         void executeOnNextDraw(ViewOnDrawExecutor executor);
         void bindDeepShortcutMap(HashMap<ComponentKey, Integer> deepShortcutMap);
 
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 79396b1..34a21fe 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -83,6 +83,7 @@
 import com.android.launcher3.shortcuts.ShortcutRequest.QueryResult;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IOUtils;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.LooperIdleLock;
 import com.android.launcher3.util.PackageManagerHelper;
 import com.android.launcher3.util.PackageUserKey;
@@ -173,8 +174,9 @@
         ArrayList<ItemInfo> allItems = mBgDataModel.getAllWorkspaceItems();
         // Screen set is never empty
         final int firstScreen = mBgDataModel.collectWorkspaceScreens().get(0);
+        // TODO(b/185515153): support two panel home.
 
-        filterCurrentWorkspaceItems(firstScreen, allItems, firstScreenItems,
+        filterCurrentWorkspaceItems(IntSet.wrap(firstScreen), allItems, firstScreenItems,
                 new ArrayList<>() /* otherScreenItems are ignored */);
         mFirstScreenBroadcast.sendBroadcasts(mApp.getContext(), firstScreenItems);
     }
diff --git a/src/com/android/launcher3/model/ModelUtils.java b/src/com/android/launcher3/model/ModelUtils.java
index 9b5fac8..58aa9e5 100644
--- a/src/com/android/launcher3/model/ModelUtils.java
+++ b/src/com/android/launcher3/model/ModelUtils.java
@@ -51,7 +51,8 @@
      * Filters the set of items who are directly or indirectly (via another container) on the
      * specified screen.
      */
-    public static <T extends ItemInfo> void filterCurrentWorkspaceItems(int currentScreenId,
+    public static <T extends ItemInfo> void filterCurrentWorkspaceItems(
+            IntSet currentScreenIds,
             ArrayList<T> allWorkspaceItems,
             ArrayList<T> currentScreenItems,
             ArrayList<T> otherScreenItems) {
@@ -65,7 +66,7 @@
                 (lhs, rhs) -> Integer.compare(lhs.container, rhs.container));
         for (T info : allWorkspaceItems) {
             if (info.container == LauncherSettings.Favorites.CONTAINER_DESKTOP) {
-                if (info.screenId == currentScreenId) {
+                if (currentScreenIds.contains(info.screenId)) {
                     currentScreenItems.add(info);
                     itemsOnScreen.add(info.id);
                 } else {
diff --git a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
index 5999091..b271a6a 100644
--- a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
+++ b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
@@ -41,6 +41,7 @@
 import com.android.launcher3.popup.PopupDataProvider;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.ItemInfoMatcher;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.util.ViewOnDrawExecutor;
@@ -175,8 +176,8 @@
     }
 
     @Override
-    public int getPageToBindSynchronously() {
-        return 0;
+    public IntSet getPagesToBindSynchronously() {
+        return new IntSet();
     }
 
     @Override
@@ -199,7 +200,7 @@
     }
 
     @Override
-    public void finishBindingItems(int pageBoundFirst) { }
+    public void finishBindingItems(IntSet pagesBoundFirst) { }
 
     @Override
     public void preAddApps() { }
@@ -229,7 +230,7 @@
     public void bindAllWidgets(List<WidgetsListBaseEntry> widgets) { }
 
     @Override
-    public void onPageBoundSynchronously(int page) { }
+    public void onPagesBoundSynchronously(IntSet pages) { }
 
     @Override
     public void executeOnNextDraw(ViewOnDrawExecutor executor) {
diff --git a/src/com/android/launcher3/util/IntArray.java b/src/com/android/launcher3/util/IntArray.java
index 7252f7a..e7235e7 100644
--- a/src/com/android/launcher3/util/IntArray.java
+++ b/src/com/android/launcher3/util/IntArray.java
@@ -17,13 +17,14 @@
 package com.android.launcher3.util;
 
 import java.util.Arrays;
+import java.util.Iterator;
 import java.util.StringTokenizer;
 
 /**
  * Copy of the platform hidden implementation of android.util.IntArray.
  * Implements a growing array of int primitives.
  */
-public class IntArray implements Cloneable {
+public class IntArray implements Cloneable, Iterable<Integer> {
     private static final int MIN_CAPACITY_INCREMENT = 12;
 
     private static final int[] EMPTY_INT = new int[0];
@@ -272,4 +273,30 @@
             throw new ArrayIndexOutOfBoundsException("length=" + len + "; index=" + index);
         }
     }
+
+    @Override
+    public Iterator<Integer> iterator() {
+        return new ValueIterator();
+    }
+
+    @Thunk
+    class ValueIterator implements Iterator<Integer> {
+
+        private int mNextIndex = 0;
+
+        @Override
+        public boolean hasNext() {
+            return mNextIndex < size();
+        }
+
+        @Override
+        public Integer next() {
+            return get(mNextIndex++);
+        }
+
+        @Override
+        public void remove() {
+            throw new UnsupportedOperationException();
+        }
+    }
 }
\ No newline at end of file
diff --git a/src/com/android/launcher3/util/IntSet.java b/src/com/android/launcher3/util/IntSet.java
index 851f129..0f4df62 100644
--- a/src/com/android/launcher3/util/IntSet.java
+++ b/src/com/android/launcher3/util/IntSet.java
@@ -16,11 +16,13 @@
 package com.android.launcher3.util;
 
 import java.util.Arrays;
+import java.util.Iterator;
 
 /**
  * A wrapper over IntArray implementing a growing set of int primitives.
+ * The elements in the array are sorted in ascending order.
  */
-public class IntSet {
+public class IntSet implements Iterable<Integer> {
 
     final IntArray mArray = new IntArray();
 
@@ -34,6 +36,16 @@
         }
     }
 
+    /**
+     * Removes the specified value from the set if it exist.
+     */
+    public void remove(int value) {
+        int index = Arrays.binarySearch(mArray.mValues, 0, mArray.mSize, value);
+        if (index >= 0) {
+            mArray.removeIndex(index);
+        }
+    }
+
     public boolean contains(int value) {
         return Arrays.binarySearch(mArray.mValues, 0, mArray.mSize, value) >= 0;
     }
@@ -61,6 +73,9 @@
         return (obj instanceof IntSet) && ((IntSet) obj).mArray.equals(mArray);
     }
 
+    /**
+     * Returns the wrapped IntArray. The elements in the array are sorted in ascending order.
+     */
     public IntArray getArray() {
         return mArray;
     }
@@ -78,4 +93,21 @@
         Arrays.sort(set.mArray.mValues, 0, set.mArray.mSize);
         return set;
     }
+
+    /**
+     * Returns an IntSet with the given values.
+     */
+    public static IntSet wrap(int... array) {
+        return wrap(IntArray.wrap(array));
+    }
+
+    @Override
+    public Iterator<Integer> iterator() {
+        return mArray.iterator();
+    }
+
+    @Override
+    public String toString() {
+        return "IntSet{" + mArray.toConcatString() + '}';
+    }
 }
diff --git a/src/com/android/launcher3/views/ActivityContext.java b/src/com/android/launcher3/views/ActivityContext.java
index 646b669..b95904e 100644
--- a/src/com/android/launcher3/views/ActivityContext.java
+++ b/src/com/android/launcher3/views/ActivityContext.java
@@ -21,9 +21,12 @@
 import android.view.LayoutInflater;
 import android.view.View.AccessibilityDelegate;
 
+import androidx.annotation.Nullable;
+
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.dot.DotInfo;
 import com.android.launcher3.dragndrop.DragController;
+import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.ViewCache;
 
@@ -100,6 +103,13 @@
     }
 
     /**
+     * Returns the FolderIcon with the given item id, if it exists.
+     */
+    default @Nullable FolderIcon findFolderIcon(final int folderIconId) {
+        return null;
+    }
+
+    /**
      * Returns the ActivityContext associated with the given Context.
      */
     static <T extends Context & ActivityContext> T lookupContext(Context context) {