Improve workspace loading times.

Updated loadWorkspace to load all required icons in a series of bulk sql queries. This reduces the cost of SQL lookups (up to two lookups per user, rather than one lookup per icon)

Bug: 195674813
Test: Added all icons to workspace, added duplicate icons, added icons for same component name from different users

Change-Id: I56afaa04e7c7701f0d3c86b31c53f578dfa73fe6
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index f091262..382f7a7 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -147,6 +147,11 @@
     public static final BooleanFlag ENABLE_THEMED_ICONS = getDebugFlag(
             "ENABLE_THEMED_ICONS", true, "Enable themed icons on workspace");
 
+    public static final BooleanFlag ENABLE_BULK_WORKSPACE_ICON_LOADING = getDebugFlag(
+            "ENABLE_BULK_WORKSPACE_ICON_LOADING",
+            false,
+            "Enable loading workspace icons in bulk.");
+
     // Keep as DeviceFlag for remote disable in emergency.
     public static final BooleanFlag ENABLE_OVERVIEW_SELECTIONS = new DeviceFlag(
             "ENABLE_OVERVIEW_SELECTIONS", true, "Show Select Mode button in Overview Actions");
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index 1a468ae..60d6e83 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -19,6 +19,8 @@
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
+import static java.util.stream.Collectors.groupingBy;
+
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
@@ -30,10 +32,15 @@
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManager.NameNotFoundException;
 import android.content.pm.ShortcutInfo;
+import android.database.Cursor;
+import android.database.sqlite.SQLiteException;
 import android.graphics.drawable.Drawable;
 import android.os.Process;
+import android.os.Trace;
 import android.os.UserHandle;
+import android.text.TextUtils;
 import android.util.Log;
+import android.util.Pair;
 
 import androidx.annotation.NonNull;
 
@@ -47,6 +54,7 @@
 import com.android.launcher3.icons.cache.CachingLogic;
 import com.android.launcher3.icons.cache.HandlerRunnable;
 import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.model.data.IconRequestInfo;
 import com.android.launcher3.model.data.ItemInfoWithIcon;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
@@ -56,8 +64,13 @@
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.Preconditions;
 
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
+import java.util.stream.Stream;
 
 /**
  * Cache of application icons.  Icons can be made from any thread.
@@ -306,6 +319,87 @@
         applyCacheEntry(entry, infoInOut);
     }
 
+    /**
+     * Creates an sql cursor for a query of a set of ItemInfoWithIcon icons and titles.
+     *
+     * @param iconRequestInfos List of IconRequestInfos representing titles and icons to query.
+     * @param user UserHandle all the given iconRequestInfos share
+     * @param useLowResIcons whether we should exclude the icon column from the sql results.
+     */
+    private <T extends ItemInfoWithIcon> Cursor createBulkQueryCursor(
+            List<IconRequestInfo<T>> iconRequestInfos, UserHandle user, boolean useLowResIcons)
+            throws SQLiteException {
+        String[] queryParams = Stream.concat(
+                iconRequestInfos.stream()
+                        .map(r -> r.itemInfo.getTargetComponent())
+                        .filter(Objects::nonNull)
+                        .distinct()
+                        .map(ComponentName::flattenToString),
+                Stream.of(Long.toString(getSerialNumberForUser(user)))).toArray(String[]::new);
+        String componentNameQuery = TextUtils.join(
+                ",", Collections.nCopies(queryParams.length - 1, "?"));
+
+        return mIconDb.query(
+                useLowResIcons ? IconDB.COLUMNS_LOW_RES : IconDB.COLUMNS_HIGH_RES,
+                IconDB.COLUMN_COMPONENT
+                        + " IN ( " + componentNameQuery + " )"
+                        + " AND " + IconDB.COLUMN_USER + " = ?",
+                queryParams);
+    }
+
+    /**
+     * Load and fill icons requested in iconRequestInfos using a single bulk sql query.
+     */
+    public synchronized <T extends ItemInfoWithIcon> void getTitlesAndIconsInBulk(
+            List<IconRequestInfo<T>> iconRequestInfos) {
+        Map<Pair<UserHandle, Boolean>, List<IconRequestInfo<T>>> iconLoadSubsectionsMap =
+                iconRequestInfos.stream()
+                        .collect(groupingBy(iconRequest ->
+                                Pair.create(iconRequest.itemInfo.user, iconRequest.useLowResIcon)));
+
+        Trace.beginSection("loadIconsInBulk");
+        iconLoadSubsectionsMap.forEach((sectionKey, filteredList) -> {
+            Map<ComponentName, List<IconRequestInfo<T>>> duplicateIconRequestsMap =
+                    filteredList.stream()
+                            .collect(groupingBy(iconRequest ->
+                                    iconRequest.itemInfo.getTargetComponent()));
+
+            Trace.beginSection("loadIconSubsectionInBulk");
+            try (Cursor c = createBulkQueryCursor(
+                    filteredList,
+                    /* user = */ sectionKey.first,
+                    /* useLowResIcons = */ sectionKey.second)) {
+                int componentNameColumnIndex = c.getColumnIndexOrThrow(IconDB.COLUMN_COMPONENT);
+                while (c.moveToNext()) {
+                    ComponentName cn = ComponentName.unflattenFromString(
+                            c.getString(componentNameColumnIndex));
+                    List<IconRequestInfo<T>> duplicateIconRequests =
+                            duplicateIconRequestsMap.get(cn);
+
+                    if (cn != null) {
+                        CacheEntry entry = cacheLocked(
+                                cn,
+                                /* user = */ sectionKey.first,
+                                () -> duplicateIconRequests.get(0).launcherActivityInfo,
+                                mLauncherActivityInfoCachingLogic,
+                                c,
+                                /* usePackageIcon= */ false,
+                                /* useLowResIcons = */ sectionKey.second);
+
+                        for (IconRequestInfo<T> iconRequest : duplicateIconRequests) {
+                            applyCacheEntry(entry, iconRequest.itemInfo);
+                        }
+                    }
+                }
+            } catch (SQLiteException e) {
+                Log.d(TAG, "Error reading icon cache", e);
+            } finally {
+                Trace.endSection();
+            }
+        });
+        Trace.endSection();
+    }
+
 
     /**
      * Fill in {@param infoInOut} with the corresponding icon and label.
diff --git a/src/com/android/launcher3/model/LoaderCursor.java b/src/com/android/launcher3/model/LoaderCursor.java
index 7e3bcee..8a5a9bf 100644
--- a/src/com/android/launcher3/model/LoaderCursor.java
+++ b/src/com/android/launcher3/model/LoaderCursor.java
@@ -16,13 +16,10 @@
 
 package com.android.launcher3.model;
 
-import static android.graphics.BitmapFactory.decodeByteArray;
-
 import android.content.ComponentName;
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.Intent;
-import android.content.Intent.ShortcutIconResource;
 import android.content.pm.LauncherActivityInfo;
 import android.content.pm.LauncherApps;
 import android.content.pm.PackageManager;
@@ -45,11 +42,10 @@
 import com.android.launcher3.Utilities;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.config.FeatureFlags;
-import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.IconCache;
-import com.android.launcher3.icons.LauncherIcons;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.model.data.IconRequestInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.shortcuts.ShortcutKey;
@@ -184,32 +180,21 @@
      * Loads the icon from the cursor and updates the {@param info} if the icon is an app resource.
      */
     protected boolean loadIcon(WorkspaceItemInfo info) {
-        try (LauncherIcons li = LauncherIcons.obtain(mContext)) {
-            if (itemType == LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT) {
-                String packageName = getString(iconPackageIndex);
-                String resourceName = getString(iconResourceIndex);
-                if (!TextUtils.isEmpty(packageName) || !TextUtils.isEmpty(resourceName)) {
-                    info.iconResource = new ShortcutIconResource();
-                    info.iconResource.packageName = packageName;
-                    info.iconResource.resourceName = resourceName;
-                    BitmapInfo iconInfo = li.createIconBitmap(info.iconResource);
-                    if (iconInfo != null) {
-                        info.bitmap = iconInfo;
-                        return true;
-                    }
-                }
-            }
+        return createIconRequestInfo(info, false).loadWorkspaceIcon(mContext);
+    }
 
-            // Failed to load from resource, try loading from DB.
-            byte[] data = getBlob(iconIndex);
-            try {
-                info.bitmap = li.createIconBitmap(decodeByteArray(data, 0, data.length));
-                return true;
-            } catch (Exception e) {
-                Log.e(TAG, "Failed to decode byte array for info " + info, e);
-                return false;
-            }
-        }
+    public IconRequestInfo<WorkspaceItemInfo> createIconRequestInfo(
+            WorkspaceItemInfo wai, boolean useLowResIcon) {
+        String packageName = itemType == LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT
+                ? getString(iconPackageIndex) : null;
+        String resourceName = itemType == LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT
+                ? getString(iconResourceIndex) : null;
+        byte[] iconBlob = itemType == LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT
+                || restoreFlag != 0
+                ? getBlob(iconIndex) : null;
+
+        return new IconRequestInfo<>(
+                wai, mActivityInfo, packageName, resourceName, iconBlob, useLowResIcon);
     }
 
     /**
@@ -262,6 +247,11 @@
      */
     public WorkspaceItemInfo getAppShortcutInfo(
             Intent intent, boolean allowMissingTarget, boolean useLowResIcon) {
+        return getAppShortcutInfo(intent, allowMissingTarget, useLowResIcon, true);
+    }
+
+    public WorkspaceItemInfo getAppShortcutInfo(
+            Intent intent, boolean allowMissingTarget, boolean useLowResIcon, boolean loadIcon) {
         if (user == null) {
             Log.d(TAG, "Null user found in getShortcutInfo");
             return null;
@@ -288,9 +278,11 @@
         info.user = user;
         info.intent = newIntent;
 
-        mIconCache.getTitleAndIcon(info, mActivityInfo, useLowResIcon);
-        if (mIconCache.isDefaultIcon(info.bitmap, user)) {
-            loadIcon(info);
+        if (loadIcon) {
+            mIconCache.getTitleAndIcon(info, mActivityInfo, useLowResIcon);
+            if (mIconCache.isDefaultIcon(info.bitmap, user)) {
+                loadIcon(info);
+            }
         }
 
         if (mActivityInfo != null) {
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index c178b02..1249606 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -72,6 +72,7 @@
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.FolderInfo;
+import com.android.launcher3.model.data.IconRequestInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.ItemInfoWithIcon;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
@@ -420,6 +421,7 @@
                 LauncherAppWidgetProviderInfo widgetProviderInfo;
                 Intent intent;
                 String targetPkg;
+                List<IconRequestInfo<WorkspaceItemInfo>> iconRequestInfos = new ArrayList<>();
 
                 while (!mStopped && c.moveToNext()) {
                     try {
@@ -542,7 +544,10 @@
                             } else if (c.itemType ==
                                     LauncherSettings.Favorites.ITEM_TYPE_APPLICATION) {
                                 info = c.getAppShortcutInfo(
-                                        intent, allowMissingTarget, useLowResIcon);
+                                        intent,
+                                        allowMissingTarget,
+                                        useLowResIcon,
+                                        !FeatureFlags.ENABLE_BULK_WORKSPACE_ICON_LOADING.get());
                             } else if (c.itemType ==
                                     LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT) {
 
@@ -594,6 +599,8 @@
                             }
 
                             if (info != null) {
+                                iconRequestInfos.add(c.createIconRequestInfo(info, useLowResIcon));
+
                                 c.applyCommonProperties(info);
 
                                 info.intent = intent;
@@ -811,6 +818,21 @@
                         Log.e(TAG, "Desktop items loading interrupted", e);
                     }
                 }
+                if (FeatureFlags.ENABLE_BULK_WORKSPACE_ICON_LOADING.get()) {
+                    Trace.beginSection("LoadWorkspaceIconsInBulk");
+                    try {
+                        mIconCache.getTitlesAndIconsInBulk(iconRequestInfos);
+                        for (IconRequestInfo<WorkspaceItemInfo> iconRequestInfo :
+                                iconRequestInfos) {
+                            WorkspaceItemInfo wai = iconRequestInfo.itemInfo;
+                            if (mIconCache.isDefaultIcon(wai.bitmap, wai.user)) {
+                                iconRequestInfo.loadWorkspaceIcon(mApp.getContext());
+                            }
+                        }
+                    } finally {
+                        Trace.endSection();
+                    }
+                }
             } finally {
                 IOUtils.closeSilently(c);
             }
diff --git a/src/com/android/launcher3/model/data/IconRequestInfo.java b/src/com/android/launcher3/model/data/IconRequestInfo.java
new file mode 100644
index 0000000..2f566f6
--- /dev/null
+++ b/src/com/android/launcher3/model/data/IconRequestInfo.java
@@ -0,0 +1,101 @@
+/*
+ * Copyright (C) 2021 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.launcher3.model.data;
+
+import static android.graphics.BitmapFactory.decodeByteArray;
+
+import android.content.Context;
+import android.content.Intent;
+import android.content.pm.LauncherActivityInfo;
+import android.text.TextUtils;
+import android.util.Log;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.icons.BitmapInfo;
+import com.android.launcher3.icons.LauncherIcons;
+
+/**
+ * Class representing one request for an icon to be queried in a sql database.
+ *
+ * @param <T> ItemInfoWithIcon subclass whose title and icon can be loaded and filled by an sql
+ *           query.
+ */
+public class IconRequestInfo<T extends ItemInfoWithIcon> {
+
+    private static final String TAG = "IconRequestInfo";
+
+    @NonNull public final T itemInfo;
+    @Nullable public final LauncherActivityInfo launcherActivityInfo;
+    @Nullable public final String packageName;
+    @Nullable public final String resourceName;
+    @Nullable public final byte[] iconBlob;
+    public final boolean useLowResIcon;
+
+    public IconRequestInfo(
+            @NonNull T itemInfo,
+            @Nullable LauncherActivityInfo launcherActivityInfo,
+            @Nullable String packageName,
+            @Nullable String resourceName,
+            @Nullable byte[] iconBlob,
+            boolean useLowResIcon) {
+        this.itemInfo = itemInfo;
+        this.launcherActivityInfo = launcherActivityInfo;
+        this.packageName = packageName;
+        this.resourceName = resourceName;
+        this.iconBlob = iconBlob;
+        this.useLowResIcon = useLowResIcon;
+    }
+
+    /** Loads  */
+    public boolean loadWorkspaceIcon(Context context) {
+        if (!(itemInfo instanceof WorkspaceItemInfo)) {
+            throw new IllegalStateException(
+                    "loadWorkspaceIcon should only be use for a WorkspaceItemInfos: " + itemInfo);
+        }
+
+        try (LauncherIcons li = LauncherIcons.obtain(context)) {
+            WorkspaceItemInfo info = (WorkspaceItemInfo) itemInfo;
+            if (itemInfo.itemType == LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT) {
+                if (!TextUtils.isEmpty(packageName) || !TextUtils.isEmpty(resourceName)) {
+                    info.iconResource = new Intent.ShortcutIconResource();
+                    info.iconResource.packageName = packageName;
+                    info.iconResource.resourceName = resourceName;
+                    BitmapInfo iconInfo = li.createIconBitmap(info.iconResource);
+                    if (iconInfo != null) {
+                        info.bitmap = iconInfo;
+                        return true;
+                    }
+                }
+            }
+
+            // Failed to load from resource, try loading from DB.
+            try {
+                if (iconBlob == null) {
+                    return false;
+                }
+                info.bitmap = li.createIconBitmap(decodeByteArray(
+                        iconBlob, 0, iconBlob.length));
+                return true;
+            } catch (Exception e) {
+                Log.e(TAG, "Failed to decode byte array for info " + info, e);
+                return false;
+            }
+        }
+    }
+}