Merge "Adding support for custom layouts in Preview" into main

7 files changed, 331 insertions, 133 deletions

diff --git a/src/com/android/launcher3/graphics/GridCustomizationsProxy.java b/src/com/android/launcher3/graphics/GridCustomizationsProxy.java
index b40e099218..7bd9a28416 100644
--- a/src/com/android/launcher3/graphics/GridCustomizationsProxy.java
+++ b/src/com/android/launcher3/graphics/GridCustomizationsProxy.java
@@ -27,6 +27,7 @@ import android.content.Context;
 import android.database.Cursor;
 import android.database.MatrixCursor;
 import android.net.Uri;
+import android.os.Binder;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.IBinder.DeathRecipient;
@@ -313,7 +314,7 @@ public class GridCustomizationsProxy implements ProxyProvider {
         RunnableList lifeCycleTracker = new RunnableList();
         try {
             PreviewSurfaceRenderer renderer = new PreviewSurfaceRenderer(
-                    mContext, lifeCycleTracker, request);
+                    mContext, lifeCycleTracker, request, Binder.getCallingPid());
             PreviewLifecycleObserver observer =
                     new PreviewLifecycleObserver(lifeCycleTracker, renderer);
 
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index e4e4b900fa..a062a54ee3 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -23,6 +23,7 @@ import static android.view.View.VISIBLE;
 import static com.android.launcher3.BubbleTextView.DISPLAY_TASKBAR;
 import static com.android.launcher3.BubbleTextView.DISPLAY_WORKSPACE;
 import static com.android.launcher3.DeviceProfile.DEFAULT_SCALE;
+import static com.android.launcher3.Flags.extendibleThemeManager;
 import static com.android.launcher3.Hotseat.ALPHA_CHANNEL_PREVIEW_RENDERER;
 import static com.android.launcher3.LauncherPrefs.GRID_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
@@ -30,6 +31,7 @@ import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
 import static com.android.launcher3.graphics.ThemeManager.PREF_ICON_SHAPE;
 import static com.android.launcher3.model.ModelUtils.currentScreenContentFilter;
+import static com.android.launcher3.widget.LauncherWidgetHolder.APPWIDGET_HOST_ID;
 
 import android.app.Fragment;
 import android.app.WallpaperColors;
@@ -44,6 +46,7 @@ import android.graphics.PointF;
 import android.graphics.Rect;
 import android.os.Handler;
 import android.os.Looper;
+import android.text.TextUtils;
 import android.util.AttributeSet;
 import android.util.Size;
 import android.util.SparseArray;
@@ -76,13 +79,20 @@ import com.android.launcher3.apppairs.AppPairIcon;
 import com.android.launcher3.celllayout.CellLayoutLayoutParams;
 import com.android.launcher3.celllayout.CellPosMapper;
 import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.dagger.ApiWrapperModule;
+import com.android.launcher3.dagger.AppModule;
 import com.android.launcher3.dagger.LauncherAppComponent;
-import com.android.launcher3.dagger.LauncherAppModule;
 import com.android.launcher3.dagger.LauncherAppSingleton;
+import com.android.launcher3.dagger.LauncherComponentProvider;
+import com.android.launcher3.dagger.PluginManagerWrapperModule;
+import com.android.launcher3.dagger.StaticObjectModule;
+import com.android.launcher3.dagger.WindowManagerProxyModule;
 import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.model.BaseLauncherBinder.BaseLauncherBinderFactory;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
+import com.android.launcher3.model.LayoutParserFactory;
+import com.android.launcher3.model.LayoutParserFactory.XmlLayoutParserFactory;
 import com.android.launcher3.model.LoaderTask.LoaderTaskFactory;
 import com.android.launcher3.model.data.AppPairInfo;
 import com.android.launcher3.model.data.CollectionInfo;
@@ -104,6 +114,7 @@ import com.android.launcher3.views.BaseDragLayer;
 import com.android.launcher3.widget.BaseLauncherAppWidgetHostView;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.LauncherWidgetHolder;
+import com.android.launcher3.widget.LauncherWidgetHolder.WidgetHolderFactory;
 import com.android.launcher3.widget.LocalColorExtractor;
 import com.android.launcher3.widget.util.WidgetSizes;
 import com.android.systemui.shared.Flags;
@@ -111,6 +122,8 @@ import com.android.systemui.shared.Flags;
 import dagger.BindsInstance;
 import dagger.Component;
 
+import java.io.File;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -139,21 +152,63 @@ public class LauncherPreviewRenderer extends BaseContext
 
         private final String mPrefName;
 
+        private final File mDbDir;
+
         public PreviewContext(Context base, String gridName, String shapeKey) {
+            this(base, gridName, shapeKey, APPWIDGET_HOST_ID, null);
+        }
+
+        public PreviewContext(Context base, String gridName, String shapeKey,
+                int widgetHostId, @Nullable String layoutXml) {
             super(base);
-            mPrefName = "preview-" + UUID.randomUUID().toString();
+            String randomUid = UUID.randomUUID().toString();
+            mPrefName = "preview-" + randomUid;
             LauncherPrefs prefs =
                     new ProxyPrefs(this, getSharedPreferences(mPrefName, MODE_PRIVATE));
             prefs.put(GRID_NAME, gridName);
             prefs.put(PREF_ICON_SHAPE, shapeKey);
-            initDaggerComponent(
-                    DaggerLauncherPreviewRenderer_PreviewAppComponent.builder().bindPrefs(prefs));
+
+            PreviewAppComponent.Builder builder =
+                    DaggerLauncherPreviewRenderer_PreviewAppComponent.builder().bindPrefs(prefs);
+            if (TextUtils.isEmpty(layoutXml) || !extendibleThemeManager()) {
+                mDbDir = null;
+                builder.bindParserFactory(new LayoutParserFactory(this))
+                        .bindWidgetsFactory(
+                                LauncherComponentProvider.get(base).getWidgetHolderFactory());
+            } else {
+                mDbDir = new File(base.getFilesDir(), randomUid);
+                emptyDbDir();
+                mDbDir.mkdirs();
+                builder.bindParserFactory(new XmlLayoutParserFactory(this, layoutXml))
+                        .bindWidgetsFactory(c -> new LauncherWidgetHolder(c, widgetHostId));
+            }
+            initDaggerComponent(builder);
+
+            if (!TextUtils.isEmpty(layoutXml)) {
+                // Use null the DB file so that we use a new in-memory DB
+                InvariantDeviceProfile.INSTANCE.get(this).dbFile = null;
+            }
+        }
+
+        private void emptyDbDir() {
+            if (mDbDir != null && mDbDir.exists()) {
+                Arrays.stream(mDbDir.listFiles()).forEach(File::delete);
+            }
         }
 
         @Override
         protected void cleanUpObjects() {
             super.cleanUpObjects();
             deleteSharedPreferences(mPrefName);
+            if (mDbDir != null) {
+                emptyDbDir();
+                mDbDir.delete();
+            }
+        }
+
+        @Override
+        public File getDatabasePath(String name) {
+            return mDbDir != null ? new File(mDbDir, name) :  super.getDatabasePath(name);
         }
     }
 
@@ -173,13 +228,15 @@ public class LauncherPreviewRenderer extends BaseContext
 
     public LauncherPreviewRenderer(Context context,
             InvariantDeviceProfile idp,
+            int widgetHostId,
             WallpaperColors wallpaperColorsOverride,
             @Nullable final SparseArray<Size> launcherWidgetSpanInfo) {
-        this(context, idp, null, wallpaperColorsOverride, launcherWidgetSpanInfo);
+        this(context, idp, widgetHostId, null, wallpaperColorsOverride, launcherWidgetSpanInfo);
     }
 
     public LauncherPreviewRenderer(Context context,
             InvariantDeviceProfile idp,
+            int widgetHostId,
             SparseIntArray previewColorOverride,
             WallpaperColors wallpaperColorsOverride,
             @Nullable final SparseArray<Size> launcherWidgetSpanInfo) {
@@ -262,9 +319,13 @@ public class LauncherPreviewRenderer extends BaseContext
                     wallpaperColors)
                     : null;
         }
-        mAppWidgetHost = new LauncherPreviewAppWidgetHost(context);
+        mAppWidgetHost = new LauncherPreviewAppWidgetHost(context, widgetHostId);
 
         onViewCreated();
+
+        if (widgetHostId != APPWIDGET_HOST_ID) {
+            mAppWidgetHost.stopListening();
+        }
     }
 
     @Override
@@ -567,8 +628,8 @@ public class LauncherPreviewRenderer extends BaseContext
 
     private class LauncherPreviewAppWidgetHost extends AppWidgetHost {
 
-        private LauncherPreviewAppWidgetHost(Context context) {
-            super(context, LauncherWidgetHolder.APPWIDGET_HOST_ID);
+        private LauncherPreviewAppWidgetHost(Context context, int hostId) {
+            super(context, hostId);
         }
 
         @Override
@@ -604,7 +665,12 @@ public class LauncherPreviewRenderer extends BaseContext
     }
 
     @LauncherAppSingleton
-    @Component(modules = LauncherAppModule.class)
+    // Exclude widget module since we bind widget holder separately
+    @Component(modules = {WindowManagerProxyModule.class,
+            ApiWrapperModule.class,
+            PluginManagerWrapperModule.class,
+            StaticObjectModule.class,
+            AppModule.class})
     public interface PreviewAppComponent extends LauncherAppComponent {
 
         LoaderTaskFactory getLoaderTaskFactory();
@@ -615,6 +681,8 @@ public class LauncherPreviewRenderer extends BaseContext
         @Component.Builder
         interface Builder extends LauncherAppComponent.Builder {
             @BindsInstance Builder bindPrefs(LauncherPrefs prefs);
+            @BindsInstance Builder bindParserFactory(LayoutParserFactory parserFactory);
+            @BindsInstance Builder bindWidgetsFactory(WidgetHolderFactory holderFactory);
             PreviewAppComponent build();
         }
     }
diff --git a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
index 457d12e8ca..93e5f325bf 100644
--- a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
+++ b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
@@ -20,12 +20,18 @@ import static android.content.res.Configuration.UI_MODE_NIGHT_NO;
 import static android.content.res.Configuration.UI_MODE_NIGHT_YES;
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static com.android.launcher3.Flags.extendibleThemeManager;
 import static com.android.launcher3.LauncherPrefs.GRID_NAME;
+import static com.android.launcher3.WorkspaceLayoutManager.FIRST_SCREEN_ID;
+import static com.android.launcher3.WorkspaceLayoutManager.SECOND_SCREEN_ID;
 import static com.android.launcher3.graphics.ThemeManager.PREF_ICON_SHAPE;
+import static com.android.launcher3.provider.LauncherDbUtils.selectionForWorkspaceScreen;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+import static com.android.launcher3.widget.LauncherWidgetHolder.APPWIDGET_HOST_ID;
 
 import android.app.WallpaperColors;
+import android.appwidget.AppWidgetHost;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
 import android.content.res.Configuration;
@@ -33,6 +39,7 @@ import android.database.Cursor;
 import android.hardware.display.DisplayManager;
 import android.os.Bundle;
 import android.os.IBinder;
+import android.text.TextUtils;
 import android.util.Log;
 import android.util.Size;
 import android.util.SparseArray;
@@ -56,7 +63,6 @@ import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings;
-import com.android.launcher3.Workspace;
 import com.android.launcher3.dagger.LauncherComponentProvider;
 import com.android.launcher3.graphics.LauncherPreviewRenderer.PreviewAppComponent;
 import com.android.launcher3.graphics.LauncherPreviewRenderer.PreviewContext;
@@ -89,18 +95,22 @@ public class PreviewSurfaceRenderer {
     private static final String KEY_COLOR_RESOURCE_IDS = "color_resource_ids";
     private static final String KEY_COLOR_VALUES = "color_values";
     private static final String KEY_DARK_MODE = "use_dark_mode";
+    private static final String KEY_LAYOUT_XML = "layout_xml";
     public static final String KEY_SKIP_ANIMATIONS = "skip_animations";
 
     private final Context mContext;
     private SparseIntArray mPreviewColorOverride;
     private String mGridName;
     private String mShapeKey;
+    private String mLayoutXml;
 
     @Nullable private Boolean mDarkMode;
     private boolean mDestroyed = false;
     private boolean mHideQsb;
     @Nullable private FrameLayout mViewRoot = null;
+    private boolean mDeletingHostOnExit = false;
 
+    private final int mCallingPid;
     private final IBinder mHostToken;
     private final int mWidth;
     private final int mHeight;
@@ -111,11 +121,11 @@ public class PreviewSurfaceRenderer {
     private final RunnableList mLifeCycleTracker;
     private final SurfaceControlViewHost mSurfaceControlViewHost;
 
-
-    public PreviewSurfaceRenderer(
-            Context context, RunnableList lifecycleTracker, Bundle bundle) throws Exception {
+    public PreviewSurfaceRenderer(Context context, RunnableList lifecycleTracker, Bundle bundle,
+            int callingPid) throws Exception {
         mContext = context;
         mLifeCycleTracker = lifecycleTracker;
+        mCallingPid = callingPid;
         mGridName = bundle.getString("name");
         bundle.remove("name");
         if (mGridName == null) {
@@ -135,6 +145,7 @@ public class PreviewSurfaceRenderer {
         mDisplayId = bundle.getInt(KEY_DISPLAY_ID);
         mDisplay = context.getSystemService(DisplayManager.class)
                 .getDisplay(mDisplayId);
+        mLayoutXml = bundle.getString(KEY_LAYOUT_XML);
         if (mDisplay == null) {
             throw new IllegalArgumentException("Display ID does not match any displays.");
         }
@@ -335,41 +346,52 @@ public class PreviewSurfaceRenderer {
     private void loadModelData() {
         final Context inflationContext = getPreviewContext();
         if (!mGridName.equals(LauncherPrefs.INSTANCE.get(mContext).get(GRID_NAME))
-                || !mShapeKey.equals(LauncherPrefs.INSTANCE.get(mContext).get(PREF_ICON_SHAPE))) {
+                || !mShapeKey.equals(LauncherPrefs.INSTANCE.get(mContext).get(PREF_ICON_SHAPE))
+                || !TextUtils.isEmpty(mLayoutXml)) {
+
+            boolean isCustomLayout = extendibleThemeManager() &&  !TextUtils.isEmpty(mLayoutXml);
+            int widgetHostId = isCustomLayout ? APPWIDGET_HOST_ID + mCallingPid : APPWIDGET_HOST_ID;
+
             // Start the migration
-            PreviewContext previewContext =
-                    new PreviewContext(inflationContext, mGridName, mShapeKey);
+            PreviewContext previewContext = new PreviewContext(
+                    inflationContext, mGridName, mShapeKey, widgetHostId, mLayoutXml);
             PreviewAppComponent appComponent =
                     (PreviewAppComponent) LauncherComponentProvider.get(previewContext);
 
+            if (extendibleThemeManager() && isCustomLayout && !mDeletingHostOnExit) {
+                mDeletingHostOnExit = true;
+                mLifeCycleTracker.add(() -> {
+                    AppWidgetHost host = new AppWidgetHost(mContext, widgetHostId);
+                    // Start listening here, so that any previous active host is disabled
+                    host.startListening();
+                    host.stopListening();
+                    host.deleteHost();
+                });
+            }
+
             LoaderTask task = appComponent.getLoaderTaskFactory().newLoaderTask(
                     appComponent.getBaseLauncherBinderFactory().createBinder(new Callbacks[0]),
                     new UserManagerState());
 
             InvariantDeviceProfile idp = appComponent.getIDP();
             DeviceProfile deviceProfile = idp.getDeviceProfile(previewContext);
-            String query =
-                    LauncherSettings.Favorites.SCREEN + " = " + Workspace.FIRST_SCREEN_ID
-                            + " or " + LauncherSettings.Favorites.CONTAINER + " = "
-                            + LauncherSettings.Favorites.CONTAINER_HOTSEAT;
-            if (deviceProfile.isTwoPanels) {
-                query += " or " + LauncherSettings.Favorites.SCREEN + " = "
-                        + Workspace.SECOND_SCREEN_ID;
-            }
-
+            String query = deviceProfile.isTwoPanels
+                    ? selectionForWorkspaceScreen(FIRST_SCREEN_ID, SECOND_SCREEN_ID)
+                    : selectionForWorkspaceScreen(FIRST_SCREEN_ID);
             Map<ComponentKey, AppWidgetProviderInfo> widgetProviderInfoMap = new HashMap<>();
             task.loadWorkspaceForPreview(query, widgetProviderInfoMap);
             final SparseArray<Size> spanInfo = getLoadedLauncherWidgetInfo();
             MAIN_EXECUTOR.execute(() -> {
-                renderView(previewContext, appComponent.getDataModel(), widgetProviderInfoMap,
-                        spanInfo, idp);
+                renderView(previewContext, appComponent.getDataModel(), widgetHostId,
+                        widgetProviderInfoMap, spanInfo, idp);
                 mLifeCycleTracker.add(previewContext::onDestroy);
             });
         } else {
             LauncherAppState.getInstance(inflationContext).getModel().loadAsync(dataModel -> {
                 if (dataModel != null) {
-                    MAIN_EXECUTOR.execute(() -> renderView(inflationContext, dataModel, null,
-                            null, LauncherAppState.getIDP(inflationContext)));
+                    MAIN_EXECUTOR.execute(() -> renderView(inflationContext, dataModel,
+                            APPWIDGET_HOST_ID, null, null,
+                            LauncherAppState.getIDP(inflationContext)));
                 } else {
                     Log.e(TAG, "Model loading failed");
                 }
@@ -378,7 +400,7 @@ public class PreviewSurfaceRenderer {
     }
 
     @UiThread
-    private void renderView(Context inflationContext, BgDataModel dataModel,
+    private void renderView(Context inflationContext, BgDataModel dataModel, int widgetHostId,
             Map<ComponentKey, AppWidgetProviderInfo> widgetProviderInfoMap,
             @Nullable final SparseArray<Size> launcherWidgetSpanInfo, InvariantDeviceProfile idp) {
         if (mDestroyed) {
@@ -386,10 +408,10 @@ public class PreviewSurfaceRenderer {
         }
         LauncherPreviewRenderer renderer;
         if (Flags.newCustomizationPickerUi()) {
-            renderer = new LauncherPreviewRenderer(inflationContext, idp, mPreviewColorOverride,
-                    mWallpaperColors, launcherWidgetSpanInfo);
+            renderer = new LauncherPreviewRenderer(inflationContext, idp, widgetHostId,
+                    mPreviewColorOverride, mWallpaperColors, launcherWidgetSpanInfo);
         } else {
-            renderer = new LauncherPreviewRenderer(inflationContext, idp,
+            renderer = new LauncherPreviewRenderer(inflationContext, idp, widgetHostId,
                     mWallpaperColors, launcherWidgetSpanInfo);
         }
         renderer.hideBottomRow(mHideQsb);
diff --git a/src/com/android/launcher3/model/LayoutParserFactory.kt b/src/com/android/launcher3/model/LayoutParserFactory.kt
new file mode 100644
index 0000000000..32aa58a03b
--- /dev/null
+++ b/src/com/android/launcher3/model/LayoutParserFactory.kt
@@ -0,0 +1,181 @@
+/*
+ * Copyright (C) 2025 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
+
+import android.app.blob.BlobHandle
+import android.app.blob.BlobStoreManager
+import android.content.Context
+import android.os.ParcelFileDescriptor.AutoCloseInputStream
+import android.provider.Settings.Secure
+import android.text.TextUtils
+import android.util.Base64
+import android.util.Log
+import android.util.Xml
+import com.android.launcher3.AutoInstallsLayout
+import com.android.launcher3.AutoInstallsLayout.SourceResources
+import com.android.launcher3.DefaultLayoutParser
+import com.android.launcher3.DefaultLayoutParser.RES_PARTNER_DEFAULT_LAYOUT
+import com.android.launcher3.LauncherSettings.Settings
+import com.android.launcher3.dagger.ApplicationContext
+import com.android.launcher3.util.IOUtils
+import com.android.launcher3.util.Partner
+import com.android.launcher3.widget.LauncherWidgetHolder
+import java.io.StringReader
+import javax.inject.Inject
+
+private const val TAG = "LayoutParserFactory"
+
+/** Utility class for providing default layout parsers */
+open class LayoutParserFactory
+@Inject
+constructor(@ApplicationContext private val context: Context) {
+
+    open fun createExternalLayoutParser(
+        widgetHolder: LauncherWidgetHolder,
+        openHelper: DatabaseHelper,
+    ): AutoInstallsLayout? {
+
+        createWorkspaceLoaderFromAppRestriction(widgetHolder, openHelper)?.let {
+            return it
+        }
+        AutoInstallsLayout.get(context, widgetHolder, openHelper)?.let {
+            return it
+        }
+
+        val partner = Partner.get(context.packageManager)
+        if (partner != null) {
+            val workspaceResId = partner.getXmlResId(RES_PARTNER_DEFAULT_LAYOUT)
+            if (workspaceResId != 0) {
+                return DefaultLayoutParser(
+                    context,
+                    widgetHolder,
+                    openHelper,
+                    partner.resources,
+                    workspaceResId,
+                )
+            }
+        }
+        return null
+    }
+
+    /**
+     * Creates workspace loader from an XML resource listed in the app restrictions.
+     *
+     * @return the loader if the restrictions are set and the resource exists; null otherwise.
+     */
+    private fun createWorkspaceLoaderFromAppRestriction(
+        widgetHolder: LauncherWidgetHolder,
+        openHelper: DatabaseHelper,
+    ): AutoInstallsLayout? {
+        val systemLayoutProvider =
+            Secure.getString(context.contentResolver, Settings.LAYOUT_PROVIDER_KEY)
+        if (TextUtils.isEmpty(systemLayoutProvider)) {
+            return null
+        }
+
+        // Try the blob store first
+        val blobManager = context.getSystemService(BlobStoreManager::class.java)
+        if (systemLayoutProvider.startsWith(Settings.BLOB_KEY_PREFIX) && blobManager != null) {
+            val blobHandlerDigest = systemLayoutProvider.substring(Settings.BLOB_KEY_PREFIX.length)
+            try {
+                AutoCloseInputStream(
+                        blobManager.openBlob(
+                            BlobHandle.createWithSha256(
+                                Base64.decode(
+                                    blobHandlerDigest,
+                                    Base64.NO_WRAP or Base64.NO_PADDING,
+                                ),
+                                Settings.LAYOUT_DIGEST_LABEL,
+                                0,
+                                Settings.LAYOUT_DIGEST_TAG,
+                            )
+                        )
+                    )
+                    .use {
+                        return getAutoInstallsLayoutFromIS(
+                            widgetHolder,
+                            openHelper,
+                            String(IOUtils.toByteArray(it)),
+                        )
+                    }
+            } catch (e: Exception) {
+                Log.e(TAG, "Error getting layout from blob handle", e)
+                return null
+            }
+        }
+
+        // Try contentProvider based provider
+        val pm = context.packageManager
+        val pi = pm.resolveContentProvider(systemLayoutProvider, 0)
+        if (pi == null) {
+            Log.e(TAG, "No provider found for authority $systemLayoutProvider")
+            return null
+        }
+        val uri = ModelDbController.getLayoutUri(systemLayoutProvider, context)
+        try {
+            context.contentResolver.openInputStream(uri)?.use {
+                Log.d(TAG, "Loading layout from $systemLayoutProvider")
+                val res = pm.getResourcesForApplication(pi.applicationInfo)
+                return getAutoInstallsLayoutFromIS(
+                    widgetHolder,
+                    openHelper,
+                    String(IOUtils.toByteArray(it)),
+                    SourceResources.wrap(res),
+                )
+            }
+        } catch (e: Exception) {
+            Log.e(TAG, "Error getting layout stream from: $systemLayoutProvider", e)
+        }
+        return null
+    }
+
+    @Throws(Exception::class)
+    protected fun getAutoInstallsLayoutFromIS(
+        widgetHolder: LauncherWidgetHolder,
+        openHelper: DatabaseHelper,
+        xml: String,
+        res: SourceResources = object : SourceResources {},
+    ): AutoInstallsLayout {
+        val parser = Xml.newPullParser()
+        parser.setInput(StringReader(xml))
+
+        return AutoInstallsLayout(
+            context,
+            widgetHolder,
+            openHelper,
+            res,
+            { parser },
+            AutoInstallsLayout.TAG_WORKSPACE,
+        )
+    }
+
+    /** Layout parser factory with fixed xml */
+    class XmlLayoutParserFactory(ctx: Context, private val xml: String) : LayoutParserFactory(ctx) {
+
+        override fun createExternalLayoutParser(
+            widgetHolder: LauncherWidgetHolder,
+            openHelper: DatabaseHelper,
+        ): AutoInstallsLayout? {
+            try {
+                return getAutoInstallsLayoutFromIS(widgetHolder, openHelper, xml)
+            } catch (e: Exception) {
+                Log.e(TAG, "Error getting layout from provided xml", e)
+                return super.createExternalLayoutParser(widgetHolder, openHelper)
+            }
+        }
+    }
+}
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
index 64b9c1c9cd..f48d8488c2 100644
--- a/src/com/android/launcher3/model/ModelDbController.java
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -16,10 +16,7 @@
 package com.android.launcher3.model;
 
 import static android.provider.BaseColumns._ID;
-import static android.util.Base64.NO_PADDING;
-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;
@@ -27,40 +24,25 @@ import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.addTableToDb;
-import static com.android.launcher3.LauncherSettings.Settings.BLOB_KEY_PREFIX;
-import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_LABEL;
-import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_TAG;
-import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_PROVIDER_KEY;
 import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
 
-import android.app.blob.BlobHandle;
-import android.app.blob.BlobStoreManager;
-import android.content.ContentResolver;
 import android.content.ContentValues;
 import android.content.Context;
-import android.content.pm.PackageManager;
-import android.content.pm.ProviderInfo;
-import android.content.res.Resources;
 import android.database.Cursor;
 import android.database.SQLException;
 import android.database.sqlite.SQLiteDatabase;
 import android.net.Uri;
 import android.os.Bundle;
-import android.os.ParcelFileDescriptor;
 import android.os.Process;
 import android.os.UserHandle;
 import android.os.UserManager;
-import android.provider.Settings;
 import android.text.TextUtils;
-import android.util.Base64;
 import android.util.Log;
-import android.util.Xml;
 
 import androidx.annotation.Nullable;
 import androidx.annotation.WorkerThread;
 
 import com.android.launcher3.AutoInstallsLayout;
-import com.android.launcher3.AutoInstallsLayout.SourceResources;
 import com.android.launcher3.ConstantItem;
 import com.android.launcher3.DefaultLayoutParser;
 import com.android.launcher3.EncryptionType;
@@ -80,16 +62,10 @@ import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.provider.LauncherDbUtils;
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
 import com.android.launcher3.provider.RestoreDbTask;
-import com.android.launcher3.util.IOUtils;
 import com.android.launcher3.util.IntArray;
-import com.android.launcher3.util.Partner;
 import com.android.launcher3.widget.LauncherWidgetHolder;
 
-import org.xmlpull.v1.XmlPullParser;
-
 import java.io.File;
-import java.io.InputStream;
-import java.io.StringReader;
 import java.util.List;
 import java.util.stream.Collectors;
 
@@ -113,17 +89,20 @@ public class ModelDbController {
     private final InvariantDeviceProfile mIdp;
     private final LauncherPrefs mPrefs;
     private final UserCache mUserCache;
+    private final LayoutParserFactory mLayoutParserFactory;
 
     @Inject
     ModelDbController(
             @ApplicationContext Context context,
             InvariantDeviceProfile idp,
             LauncherPrefs prefs,
-            UserCache userCache) {
+            UserCache userCache,
+            LayoutParserFactory layoutParserFactory) {
         mContext = context;
         mIdp = idp;
         mPrefs = prefs;
         mUserCache = userCache;
+        mLayoutParserFactory = layoutParserFactory;
     }
 
     private void printDBs(String prefix) {
@@ -368,7 +347,6 @@ public class ModelDbController {
     public void attemptMigrateDb(LauncherRestoreEventLogger restoreEventLogger,
             ModelDelegate modelDelegate) throws Exception {
         createDbIfNotExists();
-
         if (shouldResetDb()) {
             resetLauncherDb(restoreEventLogger);
             return;
@@ -634,20 +612,8 @@ public class ModelDbController {
 
             LauncherWidgetHolder widgetHolder = mOpenHelper.newLauncherWidgetHolder();
             try {
-                AutoInstallsLayout loader = createWorkspaceLoaderFromAppRestriction(widgetHolder);
-                if (loader == null) {
-                    loader = AutoInstallsLayout.get(mContext, widgetHolder, mOpenHelper);
-                }
-                if (loader == null) {
-                    final Partner partner = Partner.get(mContext.getPackageManager());
-                    if (partner != null) {
-                        int workspaceResId = partner.getXmlResId(RES_PARTNER_DEFAULT_LAYOUT);
-                        if (workspaceResId != 0) {
-                            loader = new DefaultLayoutParser(mContext, widgetHolder,
-                                    mOpenHelper, partner.getResources(), workspaceResId);
-                        }
-                    }
-                }
+                AutoInstallsLayout loader =
+                        mLayoutParserFactory.createExternalLayoutParser(widgetHolder, mOpenHelper);
 
                 final boolean usingExternallyProvidedLayout = loader != null;
                 if (loader == null) {
@@ -672,64 +638,6 @@ public class ModelDbController {
         }
     }
 
-    /**
-     * Creates workspace loader from an XML resource listed in the app restrictions.
-     *
-     * @return the loader if the restrictions are set and the resource exists; null otherwise.
-     */
-    private AutoInstallsLayout createWorkspaceLoaderFromAppRestriction(
-            LauncherWidgetHolder widgetHolder) {
-        ContentResolver cr = mContext.getContentResolver();
-        String systemLayoutProvider = Settings.Secure.getString(cr, LAYOUT_PROVIDER_KEY);
-        if (TextUtils.isEmpty(systemLayoutProvider)) {
-            return null;
-        }
-
-        // Try the blob store first
-        if (systemLayoutProvider.startsWith(BLOB_KEY_PREFIX)) {
-            BlobStoreManager blobManager = mContext.getSystemService(BlobStoreManager.class);
-            String blobHandlerDigest = systemLayoutProvider.substring(BLOB_KEY_PREFIX.length());
-            try (InputStream in = new ParcelFileDescriptor.AutoCloseInputStream(
-                    blobManager.openBlob(BlobHandle.createWithSha256(
-                            Base64.decode(blobHandlerDigest, NO_WRAP | NO_PADDING),
-                            LAYOUT_DIGEST_LABEL, 0, LAYOUT_DIGEST_TAG)))) {
-                return getAutoInstallsLayoutFromIS(in, widgetHolder, new SourceResources() { });
-            } catch (Exception e) {
-                Log.e(TAG, "Error getting layout from blob handle" , e);
-                return null;
-            }
-        }
-
-        // Try contentProvider based provider
-        PackageManager pm = mContext.getPackageManager();
-        ProviderInfo pi = pm.resolveContentProvider(systemLayoutProvider, 0);
-        if (pi == null) {
-            Log.e(TAG, "No provider found for authority " + systemLayoutProvider);
-            return null;
-        }
-        Uri uri = getLayoutUri(systemLayoutProvider, mContext);
-        try (InputStream in = cr.openInputStream(uri)) {
-            Log.d(TAG, "Loading layout from " + systemLayoutProvider);
-
-            Resources res = pm.getResourcesForApplication(pi.applicationInfo);
-            return getAutoInstallsLayoutFromIS(in, widgetHolder, SourceResources.wrap(res));
-        } catch (Exception e) {
-            Log.e(TAG, "Error getting layout stream from: " + systemLayoutProvider , e);
-            return null;
-        }
-    }
-
-    private AutoInstallsLayout getAutoInstallsLayoutFromIS(InputStream in,
-            LauncherWidgetHolder widgetHolder, SourceResources res) throws Exception {
-        // Read the full xml so that we fail early in case of any IO error.
-        String layout = new String(IOUtils.toByteArray(in));
-        XmlPullParser parser = Xml.newPullParser();
-        parser.setInput(new StringReader(layout));
-
-        return new AutoInstallsLayout(mContext, widgetHolder, mOpenHelper, res,
-                () -> parser, AutoInstallsLayout.TAG_WORKSPACE);
-    }
-
     public static Uri getLayoutUri(String authority, Context ctx) {
         InvariantDeviceProfile grid = LauncherAppState.getIDP(ctx);
         return new Uri.Builder().scheme("content").authority(authority).path("launcher_layout")
diff --git a/src/com/android/launcher3/provider/LauncherDbUtils.kt b/src/com/android/launcher3/provider/LauncherDbUtils.kt
index 36413715e0..7ee49880be 100644
--- a/src/com/android/launcher3/provider/LauncherDbUtils.kt
+++ b/src/com/android/launcher3/provider/LauncherDbUtils.kt
@@ -27,7 +27,13 @@ import android.os.Process
 import android.os.UserManager
 import android.text.TextUtils
 import com.android.launcher3.LauncherSettings
+import com.android.launcher3.LauncherSettings.Favorites
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER
 import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT
+import com.android.launcher3.LauncherSettings.Favorites.SCREEN
+import com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME
+import com.android.launcher3.LauncherSettings.Favorites._ID
 import com.android.launcher3.Utilities
 import com.android.launcher3.dagger.LauncherComponentProvider.appComponent
 import com.android.launcher3.icons.IconCache
@@ -45,6 +51,14 @@ object LauncherDbUtils {
      */
     @JvmStatic fun itemIdMatch(itemId: Int): String = "_id=$itemId"
 
+    /**
+     * Returns a string which can be used as a where clause for DB query to match the given
+     * workspace screens or hotseat or a collection in workspace screens or hotseat
+     */
+    @JvmStatic
+    fun selectionForWorkspaceScreen(vararg screens: Int) =
+        "$SCREEN in (${screens.joinToString()}) or $CONTAINER = $CONTAINER_HOTSEAT or $CONTAINER in (select $_ID from $TABLE_NAME where $SCREEN in (${screens.joinToString()}) or $CONTAINER = $CONTAINER_HOTSEAT)"
+
     @JvmStatic
     fun queryIntArray(
         distinct: Boolean,
diff --git a/src/com/android/launcher3/widget/LauncherWidgetHolder.java b/src/com/android/launcher3/widget/LauncherWidgetHolder.java
index 642f35abcd..f7c34f6fad 100644
--- a/src/com/android/launcher3/widget/LauncherWidgetHolder.java
+++ b/src/com/android/launcher3/widget/LauncherWidgetHolder.java
@@ -107,7 +107,11 @@ public class LauncherWidgetHolder {
 
     @AssistedInject
     protected LauncherWidgetHolder(@Assisted("UI_CONTEXT") @NonNull Context context) {
-        this(context, new LauncherAppWidgetHost(context, APPWIDGET_HOST_ID));
+        this(context, APPWIDGET_HOST_ID);
+    }
+
+    public LauncherWidgetHolder(@NonNull Context context, int hostId) {
+        this(context, new LauncherAppWidgetHost(context, hostId));
     }
 
     protected LauncherWidgetHolder(