2D Recents: use correct layout logic for drawing transition headers

Bug: 33752154
Test: Checked/screencasted app -> recents -> app transitions on sw60dp device
Change-Id: Ib29e5feaadb4bdacae6c60d25526d0c7af95361f
diff --git a/packages/SystemUI/src/com/android/systemui/recents/RecentsImpl.java b/packages/SystemUI/src/com/android/systemui/recents/RecentsImpl.java
index 0e42a54..7ef82ff 100644
--- a/packages/SystemUI/src/com/android/systemui/recents/RecentsImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/recents/RecentsImpl.java
@@ -71,6 +71,7 @@
 import com.android.systemui.recents.views.TaskStackViewScroller;
 import com.android.systemui.recents.views.TaskViewHeader;
 import com.android.systemui.recents.views.TaskViewTransform;
+import com.android.systemui.recents.views.grid.TaskGridLayoutAlgorithm;
 import com.android.systemui.stackdivider.DividerView;
 import com.android.systemui.statusbar.BaseStatusBar;
 import com.android.systemui.statusbar.phone.NavigationBarGestureHelper;
@@ -618,10 +619,21 @@
             stackLayout.initialize(displayRect, windowRect, mTaskStackBounds,
                     TaskStackLayoutAlgorithm.StackState.getStackStateForStack(stack));
             mDummyStackView.setTasks(stack, false /* allowNotifyStackChanges */);
+            // Get the width of a task view so that we know how wide to draw the header bar.
+            int taskViewWidth = 0;
+            if (mDummyStackView.useGridLayout()) {
+                TaskGridLayoutAlgorithm gridLayout = mDummyStackView.getGridAlgorithm();
+                gridLayout.initialize(windowRect);
+                taskViewWidth = (int) gridLayout.getTransform(0 /* taskIndex */,
+                        stack.getTaskCount(), new TaskViewTransform(), stackLayout).rect.width();
+            } else {
+                Rect taskViewBounds = stackLayout.getUntransformedTaskViewBounds();
+                if (!taskViewBounds.isEmpty()) {
+                    taskViewWidth = taskViewBounds.width();
+                }
+            }
 
-            Rect taskViewBounds = stackLayout.getUntransformedTaskViewBounds();
-            if (!taskViewBounds.isEmpty()) {
-                int taskViewWidth = taskViewBounds.width();
+            if (taskViewWidth > 0) {
                 synchronized (mHeaderBarLock) {
                     if (mHeaderBar.getMeasuredWidth() != taskViewWidth ||
                             mHeaderBar.getMeasuredHeight() != mTaskBarHeight) {