Fix layout for multi-element rows (e.g. Screenshots).

Demo with standard 3 element row and overflowing 2 element row:
https://drive.google.com/drive/folders/1o1gvD3mOxFmVW2vqjn3D1Gukz-ie_kmW?resourcekey=0-bHRsdqpUWITEUo39tAMJDg&usp=sharing

Test: Manually with 1, 2, 3, and 4 length rows of 3 screenshots
and putting web suggestions into 2 columns.
Fix: 247837336
Bug: 239927522

Change-Id: I391beb7a24f189b5e55cb7b65c5ffc8349c83cfc
diff --git a/src/com/android/launcher3/allapps/AllAppsGridAdapter.java b/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
index e0b4951..368a373 100644
--- a/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
+++ b/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
@@ -42,25 +42,29 @@
         BaseAllAppsAdapter<T> {
 
     public static final String TAG = "AppsGridAdapter";
-    private final GridLayoutManager mGridLayoutMgr;
-    private final GridSpanSizer mGridSizer;
+    private final AppsGridLayoutManager mGridLayoutMgr;
 
     public AllAppsGridAdapter(T activityContext, LayoutInflater inflater,
             AlphabeticalAppsList apps, BaseAdapterProvider[] adapterProviders) {
         super(activityContext, inflater, apps, adapterProviders);
-        mGridSizer = new GridSpanSizer();
         mGridLayoutMgr = new AppsGridLayoutManager(mActivityContext);
-        mGridLayoutMgr.setSpanSizeLookup(mGridSizer);
+        mGridLayoutMgr.setSpanSizeLookup(new GridSpanSizer());
         setAppsPerRow(activityContext.getDeviceProfile().numShownAllAppsColumns);
     }
 
     /**
      * Returns the grid layout manager.
      */
-    public RecyclerView.LayoutManager getLayoutManager() {
+    public AppsGridLayoutManager getLayoutManager() {
         return mGridLayoutMgr;
     }
 
+    /** @return the column index that the given adapter index falls. */
+    public int getSpanIndex(int adapterIndex) {
+        AppsGridLayoutManager lm = getLayoutManager();
+        return lm.getSpanSizeLookup().getSpanIndex(adapterIndex, lm.getSpanCount());
+    }
+
     /**
      * A subclass of GridLayoutManager that overrides accessibility values during app search.
      */
diff --git a/src/com/android/launcher3/allapps/SearchTransitionController.java b/src/com/android/launcher3/allapps/SearchTransitionController.java
index 8fc7965..a1f5bc6 100644
--- a/src/com/android/launcher3/allapps/SearchTransitionController.java
+++ b/src/com/android/launcher3/allapps/SearchTransitionController.java
@@ -29,6 +29,7 @@
 import android.animation.ObjectAnimator;
 import android.graphics.drawable.Drawable;
 import android.util.FloatProperty;
+import android.util.Log;
 import android.view.View;
 import android.view.animation.Interpolator;
 
@@ -36,10 +37,13 @@
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
 
 /** Coordinates the transition between Search and A-Z in All Apps. */
 public class SearchTransitionController {
 
+    private static final String LOG_TAG = "SearchTransitionCtrl";
+
     // Interpolator when the user taps the QSB while already in All Apps.
     private static final Interpolator INTERPOLATOR_WITHIN_ALL_APPS = DEACCEL_1_7;
     // Interpolator when the user taps the QSB from home screen, so transition to all apps is
@@ -171,6 +175,7 @@
         int appRowHeight = 0;
         Integer top = null;
         SearchRecyclerView searchRecyclerView = getSearchRecyclerView();
+
         for (int i = 0; i < searchRecyclerView.getChildCount(); i++) {
             View searchResultView = searchRecyclerView.getChildAt(i);
             if (searchResultView == null) {
@@ -226,15 +231,43 @@
             float scaleY = 1 - mSearchToAzProgress;
             int scaledHeight = (int) (searchResultView.getHeight() * scaleY);
             searchResultView.setScaleY(scaleY);
-            searchResultView.setY(top + totalHeight);
 
-            numSearchResultsAnimated++;
-            totalHeight += scaledHeight;
+            // For rows with multiple elements, only count the height once and translate elements to
+            // the same y position.
+            int y = top + totalHeight;
+            int spanIndex = getSpanIndex(searchRecyclerView, adapterPosition);
+            if (spanIndex > 0) {
+                // Continuation of an existing row; move this item into the row.
+                y -= scaledHeight;
+            } else {
+                // Start of a new row contributes to total height and animation stagger.
+                numSearchResultsAnimated++;
+                totalHeight += scaledHeight;
+            }
+            searchResultView.setY(y);
         }
 
         return totalHeight - appRowHeight;
     }
 
+    /** @return the column that the view at this position is found (0 assumed if indeterminate). */
+    private int getSpanIndex(SearchRecyclerView searchRecyclerView, int adapterPosition) {
+        if (adapterPosition == NO_POSITION) {
+            Log.w(LOG_TAG, "Can't determine span index - child not found in adapter");
+            return 0;
+        }
+        if (!(searchRecyclerView.getAdapter() instanceof AllAppsGridAdapter<?>)) {
+            Log.e(LOG_TAG, "Search RV doesn't have an AllAppsGridAdapter?");
+            // This case shouldn't happen, but for debug devices we will continue to create a more
+            // visible crash.
+            if (!Utilities.IS_DEBUG_DEVICE) {
+                return 0;
+            }
+        }
+        AllAppsGridAdapter<?> adapter = (AllAppsGridAdapter<?>) searchRecyclerView.getAdapter();
+        return adapter.getSpanIndex(adapterPosition);
+    }
+
     /** Called just before a child is attached to the SearchRecyclerView. */
     private void onSearchChildAttached(View child) {
         // Avoid allocating hardware layers for alpha changes.