Merge "Enable currently functional multivalent tests" into main
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index 7ff887c..6279903 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -245,7 +245,7 @@
 
     private Animator mTaskbarBackgroundAlphaAnimator;
     private long mTaskbarBackgroundDuration;
-    private boolean mIsGoingHome;
+    private boolean mUserIsNotGoingHome = false;
 
     // Evaluate whether the handle should be stashed
     private final LongPredicate mIsStashedPredicate = flags -> {
@@ -828,17 +828,13 @@
             private boolean mTaskbarBgAlphaAnimationStarted = false;
             @Override
             public void onAnimationUpdate(ValueAnimator valueAnimator) {
-                if (mIsGoingHome) {
-                    mTaskbarBgAlphaAnimationStarted = true;
-                }
                 if (mTaskbarBgAlphaAnimationStarted) {
                     return;
                 }
 
                 if (valueAnimator.getAnimatedFraction() >= ANIMATED_FRACTION_THRESHOLD) {
-                    if (!mIsGoingHome) {
+                    if (mUserIsNotGoingHome) {
                         playTaskbarBackgroundAlphaAnimation();
-                        setUserIsGoingHome(false);
                         mTaskbarBgAlphaAnimationStarted = true;
                     }
                 }
@@ -850,8 +846,8 @@
     /**
      * Sets whether the user is going home based on the current gesture.
      */
-    public void setUserIsGoingHome(boolean isGoingHome) {
-        mIsGoingHome = isGoingHome;
+    public void setUserIsNotGoingHome(boolean userIsNotGoingHome) {
+        mUserIsNotGoingHome = userIsNotGoingHome;
     }
 
     /**
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
index 593285f..ce281c3 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
@@ -415,7 +415,7 @@
     /**
      * Sets whether the user is going home based on the current gesture.
      */
-    public void setUserIsGoingHome(boolean isGoingHome) {
-        mControllers.taskbarStashController.setUserIsGoingHome(isGoingHome);
+    public void setUserIsNotGoingHome(boolean isNotGoingHome) {
+        mControllers.taskbarStashController.setUserIsNotGoingHome(isNotGoingHome);
     }
 }
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 1acafab..93f72fc 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -1196,7 +1196,7 @@
         }
         if (mContainerInterface.getTaskbarController() != null) {
             // Resets this value as the gesture is now complete.
-            mContainerInterface.getTaskbarController().setUserIsGoingHome(false);
+            mContainerInterface.getTaskbarController().setUserIsNotGoingHome(false);
         }
         ActiveGestureLog.INSTANCE.addLog(
                 new ActiveGestureLog.CompoundString("onSettledOnEndTarget ")
@@ -1350,7 +1350,7 @@
                 && mIsTransientTaskbar
                 && mContainerInterface.getTaskbarController() != null) {
             mContainerInterface.getTaskbarController()
-                    .setUserIsGoingHome(endTarget == GestureState.GestureEndTarget.HOME);
+                    .setUserIsNotGoingHome(endTarget != GestureState.GestureEndTarget.HOME);
         }
 
         float endShift = endTarget == ALL_APPS ? mDragLengthFactor
diff --git a/quickstep/src/com/android/quickstep/OverviewCommandHelper.java b/quickstep/src/com/android/quickstep/OverviewCommandHelper.java
index 7da92bc..8f533a3 100644
--- a/quickstep/src/com/android/quickstep/OverviewCommandHelper.java
+++ b/quickstep/src/com/android/quickstep/OverviewCommandHelper.java
@@ -265,6 +265,10 @@
                 case TYPE_HOME:
                     ActiveGestureLog.INSTANCE.addLog(
                             "OverviewCommandHelper.executeCommand(TYPE_HOME)");
+                    // Although IActivityTaskManager$Stub$Proxy.startActivity is a slow binder call,
+                    // we should still call it on main thread because launcher is waiting for
+                    // ActivityTaskManager to resume it. Also calling startActivity() on bg thread
+                    // could potentially delay resuming launcher. See b/348668521 for more details.
                     mService.startActivity(mOverviewComponentObserver.getHomeIntent());
                     return true;
                 case TYPE_SHOW:
diff --git a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
index fd141c3..f7e1b4e 100644
--- a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
@@ -228,6 +228,7 @@
 
                 // Take the thumbnail of the task without a scrim and apply it back after
                 float alpha = mThumbnailView.getDimAlpha();
+                // TODO(b/348643341) add ability to get override the scrim for this Bitmap retrieval
                 mThumbnailView.setDimAlpha(0);
                 Bitmap thumbnail = RecentsTransition.drawViewIntoHardwareBitmap(
                         taskBounds.width(), taskBounds.height(), mThumbnailView, 1f,
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
index 2836c89..dbe2b19 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
@@ -53,25 +53,31 @@
         TaskThumbnailViewModel(
             recentsView.mRecentsViewData,
             (parent as TaskView).taskViewData,
+            (parent as TaskView).getTaskContainerForTaskThumbnailView(this)!!.taskContainerData,
             recentsView.mTasksRepository,
         )
     }
 
     private var uiState: TaskThumbnailUiState = Uninitialized
     private var inheritedScale: Float = 1f
+    private var dimProgress: Float = 0f
 
     private val backgroundPaint = Paint(Paint.ANTI_ALIAS_FLAG)
+    private val scrimPaint = Paint().apply { color = Color.BLACK }
     private val _measuredBounds = Rect()
     private val measuredBounds: Rect
         get() {
             _measuredBounds.set(0, 0, measuredWidth, measuredHeight)
             return _measuredBounds
         }
+
     private var cornerRadius: Float = TaskCornerRadius.get(context)
     private var fullscreenCornerRadius: Float = QuickStepContract.getWindowCornerRadius(context)
 
     constructor(context: Context?) : super(context)
+
     constructor(context: Context?, attrs: AttributeSet?) : super(context, attrs)
+
     constructor(
         context: Context?,
         attrs: AttributeSet?,
@@ -87,6 +93,13 @@
                 invalidate()
             }
         }
+        MainScope().launch {
+            viewModel.dimProgress.collect { dimProgress ->
+                // TODO(b/348195366) Add fade in/out for scrim
+                this@TaskThumbnailView.dimProgress = dimProgress
+                invalidate()
+            }
+        }
         MainScope().launch { viewModel.recentsFullscreenProgress.collect { invalidateOutline() } }
         MainScope().launch {
             viewModel.inheritedScale.collect { viewModelInheritedScale ->
@@ -111,6 +124,10 @@
             is Snapshot -> drawSnapshotState(canvas, uiStateVal)
             is BackgroundOnly -> drawBackgroundOnly(canvas, uiStateVal.backgroundColor)
         }
+
+        if (dimProgress > 0) {
+            drawScrim(canvas)
+        }
     }
 
     private fun drawBackgroundOnly(canvas: Canvas, @ColorInt backgroundColor: Int) {
@@ -135,6 +152,11 @@
         canvas.drawBitmap(snapshot.bitmap, snapshot.drawnRect, measuredBounds, null)
     }
 
+    private fun drawScrim(canvas: Canvas) {
+        scrimPaint.alpha = (dimProgress * MAX_SCRIM_ALPHA).toInt()
+        canvas.drawRect(measuredBounds, scrimPaint)
+    }
+
     private fun getCurrentCornerRadius() =
         Utilities.mapRange(
             viewModel.recentsFullscreenProgress.value,
@@ -145,5 +167,6 @@
     companion object {
         private val CLEAR_PAINT =
             Paint().apply { xfermode = PorterDuffXfermode(PorterDuff.Mode.CLEAR) }
+        private const val MAX_SCRIM_ALPHA = (0.4f * 255).toInt()
     }
 }
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModel.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModel.kt
index 4511ea7..fe21174 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModel.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModel.kt
@@ -25,6 +25,7 @@
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
+import com.android.quickstep.task.viewmodel.TaskContainerData
 import com.android.quickstep.task.viewmodel.TaskViewData
 import com.android.systemui.shared.recents.model.Task
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -40,6 +41,7 @@
 class TaskThumbnailViewModel(
     recentsViewData: RecentsViewData,
     taskViewData: TaskViewData,
+    taskContainerData: TaskContainerData,
     private val tasksRepository: RecentTasksRepository,
 ) {
     private val task = MutableStateFlow<Flow<Task?>>(flowOf(null))
@@ -50,6 +52,7 @@
         combine(recentsViewData.scale, taskViewData.scale) { recentsScale, taskScale ->
             recentsScale * taskScale
         }
+    val dimProgress: Flow<Float> = taskContainerData.taskMenuOpenProgress
     val uiState: Flow<TaskThumbnailUiState> =
         task
             .flatMapLatest { taskFlow ->
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskContainerData.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskContainerData.kt
new file mode 100644
index 0000000..769424c
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskContainerData.kt
@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2024 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.task.viewmodel
+
+import kotlinx.coroutines.flow.MutableStateFlow
+
+class TaskContainerData {
+    val taskMenuOpenProgress = MutableStateFlow(0f)
+}
diff --git a/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt b/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
index 936f6a1..4c78e21 100644
--- a/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
@@ -193,23 +193,23 @@
             }
             val taskContainer =
                 TaskContainer(
-                        task,
-                        // TODO(b/338360089): Support new TTV for DesktopTaskView
-                        thumbnailView = null,
-                        thumbnailViewDeprecated,
-                        iconView,
-                        TransformingTouchDelegate(iconView.asView()),
-                        SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
-                        digitalWellBeingToast = null,
-                        showWindowsView = null,
-                        taskOverlayFactory
-                    )
-                    .apply { thumbnailViewDeprecated.bind(task, overlay) }
+                    task,
+                    // TODO(b/338360089): Support new TTV for DesktopTaskView
+                    thumbnailView = null,
+                    thumbnailViewDeprecated,
+                    iconView,
+                    TransformingTouchDelegate(iconView.asView()),
+                    SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
+                    digitalWellBeingToast = null,
+                    showWindowsView = null,
+                    taskOverlayFactory
+                )
             if (index >= taskContainers.size) {
                 taskContainers.add(taskContainer)
             } else {
                 taskContainers[index] = taskContainer
             }
+            taskContainer.bind()
         }
         repeat(taskContainers.size - tasks.size) {
             with(taskContainers.removeLast()) {
diff --git a/quickstep/src/com/android/quickstep/views/GroupedTaskView.kt b/quickstep/src/com/android/quickstep/views/GroupedTaskView.kt
index d6a3376..6296b0e 100644
--- a/quickstep/src/com/android/quickstep/views/GroupedTaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/GroupedTaskView.kt
@@ -145,6 +145,8 @@
                     taskOverlayFactory
                 )
             )
+        taskContainers.forEach { it.bind() }
+
         this.splitBoundsConfig =
             splitBoundsConfig?.also {
                 taskContainers[0]
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 43a3eda..d806e3d 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -316,7 +316,6 @@
     /**
      * Can be used to tint the color of the RecentsView to simulate a scrim that can views
      * excluded from. Really should be a proper scrim.
-     * TODO(b/187528071): Remove this and replace with a real scrim.
      */
     private static final FloatProperty<RecentsView> COLOR_TINT =
             new FloatProperty<RecentsView>("colorTint") {
@@ -557,7 +556,6 @@
     @Nullable
     protected GestureState.GestureEndTarget mCurrentGestureEndTarget;
 
-    // TODO(b/187528071): Remove these and replace with a real scrim.
     private float mColorTint;
     private final int mTintingColor;
     @Nullable
@@ -6032,6 +6030,7 @@
      * tasks to be dimmed while other elements in the recents view are left alone.
      */
     public void showForegroundScrim(boolean show) {
+        // TODO(b/335606129) Add scrim response into new TTV - this is called from overlay
         if (!show && mColorTint == 0) {
             if (mTintingAnimator != null) {
                 mTintingAnimator.cancel();
@@ -6047,7 +6046,6 @@
     }
 
     /** Tint the RecentsView and TaskViews in to simulate a scrim. */
-    // TODO(b/187528071): Replace this tinting with a scrim on top of RecentsView
     private void setColorTint(float tintAmount) {
         mColorTint = tintAmount;
 
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuView.java b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
index eda58c5..4f446b2 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
@@ -18,6 +18,7 @@
 
 import static com.android.app.animation.Interpolators.EMPHASIZED;
 import static com.android.launcher3.Flags.enableOverviewIconMenu;
+import static com.android.launcher3.Flags.enableRefactorTaskThumbnail;
 import static com.android.launcher3.util.MultiPropertyFactory.MULTI_PROPERTY_VALUE;
 import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.quickstep.views.TaskThumbnailViewDeprecated.DIM_ALPHA;
@@ -367,6 +368,14 @@
                         mTaskContainer.getThumbnailViewDeprecated(), DIM_ALPHA,
                         closing ? 0 : TaskView.MAX_PAGE_SCRIM_ALPHA),
                 ObjectAnimator.ofFloat(this, ALPHA, closing ? 0 : 1));
+        if (enableRefactorTaskThumbnail()) {
+            mRevealAnimator.addUpdateListener(animation -> {
+                float animatedFraction = animation.getAnimatedFraction();
+                float openProgress = closing ? (1 - animatedFraction) : animatedFraction;
+                mTaskContainer.getTaskContainerData()
+                        .getTaskMenuOpenProgress().setValue(openProgress);
+            });
+        }
         mOpenCloseAnimator.addListener(new AnimationSuccessListener() {
             @Override
             public void onAnimationStart(Animator animation) {
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index 9c1aaa6..7a3b00f 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -88,6 +88,7 @@
 import com.android.quickstep.orientation.RecentsPagedOrientationHandler
 import com.android.quickstep.task.thumbnail.TaskThumbnail
 import com.android.quickstep.task.thumbnail.TaskThumbnailView
+import com.android.quickstep.task.viewmodel.TaskContainerData
 import com.android.quickstep.task.viewmodel.TaskViewData
 import com.android.quickstep.util.ActiveGestureErrorDetector
 import com.android.quickstep.util.ActiveGestureLog
@@ -667,6 +668,7 @@
                     taskOverlayFactory
                 )
             )
+        taskContainers.forEach { it.bind() }
         setOrientationState(orientedState)
     }
 
@@ -693,24 +695,16 @@
         }
         val iconView = getOrInflateIconView(iconViewId)
         return TaskContainer(
-                task,
-                thumbnailView,
-                thumbnailViewDeprecated,
-                iconView,
-                TransformingTouchDelegate(iconView.asView()),
-                stagePosition,
-                DigitalWellBeingToast(container, this),
-                findViewById(showWindowViewId)!!,
-                taskOverlayFactory
-            )
-            .apply {
-                if (enableRefactorTaskThumbnail()) {
-                    thumbnailViewDeprecated.setTaskOverlay(overlay)
-                    bindThumbnailView()
-                } else {
-                    thumbnailViewDeprecated.bind(task, overlay)
-                }
-            }
+            task,
+            thumbnailView,
+            thumbnailViewDeprecated,
+            iconView,
+            TransformingTouchDelegate(iconView.asView()),
+            stagePosition,
+            DigitalWellBeingToast(container, this),
+            findViewById(showWindowViewId)!!,
+            taskOverlayFactory
+        )
     }
 
     protected fun getOrInflateIconView(@IdRes iconViewId: Int): TaskViewIcon {
@@ -1379,7 +1373,6 @@
     open fun setColorTint(amount: Float, tintColor: Int) {
         taskContainers.forEach {
             if (!enableRefactorTaskThumbnail()) {
-                // TODO(b/334832108) Add scrim to new TTV
                 it.thumbnailViewDeprecated.dimAlpha = amount
             }
             it.iconView.setIconColorTint(tintColor, amount)
@@ -1522,6 +1515,9 @@
         resetViewTransforms()
     }
 
+    fun getTaskContainerForTaskThumbnailView(taskThumbnailView: TaskThumbnailView): TaskContainer? =
+        taskContainers.firstOrNull { it.thumbnailView == taskThumbnailView }
+
     open fun resetViewTransforms() {
         // fullscreenTranslation and accumulatedTranslation should not be reset, as
         // resetViewTransforms is called during QuickSwitch scrolling.
@@ -1623,6 +1619,7 @@
         taskOverlayFactory: TaskOverlayFactory
     ) {
         val overlay: TaskOverlay<*> = taskOverlayFactory.createOverlay(this)
+        val taskContainerData = TaskContainerData()
 
         val snapshotView: View
             get() = thumbnailView ?: thumbnailViewDeprecated
@@ -1656,6 +1653,15 @@
             thumbnailView?.let { taskView.removeView(it) }
         }
 
+        fun bind() {
+            if (enableRefactorTaskThumbnail() && thumbnailView != null) {
+                thumbnailViewDeprecated.setTaskOverlay(overlay)
+                bindThumbnailView()
+            } else {
+                thumbnailViewDeprecated.bind(task, overlay)
+            }
+        }
+
         // TODO(b/335649589): TaskView's VM will already have access to TaskThumbnailView's VM
         //  so there will be no need to access TaskThumbnailView's VM through the TaskThumbnailView
         fun bindThumbnailView() {
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
index 3b8754c..a394b65 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
@@ -28,6 +28,7 @@
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
+import com.android.quickstep.task.viewmodel.TaskContainerData
 import com.android.quickstep.task.viewmodel.TaskViewData
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
@@ -43,9 +44,10 @@
 class TaskThumbnailViewModelTest {
     private val recentsViewData = RecentsViewData()
     private val taskViewData = TaskViewData()
+    private val taskContainerData = TaskContainerData()
     private val tasksRepository = FakeTasksRepository()
     private val systemUnderTest =
-        TaskThumbnailViewModel(recentsViewData, taskViewData, tasksRepository)
+        TaskThumbnailViewModel(recentsViewData, taskViewData, taskContainerData, tasksRepository)
 
     private val tasks = (0..5).map(::createTaskWithId)
 
diff --git a/src/com/android/launcher3/Hotseat.java b/src/com/android/launcher3/Hotseat.java
index 8546454..f775673 100644
--- a/src/com/android/launcher3/Hotseat.java
+++ b/src/com/android/launcher3/Hotseat.java
@@ -97,10 +97,9 @@
         if (bubbleBarEnabled) {
             float adjustedBorderSpace = dp.getHotseatAdjustedBorderSpaceForBubbleBar(getContext());
             if (hasBubbles && Float.compare(adjustedBorderSpace, 0f) != 0) {
-                getShortcutsAndWidgets().setTranslationProvider(child -> {
-                    int index = getShortcutsAndWidgets().indexOfChild(child);
+                getShortcutsAndWidgets().setTranslationProvider(cellX -> {
                     float borderSpaceDelta = adjustedBorderSpace - dp.hotseatBorderSpace;
-                    return dp.iconSizePx + index * borderSpaceDelta;
+                    return dp.iconSizePx + cellX * borderSpaceDelta;
                 });
                 if (mQsb instanceof HorizontalInsettableView) {
                     HorizontalInsettableView insettableQsb = (HorizontalInsettableView) mQsb;
@@ -147,10 +146,7 @@
 
         // update the translation provider for future layout passes of hotseat icons.
         if (isBubbleBarVisible) {
-            icons.setTranslationProvider(child -> {
-                int index = icons.indexOfChild(child);
-                return dp.iconSizePx + index * borderSpaceDelta;
-            });
+            icons.setTranslationProvider(cellX -> dp.iconSizePx + cellX * borderSpaceDelta);
         } else {
             icons.setTranslationProvider(null);
         }
diff --git a/src/com/android/launcher3/ShortcutAndWidgetContainer.java b/src/com/android/launcher3/ShortcutAndWidgetContainer.java
index 7484b64..d2c3c78 100644
--- a/src/com/android/launcher3/ShortcutAndWidgetContainer.java
+++ b/src/com/android/launcher3/ShortcutAndWidgetContainer.java
@@ -245,7 +245,7 @@
         }
         child.layout(childLeft, childTop, childLeft + lp.width, childTop + lp.height);
         if (mTranslationProvider != null) {
-            final float tx = mTranslationProvider.getTranslationX(child);
+            final float tx = mTranslationProvider.getTranslationX(lp.getCellX());
             if (child instanceof Reorderable) {
                 ((Reorderable) child).getTranslateDelegate()
                         .getTranslationX(INDEX_BUBBLE_ADJUSTMENT_ANIM)
@@ -330,6 +330,6 @@
 
     /** Provides translation values to apply when laying out child views. */
     interface TranslationProvider {
-        float getTranslationX(View child);
+        float getTranslationX(int cellX);
     }
 }