Fix back button sometimes not showing up in the initial setup screen

A few things to note
- We use gesture nav layoutter because it's the mode in which we show back button rendered by IME. There might be other historical reasons but it is what it is.
- The back button isn't necessarily laid out correctly with this layoutter (currently we don't adjust anything in layoutButtons.

Fixes: 338876161
Test: Repeatedly get in setup mode, and make sure the back button is visible
Change-Id: I402273c2868f3f46a12b00918dab3f36fbe82204
(cherry picked from commit 9d602e684c)

quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java

@@ -951,7 +951,9 @@ public class TaskbarActivityContext extends BaseTaskbarContext {
                 mWindowLayoutParams.paramsForRotation[rot].height = size;
             }
         }
-        mControllers.taskbarInsetsController.onTaskbarOrBubblebarWindowHeightOrInsetsChanged();
+        mControllers.runAfterInit(
+                mControllers.taskbarInsetsController
+                        ::onTaskbarOrBubblebarWindowHeightOrInsetsChanged);
         notifyUpdateLayoutParams();
     }
 
@@ -1416,7 +1418,7 @@ public class TaskbarActivityContext extends BaseTaskbarContext {
         });
     }
 
-    protected boolean isUserSetupComplete() {
+    public boolean isUserSetupComplete() {
         return mIsUserSetupComplete;
     }
 

quickstep/src/com/android/launcher3/taskbar/navbutton/AbstractNavButtonLayoutter.kt

@@ -24,8 +24,10 @@ import android.widget.FrameLayout
 import android.widget.ImageView
 import android.widget.LinearLayout
 import android.widget.Space
+import com.android.launcher3.DeviceProfile
 import com.android.launcher3.R
 import com.android.launcher3.Utilities
+import com.android.launcher3.taskbar.TaskbarActivityContext
 import com.android.launcher3.taskbar.navbutton.NavButtonLayoutFactory.NavButtonLayoutter
 
 /**
@@ -73,6 +75,23 @@ abstract class AbstractNavButtonLayoutter(
         return params
     }
 
+    fun adjustForSetupInPhoneMode(
+        navButtonsLayoutParams: FrameLayout.LayoutParams,
+        navButtonsViewLayoutParams: FrameLayout.LayoutParams,
+        deviceProfile: DeviceProfile
+    ) {
+        val phoneOrPortraitSetupMargin =
+            resources.getDimensionPixelSize(R.dimen.taskbar_contextual_button_suw_margin)
+        navButtonsLayoutParams.marginStart = phoneOrPortraitSetupMargin
+        navButtonsLayoutParams.bottomMargin =
+            if (!deviceProfile.isLandscape) 0
+            else
+                phoneOrPortraitSetupMargin -
+                        resources.getDimensionPixelSize(R.dimen.taskbar_nav_buttons_size) / 2
+        navButtonsViewLayoutParams.height =
+            resources.getDimensionPixelSize(R.dimen.taskbar_contextual_button_suw_height)
+    }
+
     open fun repositionContextualContainer(
         contextualContainer: ViewGroup,
         buttonSize: Int,

quickstep/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactory.kt

@@ -116,6 +116,7 @@ class NavButtonLayoutFactory {
                 isPhoneGestureMode -> {
                     PhoneGestureLayoutter(
                         resources,
+                        navButtonsView,
                         navButtonContainer,
                         endContextualContainer,
                         startContextualContainer,

quickstep/src/com/android/launcher3/taskbar/navbutton/PhoneGestureLayoutter.kt

@@ -17,15 +17,19 @@
 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 android.widget.Space
+import com.android.launcher3.DeviceProfile
 import com.android.launcher3.taskbar.TaskbarActivityContext
 
 /** Layoutter for showing gesture navigation on phone screen. No buttons here, no-op container */
 class PhoneGestureLayoutter(
     resources: Resources,
+    navButtonsView: NearestTouchFrame,
     navBarContainer: LinearLayout,
     endContextualContainer: ViewGroup,
     startContextualContainer: ViewGroup,
@@ -42,8 +46,31 @@ class PhoneGestureLayoutter(
         a11yButton,
         space
     ) {
+    private val mNavButtonsView = navButtonsView
 
     override fun layoutButtons(context: TaskbarActivityContext, isA11yButtonPersistent: Boolean) {
+        // TODO: look into if we should use SetupNavLayoutter instead.
+        if (!context.isUserSetupComplete) {
+            // Since setup wizard only has back button enabled, it looks strange to be
+            // end-aligned, so start-align instead.
+            val navButtonsLayoutParams = navButtonContainer.layoutParams as FrameLayout.LayoutParams
+            val navButtonsViewLayoutParams =
+                mNavButtonsView.layoutParams as FrameLayout.LayoutParams
+            val deviceProfile: DeviceProfile = context.deviceProfile
+
+            navButtonsLayoutParams.marginEnd = 0
+            navButtonsLayoutParams.gravity = Gravity.START
+            context.setTaskbarWindowSize(context.setupWindowSize)
+
+            adjustForSetupInPhoneMode(
+                navButtonsLayoutParams,
+                navButtonsViewLayoutParams,
+                deviceProfile
+            )
+            mNavButtonsView.layoutParams = navButtonsViewLayoutParams
+            navButtonContainer.layoutParams = navButtonsLayoutParams
+        }
+
         endContextualContainer.removeAllViews()
         startContextualContainer.removeAllViews()
     }

quickstep/src/com/android/launcher3/taskbar/navbutton/SetupNavLayoutter.kt

@@ -77,16 +77,11 @@ class SetupNavLayoutter(
             navButtonsLayoutParams.height =
                 resources.getDimensionPixelSize(R.dimen.taskbar_back_button_suw_height)
         } else {
-            val phoneOrPortraitSetupMargin =
-                resources.getDimensionPixelSize(R.dimen.taskbar_contextual_button_suw_margin)
-            navButtonsLayoutParams.marginStart = phoneOrPortraitSetupMargin
-            navButtonsLayoutParams.bottomMargin =
-                if (!deviceProfile.isLandscape) 0
-                else
-                    phoneOrPortraitSetupMargin -
-                        resources.getDimensionPixelSize(R.dimen.taskbar_nav_buttons_size) / 2
-            navButtonsViewLayoutParams.height =
-                resources.getDimensionPixelSize(R.dimen.taskbar_contextual_button_suw_height)
+            adjustForSetupInPhoneMode(
+                navButtonsLayoutParams,
+                navButtonsViewLayoutParams,
+                deviceProfile
+            )
         }
         mNavButtonsView.layoutParams = navButtonsViewLayoutParams
         navButtonContainer.layoutParams = navButtonsLayoutParams