OneGrid Grid Option Updates

Bug: 330900048
Flag: com.android.launcher3.one_grid_specs
Test: n/a
Change-Id: I919195dbc7ac78c3be42f0f9d7620193a24d7e99
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index ece6540..7acba75 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -58,9 +58,9 @@
 import com.android.launcher3.testing.shared.ResourceUtils;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.DisplayController.Info;
-import com.android.launcher3.util.LockedUserState;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.Partner;
+import com.android.launcher3.util.ResourceHelper;
 import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.WindowBounds;
 import com.android.launcher3.util.window.WindowManagerProxy;
@@ -186,6 +186,8 @@
     @XmlRes
     public int workspaceSpecsId = INVALID_RESOURCE_HANDLE;
     @XmlRes
+    public int rowCountSpecsId = INVALID_RESOURCE_HANDLE;;
+    @XmlRes
     public int workspaceSpecsTwoPanelId = INVALID_RESOURCE_HANDLE;
     @XmlRes
     public int allAppsSpecsId = INVALID_RESOURCE_HANDLE;
@@ -232,8 +234,6 @@
         if (!newGridName.equals(gridName)) {
             LauncherPrefs.get(context).put(GRID_NAME, newGridName);
         }
-        LockedUserState.get(context).runOnUserUnlocked(() ->
-            new DeviceGridState(this).writeToPrefs(context));
 
         DisplayController.INSTANCE.get(context).setPriorityListener(
                 (displayContext, info, flags) -> {
@@ -340,15 +340,29 @@
         Info displayInfo = DisplayController.INSTANCE.get(context).getInfo();
         @DeviceType int deviceType = displayInfo.getDeviceType();
 
-        ArrayList<DisplayOption> allOptions =
+        List<DisplayOption> allOptions =
                 getPredefinedDeviceProfiles(context, gridName, deviceType,
                         RestoreDbTask.isPending(context));
+
+        // Filter out options that don't have the same number of columns as the grid
+        DeviceGridState deviceGridState = new DeviceGridState(context);
+        List<DisplayOption> allOptionsFilteredByColCount =
+                filterByColumnCount(allOptions, deviceGridState.getColumns());
+
         DisplayOption displayOption =
-                invDistWeightedInterpolate(displayInfo, allOptions, deviceType);
+                invDistWeightedInterpolate(displayInfo, allOptionsFilteredByColCount.isEmpty()
+                        ? new ArrayList<>(allOptions)
+                        : new ArrayList<>(allOptionsFilteredByColCount), deviceType);
         initGrid(context, displayInfo, displayOption, deviceType);
         return displayOption.grid.name;
     }
 
+    private List<DisplayOption> filterByColumnCount(
+            List<DisplayOption> allOptions, int numColumns) {
+        return allOptions.stream().filter(
+                option -> option.grid.numColumns == numColumns).toList();
+    }
+
     /**
      * @deprecated This is a temporary solution because on the backup and restore case we modify the
      * IDP, this resets it. b/332974074
@@ -383,6 +397,7 @@
         isScalable = closestProfile.isScalable;
         devicePaddingId = closestProfile.devicePaddingId;
         workspaceSpecsId = closestProfile.mWorkspaceSpecsId;
+        rowCountSpecsId = closestProfile.mRowCountSpecsId;
         workspaceSpecsTwoPanelId = closestProfile.mWorkspaceSpecsTwoPanelId;
         allAppsSpecsId = closestProfile.mAllAppsSpecsId;
         allAppsSpecsTwoPanelId = closestProfile.mAllAppsSpecsTwoPanelId;
@@ -495,7 +510,6 @@
         mChangeListeners.remove(listener);
     }
 
-
     public void setCurrentGrid(Context context, String gridName) {
         LauncherPrefs.get(context).put(GRID_NAME, gridName);
         MAIN_EXECUTOR.execute(() -> {
@@ -526,7 +540,7 @@
         }
     }
 
-    private static ArrayList<DisplayOption> getPredefinedDeviceProfiles(Context context,
+    private List<DisplayOption> getPredefinedDeviceProfiles(Context context,
             String gridName, @DeviceType int deviceType, boolean allowDisabledGrid) {
         ArrayList<DisplayOption> profiles = new ArrayList<>();
 
@@ -539,7 +553,8 @@
                         && GridOption.TAG_NAME.equals(parser.getName())) {
 
                     GridOption gridOption = new GridOption(context, Xml.asAttributeSet(parser));
-                    if (gridOption.isEnabled(deviceType) || allowDisabledGrid) {
+                    if ((gridOption.isEnabled(deviceType) || allowDisabledGrid)
+                            && (Flags.oneGridSpecs() == gridOption.isNewGridOption())) {
                         final int displayDepth = parser.getDepth();
                         while (((type = parser.next()) != XmlPullParser.END_TAG
                                 || parser.getDepth() > displayDepth)
@@ -566,13 +581,16 @@
                 }
             }
         }
-        if (filteredProfiles.isEmpty()) {
-            // No grid found, use the default options
+        if (filteredProfiles.isEmpty() && TextUtils.isEmpty(gridName)) {
+            // Use the default options since gridName is empty and there's no valid grids.
             for (DisplayOption option : profiles) {
                 if (option.canBeDefault) {
                     filteredProfiles.add(option);
                 }
             }
+        } else if (filteredProfiles.isEmpty()) {
+            // In this case we had a grid selected but we couldn't find it.
+            filteredProfiles.addAll(profiles);
         }
         if (filteredProfiles.isEmpty()) {
             throw new RuntimeException("No display option with canBeDefault=true");
@@ -581,6 +599,72 @@
     }
 
     /**
+     * Parses through the xml to find NumRows specs. Then calls findBestRowCount to get the correct
+     * row count for this GridOption.
+     *
+     * @return the result of {@link #findBestRowCount(List, Context, int)}.
+     */
+    public static NumRows getRowCount(ResourceHelper resourceHelper, Context context,
+            int deviceType) {
+        ArrayList<NumRows> rowCounts = new ArrayList<>();
+
+        try (XmlResourceParser parser = resourceHelper.getXml()) {
+            final int depth = parser.getDepth();
+            int type;
+            while (((type = parser.next()) != XmlPullParser.END_TAG
+                    || parser.getDepth() > depth) && type != XmlPullParser.END_DOCUMENT) {
+                if ((type == XmlPullParser.START_TAG)
+                        && "NumRows".equals(parser.getName())) {
+                    rowCounts.add(new NumRows(context, Xml.asAttributeSet(parser)));
+                }
+            }
+        } catch (IOException | XmlPullParserException e) {
+            throw new RuntimeException(e);
+        }
+
+        return findBestRowCount(rowCounts, context, deviceType);
+    }
+
+    /**
+     * @return the biggest row count that fits the display dimensions spec using NumRows to
+     * determine that. If no best row count is found, return -1.
+     */
+    public static NumRows findBestRowCount(List<NumRows> list, Context context,
+            @DeviceType int deviceType) {
+        Info displayInfo = DisplayController.INSTANCE.get(context).getInfo();
+        int minWidthPx = Integer.MAX_VALUE;
+        int minHeightPx = Integer.MAX_VALUE;
+        for (WindowBounds bounds : displayInfo.supportedBounds) {
+            boolean isTablet = displayInfo.isTablet(bounds);
+            if (isTablet && deviceType == TYPE_MULTI_DISPLAY) {
+                // For split displays, take half width per page
+                minWidthPx = Math.min(minWidthPx, bounds.availableSize.x / 2);
+                minHeightPx = Math.min(minHeightPx, bounds.availableSize.y);
+            } else if (!isTablet && bounds.isLandscape()) {
+                // We will use transposed layout in this case
+                minWidthPx = Math.min(minWidthPx, bounds.availableSize.y);
+                minHeightPx = Math.min(minHeightPx, bounds.availableSize.x);
+            } else {
+                minWidthPx = Math.min(minWidthPx, bounds.availableSize.x);
+                minHeightPx = Math.min(minHeightPx, bounds.availableSize.y);
+            }
+        }
+
+        NumRows selectedRow = null;
+        for (NumRows item: list) {
+            if (minWidthPx >= item.mMinDeviceWidthPx && minHeightPx >= item.mMinDeviceHeightPx) {
+                if (selectedRow == null || selectedRow.mNumRowsNew < item.mNumRowsNew) {
+                    selectedRow = item;
+                }
+            }
+        }
+        if (selectedRow != null) {
+            return selectedRow;
+        }
+        return null;
+    }
+
+    /**
      * Returns the GridOption associated to the given file name or null if the fileName is not
      * supported.
      * Ej, launcher.db -> "normal grid", launcher_4_by_4.db -> "practical grid"
@@ -626,6 +710,7 @@
         return parseAllDefinedGridOptions(context)
                 .stream()
                 .filter(go -> go.isEnabled(deviceType))
+                .filter(go -> (Flags.oneGridSpecs() == go.isNewGridOption()))
                 .collect(Collectors.toList());
     }
 
@@ -709,7 +794,7 @@
     }
 
     private static DisplayOption invDistWeightedInterpolate(
-            Info displayInfo, ArrayList<DisplayOption> points, @DeviceType int deviceType) {
+            Info displayInfo, List<DisplayOption> points, @DeviceType int deviceType) {
         int minWidthPx = Integer.MAX_VALUE;
         int minHeightPx = Integer.MAX_VALUE;
         for (WindowBounds bounds : displayInfo.supportedBounds) {
@@ -733,7 +818,7 @@
         float height = dpiFromPx(minHeightPx, displayInfo.getDensityDpi());
 
         // Sort the profiles based on the closeness to the device size
-        Collections.sort(points, (a, b) ->
+        points.sort((a, b) ->
                 Float.compare(dist(width, height, a.minWidthDps, a.minHeightDps),
                         dist(width, height, b.minWidthDps, b.minHeightDps)));
 
@@ -855,6 +940,7 @@
         private static final int DONT_INLINE_QSB = 0;
 
         public final String name;
+        public final String title;
         public final int numRows;
         public final int numColumns;
         public final int numSearchContainerColumns;
@@ -894,17 +980,30 @@
         private final int mWorkspaceCellSpecsTwoPanelId;
         private final int mAllAppsCellSpecsId;
         private final int mAllAppsCellSpecsTwoPanelId;
+        private final int mRowCountSpecsId;
 
         public GridOption(Context context, AttributeSet attrs) {
             TypedArray a = context.obtainStyledAttributes(
                     attrs, R.styleable.GridDisplayOption);
             name = a.getString(R.styleable.GridDisplayOption_name);
-            numRows = a.getInt(R.styleable.GridDisplayOption_numRows, 0);
+            title = a.getString(R.styleable.GridDisplayOption_title);
+            deviceCategory = a.getInt(R.styleable.GridDisplayOption_deviceCategory,
+                    DEVICE_CATEGORY_ALL);
+            mRowCountSpecsId = a.getResourceId(
+                    R.styleable.GridDisplayOption_rowCountSpecsId, INVALID_RESOURCE_HANDLE);
+            if (mRowCountSpecsId != INVALID_RESOURCE_HANDLE) {
+                ResourceHelper resourceHelper = new ResourceHelper(context, mRowCountSpecsId);
+                NumRows numR = getRowCount(resourceHelper, context, deviceCategory);
+                numRows = numR.mNumRowsNew;
+                dbFile = numR.mDbFile;
+            } else {
+                numRows = a.getInt(R.styleable.GridDisplayOption_numRows, 0);
+                dbFile = a.getString(R.styleable.GridDisplayOption_dbFile);
+            }
+
             numColumns = a.getInt(R.styleable.GridDisplayOption_numColumns, 0);
             numSearchContainerColumns = a.getInt(
                     R.styleable.GridDisplayOption_numSearchContainerColumns, numColumns);
-
-            dbFile = a.getString(R.styleable.GridDisplayOption_dbFile);
             defaultLayoutId = a.getResourceId(
                     R.styleable.GridDisplayOption_defaultLayoutId, 0);
             demoModeLayoutId = a.getResourceId(
@@ -969,8 +1068,6 @@
                     R.styleable.GridDisplayOption_isScalable, false);
             devicePaddingId = a.getResourceId(
                     R.styleable.GridDisplayOption_devicePaddingId, INVALID_RESOURCE_HANDLE);
-            deviceCategory = a.getInt(R.styleable.GridDisplayOption_deviceCategory,
-                    DEVICE_CATEGORY_ALL);
 
             if (FeatureFlags.enableResponsiveWorkspace()) {
                 mWorkspaceSpecsId = a.getResourceId(
@@ -1053,6 +1150,28 @@
                     return false;
             }
         }
+
+        public boolean isNewGridOption() {
+            return mRowCountSpecsId != INVALID_RESOURCE_HANDLE;
+        }
+    }
+
+    public static final class NumRows {
+        final int mNumRowsNew;
+        final float mMinDeviceWidthPx;
+        final float mMinDeviceHeightPx;
+        final String mDbFile;
+
+        NumRows(Context context, AttributeSet attrs) {
+            TypedArray a = context.obtainStyledAttributes(attrs, R.styleable.NumRows);
+
+            mNumRowsNew = (int) a.getFloat(R.styleable.NumRows_numRowsNew, 0);
+            mMinDeviceWidthPx = a.getFloat(R.styleable.NumRows_minDeviceWidthPx, 0);
+            mMinDeviceHeightPx = a.getFloat(R.styleable.NumRows_minDeviceHeightPx, 0);
+            mDbFile = a.getString(R.styleable.NumRows_dbFile);
+
+            a.recycle();
+        }
     }
 
     @VisibleForTesting
diff --git a/src/com/android/launcher3/LauncherFiles.java b/src/com/android/launcher3/LauncherFiles.java
index d730cea..1148f79 100644
--- a/src/com/android/launcher3/LauncherFiles.java
+++ b/src/com/android/launcher3/LauncherFiles.java
@@ -18,6 +18,8 @@
     public static final String LAUNCHER_DB = "launcher.db";
     public static final String LAUNCHER_6_BY_5_DB = "launcher_6_by_5.db";
     public static final String LAUNCHER_4_BY_5_DB = "launcher_4_by_5.db";
+    public static final String LAUNCHER_4_BY_6_DB = "launcher_4_by_6.db";
+    public static final String LAUNCHER_5_BY_6_DB = "launcher_5_by_6.db";
     public static final String LAUNCHER_4_BY_4_DB = "launcher_4_by_4.db";
     public static final String LAUNCHER_3_BY_3_DB = "launcher_3_by_3.db";
     public static final String LAUNCHER_2_BY_2_DB = "launcher_2_by_2.db";
@@ -35,6 +37,8 @@
             LAUNCHER_DB,
             LAUNCHER_6_BY_5_DB,
             LAUNCHER_4_BY_5_DB,
+            LAUNCHER_4_BY_6_DB,
+            LAUNCHER_5_BY_6_DB,
             LAUNCHER_4_BY_4_DB,
             LAUNCHER_3_BY_3_DB,
             LAUNCHER_2_BY_2_DB));
diff --git a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
index 259e543..a5bcd0f 100644
--- a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
+++ b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
@@ -73,6 +73,7 @@
     private static final String TAG = "GridCustomizationsProvider";
 
     private static final String KEY_NAME = "name";
+    private static final String KEY_GRID_TITLE = "grid_title";
     private static final String KEY_ROWS = "rows";
     private static final String KEY_COLS = "cols";
     private static final String KEY_PREVIEW_COUNT = "preview_count";
@@ -117,6 +118,7 @@
                 for (GridOption gridOption : idp.parseAllGridOptions(getContext())) {
                     cursor.newRow()
                             .add(KEY_NAME, gridOption.name)
+                            .add(KEY_GRID_TITLE, gridOption.title)
                             .add(KEY_ROWS, gridOption.numRows)
                             .add(KEY_COLS, gridOption.numColumns)
                             .add(KEY_PREVIEW_COUNT, 1)
diff --git a/src/com/android/launcher3/model/DeviceGridState.java b/src/com/android/launcher3/model/DeviceGridState.java
index 729b381..90af215 100644
--- a/src/com/android/launcher3/model/DeviceGridState.java
+++ b/src/com/android/launcher3/model/DeviceGridState.java
@@ -156,10 +156,16 @@
     }
 
     public Integer getColumns() {
+        if (TextUtils.isEmpty(mGridSizeString)) {
+            return -1;
+        }
         return Integer.parseInt(String.valueOf(mGridSizeString.split(",")[0]));
     }
 
     public Integer getRows() {
+        if (TextUtils.isEmpty(mGridSizeString)) {
+            return -1;
+        }
         return Integer.parseInt(String.valueOf(mGridSizeString.split(",")[1]));
     }
 
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
index 8c3e860..094798b 100644
--- a/src/com/android/launcher3/model/ModelDbController.java
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -20,6 +20,7 @@
 import static android.util.Base64.NO_WRAP;
 
 import static com.android.launcher3.DefaultLayoutParser.RES_PARTNER_DEFAULT_LAYOUT;
+import static com.android.launcher3.LauncherPrefs.DB_FILE;
 import static com.android.launcher3.LauncherPrefs.NO_DB_FILES_RESTORED;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE;
@@ -128,16 +129,20 @@
 
     private synchronized void createDbIfNotExists() {
         if (mOpenHelper == null) {
-            mOpenHelper = createDatabaseHelper(false /* forMigration */);
+            String dbFile = LauncherPrefs.get(mContext).get(DB_FILE);
+            if (dbFile.isEmpty()) {
+                dbFile = InvariantDeviceProfile.INSTANCE.get(mContext).dbFile;
+            }
+            mOpenHelper = createDatabaseHelper(false /* forMigration */, dbFile);
             printDBs("before: ");
             RestoreDbTask.restoreIfNeeded(mContext, this);
             printDBs("after: ");
         }
     }
 
-    protected DatabaseHelper createDatabaseHelper(boolean forMigration) {
+    protected DatabaseHelper createDatabaseHelper(boolean forMigration, String dbFile) {
         boolean isSandbox = mContext instanceof SandboxContext;
-        String dbName = isSandbox ? null : InvariantDeviceProfile.INSTANCE.get(mContext).dbFile;
+        String dbName = isSandbox ? null : dbFile;
 
         // Set the flag for empty DB
         Runnable onEmptyDbCreateCallback = forMigration ? () -> { }
@@ -364,7 +369,7 @@
         InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
         DatabaseHelper oldHelper = mOpenHelper;
         mOpenHelper = (mContext instanceof SandboxContext) ? oldHelper
-                : createDatabaseHelper(true /* forMigration */);
+                : createDatabaseHelper(true, new DeviceGridState(idp).getDbFile());
         try {
             // This is the current grid we have, given by the mContext
             DeviceGridState srcDeviceState = new DeviceGridState(mContext);
@@ -388,7 +393,6 @@
      * Migrates the DB if needed. If the migration failed, it clears the DB.
      */
     public void tryMigrateDB(@Nullable LauncherRestoreEventLogger restoreEventLogger) {
-
         if (!migrateGridIfNeeded()) {
             if (restoreEventLogger != null) {
                 if (LauncherPrefs.get(mContext).get(NO_DB_FILES_RESTORED)) {
@@ -438,7 +442,7 @@
         }
         DatabaseHelper oldHelper = mOpenHelper;
         mOpenHelper = (mContext instanceof SandboxContext) ? oldHelper
-                : createDatabaseHelper(true /* forMigration */);
+                : createDatabaseHelper(true /* forMigration */, targetDbName);
         try {
             // This is the current grid we have, given by the mContext
             DeviceGridState srcDeviceState = new DeviceGridState(mContext);
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 775d248..59c27af 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -125,7 +125,8 @@
         LauncherPrefs.get(context).removeSync(RESTORE_DEVICE);
 
         if (Flags.enableNarrowGridRestore()) {
-            String oldPhoneFileName = idp.dbFile;
+            DeviceGridState deviceGridState = new DeviceGridState(context);
+            String oldPhoneFileName = deviceGridState.getDbFile();
             List<String> previousDbs = existingDbs(context);
             removeOldDBs(context, oldPhoneFileName);
             // The idp before this contains data about the old phone, after this it becomes the idp
@@ -148,6 +149,7 @@
                 context, oldPhoneDbFileName);
         // The grid option could be null if current phone doesn't support the previous db.
         if (oldPhoneGridOption != null) {
+
             /* If the user only used the default db on the previous phone and the new default db is
              * bigger than or equal to the previous one, then keep the new default db */
             if (previousDbs.size() == 1 && oldPhoneGridOption.numColumns <= idp.numColumns