Merge "Parsing test boards from text files to be able to add more and bigger tests" into udc-dev
diff --git a/OWNERS b/OWNERS
index 4fd6a50..76644b3 100644
--- a/OWNERS
+++ b/OWNERS
@@ -11,6 +11,7 @@
 twickham@google.com
 vadimt@google.com
 winsonc@google.com
+jonmiranda@google.com
 
 per-file FeatureFlags.java, globs = set noparent
 per-file FeatureFlags.java = sunnygoyal@google.com, winsonc@google.com, adamcohen@google.com, hyunyoungs@google.com, captaincole@google.com
diff --git a/quickstep/res/layout/redesigned_gesture_tutorial_fragment.xml b/quickstep/res/layout/redesigned_gesture_tutorial_fragment.xml
index 2887518..b2957aa 100644
--- a/quickstep/res/layout/redesigned_gesture_tutorial_fragment.xml
+++ b/quickstep/res/layout/redesigned_gesture_tutorial_fragment.xml
@@ -42,8 +42,6 @@
         android:id="@+id/gesture_tutorial_fake_previous_task_view"
         android:layout_width="match_parent"
         android:layout_height="match_parent"
-        android:scaleX="0.98"
-        android:scaleY="0.98"
         android:visibility="invisible">
 
         <View
diff --git a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatRestoreHelper.java b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatRestoreHelper.java
index 4956fa1..726abff 100644
--- a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatRestoreHelper.java
+++ b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatRestoreHelper.java
@@ -21,7 +21,6 @@
 
 import android.content.Context;
 
-import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.model.GridBackupTable;
@@ -42,10 +41,7 @@
                             context.getContentResolver(),
                             LauncherSettings.Settings.METHOD_NEW_TRANSACTION)
                             .getBinder(LauncherSettings.Settings.EXTRA_VALUE)) {
-                InvariantDeviceProfile idp = LauncherAppState.getIDP(context);
-                GridBackupTable backupTable = new GridBackupTable(context,
-                        transaction.getDb(), idp.numDatabaseHotseatIcons, idp.numColumns,
-                        idp.numRows);
+                GridBackupTable backupTable = new GridBackupTable(context, transaction.getDb());
                 backupTable.createCustomBackupTable(HYBRID_HOTSEAT_BACKUP_TABLE);
                 transaction.commit();
                 LauncherSettings.Settings.call(context.getContentResolver(),
@@ -67,10 +63,7 @@
                 if (!tableExists(transaction.getDb(), HYBRID_HOTSEAT_BACKUP_TABLE)) {
                     return;
                 }
-                InvariantDeviceProfile idp = LauncherAppState.getIDP(context);
-                GridBackupTable backupTable = new GridBackupTable(context,
-                        transaction.getDb(), idp.numDatabaseHotseatIcons, idp.numColumns,
-                        idp.numRows);
+                GridBackupTable backupTable = new GridBackupTable(context, transaction.getDb());
                 backupTable.restoreFromCustomBackupTable(HYBRID_HOTSEAT_BACKUP_TABLE, true);
                 transaction.commit();
                 LauncherAppState.getInstance(context).getModel().forceReload();
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
index c9e7df4..9bc8cdd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
@@ -642,14 +642,8 @@
         long resetDuration = mControllers.taskbarStashController.isInApp()
                 ? duration
                 : duration / 2;
-        boolean shouldReset =
-                mControllers.taskbarTranslationController.shouldResetBackToZero(resetDuration);
-        boolean goingToLauncher = isAnimatingToLauncher();
-        boolean isNormalState = mLauncherState == LauncherState.NORMAL;
-        // Taskbar should always reset when animating to launcher in normal state to ensure there
-        // is no jump during the handoff to the hotseat.
-        if ((goingToLauncher && isNormalState)
-                || (shouldReset && (goingToLauncher || isNormalState))) {
+        if (!mControllers.taskbarTranslationController.willAnimateToZeroBefore(resetDuration)
+                && (isAnimatingToLauncher() || mLauncherState == LauncherState.NORMAL)) {
             animatorSet.play(mControllers.taskbarTranslationController
                     .createAnimToResetTranslation(resetDuration));
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarTranslationController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarTranslationController.java
index 4b18bb6..065d111 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarTranslationController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarTranslationController.java
@@ -54,7 +54,6 @@
     private boolean mHasSprungOnceThisGesture;
     private @Nullable ValueAnimator mSpringBounce;
     private boolean mGestureEnded;
-    private boolean mGestureInProgress;
     private boolean mAnimationToHomeRunning;
 
     private final boolean mIsTransientTaskbar;
@@ -124,7 +123,6 @@
 
     private void reset() {
         mGestureEnded = false;
-        mGestureInProgress = false;
         mHasSprungOnceThisGesture = false;
     }
 
@@ -136,24 +134,18 @@
     }
 
     /**
-     * Returns {@code true} if we should reset the animation back to zero.
-     *
-     * Returns {@code false} if there is a gesture in progress, or if we are already animating
-     * to 0 within the specified duration.
+     * Returns true if we will animate to zero before the input duration.
      */
-    public boolean shouldResetBackToZero(long duration) {
-        if (mGestureInProgress) {
-            return false;
-        }
+    public boolean willAnimateToZeroBefore(long duration) {
         if (mSpringBounce != null && mSpringBounce.isRunning()) {
             long springDuration = mSpringBounce.getDuration();
             long current = mSpringBounce.getCurrentPlayTime();
-            return (springDuration - current >= duration);
+            return (springDuration - current < duration);
         }
         if (mTranslationYForSwipe.isAnimatingToValue(0)) {
-            return mTranslationYForSwipe.getRemainingTime() >= duration;
+            return mTranslationYForSwipe.getRemainingTime() < duration;
         }
-        return true;
+        return false;
     }
 
     /**
@@ -196,7 +188,6 @@
             mAnimationToHomeRunning = false;
             cancelSpringIfExists();
             reset();
-            mGestureInProgress = true;
         }
         /**
          * Called when there is movement to move the taskbar.
@@ -220,7 +211,6 @@
                 mGestureEnded = true;
                 startSpring();
             }
-            mGestureInProgress = false;
         }
     }
 
@@ -232,7 +222,6 @@
         pw.println(prefix + "\tmHasSprungOnceThisGesture=" + mHasSprungOnceThisGesture);
         pw.println(prefix + "\tmAnimationToHomeRunning=" + mAnimationToHomeRunning);
         pw.println(prefix + "\tmGestureEnded=" + mGestureEnded);
-        pw.println(prefix + "\tmGestureInProgress=" + mGestureInProgress);
         pw.println(prefix + "\tmSpringBounce is running=" + (mSpringBounce != null
                 && mSpringBounce.isRunning()));
     }
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index edd8823..9795670 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -62,6 +62,7 @@
 import android.annotation.TargetApi;
 import android.app.Activity;
 import android.app.ActivityManager;
+import android.app.TaskInfo;
 import android.app.WindowConfiguration;
 import android.content.Context;
 import android.content.Intent;
@@ -2079,13 +2080,16 @@
         if (!mCanceled) {
             TaskView nextTask = mRecentsView.getNextPageTaskView();
             if (nextTask != null) {
-                int taskId = nextTask.getTask().key.id;
+                Task.TaskKey nextTaskKey = nextTask.getTask().key;
+                int taskId = nextTaskKey.id;
                 mGestureState.updateLastStartedTaskId(taskId);
                 boolean hasTaskPreviouslyAppeared = mGestureState.getPreviouslyAppearedTaskIds()
                         .contains(taskId);
                 if (!hasTaskPreviouslyAppeared) {
                     ActiveGestureLog.INSTANCE.trackEvent(EXPECTING_TASK_APPEARED);
                 }
+                ActiveGestureLog.INSTANCE.addLog("Launching task: id=" + taskId
+                        + " pkg=" + nextTaskKey.getPackageName());
                 nextTask.launchTask(success -> {
                     resultCallback.accept(success);
                     if (success) {
@@ -2154,7 +2158,18 @@
     @Override
     public void onTasksAppeared(RemoteAnimationTarget[] appearedTaskTargets) {
         if (mRecentsAnimationController != null) {
-            if (handleTaskAppeared(appearedTaskTargets)) {
+            boolean hasStartedTaskBefore = Arrays.stream(appearedTaskTargets).anyMatch(
+                    targetCompat -> targetCompat.taskId == mGestureState.getLastStartedTaskId());
+            if (!mStateCallback.hasStates(STATE_GESTURE_COMPLETED) && !hasStartedTaskBefore) {
+                // This is a special case, if a task is started mid-gesture that wasn't a part of a
+                // previous quickswitch task launch, then cancel the animation back to the app
+                RemoteAnimationTarget appearedTaskTarget = appearedTaskTargets[0];
+                TaskInfo taskInfo = appearedTaskTarget.taskInfo;
+                ActiveGestureLog.INSTANCE.addLog("Unexpected task appeared"
+                        + " id=" + taskInfo.taskId
+                        + " pkg=" + taskInfo.baseIntent.getComponent().getPackageName());
+                finishRecentsAnimationOnTasksAppeared();
+            } else if (handleTaskAppeared(appearedTaskTargets)) {
                 Optional<RemoteAnimationTarget> taskTargetOptional =
                         Arrays.stream(appearedTaskTargets)
                                 .filter(targetCompat ->
@@ -2202,7 +2217,7 @@
         if (mRecentsAnimationController != null) {
             mRecentsAnimationController.finish(false /* toRecents */, null /* onFinishComplete */);
         }
-        ActiveGestureLog.INSTANCE.addLog("finishRecentsAnimation", false);
+        ActiveGestureLog.INSTANCE.addLog("finishRecentsAnimationOnTasksAppeared");
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/InputConsumer.java b/quickstep/src/com/android/quickstep/InputConsumer.java
index c455dc7..64c9295 100644
--- a/quickstep/src/com/android/quickstep/InputConsumer.java
+++ b/quickstep/src/com/android/quickstep/InputConsumer.java
@@ -40,6 +40,7 @@
     int TYPE_SYSUI_OVERLAY = 1 << 10;
     int TYPE_ONE_HANDED = 1 << 11;
     int TYPE_TASKBAR_STASH = 1 << 12;
+    int TYPE_STATUS_BAR = 1 << 13;
 
     String[] NAMES = new String[] {
            "TYPE_NO_OP",                    // 0
@@ -55,6 +56,7 @@
             "TYPE_SYSUI_OVERLAY",           // 10
             "TYPE_ONE_HANDED",              // 11
             "TYPE_TASKBAR_STASH",           // 12
+            "TYPE_STATUS_BAR",              // 13
     };
 
     InputConsumer NO_OP = () -> TYPE_NO_OP;
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 93363a0..1cf682b 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -26,6 +26,7 @@
 import static com.android.launcher3.Launcher.INTENT_ACTION_ALL_APPS_TOGGLE;
 import static com.android.launcher3.MotionEventsUtils.isTrackpadMultiFingerSwipe;
 import static com.android.launcher3.config.FeatureFlags.ASSISTANT_GIVES_LAUNCHER_FOCUS;
+import static com.android.launcher3.config.FeatureFlags.ENABLE_TRACKPAD_GESTURE;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.quickstep.GestureState.DEFAULT_STATE;
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.FLAG_USING_OTHER_ACTIVITY_INPUT_CONSUMER;
@@ -106,6 +107,7 @@
 import com.android.quickstep.inputconsumers.ProgressDelegateInputConsumer;
 import com.android.quickstep.inputconsumers.ResetGestureInputConsumer;
 import com.android.quickstep.inputconsumers.ScreenPinnedInputConsumer;
+import com.android.quickstep.inputconsumers.StatusBarInputConsumer;
 import com.android.quickstep.inputconsumers.SysUiOverlayInputConsumer;
 import com.android.quickstep.inputconsumers.TaskbarStashInputConsumer;
 import com.android.quickstep.util.ActiveGestureLog;
@@ -857,7 +859,14 @@
                         getBaseContext(), mDeviceState, mInputMonitorCompat);
             }
 
-
+            if (ENABLE_TRACKPAD_GESTURE.get() && mGestureState.isTrackpadGesture()
+                    && mGestureState.getActivityInterface().isResumed()
+                    && !previousGestureState.isRecentsAnimationRunning()) {
+                reasonString = newCompoundString(reasonPrefix)
+                        .append(SUBSTRING_PREFIX)
+                        .append("Trackpad 3-finger gesture, using StatusBarInputConsumer");
+                base = new StatusBarInputConsumer(getBaseContext(), base, mInputMonitorCompat);
+            }
 
             if (mDeviceState.isScreenPinningActive()) {
                 reasonString = newCompoundString(reasonPrefix)
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/StatusBarInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/StatusBarInputConsumer.java
new file mode 100644
index 0000000..f3d2a60
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/inputconsumers/StatusBarInputConsumer.java
@@ -0,0 +1,76 @@
+/*
+ * Copyright (C) 2023 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.quickstep.inputconsumers;
+
+import static android.view.MotionEvent.ACTION_DOWN;
+import static android.view.MotionEvent.ACTION_MOVE;
+
+import android.content.Context;
+import android.graphics.PointF;
+import android.view.MotionEvent;
+import android.view.ViewConfiguration;
+
+import com.android.quickstep.InputConsumer;
+import com.android.quickstep.SystemUiProxy;
+import com.android.systemui.shared.system.InputMonitorCompat;
+
+/** Allows the status bar to be pull down for notification shade */
+public class StatusBarInputConsumer extends DelegateInputConsumer {
+
+    private final SystemUiProxy mSystemUiProxy;
+    private final float mTouchSlop;
+    private final PointF mDown = new PointF();
+
+    public StatusBarInputConsumer(Context context, InputConsumer delegate,
+            InputMonitorCompat inputMonitor) {
+        super(delegate, inputMonitor);
+
+        mSystemUiProxy = SystemUiProxy.INSTANCE.get(context);
+        mTouchSlop = 2 * ViewConfiguration.get(context).getScaledTouchSlop();
+    }
+
+    @Override
+    public int getType() {
+        return TYPE_STATUS_BAR | mDelegate.getType();
+    }
+
+    @Override
+    public void onMotionEvent(MotionEvent ev) {
+        if (mState != STATE_ACTIVE) {
+            mDelegate.onMotionEvent(ev);
+
+            switch (ev.getActionMasked()) {
+                case ACTION_DOWN -> mDown.set(ev.getX(), ev.getY());
+                case ACTION_MOVE -> {
+                    float displacementY = ev.getY() - mDown.y;
+                    if (displacementY > mTouchSlop) {
+                        setActive(ev);
+                        ev.setAction(ACTION_DOWN);
+                        dispatchTouchEvent(ev);
+                    }
+                }
+            }
+        } else {
+            dispatchTouchEvent(ev);
+        }
+    }
+
+    private void dispatchTouchEvent(MotionEvent ev) {
+        if (mSystemUiProxy.isActive()) {
+            mSystemUiProxy.onStatusBarMotionEvent(ev);
+        }
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/interaction/AnimatedTaskView.java b/quickstep/src/com/android/quickstep/interaction/AnimatedTaskView.java
index 73937f5..44b3d62 100644
--- a/quickstep/src/com/android/quickstep/interaction/AnimatedTaskView.java
+++ b/quickstep/src/com/android/quickstep/interaction/AnimatedTaskView.java
@@ -15,6 +15,10 @@
  */
 package com.android.quickstep.interaction;
 
+import static com.android.launcher3.QuickstepTransitionManager.ANIMATION_NAV_FADE_OUT_DURATION;
+import static com.android.launcher3.QuickstepTransitionManager.NAV_FADE_OUT_INTERPOLATOR;
+import static com.android.launcher3.anim.Interpolators.FAST_OUT_SLOW_IN;
+import static com.android.launcher3.anim.Interpolators.LINEAR;
 import static com.android.launcher3.config.FeatureFlags.ENABLE_NEW_GESTURE_NAV_TUTORIAL;
 
 import android.animation.Animator;
@@ -29,6 +33,8 @@
 import android.graphics.Rect;
 import android.util.AttributeSet;
 import android.util.Log;
+import android.view.Display;
+import android.view.RoundedCorner;
 import android.view.View;
 import android.view.ViewOutlineProvider;
 import android.view.animation.ScaleAnimation;
@@ -40,8 +46,12 @@
 import androidx.constraintlayout.widget.ConstraintLayout;
 
 import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
+import com.android.launcher3.anim.Interpolators;
+import com.android.quickstep.util.MultiValueUpdateListener;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 
 /**
  * Helper View for the gesture tutorial mock previous app task view.
@@ -51,6 +61,8 @@
  */
 public class AnimatedTaskView extends ConstraintLayout {
 
+    private static final long ANIMATE_TO_FULL_SCREEN_DURATION = 300;
+
     private View mFullTaskView;
     private View mTopTaskView;
     private View mBottomTaskView;
@@ -92,34 +104,87 @@
         setToSingleRowLayout(false);
     }
 
-    void animateToFillScreen(@Nullable  Runnable onAnimationEndCallback) {
+    void animateToFillScreen(@Nullable Runnable onAnimationEndCallback) {
+        if (mTaskViewOutlineProvider == null) {
+            // This is an illegal state.
+            return;
+        }
+        // calculate start and end corner radius
+        Outline startOutline = new Outline();
+        mTaskViewOutlineProvider.getOutline(this, startOutline);
+        Rect outlineStartRect = new Rect();
+        startOutline.getRect(outlineStartRect);
+        float outlineStartRadius = startOutline.getRadius();
 
+        final Display display = mContext.getDisplay();;
+        RoundedCorner corner = display.getRoundedCorner(RoundedCorner.POSITION_TOP_LEFT);
+        float outlineEndRadius = corner.getRadius();
+
+        // create animation
         AnimatorSet set = new AnimatorSet();
         ArrayList<Animator> animations = new ArrayList<>();
 
         // center view
         animations.add(ObjectAnimator.ofFloat(this, TRANSLATION_X, 0));
 
-        // calculate full screen scaling, scale should be 1:1 for x and y
+        // retrieve start animation matrix to scale off of
         Matrix matrix = getAnimationMatrix();
+        if (matrix == null) {
+            // This is an illegal state.
+            return;
+        }
+
         float[] matrixValues = new float[9];
         matrix.getValues(matrixValues);
-        float scaleX = matrixValues[Matrix.MSCALE_X];
-        float scaleToFullScreen = 1 / scaleX;
+        float[] newValues = matrixValues.clone();
 
-        // scale view to full screen
-        ValueAnimator scale = ValueAnimator.ofFloat(1f, scaleToFullScreen);
-        scale.addUpdateListener(animation -> {
-            float value = (float) animation.getAnimatedValue();
-            mFullTaskView.setScaleX(value);
-            mFullTaskView.setScaleY(value);
-        });
+        ValueAnimator transformAnimation = ValueAnimator.ofFloat(0, 1);
 
-        animations.add(scale);
+        MultiValueUpdateListener listener = new MultiValueUpdateListener() {
+            Matrix currentMatrix = new Matrix();
+
+            FloatProp mOutlineRadius = new FloatProp(outlineStartRadius, outlineEndRadius, 0,
+                    ANIMATE_TO_FULL_SCREEN_DURATION, LINEAR);
+            FloatProp mTransX = new FloatProp(matrixValues[Matrix.MTRANS_X], 0f, 0,
+                    ANIMATE_TO_FULL_SCREEN_DURATION, LINEAR);
+            FloatProp mTransY = new FloatProp(matrixValues[Matrix.MTRANS_Y], 0f, 0,
+                    ANIMATE_TO_FULL_SCREEN_DURATION, LINEAR);
+            FloatProp mScaleX = new FloatProp(matrixValues[Matrix.MSCALE_X], 1f, 0,
+                    ANIMATE_TO_FULL_SCREEN_DURATION, LINEAR);
+            FloatProp mScaleY = new FloatProp(matrixValues[Matrix.MSCALE_Y], 1f, 0,
+                    ANIMATE_TO_FULL_SCREEN_DURATION, LINEAR);
+
+            @Override
+            public void onUpdate(float percent, boolean initOnly) {
+                // scale corner radius to match display radius
+                mTaskViewAnimatedRadius = mOutlineRadius.value;
+                mFullTaskView.invalidateOutline();
+
+                // translate to center, ends at translation x:0, y:0
+                newValues[Matrix.MTRANS_X] = mTransX.value;
+                newValues[Matrix.MTRANS_Y] = mTransY.value;
+
+                // scale to full size, ends at scale 1
+                newValues[Matrix.MSCALE_X] = mScaleX.value;
+                newValues[Matrix.MSCALE_Y] = mScaleY.value;
+
+                // create and set new animation matrix
+                currentMatrix.setValues(newValues);
+                setAnimationMatrix(currentMatrix);
+            }
+        };
+
+        transformAnimation.addUpdateListener(listener);
+        animations.add(transformAnimation);
         set.playSequentially(animations);
-
         set.addListener(new AnimatorListenerAdapter() {
             @Override
+            public void onAnimationStart(Animator animation) {
+                super.onAnimationStart(animation);
+                addAnimatedOutlineProvider(mFullTaskView, outlineStartRect, outlineStartRadius);
+            }
+
+            @Override
             public void onAnimationEnd(Animator animation) {
                 super.onAnimationEnd(animation);
                 if (onAnimationEndCallback != null) {
@@ -127,7 +192,6 @@
                 }
             }
         });
-
         set.start();
     }
 
@@ -158,17 +222,7 @@
             @Override
             public void onAnimationStart(Animator animation) {
                 super.onAnimationStart(animation);
-
-                mTaskViewAnimatedRect.set(outlineStartRect);
-                mTaskViewAnimatedRadius = outlineStartRadius;
-
-                mFullTaskView.setClipToOutline(true);
-                mFullTaskView.setOutlineProvider(new ViewOutlineProvider() {
-                    @Override
-                    public void getOutline(View view, Outline outline) {
-                        outline.setRoundRect(mTaskViewAnimatedRect, mTaskViewAnimatedRadius);
-                    }
-                });
+                addAnimatedOutlineProvider(mFullTaskView, outlineStartRect, outlineStartRadius);
             }
 
             @Override
@@ -247,4 +301,17 @@
         mTaskViewOutlineProvider = provider;
         mFullTaskView.setOutlineProvider(mTaskViewOutlineProvider);
     }
+
+    private void addAnimatedOutlineProvider(View view,
+            Rect outlineStartRect, float outlineStartRadius){
+        mTaskViewAnimatedRect.set(outlineStartRect);
+        mTaskViewAnimatedRadius = outlineStartRadius;
+        view.setClipToOutline(true);
+        view.setOutlineProvider(new ViewOutlineProvider() {
+            @Override
+            public void getOutline(View view, Outline outline) {
+                outline.setRoundRect(mTaskViewAnimatedRect, mTaskViewAnimatedRadius);
+            }
+        });
+    }
 }
diff --git a/res/drawable/enter_home_gardening_icon.xml b/res/drawable/enter_home_gardening_icon.xml
new file mode 100644
index 0000000..039258e
--- /dev/null
+++ b/res/drawable/enter_home_gardening_icon.xml
@@ -0,0 +1,25 @@
+<!--
+Copyright (C) 2023 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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24"
+    android:tint="?android:attr/textColorPrimary">
+    <path
+        android:fillColor="@android:color/white"
+        android:pathData="M5,19H6.4L16.45,8.975L15.75,8.25L15.025,7.55L5,17.6ZM3,21V16.75L16.45,3.325Q17.025,2.75 17.863,2.75Q18.7,2.75 19.275,3.325L20.675,4.75Q21.25,5.325 21.25,6.15Q21.25,6.975 20.675,7.55L7.25,21ZM19.25,6.15 L17.85,4.75ZM16.45,8.975 L15.75,8.25 15.025,7.55 16.45,8.975Z"/>
+</vector>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index aebc1d0..b54d4f7 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -250,6 +250,8 @@
     <string name="wallpaper_button_text">Wallpapers</string>
     <!-- Text for wallpaper change button [CHAR LIMIT=30]-->
     <string name="styles_wallpaper_button_text">Wallpaper &amp; style</string>
+    <!-- Text for edit home screen button [CHAR LIMIT=30]-->
+    <string name="edit_home_screen">Edit Home Screen</string>
     <!-- Text for settings button [CHAR LIMIT=20]-->
     <string name="settings_button_text">Home settings</string>
     <!-- Message shown when a feature is disabled by the administrator -->
diff --git a/src/com/android/launcher3/AbstractFloatingView.java b/src/com/android/launcher3/AbstractFloatingView.java
index cd9e598..29f4a62 100644
--- a/src/com/android/launcher3/AbstractFloatingView.java
+++ b/src/com/android/launcher3/AbstractFloatingView.java
@@ -127,6 +127,10 @@
             TYPE_WIDGETS_FULL_SHEET | TYPE_WIDGET_RESIZE_FRAME | TYPE_ON_BOARD_POPUP |
             TYPE_DISCOVERY_BOUNCE | TYPE_TASK_MENU | TYPE_DRAG_DROP_POPUP;
 
+    // Floating views that are exclusive to the taskbar overlay window.
+    public static final int TYPE_TASKBAR_OVERLAYS =
+            TYPE_TASKBAR_ALL_APPS | TYPE_TASKBAR_EDUCATION_DIALOG;
+
     protected boolean mIsOpen;
 
     public AbstractFloatingView(Context context, AttributeSet attrs) {
diff --git a/src/com/android/launcher3/AppWidgetsRestoredReceiver.java b/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
index 9d5b08e..e543370 100644
--- a/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
+++ b/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
@@ -107,12 +107,6 @@
                     cursor.close();
                 }
             }
-            // attempt to update widget id in backup table as well
-            new ContentWriter(context, ContentWriter.CommitParams.backupCommitParams(
-                    "appWidgetId=? and profileId=?", args))
-                    .put(LauncherSettings.Favorites.APPWIDGET_ID, newWidgetIds[i])
-                    .put(LauncherSettings.Favorites.RESTORED, state)
-                    .commit();
         }
 
         LauncherAppState app = LauncherAppState.getInstanceNoCreate();
diff --git a/src/com/android/launcher3/AutoInstallsLayout.java b/src/com/android/launcher3/AutoInstallsLayout.java
index 5367d80..197aa5a 100644
--- a/src/com/android/launcher3/AutoInstallsLayout.java
+++ b/src/com/android/launcher3/AutoInstallsLayout.java
@@ -81,7 +81,7 @@
     private static final String FORMATTED_LAYOUT_RES = "default_layout_%dx%d";
     private static final String LAYOUT_RES = "default_layout";
 
-    static AutoInstallsLayout get(Context context, LauncherWidgetHolder appWidgetHolder,
+    public static AutoInstallsLayout get(Context context, LauncherWidgetHolder appWidgetHolder,
             LayoutParserCallback callback) {
         Partner partner = Partner.get(context.getPackageManager(), ACTION_LAUNCHER_CUSTOMIZATION);
         if (partner == null) {
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 4c34648..5e07a3c 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -183,7 +183,7 @@
 
     public String dbFile;
     public int defaultLayoutId;
-    int demoModeLayoutId;
+    public int demoModeLayoutId;
     public boolean[] inlineQsb = new boolean[COUNT_SIZES];
 
     /**
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index f4892b2..dee3205 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -16,10 +16,6 @@
 
 package com.android.launcher3;
 
-import static com.android.launcher3.DefaultLayoutParser.RES_PARTNER_DEFAULT_LAYOUT;
-import static com.android.launcher3.provider.LauncherDbUtils.copyTable;
-import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
-
 import android.annotation.TargetApi;
 import android.appwidget.AppWidgetManager;
 import android.content.ComponentName;
@@ -28,61 +24,33 @@
 import android.content.ContentProviderResult;
 import android.content.ContentUris;
 import android.content.ContentValues;
-import android.content.Context;
 import android.content.OperationApplicationException;
-import android.content.SharedPreferences;
-import android.content.pm.ProviderInfo;
 import android.database.Cursor;
-import android.database.SQLException;
-import android.database.sqlite.SQLiteDatabase;
 import android.database.sqlite.SQLiteQueryBuilder;
 import android.net.Uri;
 import android.os.Binder;
 import android.os.Build;
 import android.os.Bundle;
 import android.os.Process;
-import android.os.UserManager;
-import android.provider.Settings;
 import android.text.TextUtils;
 import android.util.Log;
-import android.util.Xml;
 
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.config.FeatureFlags;
-import com.android.launcher3.model.DatabaseHelper;
-import com.android.launcher3.provider.LauncherDbUtils;
+import com.android.launcher3.model.ModelDbController;
 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.util.Thunk;
 import com.android.launcher3.widget.LauncherWidgetHolder;
 
-import org.xmlpull.v1.XmlPullParser;
-
 import java.io.FileDescriptor;
-import java.io.InputStream;
 import java.io.PrintWriter;
-import java.io.StringReader;
 import java.util.ArrayList;
-import java.util.function.Supplier;
 
 public class LauncherProvider extends ContentProvider {
     private static final String TAG = "LauncherProvider";
 
     public static final String AUTHORITY = BuildConfig.APPLICATION_ID + ".settings";
 
-    private static final int TEST_WORKSPACE_LAYOUT_RES_XML = R.xml.default_test_workspace;
-    private static final int TEST2_WORKSPACE_LAYOUT_RES_XML = R.xml.default_test2_workspace;
-    private static final int TAPL_WORKSPACE_LAYOUT_RES_XML = R.xml.default_tapl_test_workspace;
-
-    public static final String EMPTY_DATABASE_CREATED = "EMPTY_DATABASE_CREATED";
-
-    protected DatabaseHelper mOpenHelper;
-    protected String mProviderAuthority;
-
-    private int mDefaultWorkspaceLayoutOverride = 0;
+    protected ModelDbController mModelDbController;
 
     /**
      * $ adb shell dumpsys activity provider com.android.launcher3
@@ -101,6 +69,7 @@
         if (FeatureFlags.IS_STUDIO_BUILD) {
             Log.d(TAG, "Launcher process started");
         }
+        mModelDbController = new ModelDbController(getContext());
 
         // The content provider exists for the entire duration of the launcher main process and
         // is the first component to get created.
@@ -118,49 +87,17 @@
         }
     }
 
-    /**
-     * Overridden in tests
-     */
-    protected synchronized void createDbIfNotExists() {
-        if (mOpenHelper == null) {
-            mOpenHelper = DatabaseHelper.createDatabaseHelper(
-                    getContext(), false /* forMigration */);
-
-            RestoreDbTask.restoreIfNeeded(getContext(), mOpenHelper);
-        }
-    }
-
-    private synchronized boolean prepForMigration(String dbFile, String targetTableName,
-            Supplier<DatabaseHelper> src, Supplier<DatabaseHelper> dst) {
-        if (TextUtils.equals(dbFile, mOpenHelper.getDatabaseName())) {
-            Log.e(TAG, "prepForMigration - target db is same as current: " + dbFile);
-            return false;
-        }
-
-        final DatabaseHelper helper = src.get();
-        mOpenHelper = dst.get();
-        copyTable(helper.getReadableDatabase(), Favorites.TABLE_NAME,
-                mOpenHelper.getWritableDatabase(), targetTableName, getContext());
-        helper.close();
-        return true;
-    }
-
     @Override
     public Cursor query(Uri uri, String[] projection, String selection,
             String[] selectionArgs, String sortOrder) {
-        createDbIfNotExists();
 
         SqlArguments args = new SqlArguments(uri, selection, selectionArgs);
         SQLiteQueryBuilder qb = new SQLiteQueryBuilder();
         qb.setTables(args.table);
 
-        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-        Cursor result = qb.query(db, projection, args.where, args.args, null, null, sortOrder);
-        final Bundle extra = new Bundle();
-        extra.putString(LauncherSettings.Settings.EXTRA_DB_NAME, mOpenHelper.getDatabaseName());
-        result.setExtras(extra);
+        Cursor result = mModelDbController.query(
+                args.table, projection, args.where, args.args, sortOrder);
         result.setNotificationUri(getContext().getContentResolver(), uri);
-
         return result;
     }
 
@@ -175,9 +112,6 @@
 
     @Override
     public Uri insert(Uri uri, ContentValues initialValues) {
-        createDbIfNotExists();
-        SqlArguments args = new SqlArguments(uri);
-
         // In very limited cases, we support system|signature permission apps to modify the db.
         if (Binder.getCallingPid() != Process.myPid()) {
             if (!initializeExternalAdd(initialValues)) {
@@ -185,11 +119,9 @@
             }
         }
 
-        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-        addModifiedTime(initialValues);
-        final int rowId = mOpenHelper.dbInsertAndCheck(db, args.table, initialValues);
+        SqlArguments args = new SqlArguments(uri);
+        int rowId = mModelDbController.insert(args.table, initialValues);
         if (rowId < 0) return null;
-        onAddOrDeleteOp(db);
 
         uri = ContentUris.withAppendedId(uri, rowId);
         reloadLauncherIfExternal();
@@ -198,7 +130,7 @@
 
     private boolean initializeExternalAdd(ContentValues values) {
         // 1. Ensure that externally added items have a valid item id
-        int id = mOpenHelper.generateNewItemId();
+        int id = mModelDbController.generateNewItemId();
         values.put(LauncherSettings.Favorites._ID, id);
 
         // 2. In the case of an app widget, and if no app widget id is specified, we
@@ -213,7 +145,7 @@
                     values.getAsString(Favorites.APPWIDGET_PROVIDER));
 
             if (cn != null) {
-                LauncherWidgetHolder widgetHolder = mOpenHelper.newLauncherWidgetHolder();
+                LauncherWidgetHolder widgetHolder = LauncherWidgetHolder.newInstance(getContext());
                 try {
                     int appWidgetId = widgetHolder.allocateAppWidgetId();
                     values.put(LauncherSettings.Favorites.APPWIDGET_ID, appWidgetId);
@@ -238,22 +170,8 @@
 
     @Override
     public int bulkInsert(Uri uri, ContentValues[] values) {
-        createDbIfNotExists();
         SqlArguments args = new SqlArguments(uri);
-
-        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-        try (SQLiteTransaction t = new SQLiteTransaction(db)) {
-            int numValues = values.length;
-            for (int i = 0; i < numValues; i++) {
-                addModifiedTime(values[i]);
-                if (mOpenHelper.dbInsertAndCheck(db, args.table, values[i]) < 0) {
-                    return 0;
-                }
-            }
-            onAddOrDeleteOp(db);
-            t.commit();
-        }
-
+        mModelDbController.bulkInsert(args.table, values);
         reloadLauncherIfExternal();
         return values.length;
     }
@@ -262,23 +180,13 @@
     @Override
     public ContentProviderResult[] applyBatch(ArrayList<ContentProviderOperation> operations)
             throws OperationApplicationException {
-        createDbIfNotExists();
-        try (SQLiteTransaction t = new SQLiteTransaction(mOpenHelper.getWritableDatabase())) {
-            boolean isAddOrDelete = false;
-
+        try (SQLiteTransaction t = mModelDbController.newTransaction()) {
             final int numOperations = operations.size();
             final ContentProviderResult[] results = new ContentProviderResult[numOperations];
             for (int i = 0; i < numOperations; i++) {
                 ContentProviderOperation op = operations.get(i);
                 results[i] = op.apply(this, results, i);
-
-                isAddOrDelete |= (op.isInsert() || op.isDelete()) &&
-                        results[i].count != null && results[i].count > 0;
             }
-            if (isAddOrDelete) {
-                onAddOrDeleteOp(t.getDb());
-            }
-
             t.commit();
             reloadLauncherIfExternal();
             return results;
@@ -287,18 +195,9 @@
 
     @Override
     public int delete(Uri uri, String selection, String[] selectionArgs) {
-        createDbIfNotExists();
         SqlArguments args = new SqlArguments(uri, selection, selectionArgs);
-
-        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-
-        if (Binder.getCallingPid() != Process.myPid()
-                && Favorites.TABLE_NAME.equalsIgnoreCase(args.table)) {
-            mOpenHelper.removeGhostWidgets(mOpenHelper.getWritableDatabase());
-        }
-        int count = db.delete(args.table, args.where, args.args);
+        int count = mModelDbController.delete(args.table, args.where, args.args);
         if (count > 0) {
-            onAddOrDeleteOp(db);
             reloadLauncherIfExternal();
         }
         return count;
@@ -306,12 +205,8 @@
 
     @Override
     public int update(Uri uri, ContentValues values, String selection, String[] selectionArgs) {
-        createDbIfNotExists();
         SqlArguments args = new SqlArguments(uri, selection, selectionArgs);
-
-        addModifiedTime(values);
-        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-        int count = db.update(args.table, values, args.where, args.args);
+        int count = mModelDbController.update(args.table, values, args.where, args.args);
         reloadLauncherIfExternal();
         return count;
     }
@@ -321,260 +216,76 @@
         if (Binder.getCallingUid() != Process.myUid()) {
             return null;
         }
-        createDbIfNotExists();
 
         switch (method) {
             case LauncherSettings.Settings.METHOD_CLEAR_EMPTY_DB_FLAG: {
-                clearFlagEmptyDbCreated();
+                mModelDbController.clearEmptyDbFlag();
                 return null;
             }
             case LauncherSettings.Settings.METHOD_DELETE_EMPTY_FOLDERS: {
                 Bundle result = new Bundle();
-                result.putIntArray(LauncherSettings.Settings.EXTRA_VALUE, deleteEmptyFolders()
-                        .toArray());
+                result.putIntArray(LauncherSettings.Settings.EXTRA_VALUE,
+                        mModelDbController.deleteEmptyFolders().toArray());
                 return result;
             }
             case LauncherSettings.Settings.METHOD_NEW_ITEM_ID: {
                 Bundle result = new Bundle();
                 result.putInt(LauncherSettings.Settings.EXTRA_VALUE,
-                        mOpenHelper.generateNewItemId());
+                        mModelDbController.generateNewItemId());
                 return result;
             }
             case LauncherSettings.Settings.METHOD_NEW_SCREEN_ID: {
                 Bundle result = new Bundle();
                 result.putInt(LauncherSettings.Settings.EXTRA_VALUE,
-                        mOpenHelper.getNewScreenId());
+                        mModelDbController.getNewScreenId());
                 return result;
             }
             case LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB: {
-                mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
+                mModelDbController.createEmptyDB();
                 return null;
             }
             case LauncherSettings.Settings.METHOD_SET_USE_TEST_WORKSPACE_LAYOUT_FLAG: {
-                switch (arg) {
-                    case LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TEST:
-                        mDefaultWorkspaceLayoutOverride = TEST_WORKSPACE_LAYOUT_RES_XML;
-                        break;
-                    case LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TEST2:
-                        mDefaultWorkspaceLayoutOverride = TEST2_WORKSPACE_LAYOUT_RES_XML;
-                        break;
-                    case LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TAPL:
-                        mDefaultWorkspaceLayoutOverride = TAPL_WORKSPACE_LAYOUT_RES_XML;
-                        break;
-                    default:
-                        mDefaultWorkspaceLayoutOverride = 0;
-                        break;
-                }
+                mModelDbController.setUseTestWorkspaceLayout(arg);
                 return null;
             }
             case LauncherSettings.Settings.METHOD_CLEAR_USE_TEST_WORKSPACE_LAYOUT_FLAG: {
-                mDefaultWorkspaceLayoutOverride = 0;
+                mModelDbController.setUseTestWorkspaceLayout(null);
                 return null;
             }
             case LauncherSettings.Settings.METHOD_LOAD_DEFAULT_FAVORITES: {
-                loadDefaultFavoritesIfNecessary();
+                mModelDbController.loadDefaultFavoritesIfNecessary();
                 return null;
             }
             case LauncherSettings.Settings.METHOD_REMOVE_GHOST_WIDGETS: {
-                mOpenHelper.removeGhostWidgets(mOpenHelper.getWritableDatabase());
+                mModelDbController.removeGhostWidgets();
                 return null;
             }
             case LauncherSettings.Settings.METHOD_NEW_TRANSACTION: {
                 Bundle result = new Bundle();
                 result.putBinder(LauncherSettings.Settings.EXTRA_VALUE,
-                        new SQLiteTransaction(mOpenHelper.getWritableDatabase()));
+                        mModelDbController.newTransaction());
                 return result;
             }
             case LauncherSettings.Settings.METHOD_REFRESH_HOTSEAT_RESTORE_TABLE: {
-                mOpenHelper.mHotseatRestoreTableExists = tableExists(
-                        mOpenHelper.getReadableDatabase(), Favorites.HYBRID_HOTSEAT_BACKUP_TABLE);
+                mModelDbController.refreshHotseatRestoreTable();
                 return null;
             }
             case LauncherSettings.Settings.METHOD_UPDATE_CURRENT_OPEN_HELPER: {
                 Bundle result = new Bundle();
                 result.putBoolean(LauncherSettings.Settings.EXTRA_VALUE,
-                        prepForMigration(
-                                arg /* dbFile */,
-                                Favorites.TMP_TABLE,
-                                () -> mOpenHelper,
-                                () -> DatabaseHelper.createDatabaseHelper(
-                                        getContext(), true /* forMigration */)));
+                        mModelDbController.updateCurrentOpenHelper(arg /* dbFile */));
                 return result;
             }
             case LauncherSettings.Settings.METHOD_PREP_FOR_PREVIEW: {
                 Bundle result = new Bundle();
                 result.putBoolean(LauncherSettings.Settings.EXTRA_VALUE,
-                        prepForMigration(
-                                arg /* dbFile */,
-                                Favorites.PREVIEW_TABLE_NAME,
-                                () -> DatabaseHelper.createDatabaseHelper(
-                                        getContext(), arg, true /* forMigration */),
-                                () -> mOpenHelper));
+                        mModelDbController.prepareForPreview(arg /* dbFile */));
                 return result;
             }
         }
         return null;
     }
 
-    private void onAddOrDeleteOp(SQLiteDatabase db) {
-        mOpenHelper.onAddOrDeleteOp(db);
-    }
-
-    /**
-     * Deletes any empty folder from the DB.
-     * @return Ids of deleted folders.
-     */
-    private IntArray deleteEmptyFolders() {
-        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-        try (SQLiteTransaction t = new SQLiteTransaction(db)) {
-            // Select folders whose id do not match any container value.
-            String selection = LauncherSettings.Favorites.ITEM_TYPE + " = "
-                    + LauncherSettings.Favorites.ITEM_TYPE_FOLDER + " AND "
-                    + LauncherSettings.Favorites._ID +  " NOT IN (SELECT " +
-                            LauncherSettings.Favorites.CONTAINER + " FROM "
-                                + Favorites.TABLE_NAME + ")";
-
-            IntArray folderIds = LauncherDbUtils.queryIntArray(false, db, Favorites.TABLE_NAME,
-                    Favorites._ID, selection, null, null);
-            if (!folderIds.isEmpty()) {
-                db.delete(Favorites.TABLE_NAME, Utilities.createDbSelectionQuery(
-                        LauncherSettings.Favorites._ID, folderIds), null);
-            }
-            t.commit();
-            return folderIds;
-        } catch (SQLException ex) {
-            Log.e(TAG, ex.getMessage(), ex);
-            return new IntArray();
-        }
-    }
-
-    @Thunk static void addModifiedTime(ContentValues values) {
-        values.put(LauncherSettings.Favorites.MODIFIED, System.currentTimeMillis());
-    }
-
-    private void clearFlagEmptyDbCreated() {
-        LauncherPrefs.getPrefs(getContext()).edit()
-                .remove(mOpenHelper.getKey(EMPTY_DATABASE_CREATED)).commit();
-    }
-
-    /**
-     * Loads the default workspace based on the following priority scheme:
-     *   1) From the app restrictions
-     *   2) From a package provided by play store
-     *   3) From a partner configuration APK, already in the system image
-     *   4) The default configuration for the particular device
-     */
-    synchronized private void loadDefaultFavoritesIfNecessary() {
-        SharedPreferences sp = LauncherPrefs.getPrefs(getContext());
-
-        if (sp.getBoolean(mOpenHelper.getKey(EMPTY_DATABASE_CREATED), false)) {
-            Log.d(TAG, "loading default workspace");
-
-            LauncherWidgetHolder widgetHolder = mOpenHelper.newLauncherWidgetHolder();
-            try {
-                AutoInstallsLayout loader = createWorkspaceLoaderFromAppRestriction(widgetHolder);
-                if (loader == null) {
-                    loader = AutoInstallsLayout.get(getContext(), widgetHolder, mOpenHelper);
-                }
-                if (loader == null) {
-                    final Partner partner = Partner.get(getContext().getPackageManager());
-                    if (partner != null) {
-                        int workspaceResId = partner.getXmlResId(RES_PARTNER_DEFAULT_LAYOUT);
-                        if (workspaceResId != 0) {
-                            loader = new DefaultLayoutParser(getContext(), widgetHolder,
-                                    mOpenHelper, partner.getResources(), workspaceResId);
-                        }
-                    }
-                }
-
-                final boolean usingExternallyProvidedLayout = loader != null;
-                if (loader == null) {
-                    loader = getDefaultLayoutParser(widgetHolder);
-                }
-
-                // There might be some partially restored DB items, due to buggy restore logic in
-                // previous versions of launcher.
-                mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
-                // Populate favorites table with initial favorites
-                if ((mOpenHelper.loadFavorites(mOpenHelper.getWritableDatabase(), loader) <= 0)
-                        && usingExternallyProvidedLayout) {
-                    // Unable to load external layout. Cleanup and load the internal layout.
-                    mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
-                    mOpenHelper.loadFavorites(mOpenHelper.getWritableDatabase(),
-                            getDefaultLayoutParser(widgetHolder));
-                }
-                clearFlagEmptyDbCreated();
-            } finally {
-                widgetHolder.destroy();
-            }
-        }
-    }
-
-    /**
-     * 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) {
-        Context ctx = getContext();
-        final String authority;
-        if (!TextUtils.isEmpty(mProviderAuthority)) {
-            authority = mProviderAuthority;
-        } else {
-            authority = Settings.Secure.getString(ctx.getContentResolver(),
-                    "launcher3.layout.provider");
-        }
-        if (TextUtils.isEmpty(authority)) {
-            return null;
-        }
-
-        ProviderInfo pi = ctx.getPackageManager().resolveContentProvider(authority, 0);
-        if (pi == null) {
-            Log.e(TAG, "No provider found for authority " + authority);
-            return null;
-        }
-        Uri uri = getLayoutUri(authority, ctx);
-        try (InputStream in = ctx.getContentResolver().openInputStream(uri)) {
-            // 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));
-
-            Log.d(TAG, "Loading layout from " + authority);
-            return new AutoInstallsLayout(ctx, widgetHolder, mOpenHelper,
-                    ctx.getPackageManager().getResourcesForApplication(pi.applicationInfo),
-                    () -> parser, AutoInstallsLayout.TAG_WORKSPACE);
-        } catch (Exception e) {
-            Log.e(TAG, "Error getting layout stream from: " + authority , e);
-            return null;
-        }
-    }
-
-    public static Uri getLayoutUri(String authority, Context ctx) {
-        InvariantDeviceProfile grid = LauncherAppState.getIDP(ctx);
-        return new Uri.Builder().scheme("content").authority(authority).path("launcher_layout")
-                .appendQueryParameter("version", "1")
-                .appendQueryParameter("gridWidth", Integer.toString(grid.numColumns))
-                .appendQueryParameter("gridHeight", Integer.toString(grid.numRows))
-                .appendQueryParameter("hotseatSize", Integer.toString(grid.numDatabaseHotseatIcons))
-                .build();
-    }
-
-    private DefaultLayoutParser getDefaultLayoutParser(LauncherWidgetHolder widgetHolder) {
-        InvariantDeviceProfile idp = LauncherAppState.getIDP(getContext());
-        int defaultLayout = mDefaultWorkspaceLayoutOverride > 0
-                ? mDefaultWorkspaceLayoutOverride : idp.defaultLayoutId;
-
-        if (getContext().getSystemService(UserManager.class).isDemoUser()
-                && idp.demoModeLayoutId != 0) {
-            defaultLayout = idp.demoModeLayoutId;
-        }
-
-        return new DefaultLayoutParser(getContext(), widgetHolder,
-                mOpenHelper, getContext().getResources(), defaultLayout);
-    }
-
     static class SqlArguments {
         public final String table;
         public final String where;
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index 1bbb09a..82dbe53 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -139,11 +139,6 @@
         public static final String TABLE_NAME = "favorites";
 
         /**
-         * Backup table created when the favorites table is modified during grid migration
-         */
-        public static final String BACKUP_TABLE_NAME = "favorites_bakup";
-
-        /**
          * Backup table created when user hotseat is moved to workspace for hybrid hotseat
          */
         public static final String HYBRID_HOTSEAT_BACKUP_TABLE = "hotseat_restore_backup";
@@ -165,12 +160,6 @@
                 + LauncherProvider.AUTHORITY + "/" + TABLE_NAME);
 
         /**
-         * The content:// style URL for "favorites_bakup" table
-         */
-        public static final Uri BACKUP_CONTENT_URI = Uri.parse("content://"
-                + LauncherProvider.AUTHORITY + "/" + BACKUP_TABLE_NAME);
-
-        /**
          * The content:// style URL for "favorites_preview" table
          */
         public static final Uri PREVIEW_CONTENT_URI = Uri.parse("content://"
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index 7f46324..5b4a02b 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -797,6 +797,7 @@
 
         enum LatencyType {
             UNKNOWN(0),
+            // example: launcher restart that happens via daily backup and restore
             COLD(1),
             HOT(2),
             TIMEOUT(3),
@@ -804,7 +805,9 @@
             COLD_USERWAITING(5),
             ATOMIC(6),
             CONTROLLED(7),
-            CACHED(8);
+            CACHED(8),
+            // example: device is rebooting via power key or shell command `adb reboot`
+            COLD_DEVICE_REBOOTING(9);
             private final int mId;
 
             LatencyType(int id) {
diff --git a/src/com/android/launcher3/model/DatabaseHelper.java b/src/com/android/launcher3/model/DatabaseHelper.java
index 1840b75..dc5fcf7 100644
--- a/src/com/android/launcher3/model/DatabaseHelper.java
+++ b/src/com/android/launcher3/model/DatabaseHelper.java
@@ -39,7 +39,6 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherFiles;
 import com.android.launcher3.LauncherPrefs;
-import com.android.launcher3.LauncherProvider;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.Utilities;
@@ -77,6 +76,7 @@
     private static final boolean LOGD = false;
 
     private static final String DOWNGRADE_SCHEMA_FILE = "downgrade_schema.json";
+    public static final String EMPTY_DATABASE_CREATED = "EMPTY_DATABASE_CREATED";
 
     private final Context mContext;
     private final boolean mForMigration;
@@ -165,8 +165,7 @@
      */
     protected void onEmptyDbCreated() {
         // Set the flag for empty DB
-        LauncherPrefs.getPrefs(mContext).edit().putBoolean(getKey(
-                        LauncherProvider.EMPTY_DATABASE_CREATED), true)
+        LauncherPrefs.getPrefs(mContext).edit().putBoolean(getKey(EMPTY_DATABASE_CREATED), true)
                 .commit();
     }
 
diff --git a/src/com/android/launcher3/model/GridBackupTable.java b/src/com/android/launcher3/model/GridBackupTable.java
index 51cbf4b..6bd8518 100644
--- a/src/com/android/launcher3/model/GridBackupTable.java
+++ b/src/com/android/launcher3/model/GridBackupTable.java
@@ -15,18 +15,12 @@
  */
 package com.android.launcher3.model;
 
-import static com.android.launcher3.LauncherSettings.Favorites.BACKUP_TABLE_NAME;
 import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
 import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
 
-import android.content.ContentValues;
 import android.content.Context;
-import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
 import android.os.Process;
-import android.util.Log;
-
-import androidx.annotation.IntDef;
 
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.pm.UserCache;
@@ -36,50 +30,13 @@
  * within the same data base.
  */
 public class GridBackupTable {
-    private static final String TAG = "GridBackupTable";
-
-    private static final int ID_PROPERTY = -1;
-
-    private static final String KEY_HOTSEAT_SIZE = Favorites.SCREEN;
-    private static final String KEY_GRID_X_SIZE = Favorites.SPANX;
-    private static final String KEY_GRID_Y_SIZE = Favorites.SPANY;
-    private static final String KEY_DB_VERSION = Favorites.RANK;
-
-    public static final int OPTION_REQUIRES_SANITIZATION = 1;
-
-    /** STATE_NOT_FOUND indicates backup doesn't exist in the db. */
-    private static final int STATE_NOT_FOUND = 0;
-    /**
-     *  STATE_RAW indicates the backup has not yet been sanitized. This implies it might still
-     *  posses app info that doesn't exist in the workspace and needed to be sanitized before
-     *  put into use.
-     */
-    private static final int STATE_RAW = 1;
-    /** STATE_SANITIZED indicates the backup has already been sanitized, thus can be used as-is. */
-    private static final int STATE_SANITIZED = 2;
 
     private final Context mContext;
     private final SQLiteDatabase mDb;
 
-    private final int mOldHotseatSize;
-    private final int mOldGridX;
-    private final int mOldGridY;
-
-    private int mRestoredHotseatSize;
-    private int mRestoredGridX;
-    private int mRestoredGridY;
-
-    @IntDef({STATE_NOT_FOUND, STATE_RAW, STATE_SANITIZED})
-    private @interface BackupState { }
-
-    public GridBackupTable(Context context, SQLiteDatabase db, int hotseatSize, int gridX,
-            int gridY) {
+    public GridBackupTable(Context context, SQLiteDatabase db) {
         mContext = context;
         mDb = db;
-
-        mOldHotseatSize = hotseatSize;
-        mOldGridX = gridX;
-        mOldGridY = gridY;
     }
 
     /**
@@ -89,7 +46,6 @@
         long profileId = UserCache.INSTANCE.get(mContext).getSerialNumberForUser(
                 Process.myUserHandle());
         copyTable(mDb, Favorites.TABLE_NAME, tableName, profileId);
-        encodeDBProperties(0);
     }
 
     /**
@@ -114,78 +70,6 @@
     private static void copyTable(SQLiteDatabase db, String from, String to, long userSerial) {
         dropTable(db, to);
         Favorites.addTableToDb(db, userSerial, false, to);
-        db.execSQL("INSERT INTO " + to + " SELECT * FROM " + from + " where _id > " + ID_PROPERTY);
-    }
-
-    private void encodeDBProperties(int options) {
-        ContentValues values = new ContentValues();
-        values.put(Favorites._ID, ID_PROPERTY);
-        values.put(KEY_DB_VERSION, mDb.getVersion());
-        values.put(KEY_GRID_X_SIZE, mOldGridX);
-        values.put(KEY_GRID_Y_SIZE, mOldGridY);
-        values.put(KEY_HOTSEAT_SIZE, mOldHotseatSize);
-        values.put(Favorites.OPTIONS, options);
-        mDb.insert(BACKUP_TABLE_NAME, null, values);
-    }
-
-    /**
-     * Load DB properties from grid backup table.
-     */
-    public @BackupState int loadDBProperties() {
-        try (Cursor c = mDb.query(BACKUP_TABLE_NAME, new String[] {
-                KEY_DB_VERSION,     // 0
-                KEY_GRID_X_SIZE,    // 1
-                KEY_GRID_Y_SIZE,    // 2
-                KEY_HOTSEAT_SIZE,   // 3
-                Favorites.OPTIONS}, // 4
-                "_id=" + ID_PROPERTY, null, null, null, null)) {
-            if (!c.moveToNext()) {
-                Log.e(TAG, "Meta data not found in backup table");
-                return STATE_NOT_FOUND;
-            }
-            if (!validateDBVersion(mDb.getVersion(), c.getInt(0))) {
-                return STATE_NOT_FOUND;
-            }
-
-            mRestoredGridX = c.getInt(1);
-            mRestoredGridY = c.getInt(2);
-            mRestoredHotseatSize = c.getInt(3);
-            boolean isSanitized = (c.getInt(4) & OPTION_REQUIRES_SANITIZATION) == 0;
-            return isSanitized ? STATE_SANITIZED : STATE_RAW;
-        }
-    }
-
-    /**
-     * Restore workspace from raw backup if available.
-     */
-    public boolean restoreFromRawBackupIfAvailable(long oldProfileId) {
-        if (!tableExists(mDb, Favorites.BACKUP_TABLE_NAME)
-                || loadDBProperties() != STATE_RAW
-                || mOldHotseatSize != mRestoredHotseatSize
-                || mOldGridX != mRestoredGridX
-                || mOldGridY != mRestoredGridY) {
-            // skip restore if dimensions in backup table differs from current setup.
-            return false;
-        }
-        copyTable(mDb, Favorites.BACKUP_TABLE_NAME, Favorites.TABLE_NAME, oldProfileId);
-        Log.d(TAG, "Backup restored");
-        return true;
-    }
-
-    /**
-     * Performs a backup on the workspace layout.
-     */
-    public void doBackup(long profileId, int options) {
-        copyTable(mDb, Favorites.TABLE_NAME, Favorites.BACKUP_TABLE_NAME, profileId);
-        encodeDBProperties(options);
-    }
-
-    private static boolean validateDBVersion(int expected, int actual) {
-        if (expected != actual) {
-            Log.e(TAG, String.format("Launcher.db version mismatch, expecting %d but %d was found",
-                    expected, actual));
-            return false;
-        }
-        return true;
+        db.execSQL("INSERT INTO " + to + " SELECT * FROM " + from);
     }
 }
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
new file mode 100644
index 0000000..7452bcd
--- /dev/null
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -0,0 +1,462 @@
+/*
+ * Copyright (C) 2023 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 static com.android.launcher3.DefaultLayoutParser.RES_PARTNER_DEFAULT_LAYOUT;
+import static com.android.launcher3.model.DatabaseHelper.EMPTY_DATABASE_CREATED;
+import static com.android.launcher3.provider.LauncherDbUtils.copyTable;
+import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
+
+import android.content.ContentValues;
+import android.content.Context;
+import android.content.SharedPreferences;
+import android.content.pm.ProviderInfo;
+import android.database.Cursor;
+import android.database.SQLException;
+import android.database.sqlite.SQLiteDatabase;
+import android.net.Uri;
+import android.os.Binder;
+import android.os.Bundle;
+import android.os.Process;
+import android.os.UserManager;
+import android.provider.Settings;
+import android.text.TextUtils;
+import android.util.Log;
+import android.util.Xml;
+
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.AutoInstallsLayout;
+import com.android.launcher3.DefaultLayoutParser;
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherPrefs;
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
+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.InputStream;
+import java.io.StringReader;
+import java.util.function.Supplier;
+
+/**
+ * Utility class which maintains an instance of Launcher database and provides utility methods
+ * around it.
+ */
+public class ModelDbController {
+    private static final String TAG = "LauncherProvider";
+
+    private static final int TEST_WORKSPACE_LAYOUT_RES_XML = R.xml.default_test_workspace;
+    private static final int TEST2_WORKSPACE_LAYOUT_RES_XML = R.xml.default_test2_workspace;
+    private static final int TAPL_WORKSPACE_LAYOUT_RES_XML = R.xml.default_tapl_test_workspace;
+
+    protected DatabaseHelper mOpenHelper;
+    protected String mProviderAuthority;
+
+    private int mDefaultWorkspaceLayoutOverride = 0;
+
+    private final Context mContext;
+
+    public ModelDbController(Context context) {
+        mContext = context;
+    }
+
+    private synchronized void createDbIfNotExists() {
+        if (mOpenHelper == null) {
+            mOpenHelper = DatabaseHelper.createDatabaseHelper(
+                    mContext, false /* forMigration */);
+
+            RestoreDbTask.restoreIfNeeded(mContext, mOpenHelper);
+        }
+    }
+
+    private synchronized boolean prepForMigration(String dbFile, String targetTableName,
+            Supplier<DatabaseHelper> src, Supplier<DatabaseHelper> dst) {
+        if (TextUtils.equals(dbFile, mOpenHelper.getDatabaseName())) {
+            Log.e(TAG, "prepForMigration - target db is same as current: " + dbFile);
+            return false;
+        }
+
+        final DatabaseHelper helper = src.get();
+        mOpenHelper = dst.get();
+        copyTable(helper.getReadableDatabase(), Favorites.TABLE_NAME,
+                mOpenHelper.getWritableDatabase(), targetTableName, mContext);
+        helper.close();
+        return true;
+    }
+
+    /**
+     * Refer {@link SQLiteDatabase#query}
+     */
+    public Cursor query(String table, String[] projection, String selection,
+            String[] selectionArgs, String sortOrder) {
+        createDbIfNotExists();
+        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
+        Cursor result = db.query(
+                table, projection, selection, selectionArgs, null, null, sortOrder);
+
+        final Bundle extra = new Bundle();
+        extra.putString(LauncherSettings.Settings.EXTRA_DB_NAME, mOpenHelper.getDatabaseName());
+        result.setExtras(extra);
+        return result;
+    }
+
+    /**
+     * Refer {@link SQLiteDatabase#insert(String, String, ContentValues)}
+     */
+    public int insert(String table, ContentValues initialValues) {
+        createDbIfNotExists();
+
+        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
+        addModifiedTime(initialValues);
+        int rowId = mOpenHelper.dbInsertAndCheck(db, table, initialValues);
+        if (rowId >= 0) {
+            onAddOrDeleteOp(db);
+        }
+        return rowId;
+    }
+
+    /**
+     * Similar to insert but for adding multiple values in a transaction.
+     */
+    public int bulkInsert(String table, ContentValues[] values) {
+        createDbIfNotExists();
+
+        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
+        try (SQLiteTransaction t = new SQLiteTransaction(db)) {
+            int numValues = values.length;
+            for (int i = 0; i < numValues; i++) {
+                addModifiedTime(values[i]);
+                if (mOpenHelper.dbInsertAndCheck(db, table, values[i]) < 0) {
+                    return 0;
+                }
+            }
+            onAddOrDeleteOp(db);
+            t.commit();
+        }
+        return values.length;
+    }
+
+    /**
+     * Refer {@link SQLiteDatabase#delete(String, String, String[])}
+     */
+    public int delete(String table, String selection, String[] selectionArgs) {
+        createDbIfNotExists();
+        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
+
+        if (Binder.getCallingPid() != Process.myPid()
+                && Favorites.TABLE_NAME.equalsIgnoreCase(table)) {
+            mOpenHelper.removeGhostWidgets(mOpenHelper.getWritableDatabase());
+        }
+        int count = db.delete(table, selection, selectionArgs);
+        if (count > 0) {
+            onAddOrDeleteOp(db);
+        }
+        return count;
+    }
+
+    /**
+     * Refer {@link SQLiteDatabase#update(String, ContentValues, String, String[])}
+     */
+    public int update(String table, ContentValues values,
+            String selection, String[] selectionArgs) {
+        createDbIfNotExists();
+
+        addModifiedTime(values);
+        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
+        int count = db.update(table, values, selection, selectionArgs);
+        return count;
+    }
+
+    /**
+     * Clears a previously set flag corresponding to empty db creation
+     */
+    public void clearEmptyDbFlag() {
+        createDbIfNotExists();
+        clearFlagEmptyDbCreated();
+    }
+
+    /**
+     * Generates an id to be used for new item in the favorites table
+     */
+    public int generateNewItemId() {
+        createDbIfNotExists();
+        return mOpenHelper.generateNewItemId();
+    }
+
+    /**
+     * Generates an id to be used for new workspace screen
+     */
+    public int getNewScreenId() {
+        createDbIfNotExists();
+        return mOpenHelper.getNewScreenId();
+    }
+
+    /**
+     * Creates an empty DB clearing all existing data
+     */
+    public void createEmptyDB() {
+        createDbIfNotExists();
+        mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
+    }
+
+    /**
+     * Overrides the default xml to be used for setting up workspace
+     */
+    public void setUseTestWorkspaceLayout(@Nullable String layout) {
+        if (LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TEST.equals(layout)) {
+            mDefaultWorkspaceLayoutOverride = TEST_WORKSPACE_LAYOUT_RES_XML;
+        } else if (LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TEST2.equals(layout)) {
+            mDefaultWorkspaceLayoutOverride = TEST2_WORKSPACE_LAYOUT_RES_XML;
+        } else if (LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TAPL.equals(layout)) {
+            mDefaultWorkspaceLayoutOverride = TAPL_WORKSPACE_LAYOUT_RES_XML;
+        } else {
+            mDefaultWorkspaceLayoutOverride = 0;
+        }
+    }
+
+    /**
+     * Removes any widget which are present in the framework, but not in out internal DB
+     */
+    public void removeGhostWidgets() {
+        createDbIfNotExists();
+        mOpenHelper.removeGhostWidgets(mOpenHelper.getWritableDatabase());
+    }
+
+    /**
+     * Returns a new {@link SQLiteTransaction}
+     */
+    public SQLiteTransaction newTransaction() {
+        createDbIfNotExists();
+        return new SQLiteTransaction(mOpenHelper.getWritableDatabase());
+    }
+
+    /**
+     * Refreshes the internal state corresponding to presence of hotseat table
+     */
+    public void refreshHotseatRestoreTable() {
+        createDbIfNotExists();
+        mOpenHelper.mHotseatRestoreTableExists = tableExists(
+                mOpenHelper.getReadableDatabase(), Favorites.HYBRID_HOTSEAT_BACKUP_TABLE);
+    }
+
+    /**
+     * Updates the current DB and copies all the existing data to the temp table
+     * @param dbFile name of the target db file name
+     */
+    public boolean updateCurrentOpenHelper(String dbFile) {
+        createDbIfNotExists();
+        return prepForMigration(
+                dbFile,
+                Favorites.TMP_TABLE,
+                () -> mOpenHelper,
+                () -> DatabaseHelper.createDatabaseHelper(
+                        mContext, true /* forMigration */));
+    }
+
+    /**
+     * Returns the current DatabaseHelper.
+     * Only for tests
+     */
+    public DatabaseHelper getDatabaseHelper() {
+        createDbIfNotExists();
+        return mOpenHelper;
+    }
+
+    /**
+     * Prepares the DB for preview by copying all existing data to preview table
+     */
+    public boolean prepareForPreview(String dbFile) {
+        createDbIfNotExists();
+        return prepForMigration(
+                dbFile,
+                Favorites.PREVIEW_TABLE_NAME,
+                () -> DatabaseHelper.createDatabaseHelper(
+                        mContext, dbFile, true /* forMigration */),
+                () -> mOpenHelper);
+    }
+
+    private void onAddOrDeleteOp(SQLiteDatabase db) {
+        mOpenHelper.onAddOrDeleteOp(db);
+    }
+
+    /**
+     * Deletes any empty folder from the DB.
+     * @return Ids of deleted folders.
+     */
+    public IntArray deleteEmptyFolders() {
+        createDbIfNotExists();
+
+        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
+        try (SQLiteTransaction t = new SQLiteTransaction(db)) {
+            // Select folders whose id do not match any container value.
+            String selection = LauncherSettings.Favorites.ITEM_TYPE + " = "
+                    + LauncherSettings.Favorites.ITEM_TYPE_FOLDER + " AND "
+                    + LauncherSettings.Favorites._ID +  " NOT IN (SELECT "
+                    + LauncherSettings.Favorites.CONTAINER + " FROM "
+                    + Favorites.TABLE_NAME + ")";
+
+            IntArray folderIds = LauncherDbUtils.queryIntArray(false, db, Favorites.TABLE_NAME,
+                    Favorites._ID, selection, null, null);
+            if (!folderIds.isEmpty()) {
+                db.delete(Favorites.TABLE_NAME, Utilities.createDbSelectionQuery(
+                        LauncherSettings.Favorites._ID, folderIds), null);
+            }
+            t.commit();
+            return folderIds;
+        } catch (SQLException ex) {
+            Log.e(TAG, ex.getMessage(), ex);
+            return new IntArray();
+        }
+    }
+
+    private static void addModifiedTime(ContentValues values) {
+        values.put(LauncherSettings.Favorites.MODIFIED, System.currentTimeMillis());
+    }
+
+    private void clearFlagEmptyDbCreated() {
+        LauncherPrefs.getPrefs(mContext).edit()
+                .remove(mOpenHelper.getKey(EMPTY_DATABASE_CREATED)).commit();
+    }
+
+    /**
+     * Loads the default workspace based on the following priority scheme:
+     *   1) From the app restrictions
+     *   2) From a package provided by play store
+     *   3) From a partner configuration APK, already in the system image
+     *   4) The default configuration for the particular device
+     */
+    public synchronized void loadDefaultFavoritesIfNecessary() {
+        createDbIfNotExists();
+        SharedPreferences sp = LauncherPrefs.getPrefs(mContext);
+
+        if (sp.getBoolean(mOpenHelper.getKey(EMPTY_DATABASE_CREATED), false)) {
+            Log.d(TAG, "loading default workspace");
+
+            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);
+                        }
+                    }
+                }
+
+                final boolean usingExternallyProvidedLayout = loader != null;
+                if (loader == null) {
+                    loader = getDefaultLayoutParser(widgetHolder);
+                }
+
+                // There might be some partially restored DB items, due to buggy restore logic in
+                // previous versions of launcher.
+                mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
+                // Populate favorites table with initial favorites
+                if ((mOpenHelper.loadFavorites(mOpenHelper.getWritableDatabase(), loader) <= 0)
+                        && usingExternallyProvidedLayout) {
+                    // Unable to load external layout. Cleanup and load the internal layout.
+                    mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
+                    mOpenHelper.loadFavorites(mOpenHelper.getWritableDatabase(),
+                            getDefaultLayoutParser(widgetHolder));
+                }
+                clearFlagEmptyDbCreated();
+            } finally {
+                widgetHolder.destroy();
+            }
+        }
+    }
+
+    /**
+     * 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) {
+        final String authority;
+        if (!TextUtils.isEmpty(mProviderAuthority)) {
+            authority = mProviderAuthority;
+        } else {
+            authority = Settings.Secure.getString(mContext.getContentResolver(),
+                    "launcher3.layout.provider");
+        }
+        if (TextUtils.isEmpty(authority)) {
+            return null;
+        }
+
+        ProviderInfo pi = mContext.getPackageManager().resolveContentProvider(authority, 0);
+        if (pi == null) {
+            Log.e(TAG, "No provider found for authority " + authority);
+            return null;
+        }
+        Uri uri = getLayoutUri(authority, mContext);
+        try (InputStream in = mContext.getContentResolver().openInputStream(uri)) {
+            // 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));
+
+            Log.d(TAG, "Loading layout from " + authority);
+            return new AutoInstallsLayout(mContext, widgetHolder, mOpenHelper,
+                    mContext.getPackageManager().getResourcesForApplication(pi.applicationInfo),
+                    () -> parser, AutoInstallsLayout.TAG_WORKSPACE);
+        } catch (Exception e) {
+            Log.e(TAG, "Error getting layout stream from: " + authority , e);
+            return null;
+        }
+    }
+
+    private static Uri getLayoutUri(String authority, Context ctx) {
+        InvariantDeviceProfile grid = LauncherAppState.getIDP(ctx);
+        return new Uri.Builder().scheme("content").authority(authority).path("launcher_layout")
+                .appendQueryParameter("version", "1")
+                .appendQueryParameter("gridWidth", Integer.toString(grid.numColumns))
+                .appendQueryParameter("gridHeight", Integer.toString(grid.numRows))
+                .appendQueryParameter("hotseatSize", Integer.toString(grid.numDatabaseHotseatIcons))
+                .build();
+    }
+
+    private DefaultLayoutParser getDefaultLayoutParser(LauncherWidgetHolder widgetHolder) {
+        InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
+        int defaultLayout = mDefaultWorkspaceLayoutOverride > 0
+                ? mDefaultWorkspaceLayoutOverride : idp.defaultLayoutId;
+
+        if (mContext.getSystemService(UserManager.class).isDemoUser()
+                && idp.demoModeLayoutId != 0) {
+            defaultLayout = idp.demoModeLayoutId;
+        }
+
+        return new DefaultLayoutParser(mContext, widgetHolder,
+                mOpenHelper, mContext.getResources(), defaultLayout);
+    }
+}
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index c4eb14f..ba5249c 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -39,21 +39,19 @@
 import android.util.SparseLongArray;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.AppWidgetsRestoredReceiver;
 import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.DatabaseHelper;
 import com.android.launcher3.model.DeviceGridState;
-import com.android.launcher3.model.GridBackupTable;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
-import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
 import com.android.launcher3.uioverrides.ApiWrapper;
 import com.android.launcher3.util.IntArray;
@@ -108,7 +106,6 @@
         SQLiteDatabase db = helper.getWritableDatabase();
         try (SQLiteTransaction t = new SQLiteTransaction(db)) {
             RestoreDbTask task = new RestoreDbTask();
-            task.backupWorkspace(context, db);
             task.sanitizeDB(context, helper, db, new BackupManager(context));
             task.restoreAppWidgetIdsIfExists(context);
             t.commit();
@@ -120,49 +117,6 @@
     }
 
     /**
-     * Restore the workspace if backup is available.
-     */
-    public static boolean restoreIfPossible(@NonNull Context context,
-            @NonNull DatabaseHelper helper, @NonNull BackupManager backupManager) {
-        final SQLiteDatabase db = helper.getWritableDatabase();
-        try (SQLiteTransaction t = new SQLiteTransaction(db)) {
-            RestoreDbTask task = new RestoreDbTask();
-            task.restoreWorkspace(context, db, helper, backupManager);
-            t.commit();
-            return true;
-        } catch (Exception e) {
-            FileLog.e(TAG, "Failed to restore db", e);
-            return false;
-        }
-    }
-
-    /**
-     * Backup the workspace so that if things go south in restore, we can recover these entries.
-     */
-    private void backupWorkspace(Context context, SQLiteDatabase db) throws Exception {
-        InvariantDeviceProfile idp = LauncherAppState.getIDP(context);
-        new GridBackupTable(context, db, idp.numDatabaseHotseatIcons, idp.numColumns, idp.numRows)
-                .doBackup(getDefaultProfileId(db), GridBackupTable.OPTION_REQUIRES_SANITIZATION);
-    }
-
-    private void restoreWorkspace(@NonNull Context context, @NonNull SQLiteDatabase db,
-            @NonNull DatabaseHelper helper, @NonNull BackupManager backupManager)
-            throws Exception {
-        final InvariantDeviceProfile idp = LauncherAppState.getIDP(context);
-        GridBackupTable backupTable = new GridBackupTable(context, db, idp.numDatabaseHotseatIcons,
-                idp.numColumns, idp.numRows);
-        if (backupTable.restoreFromRawBackupIfAvailable(getDefaultProfileId(db))) {
-            int itemsDeleted = sanitizeDB(context, helper, db, backupManager);
-            LauncherAppState.getInstance(context).getModel().forceReload();
-            restoreAppWidgetIdsIfExists(context);
-            if (itemsDeleted == 0) {
-                // all the items are restored, we no longer need the backup table
-                dropTable(db, Favorites.BACKUP_TABLE_NAME);
-            }
-        }
-    }
-
-    /**
      * Makes the following changes in the provider DB.
      *   1. Removes all entries belonging to any profiles that were not restored.
      *   2. Marks all entries as restored. The flags are updated during first load or as
@@ -174,7 +128,8 @@
      *
      * @return number of items deleted.
      */
-    private int sanitizeDB(Context context, DatabaseHelper helper, SQLiteDatabase db,
+    @VisibleForTesting
+    protected int sanitizeDB(Context context, DatabaseHelper helper, SQLiteDatabase db,
             BackupManager backupManager) throws Exception {
         // Primary user ids
         long myProfileId = helper.getDefaultUserSerial();
@@ -258,7 +213,7 @@
         }
 
         // Override shortcuts
-        maybeOverrideShortcuts(context, db, myProfileId);
+        maybeOverrideShortcuts(context, helper, db, myProfileId);
 
         return itemsDeleted;
     }
@@ -388,8 +343,8 @@
                 APP_WIDGET_IDS.to(IntArray.wrap(newIds).toConcatString()));
     }
 
-    protected static void maybeOverrideShortcuts(Context context, SQLiteDatabase db,
-            long currentUser) {
+    protected static void maybeOverrideShortcuts(Context context, DatabaseHelper helper,
+            SQLiteDatabase db, long currentUser) {
         Map<String, LauncherActivityInfo> activityOverrides = ApiWrapper.getActivityOverrides(
                 context);
 
@@ -412,8 +367,7 @@
                 if (override != null) {
                     ContentValues values = new ContentValues();
                     values.put(Favorites.PROFILE_ID,
-                            UserCache.INSTANCE.get(context).getSerialNumberForUser(
-                                    override.getUser()));
+                            helper.getSerialNumberForUser(override.getUser()));
                     values.put(Favorites.INTENT, AppInfo.makeLaunchIntent(override).toUri(0));
                     db.update(Favorites.TABLE_NAME, values, String.format("%s=?", Favorites._ID),
                             new String[]{String.valueOf(c.getInt(idIndex))});
diff --git a/src/com/android/launcher3/util/ContentWriter.java b/src/com/android/launcher3/util/ContentWriter.java
index ee64e98..55c2585 100644
--- a/src/com/android/launcher3/util/ContentWriter.java
+++ b/src/com/android/launcher3/util/ContentWriter.java
@@ -118,21 +118,9 @@
         final String[] mSelectionArgs;
 
         public CommitParams(String where, String[] selectionArgs) {
-            this(LauncherSettings.Favorites.CONTENT_URI, where, selectionArgs);
-        }
-
-        private CommitParams(Uri uri, String where, String[] selectionArgs) {
-            mUri = uri;
+            mUri = LauncherSettings.Favorites.CONTENT_URI;
             mWhere = where;
             mSelectionArgs = selectionArgs;
         }
-
-        /**
-         * Creates commit params for backup table.
-         */
-        public static CommitParams backupCommitParams(String where, String[] selectionArgs) {
-            return new CommitParams(
-                    LauncherSettings.Favorites.BACKUP_CONTENT_URI, where, selectionArgs);
-        }
     }
 }
diff --git a/src/com/android/launcher3/views/OptionsPopupView.java b/src/com/android/launcher3/views/OptionsPopupView.java
index 315cbad..64ad390 100644
--- a/src/com/android/launcher3/views/OptionsPopupView.java
+++ b/src/com/android/launcher3/views/OptionsPopupView.java
@@ -18,6 +18,7 @@
 import static androidx.core.content.ContextCompat.getColorStateList;
 
 import static com.android.launcher3.config.FeatureFlags.ENABLE_MATERIAL_U_POPUP;
+import static com.android.launcher3.config.FeatureFlags.MULTI_SELECT_EDIT_MODE;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.IGNORE;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_SETTINGS_BUTTON_TAP_OR_LONGPRESS;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_WIDGETSTRAY_BUTTON_TAP_OR_LONGPRESS;
@@ -205,6 +206,13 @@
                     LAUNCHER_WIDGETSTRAY_BUTTON_TAP_OR_LONGPRESS,
                     OptionsPopupView::onWidgetsClicked));
         }
+        if (MULTI_SELECT_EDIT_MODE.get()) {
+            options.add(new OptionItem(launcher,
+                    R.string.edit_home_screen,
+                    R.drawable.enter_home_gardening_icon,
+                    LAUNCHER_SETTINGS_BUTTON_TAP_OR_LONGPRESS,
+                    OptionsPopupView::enterHomeGardening));
+        }
         options.add(new OptionItem(launcher,
                 R.string.settings_button_text,
                 R.drawable.ic_setting,
@@ -213,6 +221,10 @@
         return options;
     }
 
+    private static boolean enterHomeGardening(View view) {
+        return true;
+    }
+
     private static boolean onWidgetsClicked(View view) {
         return openWidgets(Launcher.getLauncher(view.getContext())) != null;
     }
diff --git a/src/com/android/launcher3/widget/WidgetManagerHelper.java b/src/com/android/launcher3/widget/WidgetManagerHelper.java
index 15fa844..737cdbd 100644
--- a/src/com/android/launcher3/widget/WidgetManagerHelper.java
+++ b/src/com/android/launcher3/widget/WidgetManagerHelper.java
@@ -78,8 +78,16 @@
             return allWidgetsSteam(mContext).collect(Collectors.toList());
         }
 
-        return mAppWidgetManager.getInstalledProvidersForPackage(
-                packageUser.mPackageName, packageUser.mUser);
+        try {
+            return mAppWidgetManager.getInstalledProvidersForPackage(
+                    packageUser.mPackageName, packageUser.mUser);
+        } catch (IllegalStateException e) {
+            // b/277189566: Launcher will load the widget when it gets the user-unlock event.
+            // If exception is thrown because of device is locked, it means a race condition occurs
+            // that the user got locked again while launcher is processing the event. In this case
+            // we should return empty list.
+            return Collections.emptyList();
+        }
     }
 
     /**
diff --git a/tests/src/com/android/launcher3/model/BackupRestoreTest.java b/tests/src/com/android/launcher3/model/BackupRestoreTest.java
deleted file mode 100644
index 41914de..0000000
--- a/tests/src/com/android/launcher3/model/BackupRestoreTest.java
+++ /dev/null
@@ -1,211 +0,0 @@
-/*
- * Copyright (C) 2020 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 static android.content.pm.PackageManager.INSTALL_REASON_DEVICE_RESTORE;
-import static android.os.Process.myUserHandle;
-
-import static com.android.launcher3.LauncherSettings.Favorites.BACKUP_TABLE_NAME;
-import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
-import static com.android.launcher3.LauncherSettings.Favorites.addTableToDb;
-import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
-import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
-import static com.android.launcher3.util.LauncherModelHelper.APP_ICON;
-import static com.android.launcher3.util.LauncherModelHelper.NO__ICON;
-import static com.android.launcher3.util.LauncherModelHelper.SHORTCUT;
-import static com.android.launcher3.util.ReflectionHelpers.getField;
-import static com.android.launcher3.util.ReflectionHelpers.setField;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.spy;
-
-import android.app.backup.BackupManager;
-import android.content.pm.PackageInstaller;
-import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
-import android.os.UserHandle;
-import android.util.ArrayMap;
-import android.util.LongSparseArray;
-
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
-import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.pm.UserCache;
-import com.android.launcher3.provider.RestoreDbTask;
-import com.android.launcher3.util.LauncherModelHelper;
-import com.android.launcher3.util.SafeCloseable;
-
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-
-/**
- * Tests to verify backup and restore flow.
- */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
-public class BackupRestoreTest {
-
-    private static final int PER_USER_RANGE = 200000;
-
-
-    private long mCurrentMyProfileId;
-    private long mOldMyProfileId;
-
-    private long mCurrentWorkProfileId;
-    private long mOldWorkProfileId;
-
-    private BackupManager mBackupManager;
-    private LauncherModelHelper mModelHelper;
-    private SQLiteDatabase mDb;
-    private InvariantDeviceProfile mIdp;
-
-    private UserHandle mWorkUserHandle;
-
-    private SafeCloseable mUserChangeListener;
-
-    @Before
-    public void setUp() {
-        mModelHelper = new LauncherModelHelper();
-
-        mCurrentMyProfileId = mModelHelper.defaultProfileId;
-        mOldMyProfileId = mCurrentMyProfileId + 1;
-        mCurrentWorkProfileId = mOldMyProfileId + 1;
-        mOldWorkProfileId = mCurrentWorkProfileId + 1;
-
-        mWorkUserHandle = UserHandle.getUserHandleForUid(PER_USER_RANGE);
-        mUserChangeListener = UserCache.INSTANCE.get(mModelHelper.sandboxContext)
-                .addUserChangeListener(() -> { });
-
-        setupUserManager();
-        setupBackupManager();
-        RestoreDbTask.setPending(mModelHelper.sandboxContext);
-        mDb = mModelHelper.provider.getDb();
-        mIdp = InvariantDeviceProfile.INSTANCE.get(mModelHelper.sandboxContext);
-
-    }
-
-    @After
-    public void tearDown() {
-        mUserChangeListener.close();
-        mModelHelper.destroy();
-    }
-
-    private void setupUserManager() {
-        UserCache cache = UserCache.INSTANCE.get(mModelHelper.sandboxContext);
-        synchronized (cache) {
-            LongSparseArray<UserHandle> users = getField(cache, "mUsers");
-            users.clear();
-            users.put(mCurrentMyProfileId, myUserHandle());
-            users.put(mCurrentWorkProfileId, mWorkUserHandle);
-
-            ArrayMap<UserHandle, Long> userMap = getField(cache, "mUserToSerialMap");
-            userMap.clear();
-            userMap.put(myUserHandle(), mCurrentMyProfileId);
-            userMap.put(mWorkUserHandle, mCurrentWorkProfileId);
-        }
-    }
-
-    private void setupBackupManager() {
-        mBackupManager = spy(new BackupManager(mModelHelper.sandboxContext));
-        doReturn(myUserHandle()).when(mBackupManager)
-                .getUserForAncestralSerialNumber(eq(mOldMyProfileId));
-        doReturn(mWorkUserHandle).when(mBackupManager)
-                .getUserForAncestralSerialNumber(eq(mOldWorkProfileId));
-    }
-
-    @Test
-    public void testOnCreateDbIfNotExists_CreatesBackup() {
-        assertTrue(tableExists(mDb, BACKUP_TABLE_NAME));
-    }
-
-    @Test
-    public void testOnRestoreSessionWithValidCondition_PerformsRestore() throws Exception {
-        setupBackup();
-        verifyTableIsFilled(BACKUP_TABLE_NAME, false);
-        verifyTableIsEmpty(TABLE_NAME);
-        createRestoreSession();
-        verifyTableIsFilled(TABLE_NAME, true);
-    }
-
-    private void setupBackup() {
-        createTableUsingOldProfileId();
-        // setup grid for main user on first screen
-        mModelHelper.createGrid(new int[][][]{{
-                { APP_ICON, APP_ICON, SHORTCUT, SHORTCUT},
-                { SHORTCUT, SHORTCUT, NO__ICON, NO__ICON},
-                { NO__ICON, NO__ICON, SHORTCUT, SHORTCUT},
-                { APP_ICON, SHORTCUT, SHORTCUT, APP_ICON},
-            }}, 1, mOldMyProfileId);
-        // setup grid for work profile on second screen
-        mModelHelper.createGrid(new int[][][]{{
-                { NO__ICON, APP_ICON, SHORTCUT, SHORTCUT},
-                { SHORTCUT, SHORTCUT, NO__ICON, NO__ICON},
-                { NO__ICON, NO__ICON, SHORTCUT, SHORTCUT},
-                { APP_ICON, SHORTCUT, SHORTCUT, NO__ICON},
-            }}, 2, mOldWorkProfileId);
-        // simulates the creation of backup upon restore
-        new GridBackupTable(mModelHelper.sandboxContext, mDb, mIdp.numDatabaseHotseatIcons,
-                mIdp.numColumns, mIdp.numRows).doBackup(
-                mOldMyProfileId, GridBackupTable.OPTION_REQUIRES_SANITIZATION);
-        // reset favorites table
-        createTableUsingOldProfileId();
-    }
-
-    private void verifyTableIsEmpty(String tableName) {
-        assertEquals(0, getCount(mDb, "SELECT * FROM " + tableName));
-    }
-
-    private void verifyTableIsFilled(String tableName, boolean sanitized) {
-        assertEquals(sanitized ? 12 : 13, getCount(mDb,
-                "SELECT * FROM " + tableName + " WHERE profileId = "
-                        + (sanitized ? mCurrentMyProfileId : mOldMyProfileId)));
-        assertEquals(10, getCount(mDb, "SELECT * FROM " + tableName + " WHERE profileId = "
-                + (sanitized ? mCurrentWorkProfileId : mOldWorkProfileId)));
-    }
-
-    private void createTableUsingOldProfileId() {
-        // simulates the creation of favorites table on old device
-        dropTable(mDb, TABLE_NAME);
-        addTableToDb(mDb, mOldMyProfileId, false);
-    }
-
-    private void createRestoreSession() throws Exception {
-        final PackageInstaller.SessionParams params = new PackageInstaller.SessionParams(
-                PackageInstaller.SessionParams.MODE_FULL_INSTALL);
-        final PackageInstaller installer = mModelHelper.sandboxContext.getPackageManager()
-                .getPackageInstaller();
-        final int sessionId = installer.createSession(params);
-        final PackageInstaller.SessionInfo info = installer.getSessionInfo(sessionId);
-        setField(info, "installReason", INSTALL_REASON_DEVICE_RESTORE);
-        // TODO: (b/148410677) we should verify the following call instead
-        //  InstallSessionHelper.INSTANCE.get(getContext()).restoreDbIfApplicable(info);
-        RestoreDbTask.restoreIfPossible(mModelHelper.sandboxContext,
-                mModelHelper.provider.getHelper(), mBackupManager);
-    }
-
-    private static int getCount(SQLiteDatabase db, String sql) {
-        try (Cursor c = db.rawQuery(sql, null)) {
-            return c.getCount();
-        }
-    }
-}
diff --git a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index aa091b6..67de1f5 100644
--- a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -15,17 +15,35 @@
  */
 package com.android.launcher3.provider;
 
+import static android.os.Process.myUserHandle;
+
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
+import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
+
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 
+import android.app.backup.BackupManager;
 import android.content.ContentValues;
+import android.content.Context;
+import android.content.Intent;
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
+import android.os.UserHandle;
+import android.os.UserManager;
+import android.util.LongSparseArray;
 
-import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.model.DatabaseHelper;
 
@@ -39,6 +57,10 @@
 @RunWith(AndroidJUnit4.class)
 public class RestoreDbTaskTest {
 
+    private static final int PER_USER_RANGE = 200000;
+
+    private final UserHandle mWorkUser = UserHandle.getUserHandleForUid(PER_USER_RANGE);
+
     @Test
     public void testGetProfileId() throws Exception {
         SQLiteDatabase db = new MyDatabaseHelper(23).getWritableDatabase();
@@ -89,6 +111,90 @@
     }
 
     @Test
+    public void testSanitizeDB_bothProfiles() throws Exception {
+        Context context = getInstrumentation().getTargetContext();
+        UserHandle myUser = myUserHandle();
+        long myProfileId = context.getSystemService(UserManager.class)
+                .getSerialNumberForUser(myUser);
+        long myProfileId_old = myProfileId + 1;
+        long workProfileId = myProfileId + 2;
+        long workProfileId_old = myProfileId + 3;
+
+        MyDatabaseHelper helper = new MyDatabaseHelper(myProfileId);
+        SQLiteDatabase db = helper.getWritableDatabase();
+        BackupManager bm = spy(new BackupManager(context));
+        doReturn(myUserHandle()).when(bm).getUserForAncestralSerialNumber(eq(myProfileId_old));
+        doReturn(mWorkUser).when(bm).getUserForAncestralSerialNumber(eq(workProfileId_old));
+        helper.users.put(workProfileId, mWorkUser);
+
+        addIconsBulk(helper, 10, 1, myProfileId_old);
+        addIconsBulk(helper, 6, 2, workProfileId_old);
+        assertEquals(10, getItemCountForProfile(db, myProfileId_old));
+        assertEquals(6, getItemCountForProfile(db, workProfileId_old));
+
+        RestoreDbTask task = new RestoreDbTask();
+        task.sanitizeDB(context, helper, helper.getWritableDatabase(), bm);
+
+        // All the data has been migrated to the new user ids
+        assertEquals(0, getItemCountForProfile(db, myProfileId_old));
+        assertEquals(0, getItemCountForProfile(db, workProfileId_old));
+        assertEquals(10, getItemCountForProfile(db, myProfileId));
+        assertEquals(6, getItemCountForProfile(db, workProfileId));
+    }
+
+    @Test
+    public void testSanitizeDB_workItemsRemoved() throws Exception {
+        Context context = getInstrumentation().getTargetContext();
+        UserHandle myUser = myUserHandle();
+        long myProfileId = context.getSystemService(UserManager.class)
+                .getSerialNumberForUser(myUser);
+        long myProfileId_old = myProfileId + 1;
+        long workProfileId_old = myProfileId + 3;
+
+        MyDatabaseHelper helper = new MyDatabaseHelper(myProfileId);
+        SQLiteDatabase db = helper.getWritableDatabase();
+        BackupManager bm = spy(new BackupManager(context));
+        doReturn(myUserHandle()).when(bm).getUserForAncestralSerialNumber(eq(myProfileId_old));
+        // Work profile is not migrated
+        doReturn(null).when(bm).getUserForAncestralSerialNumber(eq(workProfileId_old));
+
+        addIconsBulk(helper, 10, 1, myProfileId_old);
+        addIconsBulk(helper, 6, 2, workProfileId_old);
+        assertEquals(10, getItemCountForProfile(db, myProfileId_old));
+        assertEquals(6, getItemCountForProfile(db, workProfileId_old));
+
+        RestoreDbTask task = new RestoreDbTask();
+        task.sanitizeDB(context, helper, helper.getWritableDatabase(), bm);
+
+        // All the data has been migrated to the new user ids
+        assertEquals(0, getItemCountForProfile(db, myProfileId_old));
+        assertEquals(0, getItemCountForProfile(db, workProfileId_old));
+        assertEquals(10, getItemCountForProfile(db, myProfileId));
+        assertEquals(10, getCount(db, "select * from favorites"));
+    }
+
+    private void addIconsBulk(DatabaseHelper helper, int count, int screen, long profileId) {
+        int columns = LauncherAppState.getIDP(getInstrumentation().getTargetContext()).numColumns;
+        String packageName = getInstrumentation().getContext().getPackageName();
+        for (int i = 0; i < count; i++) {
+            ContentValues values = new ContentValues();
+            values.put(LauncherSettings.Favorites._ID, helper.generateNewItemId());
+            values.put(LauncherSettings.Favorites.CONTAINER, CONTAINER_DESKTOP);
+            values.put(LauncherSettings.Favorites.SCREEN, screen);
+            values.put(LauncherSettings.Favorites.CELLX, i % columns);
+            values.put(LauncherSettings.Favorites.CELLY, i / columns);
+            values.put(LauncherSettings.Favorites.SPANX, 1);
+            values.put(LauncherSettings.Favorites.SPANY, 1);
+            values.put(LauncherSettings.Favorites.PROFILE_ID, profileId);
+            values.put(LauncherSettings.Favorites.ITEM_TYPE, ITEM_TYPE_APPLICATION);
+            values.put(LauncherSettings.Favorites.INTENT,
+                    new Intent(Intent.ACTION_MAIN).setPackage(packageName).toUri(0));
+            helper.getWritableDatabase().insert(TABLE_NAME, null, values);
+        }
+    }
+
+
+    @Test
     public void testRemoveScreenIdGaps_firstScreenEmpty() {
         runRemoveScreenIdGapsTest(
                 new int[]{1, 2, 5, 6, 6, 7, 9, 9},
@@ -116,7 +222,7 @@
             ContentValues values = new ContentValues();
             values.put(Favorites._ID, i);
             values.put(Favorites.SCREEN, screenIds[i]);
-            values.put(Favorites.CONTAINER, Favorites.CONTAINER_DESKTOP);
+            values.put(Favorites.CONTAINER, CONTAINER_DESKTOP);
             db.insert(Favorites.TABLE_NAME, null, values);
         }
         // Verify items are added
@@ -138,6 +244,10 @@
         assertArrayEquals(expectedScreenIds, resultScreenIds);
     }
 
+    public int getItemCountForProfile(SQLiteDatabase db, long profileId) {
+        return getCount(db, "select * from favorites where profileId = " + profileId);
+    }
+
     private int getCount(SQLiteDatabase db, String sql) {
         try (Cursor c = db.rawQuery(sql, null)) {
             return c.getCount();
@@ -146,16 +256,18 @@
 
     private class MyDatabaseHelper extends DatabaseHelper {
 
-        private final long mProfileId;
+        public final LongSparseArray<UserHandle> users;
 
         MyDatabaseHelper(long profileId) {
-            super(InstrumentationRegistry.getInstrumentation().getTargetContext(), null, false);
-            mProfileId = profileId;
+            super(getInstrumentation().getTargetContext(), null, false);
+            users = new LongSparseArray<>();
+            users.put(profileId, myUserHandle());
         }
 
         @Override
-        public long getDefaultUserSerial() {
-            return mProfileId;
+        public long getSerialNumberForUser(UserHandle user) {
+            int index = users.indexOfValue(user);
+            return index >= 0 ? users.keyAt(index) : -1;
         }
 
         @Override
diff --git a/tests/src/com/android/launcher3/util/LauncherModelHelper.java b/tests/src/com/android/launcher3/util/LauncherModelHelper.java
index fdfeb7d..9e88c06 100644
--- a/tests/src/com/android/launcher3/util/LauncherModelHelper.java
+++ b/tests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -61,8 +61,8 @@
 import com.android.launcher3.model.AllAppsList;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.Callbacks;
-import com.android.launcher3.model.DatabaseHelper;
 import com.android.launcher3.model.ItemInstallQueue;
+import com.android.launcher3.model.ModelDbController;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.pm.InstallSessionHelper;
@@ -370,54 +370,6 @@
         sandboxContext.getContentResolver().delete(uri, null, null);
     }
 
-    public int[][][] createGrid(int[][][] typeArray) {
-        return createGrid(typeArray, 1);
-    }
-
-    public int[][][] createGrid(int[][][] typeArray, int startScreen) {
-        LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
-                LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB);
-        LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
-                LauncherSettings.Settings.METHOD_CLEAR_EMPTY_DB_FLAG);
-        return createGrid(typeArray, startScreen, defaultProfileId);
-    }
-
-    /**
-     * Initializes the DB with mock elements to represent the provided grid structure.
-     * @param typeArray A 3d array of item types. {@see #addItem(int, long, long, int, int)} for
-     *                  type definitions. The first dimension represents the screens and the next
-     *                  two represent the workspace grid.
-     * @param startScreen First screen id from where the icons will be added.
-     * @return the same grid representation where each entry is the corresponding item id.
-     */
-    public int[][][] createGrid(int[][][] typeArray, int startScreen, long profileId) {
-        int[][][] ids = new int[typeArray.length][][];
-        for (int i = 0; i < typeArray.length; i++) {
-            // Add screen to DB
-            int screenId = startScreen + i;
-
-            // Keep the screen id counter up to date
-            LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
-                    LauncherSettings.Settings.METHOD_NEW_SCREEN_ID);
-
-            ids[i] = new int[typeArray[i].length][];
-            for (int y = 0; y < typeArray[i].length; y++) {
-                ids[i][y] = new int[typeArray[i][y].length];
-                for (int x = 0; x < typeArray[i][y].length; x++) {
-                    if (typeArray[i][y][x] < 0) {
-                        // Empty cell
-                        ids[i][y][x] = -1;
-                    } else {
-                        ids[i][y][x] = addItem(
-                                typeArray[i][y][x], screenId, DESKTOP, x, y, profileId);
-                    }
-                }
-            }
-        }
-
-        return ids;
-    }
-
     /**
      * Sets up a mock provider to load the provided layout by default, next time the layout loads
      */
@@ -471,16 +423,12 @@
 
         @Override
         public boolean onCreate() {
+            mModelDbController = new ModelDbController(getContext());
             return true;
         }
 
         public SQLiteDatabase getDb() {
-            createDbIfNotExists();
-            return mOpenHelper.getWritableDatabase();
-        }
-
-        public DatabaseHelper getHelper() {
-            return mOpenHelper;
+            return mModelDbController.getDatabaseHelper().getWritableDatabase();
         }
     }