Repostion contextual buttons upon layout changes

Also make sure the contextual buttons are centered in its respective orientation

Fixes: 300849078
Test: Rotate around, and make sure the rotation button stays in the same corner
Change-Id: Ibc40739eff050f258b938ccf94c2056ed1ee596b
diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index df6626a..4b16019 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -191,6 +191,7 @@
     private MultiValueAlpha mBackButtonAlpha;
     private MultiValueAlpha mHomeButtonAlpha;
     private FloatingRotationButton mFloatingRotationButton;
+    private ImageView mImeSwitcherButton;
 
     // Variables for moving nav buttons to a separate window above IME
     private boolean mAreNavButtonsInSeparateWindow = false;
@@ -237,10 +238,10 @@
                 InputMethodService.canImeRenderGesturalNavButtons() && mContext.imeDrawsImeNavBar();
         if (!mIsImeRenderingNavButtons) {
             // IME switcher
-            View imeSwitcherButton = addButton(R.drawable.ic_ime_switcher, BUTTON_IME_SWITCH,
+            mImeSwitcherButton = addButton(R.drawable.ic_ime_switcher, BUTTON_IME_SWITCH,
                     isThreeButtonNav ? mStartContextualContainer : mEndContextualContainer,
                     mControllers.navButtonController, R.id.ime_switcher);
-            mPropertyHolders.add(new StatePropertyHolder(imeSwitcherButton,
+            mPropertyHolders.add(new StatePropertyHolder(mImeSwitcherButton,
                     flags -> ((flags & FLAG_SWITCHER_SHOWING) != 0)
                             && ((flags & FLAG_ROTATION_BUTTON_VISIBLE) == 0)));
         }
@@ -736,7 +737,9 @@
         if (TaskbarManager.FLAG_HIDE_NAVBAR_WINDOW) {
             NavButtonLayoutter navButtonLayoutter =
                     NavButtonLayoutFactory.Companion.getUiLayoutter(
-                            dp, mNavButtonsView, res, isInKidsMode, isInSetup, isThreeButtonNav,
+                            dp, mNavButtonsView, mImeSwitcherButton,
+                            mControllers.rotationButtonController.getRotationButton(),
+                            mA11yButton, res, isInKidsMode, isInSetup, isThreeButtonNav,
                             TaskbarManager.isPhoneMode(dp),
                             mWindowManagerProxy.getRotation(mContext));
             navButtonLayoutter.layoutButtons(dp, isContextualButtonShowing());
diff --git a/quickstep/src/com/android/launcher3/taskbar/navbutton/AbstractNavButtonLayoutter.kt b/quickstep/src/com/android/launcher3/taskbar/navbutton/AbstractNavButtonLayoutter.kt
index b682081..9758d44 100644
--- a/quickstep/src/com/android/launcher3/taskbar/navbutton/AbstractNavButtonLayoutter.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/navbutton/AbstractNavButtonLayoutter.kt
@@ -18,12 +18,15 @@
 
 import android.content.res.Resources
 import android.graphics.drawable.RotateDrawable
+import android.view.Gravity
 import android.view.ViewGroup
+import android.widget.FrameLayout
 import android.widget.ImageView
 import android.widget.LinearLayout
 import com.android.launcher3.R
 import com.android.launcher3.Utilities
 import com.android.launcher3.taskbar.navbutton.NavButtonLayoutFactory.NavButtonLayoutter
+import com.android.systemui.shared.rotation.RotationButton
 
 /**
  * Meant to be a simple container for data subclasses will need
@@ -37,10 +40,13 @@
  * @property startContextualContainer ViewGroup that holds the start contextual button (ex, A11y).
  */
 abstract class AbstractNavButtonLayoutter(
-    val resources: Resources,
-    val navButtonContainer: LinearLayout,
-    protected val endContextualContainer: ViewGroup,
-    protected val startContextualContainer: ViewGroup
+        val resources: Resources,
+        val navButtonContainer: LinearLayout,
+        protected val endContextualContainer: ViewGroup,
+        protected val startContextualContainer: ViewGroup,
+        protected val imeSwitcher: ImageView?,
+        protected val rotationButton: RotationButton?,
+        protected val a11yButton: ImageView
 ) : NavButtonLayoutter {
     protected val homeButton: ImageView? = navButtonContainer.findViewById(R.id.home)
     protected val recentsButton: ImageView? = navButtonContainer.findViewById(R.id.recent_apps)
@@ -56,4 +62,11 @@
             backButton.setImageDrawable(rotateDrawable)
         }
     }
+
+    fun getParamsToCenterView(): FrameLayout.LayoutParams {
+        val params = FrameLayout.LayoutParams(
+                ViewGroup.LayoutParams.WRAP_CONTENT, ViewGroup.LayoutParams.WRAP_CONTENT)
+        params.gravity = Gravity.CENTER
+        return params;
+    }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/navbutton/KidsNavLayoutter.kt b/quickstep/src/com/android/launcher3/taskbar/navbutton/KidsNavLayoutter.kt
index 4a53c0c..f254ee8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/navbutton/KidsNavLayoutter.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/navbutton/KidsNavLayoutter.kt
@@ -26,18 +26,25 @@
 import android.widget.LinearLayout
 import com.android.launcher3.DeviceProfile
 import com.android.launcher3.taskbar.navbutton.LayoutResourceHelper.*
+import com.android.systemui.shared.rotation.RotationButton
 
 class KidsNavLayoutter(
-    resources: Resources,
-    navBarContainer: LinearLayout,
-    endContextualContainer: ViewGroup,
-    startContextualContainer: ViewGroup
+        resources: Resources,
+        navBarContainer: LinearLayout,
+        endContextualContainer: ViewGroup,
+        startContextualContainer: ViewGroup,
+        imeSwitcher: ImageView?,
+        rotationButton: RotationButton?,
+        a11yButton: ImageView
 ) :
     AbstractNavButtonLayoutter(
-        resources,
-        navBarContainer,
-        endContextualContainer,
-        startContextualContainer
+            resources,
+            navBarContainer,
+            endContextualContainer,
+            startContextualContainer,
+            imeSwitcher,
+            rotationButton,
+            a11yButton
     ) {
 
     override fun layoutButtons(dp: DeviceProfile, isContextualButtonShowing: Boolean) {
@@ -89,5 +96,28 @@
         navButtonContainer.requestLayout()
 
         homeButton.onLongClickListener = null
+
+        endContextualContainer.removeAllViews()
+        startContextualContainer.removeAllViews()
+
+        val endContextualContainerParams = FrameLayout.LayoutParams(
+                ViewGroup.LayoutParams.WRAP_CONTENT, ViewGroup.LayoutParams.MATCH_PARENT)
+        endContextualContainerParams.gravity = Gravity.END or Gravity.CENTER_VERTICAL
+        endContextualContainer.layoutParams = endContextualContainerParams
+
+        val startContextualContainerParams = FrameLayout.LayoutParams(
+                ViewGroup.LayoutParams.WRAP_CONTENT, ViewGroup.LayoutParams.MATCH_PARENT)
+        startContextualContainerParams.gravity = Gravity.START or Gravity.CENTER_VERTICAL
+        startContextualContainer.layoutParams = startContextualContainerParams
+
+        if (imeSwitcher != null) {
+            startContextualContainer.addView(imeSwitcher)
+            imeSwitcher.layoutParams = getParamsToCenterView()
+        }
+        endContextualContainer.addView(a11yButton)
+        if (rotationButton != null) {
+            endContextualContainer.addView(rotationButton.currentView)
+            rotationButton.currentView.layoutParams = getParamsToCenterView()
+        }
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactory.kt b/quickstep/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactory.kt
index 931f692..7db1a37 100644
--- a/quickstep/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactory.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactory.kt
@@ -21,11 +21,13 @@
 import android.view.Surface.Rotation
 import android.view.ViewGroup
 import android.widget.FrameLayout
+import android.widget.ImageView
 import android.widget.LinearLayout
 import com.android.launcher3.DeviceProfile
 import com.android.launcher3.taskbar.navbutton.LayoutResourceHelper.*
 import com.android.launcher3.taskbar.navbutton.NavButtonLayoutFactory.Companion
 import com.android.launcher3.taskbar.navbutton.NavButtonLayoutFactory.NavButtonLayoutter
+import com.android.systemui.shared.rotation.RotationButton
 
 /**
  * Select the correct layout for nav buttons
@@ -52,14 +54,17 @@
          * @param isThreeButtonNav are no-ops when taskbar is present/showing
          */
         fun getUiLayoutter(
-            deviceProfile: DeviceProfile,
-            navButtonsView: FrameLayout,
-            resources: Resources,
-            isKidsMode: Boolean,
-            isInSetup: Boolean,
-            isThreeButtonNav: Boolean,
-            phoneMode: Boolean,
-            @Rotation surfaceRotation: Int
+                deviceProfile: DeviceProfile,
+                navButtonsView: FrameLayout,
+                imeSwitcher: ImageView?,
+                rotationButton: RotationButton?,
+                a11yButton: ImageView,
+                resources: Resources,
+                isKidsMode: Boolean,
+                isInSetup: Boolean,
+                isThreeButtonNav: Boolean,
+                phoneMode: Boolean,
+                @Rotation surfaceRotation: Int
         ): NavButtonLayoutter {
             val navButtonContainer =
                 navButtonsView.requireViewById<LinearLayout>(ID_END_NAV_BUTTONS)
@@ -73,24 +78,33 @@
                 isPhoneNavMode -> {
                     if (!deviceProfile.isLandscape) {
                         PhonePortraitNavLayoutter(
-                            resources,
-                            navButtonContainer,
-                            endContextualContainer,
-                            startContextualContainer
+                                resources,
+                                navButtonContainer,
+                                endContextualContainer,
+                                startContextualContainer,
+                                imeSwitcher,
+                                rotationButton,
+                                a11yButton
                         )
                     } else if (surfaceRotation == ROTATION_90) {
                         PhoneLandscapeNavLayoutter(
-                            resources,
-                            navButtonContainer,
-                            endContextualContainer,
-                            startContextualContainer
+                                resources,
+                                navButtonContainer,
+                                endContextualContainer,
+                                startContextualContainer,
+                                imeSwitcher,
+                                rotationButton,
+                                a11yButton
                         )
                     } else {
                         PhoneSeascapeNavLayoutter(
                                 resources,
                                 navButtonContainer,
                                 endContextualContainer,
-                                startContextualContainer
+                                startContextualContainer,
+                                imeSwitcher,
+                                rotationButton,
+                                a11yButton
                         )
                     }
                 }
@@ -99,33 +113,45 @@
                             resources,
                             navButtonContainer,
                             endContextualContainer,
-                            startContextualContainer
+                            startContextualContainer,
+                            imeSwitcher,
+                            rotationButton,
+                            a11yButton
                     )
                 }
                 deviceProfile.isTaskbarPresent -> {
                     return when {
                         isInSetup -> {
                             SetupNavLayoutter(
-                                resources,
-                                navButtonContainer,
-                                endContextualContainer,
-                                startContextualContainer
+                                    resources,
+                                    navButtonContainer,
+                                    endContextualContainer,
+                                    startContextualContainer,
+                                    imeSwitcher,
+                                    rotationButton,
+                                    a11yButton
                             )
                         }
                         isKidsMode -> {
                             KidsNavLayoutter(
-                                resources,
-                                navButtonContainer,
-                                endContextualContainer,
-                                startContextualContainer
+                                    resources,
+                                    navButtonContainer,
+                                    endContextualContainer,
+                                    startContextualContainer,
+                                    imeSwitcher,
+                                    rotationButton,
+                                    a11yButton
                             )
                         }
                         else ->
                             TaskbarNavLayoutter(
-                                resources,
-                                navButtonContainer,
-                                endContextualContainer,
-                                startContextualContainer
+                                    resources,
+                                    navButtonContainer,
+                                    endContextualContainer,
+                                    startContextualContainer,
+                                    imeSwitcher,
+                                    rotationButton,
+                                    a11yButton
                             )
                     }
                 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneGestureLayoutter.kt b/quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneGestureLayoutter.kt
index 8525c6c..c1dae40 100644
--- a/quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneGestureLayoutter.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneGestureLayoutter.kt
@@ -18,24 +18,33 @@
 
 import android.content.res.Resources
 import android.view.ViewGroup
+import android.widget.ImageView
 import android.widget.LinearLayout
 import com.android.launcher3.DeviceProfile
+import com.android.systemui.shared.rotation.RotationButton
 
 /** Layoutter for showing gesture navigation on phone screen. No buttons here, no-op container */
 class PhoneGestureLayoutter(
         resources: Resources,
         navBarContainer: LinearLayout,
         endContextualContainer: ViewGroup,
-        startContextualContainer: ViewGroup
+        startContextualContainer: ViewGroup,
+        imeSwitcher: ImageView?,
+        rotationButton: RotationButton?,
+        a11yButton: ImageView
 ) :
         AbstractNavButtonLayoutter(
                 resources,
                 navBarContainer,
                 endContextualContainer,
-                startContextualContainer
+                startContextualContainer,
+                imeSwitcher,
+                rotationButton,
+                a11yButton
         ) {
 
     override fun layoutButtons(dp: DeviceProfile, isContextualButtonShowing: Boolean) {
-        // no-op
+        endContextualContainer.removeAllViews()
+        startContextualContainer.removeAllViews()
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneLandscapeNavLayoutter.kt b/quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneLandscapeNavLayoutter.kt
index 13ffe6e..21bbca5 100644
--- a/quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneLandscapeNavLayoutter.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneLandscapeNavLayoutter.kt
@@ -20,24 +20,32 @@
 import android.view.Gravity
 import android.view.ViewGroup
 import android.widget.FrameLayout
+import android.widget.ImageView
 import android.widget.LinearLayout
 import androidx.core.view.children
 import com.android.launcher3.DeviceProfile
 import com.android.launcher3.R
 import com.android.launcher3.taskbar.TaskbarManager
 import com.android.launcher3.util.DimensionUtils
+import com.android.systemui.shared.rotation.RotationButton
 
 open class PhoneLandscapeNavLayoutter(
-    resources: Resources,
-    navBarContainer: LinearLayout,
-    endContextualContainer: ViewGroup,
-    startContextualContainer: ViewGroup
+        resources: Resources,
+        navBarContainer: LinearLayout,
+        endContextualContainer: ViewGroup,
+        startContextualContainer: ViewGroup,
+        imeSwitcher: ImageView?,
+        rotationButton: RotationButton?,
+        a11yButton: ImageView,
 ) :
     AbstractNavButtonLayoutter(
-        resources,
-        navBarContainer,
-        endContextualContainer,
-        startContextualContainer
+            resources,
+            navBarContainer,
+            endContextualContainer,
+            startContextualContainer,
+            imeSwitcher,
+            rotationButton,
+            a11yButton
     ) {
 
     override fun layoutButtons(dp: DeviceProfile, isContextualButtonShowing: Boolean) {
@@ -81,6 +89,24 @@
                 }
             }
         }
+
+        endContextualContainer.removeAllViews()
+        startContextualContainer.removeAllViews()
+
+        val startContextualContainerParams = FrameLayout.LayoutParams(
+                ViewGroup.LayoutParams.MATCH_PARENT, ViewGroup.LayoutParams.WRAP_CONTENT)
+        startContextualContainerParams.gravity = Gravity.TOP or Gravity.CENTER_HORIZONTAL
+        startContextualContainer.layoutParams = startContextualContainerParams
+
+        if (imeSwitcher != null) {
+            startContextualContainer.addView(imeSwitcher)
+            imeSwitcher.layoutParams = getParamsToCenterView()
+        }
+        startContextualContainer.addView(a11yButton)
+        if (rotationButton != null) {
+            startContextualContainer.addView(rotationButton.currentView)
+            rotationButton.currentView.layoutParams = getParamsToCenterView()
+        }
     }
 
     open fun addThreeButtons() {
diff --git a/quickstep/src/com/android/launcher3/taskbar/navbutton/PhonePortraitNavLayoutter.kt b/quickstep/src/com/android/launcher3/taskbar/navbutton/PhonePortraitNavLayoutter.kt
index c763115..ad03e5b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/navbutton/PhonePortraitNavLayoutter.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/navbutton/PhonePortraitNavLayoutter.kt
@@ -20,23 +20,31 @@
 import android.view.Gravity
 import android.view.ViewGroup
 import android.widget.FrameLayout
+import android.widget.ImageView
 import android.widget.LinearLayout
 import com.android.launcher3.DeviceProfile
 import com.android.launcher3.R
 import com.android.launcher3.taskbar.TaskbarManager
 import com.android.launcher3.util.DimensionUtils
+import com.android.systemui.shared.rotation.RotationButton
 
 class PhonePortraitNavLayoutter(
-    resources: Resources,
-    navBarContainer: LinearLayout,
-    endContextualContainer: ViewGroup,
-    startContextualContainer: ViewGroup
+        resources: Resources,
+        navBarContainer: LinearLayout,
+        endContextualContainer: ViewGroup,
+        startContextualContainer: ViewGroup,
+        imeSwitcher: ImageView?,
+        rotationButton: RotationButton?,
+        a11yButton: ImageView,
 ) :
     AbstractNavButtonLayoutter(
-        resources,
-        navBarContainer,
-        endContextualContainer,
-        startContextualContainer
+            resources,
+            navBarContainer,
+            endContextualContainer,
+            startContextualContainer,
+            imeSwitcher,
+            rotationButton,
+            a11yButton
     ) {
 
     override fun layoutButtons(dp: DeviceProfile, isContextualButtonShowing: Boolean) {
@@ -90,5 +98,23 @@
                 }
             }
         }
+
+        endContextualContainer.removeAllViews()
+        startContextualContainer.removeAllViews()
+
+        val endContextualContainerParams = FrameLayout.LayoutParams(
+                ViewGroup.LayoutParams.WRAP_CONTENT, ViewGroup.LayoutParams.MATCH_PARENT)
+        endContextualContainerParams.gravity = Gravity.END or Gravity.CENTER_VERTICAL
+        endContextualContainer.layoutParams = endContextualContainerParams
+
+        if (imeSwitcher != null) {
+            endContextualContainer.addView(imeSwitcher)
+            imeSwitcher.layoutParams = getParamsToCenterView()
+        }
+        endContextualContainer.addView(a11yButton)
+        if (rotationButton != null) {
+            endContextualContainer.addView(rotationButton.currentView)
+            rotationButton.currentView.layoutParams = getParamsToCenterView()
+        }
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneSeascapeNavLayoutter.kt b/quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneSeascapeNavLayoutter.kt
index 2d62c3f..cde39f3 100644
--- a/quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneSeascapeNavLayoutter.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneSeascapeNavLayoutter.kt
@@ -17,20 +17,30 @@
 package com.android.launcher3.taskbar.navbutton
 
 import android.content.res.Resources
+import android.view.Gravity
 import android.view.ViewGroup
+import android.widget.FrameLayout
+import android.widget.ImageView
 import android.widget.LinearLayout
+import com.android.systemui.shared.rotation.RotationButton
 
 class PhoneSeascapeNavLayoutter(
         resources: Resources,
         navBarContainer: LinearLayout,
         endContextualContainer: ViewGroup,
-        startContextualContainer: ViewGroup
+        startContextualContainer: ViewGroup,
+        imeSwitcher: ImageView?,
+        rotationButton: RotationButton?,
+        a11yButton: ImageView
 ) :
         PhoneLandscapeNavLayoutter(
                 resources,
                 navBarContainer,
                 endContextualContainer,
-                startContextualContainer
+                startContextualContainer,
+                imeSwitcher,
+                rotationButton,
+                a11yButton
         ) {
 
     override fun addThreeButtons() {
@@ -38,5 +48,23 @@
         navButtonContainer.addView(backButton)
         navButtonContainer.addView(homeButton)
         navButtonContainer.addView(recentsButton)
+
+        endContextualContainer.removeAllViews()
+        startContextualContainer.removeAllViews()
+
+        val endContextualContainerParams = FrameLayout.LayoutParams(
+                ViewGroup.LayoutParams.MATCH_PARENT, ViewGroup.LayoutParams.WRAP_CONTENT)
+        endContextualContainerParams.gravity = Gravity.BOTTOM or Gravity.CENTER_HORIZONTAL
+        endContextualContainer.layoutParams = endContextualContainerParams
+
+        if (imeSwitcher != null) {
+            endContextualContainer.addView(imeSwitcher)
+            imeSwitcher.layoutParams = getParamsToCenterView()
+        }
+        endContextualContainer.addView(a11yButton)
+        if (rotationButton != null) {
+            endContextualContainer.addView(rotationButton.currentView)
+            rotationButton.currentView.layoutParams = getParamsToCenterView()
+        }
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/navbutton/SetupNavLayoutter.kt b/quickstep/src/com/android/launcher3/taskbar/navbutton/SetupNavLayoutter.kt
index a24002c..db245b8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/navbutton/SetupNavLayoutter.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/navbutton/SetupNavLayoutter.kt
@@ -20,20 +20,28 @@
 import android.view.Gravity
 import android.view.ViewGroup
 import android.widget.FrameLayout
+import android.widget.ImageView
 import android.widget.LinearLayout
 import com.android.launcher3.DeviceProfile
+import com.android.systemui.shared.rotation.RotationButton
 
 class SetupNavLayoutter(
-    resources: Resources,
-    navButtonContainer: LinearLayout,
-    endContextualContainer: ViewGroup,
-    startContextualContainer: ViewGroup
+        resources: Resources,
+        navButtonContainer: LinearLayout,
+        endContextualContainer: ViewGroup,
+        startContextualContainer: ViewGroup,
+        imeSwitcher: ImageView?,
+        rotationButton: RotationButton?,
+        a11yButton: ImageView
 ) :
     AbstractNavButtonLayoutter(
-        resources,
-        navButtonContainer,
-        endContextualContainer,
-        startContextualContainer
+            resources,
+            navButtonContainer,
+            endContextualContainer,
+            startContextualContainer,
+            imeSwitcher,
+            rotationButton,
+            a11yButton
     ) {
 
     override fun layoutButtons(dp: DeviceProfile, isContextualButtonShowing: Boolean) {
@@ -46,5 +54,28 @@
             gravity = Gravity.START
         }
         navButtonContainer.requestLayout()
+
+        endContextualContainer.removeAllViews()
+        startContextualContainer.removeAllViews()
+
+        val endContextualContainerParams = FrameLayout.LayoutParams(
+                ViewGroup.LayoutParams.WRAP_CONTENT, ViewGroup.LayoutParams.MATCH_PARENT)
+        endContextualContainerParams.gravity = Gravity.END or Gravity.CENTER_VERTICAL
+        endContextualContainer.layoutParams = endContextualContainerParams
+
+        val startContextualContainerParams = FrameLayout.LayoutParams(
+                ViewGroup.LayoutParams.WRAP_CONTENT, ViewGroup.LayoutParams.MATCH_PARENT)
+        startContextualContainerParams.gravity = Gravity.START or Gravity.CENTER_VERTICAL
+        startContextualContainer.layoutParams = startContextualContainerParams
+
+        if (imeSwitcher != null) {
+            startContextualContainer.addView(imeSwitcher)
+            imeSwitcher.layoutParams = getParamsToCenterView()
+        }
+        endContextualContainer.addView(a11yButton)
+        if (rotationButton != null) {
+            endContextualContainer.addView(rotationButton.currentView)
+            rotationButton.currentView.layoutParams = getParamsToCenterView()
+        }
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/navbutton/TaskbarNavLayoutter.kt b/quickstep/src/com/android/launcher3/taskbar/navbutton/TaskbarNavLayoutter.kt
index 8332b7d..56e55bb 100644
--- a/quickstep/src/com/android/launcher3/taskbar/navbutton/TaskbarNavLayoutter.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/navbutton/TaskbarNavLayoutter.kt
@@ -20,22 +20,30 @@
 import android.view.Gravity
 import android.view.ViewGroup
 import android.widget.FrameLayout
+import android.widget.ImageView
 import android.widget.LinearLayout
 import com.android.launcher3.DeviceProfile
 import com.android.launcher3.R
+import com.android.systemui.shared.rotation.RotationButton
 
 /** Layoutter for showing 3 button navigation on large screen */
 class TaskbarNavLayoutter(
-    resources: Resources,
-    navBarContainer: LinearLayout,
-    endContextualContainer: ViewGroup,
-    startContextualContainer: ViewGroup
+        resources: Resources,
+        navBarContainer: LinearLayout,
+        endContextualContainer: ViewGroup,
+        startContextualContainer: ViewGroup,
+        imeSwitcher: ImageView?,
+        rotationButton: RotationButton?,
+        a11yButton: ImageView
 ) :
     AbstractNavButtonLayoutter(
-        resources,
-        navBarContainer,
-        endContextualContainer,
-        startContextualContainer
+            resources,
+            navBarContainer,
+            endContextualContainer,
+            startContextualContainer,
+            imeSwitcher,
+            rotationButton,
+            a11yButton
     ) {
 
     override fun layoutButtons(dp: DeviceProfile, isContextualButtonShowing: Boolean) {
@@ -77,5 +85,28 @@
                 }
             }
         }
+
+        endContextualContainer.removeAllViews()
+        startContextualContainer.removeAllViews()
+
+        val endContextualContainerParams = FrameLayout.LayoutParams(
+                ViewGroup.LayoutParams.WRAP_CONTENT, ViewGroup.LayoutParams.MATCH_PARENT)
+        endContextualContainerParams.gravity = Gravity.END or Gravity.CENTER_VERTICAL
+        endContextualContainer.layoutParams = endContextualContainerParams
+
+        val startContextualContainerParams = FrameLayout.LayoutParams(
+                ViewGroup.LayoutParams.WRAP_CONTENT, ViewGroup.LayoutParams.MATCH_PARENT)
+        startContextualContainerParams.gravity = Gravity.START or Gravity.CENTER_VERTICAL
+        startContextualContainer.layoutParams = startContextualContainerParams
+
+        if (imeSwitcher != null) {
+            startContextualContainer.addView(imeSwitcher)
+            imeSwitcher.layoutParams = getParamsToCenterView()
+        }
+        endContextualContainer.addView(a11yButton)
+        if (rotationButton != null) {
+            endContextualContainer.addView(rotationButton.currentView)
+            rotationButton.currentView.layoutParams = getParamsToCenterView()
+        }
     }
 }
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactoryTest.kt b/quickstep/tests/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactoryTest.kt
index 3920b08..16bfe70 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactoryTest.kt
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactoryTest.kt
@@ -13,6 +13,7 @@
 import com.android.launcher3.DeviceProfile
 import com.android.launcher3.R
 import com.android.launcher3.taskbar.TaskbarManager
+import com.android.systemui.shared.rotation.RotationButton
 import java.lang.IllegalStateException
 import org.junit.Assume.assumeTrue
 import org.junit.Before
@@ -34,6 +35,9 @@
     @Mock lateinit var mockBackButton: ImageView
     @Mock lateinit var mockRecentsButton: ImageView
     @Mock lateinit var mockHomeButton: ImageView
+    @Mock lateinit var mockImeSwitcher: ImageView
+    @Mock lateinit var mockRotationButton: RotationButton
+    @Mock lateinit var mockA11yButton: ImageView
 
     private var surfaceRotation = Surface.ROTATION_0
 
@@ -196,7 +200,10 @@
             isInSetup = isInSetup,
             isThreeButtonNav = isThreeButtonNav,
             phoneMode = phoneMode,
-            surfaceRotation = surfaceRotation
+            surfaceRotation = surfaceRotation,
+            imeSwitcher = mockImeSwitcher,
+            rotationButton = mockRotationButton,
+            a11yButton = mockA11yButton
         )
     }
 }