Fix launch animation from grid

- Fixed calculateGridTaskSize to algin to taskRect instead of gridRect
- Extracted end of grid's scroll diff calculation to getOffsetFromScrollPosition to accurately calculate end of grid task's offset between its scroll and visual position (End of grid task's scroll is all the same at the position where ClearAllButton is just invisible)
- Extracted isGridTask in TaskView so TaskViewSimualtor can use the correct task size depedning if it's grid task or not

Bug: 200813202
Test: manual
Change-Id: I11a980345aee2680abf1c2563baadd9718a40192
diff --git a/quickstep/src/com/android/quickstep/BaseActivityInterface.java b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
index f9b749e..e15aa92 100644
--- a/quickstep/src/com/android/quickstep/BaseActivityInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
@@ -322,11 +322,11 @@
     public final void calculateGridTaskSize(Context context, DeviceProfile dp, Rect outRect,
             PagedOrientationHandler orientedState) {
         Resources res = context.getResources();
-        Rect gridRect = new Rect();
-        calculateGridSize(context, dp, gridRect);
+        Rect taskRect = new Rect();
+        calculateTaskSize(context, dp, taskRect);
 
         float rowHeight =
-                (gridRect.height() + dp.overviewTaskThumbnailTopMarginPx - dp.overviewRowSpacing)
+                (taskRect.height() + dp.overviewTaskThumbnailTopMarginPx - dp.overviewRowSpacing)
                         / 2f;
 
         PointF taskDimension = getTaskDimension(context, dp);
@@ -336,7 +336,7 @@
 
         int gravity = Gravity.TOP;
         gravity |= orientedState.getRecentsRtlSetting(res) ? Gravity.RIGHT : Gravity.LEFT;
-        Gravity.apply(gravity, outWidth, outHeight, gridRect, outRect);
+        Gravity.apply(gravity, outWidth, outHeight, taskRect, outRect);
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/TaskViewUtils.java b/quickstep/src/com/android/quickstep/TaskViewUtils.java
index 67a94cc..ae3cc50 100644
--- a/quickstep/src/com/android/quickstep/TaskViewUtils.java
+++ b/quickstep/src/com/android/quickstep/TaskViewUtils.java
@@ -215,6 +215,7 @@
                     tvsLocal.taskSecondaryTranslation.value = gridTranslationSecondary;
                 }
                 tvsLocal.setScroll(startScroll);
+                tvsLocal.setIsGridTask(v.isGridTask());
 
                 // Fade in the task during the initial 20% of the animation
                 out.addFloat(targetHandle.getTransformParams(), TransformParams.TARGET_ALPHA, 0, 1,
diff --git a/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java b/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
index 280b5d7..146d235 100644
--- a/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
+++ b/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
@@ -101,6 +101,7 @@
     private int mOrientationStateId;
     private StagedSplitBounds mStagedSplitBounds;
     private boolean mDrawsBelowRecents;
+    private boolean mIsGridTask;
 
     public TaskViewSimulator(Context context, BaseActivityInterface sizeStrategy) {
         mContext = context;
@@ -140,18 +141,23 @@
         if (mDp == null) {
             return 1;
         }
-        Rect fullTaskSize = new Rect();
-        mSizeStrategy.calculateTaskSize(mContext, mDp, fullTaskSize);
+        if (mIsGridTask) {
+            mSizeStrategy.calculateGridTaskSize(mContext, mDp, mTaskRect,
+                    mOrientationState.getOrientationHandler());
+        } else {
+            mSizeStrategy.calculateTaskSize(mContext, mDp, mTaskRect);
+        }
 
+        Rect fullTaskSize;
         if (mStagedSplitBounds != null) {
             // The task rect changes according to the staged split task sizes, but recents
             // fullscreen scale and pivot remains the same since the task fits into the existing
             // sized task space bounds
-            mSizeStrategy.calculateTaskSize(mContext, mDp, mTaskRect);
+            fullTaskSize = new Rect(mTaskRect);
             mOrientationState.getOrientationHandler()
                     .setSplitTaskSwipeRect(mDp, mTaskRect, mStagedSplitBounds, mStagePosition);
         } else {
-            mTaskRect.set(fullTaskSize);
+            fullTaskSize = mTaskRect;
         }
         return mOrientationState.getFullScreenScaleAndPivot(fullTaskSize, mDp, mPivot);
     }
@@ -205,6 +211,13 @@
     }
 
     /**
+     * Sets whether the task is part of overview grid and not being focused.
+     */
+    public void setIsGridTask(boolean isGridTask) {
+        mIsGridTask = isGridTask;
+    }
+
+    /**
      * Adds animation for all the components corresponding to transition from an app to overview.
      */
     public void addAppToOverviewAnim(PendingAnimation pa, TimeInterpolator interpolator) {
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 72e01bb..223ce1c 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -157,6 +157,7 @@
 import com.android.quickstep.TaskThumbnailCache;
 import com.android.quickstep.TaskViewUtils;
 import com.android.quickstep.ViewUtils;
+import com.android.quickstep.util.GroupTask;
 import com.android.quickstep.util.LayoutUtils;
 import com.android.quickstep.util.RecentsOrientedState;
 import com.android.quickstep.util.SplitScreenBounds;
@@ -166,7 +167,6 @@
 import com.android.quickstep.util.TransformParams;
 import com.android.quickstep.util.VibratorWrapper;
 import com.android.systemui.plugins.ResourceProvider;
-import com.android.quickstep.util.GroupTask;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.Task.TaskKey;
 import com.android.systemui.shared.recents.model.ThumbnailData;
@@ -1018,8 +1018,10 @@
     }
 
     private TaskView getLastGridTaskView() {
-        IntArray topRowIdArray = getTopRowIdArray();
-        IntArray bottomRowIdArray = getBottomRowIdArray();
+        return getLastGridTaskView(getTopRowIdArray(), getBottomRowIdArray());
+    }
+
+    private TaskView getLastGridTaskView(IntArray topRowIdArray, IntArray bottomRowIdArray) {
         if (topRowIdArray.isEmpty() && bottomRowIdArray.isEmpty()) {
             return null;
         }
@@ -3145,38 +3147,23 @@
                                 }
                             }
 
-                            TaskView newLastGridTaskView = getLastGridTaskView();
+                            IntArray topRowIdArray = getTopRowIdArray();
+                            IntArray bottomRowIdArray = getBottomRowIdArray();
                             if (finalSnapToLastTask) {
                                 // If snapping to last task, find the last task after dismissal.
-                                pageToSnapTo = indexOfChild(newLastGridTaskView);
+                                pageToSnapTo = indexOfChild(
+                                        getLastGridTaskView(topRowIdArray, bottomRowIdArray));
                             } else if (taskViewIdToSnapTo != -1) {
                                 // If snapping to another page due to indices rearranging, find
                                 // the new index after dismissal & rearrange using the task view id.
                                 pageToSnapTo = indexOfChild(
                                         getTaskViewFromTaskViewId(taskViewIdToSnapTo));
-                                int taskViewToSnapToScroll = getScrollForPage(pageToSnapTo);
-                                int lastGridTaskScroll = getScrollForPage(
-                                        indexOfChild(newLastGridTaskView));
-                                if (!currentPageSnapsToEndOfGrid
-                                        && taskViewToSnapToScroll == lastGridTaskScroll) {
+                                if (!currentPageSnapsToEndOfGrid) {
                                     // If it wasn't snapped to one of the last pages, but is now
                                     // snapped to last pages, we'll need to compensate for the
-                                    // difference as last pages' scroll is the position where
-                                    // ClearAllButton is barely invisible, instead of aligned to
-                                    // mLastComputedTaskSize.
-                                    int normalTaskEnd = mIsRtl
-                                            ? mLastComputedTaskSize.right
-                                            : mLastComputedTaskSize.left;
-                                    int lastTaskStart = mIsRtl
-                                            ? mLastComputedGridSize.left
-                                            : mLastComputedGridSize.right;
-                                    // As snapped task is not the last task, it can only be the
-                                    // second last task.
-                                    int distanceToSnappedTaskEnd =
-                                            (mPageSpacing + mLastComputedGridTaskSize.width()) * 2;
-                                    int snappedTaskEnd = lastTaskStart + (mIsRtl
-                                            ? distanceToSnappedTaskEnd : -distanceToSnappedTaskEnd);
-                                    mCurrentPageScrollDiff += snappedTaskEnd - normalTaskEnd;
+                                    // offset from the page's scroll to its visual position.
+                                    mCurrentPageScrollDiff += getOffsetFromScrollPosition(
+                                            pageToSnapTo, topRowIdArray, bottomRowIdArray);
                                 }
                             }
                         }
@@ -4613,7 +4600,58 @@
                     overScrollShift, getUndampedOverScrollShift());
         }
         return getScrollForPage(pageIndex) - mOrientationHandler.getPrimaryScroll(this)
-                + overScrollShift;
+                + overScrollShift + getOffsetFromScrollPosition(pageIndex);
+    }
+
+    /**
+     * Returns how many pixels the page is offset from its scroll position.
+     */
+    private int getOffsetFromScrollPosition(int pageIndex) {
+        return getOffsetFromScrollPosition(pageIndex, getTopRowIdArray(), getBottomRowIdArray());
+    }
+
+    private int getOffsetFromScrollPosition(
+            int pageIndex, IntArray topRowIdArray, IntArray bottomRowIdArray) {
+        if (!showAsGrid()) {
+            return 0;
+        }
+
+        TaskView taskView = getTaskViewAt(pageIndex);
+        if (taskView == null) {
+            return 0;
+        }
+
+        TaskView lastGridTaskView = getLastGridTaskView(topRowIdArray, bottomRowIdArray);
+        if (lastGridTaskView == null) {
+            return 0;
+        }
+
+        if (getScrollForPage(pageIndex) != getScrollForPage(indexOfChild(lastGridTaskView))) {
+            return 0;
+        }
+
+        // Check distance from lastGridTaskView to taskView.
+        int lastGridTaskViewPosition =
+                getPositionInRow(lastGridTaskView, topRowIdArray, bottomRowIdArray);
+        int taskViewPosition = getPositionInRow(taskView, topRowIdArray, bottomRowIdArray);
+        int gridTaskSizeAndSpacing = mLastComputedGridTaskSize.width() + mPageSpacing;
+        int positionDiff = gridTaskSizeAndSpacing * (lastGridTaskViewPosition - taskViewPosition);
+
+        int lastTaskEnd = (mIsRtl
+                ? mLastComputedGridSize.left
+                : mLastComputedGridSize.right)
+                + (mIsRtl ? mPageSpacing : -mPageSpacing);
+        int taskEnd = lastTaskEnd + (mIsRtl ? positionDiff : -positionDiff);
+        int normalTaskEnd = mIsRtl
+                ? mLastComputedGridTaskSize.left
+                : mLastComputedGridTaskSize.right;
+        return taskEnd - normalTaskEnd;
+    }
+
+    private int getPositionInRow(
+            TaskView taskView, IntArray topRowIdArray, IntArray bottomRowIdArray) {
+        int position = topRowIdArray.indexOf(taskView.getTaskViewId());
+        return position != -1 ? position : bottomRowIdArray.indexOf(taskView.getTaskViewId());
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index d09ffa6..9f961c9 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -877,7 +877,7 @@
         LayoutParams snapshotParams = (LayoutParams) mSnapshotView.getLayoutParams();
         DeviceProfile deviceProfile = mActivity.getDeviceProfile();
         snapshotParams.topMargin = deviceProfile.overviewTaskThumbnailTopMarginPx;
-        boolean isGridTask = deviceProfile.overviewShowAsGrid && !isFocusedTask();
+        boolean isGridTask = isGridTask();
         int taskIconHeight = deviceProfile.overviewTaskIconSizePx;
         int taskMargin = isGridTask ? deviceProfile.overviewTaskMarginGridPx
                 : deviceProfile.overviewTaskMarginPx;
@@ -897,6 +897,14 @@
         mSnapshotView.getTaskOverlay().updateOrientationState(orientationState);
     }
 
+    /**
+     * Returns whether the task is part of overview grid and not being focused.
+     */
+    public boolean isGridTask() {
+        DeviceProfile deviceProfile = mActivity.getDeviceProfile();
+        return deviceProfile.overviewShowAsGrid && !isFocusedTask();
+    }
+
     private void setIconAndDimTransitionProgress(float progress, boolean invert) {
         if (invert) {
             progress = 1 - progress;