diff --git a/quickstep/src/com/android/launcher3/uioverrides/TwoStepSwipeController.java b/quickstep/src/com/android/launcher3/uioverrides/TwoStepSwipeController.java
index 2481c32..fb59946 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/TwoStepSwipeController.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/TwoStepSwipeController.java
@@ -90,7 +90,6 @@
     private static final int FLAG_OVERVIEW_DISABLED_OUT_OF_RANGE = 1 << 0;
     private static final int FLAG_OVERVIEW_DISABLED_FLING = 1 << 1;
     private static final int FLAG_OVERVIEW_DISABLED_CANCEL_STATE = 1 << 2;
-    private static final int FLAG_RECENTS_PLAN_LOADING = 1 << 3;
     private static final int FLAG_OVERVIEW_DISABLED = 1 << 4;
     private static final int FLAG_DISABLED_TWO_TARGETS = 1 << 5;
     private static final int FLAG_DISABLED_BACK_TARGET = 1 << 6;
@@ -272,16 +271,7 @@
             mCurrentAnimation = mLauncher.getStateManager().createAnimationToNewWorkspace(
                     mToState, mTaggedAnimatorSetBuilder, maxAccuracy);
 
-            if (TouchInteractionService.isConnected()) {
-                // Load recents plan
-                RecentsModel recentsModel = RecentsModel.getInstance(mLauncher);
-                if (recentsModel.getLastLoadPlan() != null) {
-                    onRecentsPlanLoaded(recentsModel.getLastLoadPlan());
-                } else {
-                    mDragPauseDetector.addDisabledFlags(FLAG_RECENTS_PLAN_LOADING);
-                    recentsModel.loadTasks(-1, this::onRecentsPlanLoaded);
-                }
-            } else {
+            if (!TouchInteractionService.isConnected()) {
                 mDragPauseDetector.addDisabledFlags(FLAG_OVERVIEW_DISABLED);
             }
 
@@ -302,14 +292,6 @@
         }
     }
 
-    private void onRecentsPlanLoaded(RecentsTaskLoadPlan plan) {
-        RecentsView recentsView = mLauncher.getOverviewPanel();
-        recentsView.update(plan);
-        recentsView.initToPage(0);
-
-        mDragPauseDetector.clearDisabledFlags(FLAG_RECENTS_PLAN_LOADING);
-    }
-
     private float getShiftRange() {
         return mLauncher.getAllAppsController().getShiftRange();
     }
diff --git a/quickstep/src/com/android/launcher3/uioverrides/UiFactory.java b/quickstep/src/com/android/launcher3/uioverrides/UiFactory.java
index 0e539ee..e848688 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/UiFactory.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/UiFactory.java
@@ -30,6 +30,7 @@
 import com.android.launcher3.graphics.BitmapRenderer;
 import com.android.launcher3.util.TouchController;
 import com.android.launcher3.widget.WidgetsFullSheet;
+import com.android.quickstep.RecentsView;
 import com.android.systemui.shared.recents.view.RecentsTransition;
 
 public class UiFactory {
@@ -96,4 +97,9 @@
             return result;
         }
     }
+
+    public static void resetOverview(Launcher launcher) {
+        RecentsView recents = launcher.getOverviewPanel();
+        recents.reset();
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/NavBarSwipeInteractionHandler.java b/quickstep/src/com/android/quickstep/NavBarSwipeInteractionHandler.java
index ce528e3..168c1fe 100644
--- a/quickstep/src/com/android/quickstep/NavBarSwipeInteractionHandler.java
+++ b/quickstep/src/com/android/quickstep/NavBarSwipeInteractionHandler.java
@@ -24,6 +24,7 @@
 import android.content.Context;
 import android.content.res.Resources;
 import android.graphics.Bitmap;
+import android.graphics.Color;
 import android.graphics.Rect;
 import android.os.Build;
 import android.os.Handler;
@@ -56,14 +57,10 @@
 public class NavBarSwipeInteractionHandler extends InternalStateHandler {
 
     private static final int STATE_LAUNCHER_READY = 1 << 0;
-    private static final int STATE_RECENTS_DELAY_COMPLETE = 1 << 1;
-    private static final int STATE_LOAD_PLAN_READY = 1 << 2;
-    private static final int STATE_RECENTS_FULLY_VISIBLE = 1 << 3;
     private static final int STATE_ACTIVITY_MULTIPLIER_COMPLETE = 1 << 4;
     private static final int STATE_SCALED_SNAPSHOT_RECENTS = 1 << 5;
     private static final int STATE_SCALED_SNAPSHOT_APP = 1 << 6;
 
-    private static final long RECENTS_VIEW_VISIBILITY_DELAY = 120;
     private static final long RECENTS_VIEW_VISIBILITY_DURATION = 150;
     private static final long MAX_SWIPE_DURATION = 200;
     private static final long MIN_SWIPE_DURATION = 80;
@@ -90,7 +87,7 @@
     // animated to 1, so allow for a smooth transition.
     private final AnimatedFloat mActivityMultiplier = new AnimatedFloat(this::updateFinalShift);
 
-    private final int mRunningTaskId;
+    private final Task mRunningTask;
     private final Context mContext;
 
     private final MultiStateCallback mStateCallback;
@@ -102,7 +99,6 @@
     private QuickScrubController mQuickScrubController;
     private Hotseat mHotseat;
     private AllAppsScrim mAllAppsScrim;
-    private RecentsTaskLoadPlan mLoadPlan;
 
     private boolean mLauncherReady;
     private boolean mTouchEndHandled;
@@ -114,7 +110,11 @@
 
     NavBarSwipeInteractionHandler(RunningTaskInfo runningTaskInfo, Context context,
             @TouchInteractionService.InteractionType int interactionType) {
-        mRunningTaskId = runningTaskInfo.id;
+        // TODO: We need a better way for this
+        TaskKey taskKey = new TaskKey(runningTaskInfo.id, 0, null, UserHandle.myUserId(), 0);
+        mRunningTask = new Task(taskKey, null, null, "", "", Color.BLACK, Color.BLACK,
+                true, false, false, false, null, 0, null, false);
+
         mContext = context;
         mInteractionType = interactionType;
         WindowManagerWrapper.getInstance().getStableInsets(mStableInsets);
@@ -131,11 +131,9 @@
         // Build the state callback
         mStateCallback = new MultiStateCallback();
         mStateCallback.addCallback(STATE_LAUNCHER_READY, this::onLauncherReady);
-        mStateCallback.addCallback(STATE_LOAD_PLAN_READY | STATE_RECENTS_DELAY_COMPLETE,
-                this::setTaskPlanToUi);
         mStateCallback.addCallback(STATE_SCALED_SNAPSHOT_APP, this::resumeLastTask);
-        mStateCallback.addCallback(STATE_RECENTS_FULLY_VISIBLE | STATE_SCALED_SNAPSHOT_RECENTS
-                | STATE_ACTIVITY_MULTIPLIER_COMPLETE,
+        mStateCallback.addCallback(
+                STATE_SCALED_SNAPSHOT_RECENTS | STATE_ACTIVITY_MULTIPLIER_COMPLETE,
                 this::onAnimationToLauncherComplete);
         mStateCallback.addCallback(STATE_LAUNCHER_READY | STATE_SCALED_SNAPSHOT_APP,
                 this::cleanupLauncher);
@@ -145,10 +143,6 @@
         mLauncherReady = true;
         executeFrameUpdate();
 
-        // Wait for some time before loading recents so that the first frame is fast
-        new Handler().postDelayed(() -> mStateCallback.setState(STATE_RECENTS_DELAY_COMPLETE),
-                RECENTS_VIEW_VISIBILITY_DELAY);
-
         long duration = Math.min(MAX_SWIPE_DURATION,
                 Math.max((long) (-mCurrentDisplacement / PIXEL_PER_MS), MIN_SWIPE_DURATION));
         if (mCurrentShift.getCurrentAnimation() != null) {
@@ -189,18 +183,20 @@
 
     @Override
     protected void init(Launcher launcher, boolean alreadyOnHome) {
-        AbstractFloatingView.closeAllOpenViews(launcher, alreadyOnHome);
-        launcher.getStateManager().goToState(LauncherState.OVERVIEW, alreadyOnHome);
-
         mLauncher = launcher;
+        mRecentsView = launcher.getOverviewPanel();
+        mRecentsView.showTask(mRunningTask);
+        mStateController = mRecentsView.getStateController();
+        mHotseat = mLauncher.getHotseat();
+        mAllAppsScrim = mLauncher.findViewById(R.id.all_apps_scrim);
+
+        AbstractFloatingView.closeAllOpenViews(mLauncher, alreadyOnHome);
+        mLauncher.getStateManager().goToState(LauncherState.OVERVIEW, alreadyOnHome);
+
         mDragView = new SnapshotDragView(mLauncher, mTaskSnapshot);
         mLauncher.getDragLayer().addView(mDragView);
         mDragView.setPivotX(0);
         mDragView.setPivotY(0);
-        mRecentsView = launcher.getOverviewPanel();
-        mStateController = mRecentsView.getStateController();
-        mHotseat = mLauncher.getHotseat();
-        mAllAppsScrim = mLauncher.findViewById(R.id.all_apps_scrim);
 
         boolean interactionIsQuick
                 = mInteractionType == TouchInteractionService.INTERACTION_QUICK_SCRUB
@@ -220,8 +216,6 @@
             // All-apps search box is visible in vertical bar layout.
             mLauncher.getAppsView().setVisibility(View.GONE);
         }
-        mStateController.setTransitionProgress(1);
-        mStateController.setVisibility(false);
         TraceHelper.partitionSection("TouchInt", "Launcher on new intent");
     }
 
@@ -273,26 +267,6 @@
     }
 
     @UiThread
-    public void setRecentsTaskLoadPlan(RecentsTaskLoadPlan loadPlan) {
-        mLoadPlan = loadPlan;
-        mStateCallback.setState(STATE_LOAD_PLAN_READY);
-    }
-
-    private void setTaskPlanToUi() {
-        mRecentsView.update(mLoadPlan);
-        mRecentsView.initToPage(mStartedQuickScrubFromHome ? 0 : mRecentsView.getFirstTaskIndex());
-        ObjectAnimator anim = mStateController.animateVisibility(true /* isVisible */)
-                .setDuration(RECENTS_VIEW_VISIBILITY_DURATION);
-        anim.addListener(new AnimationSuccessListener() {
-            @Override
-            public void onAnimationSuccess(Animator animator) {
-                mStateCallback.setState(STATE_RECENTS_FULLY_VISIBLE);
-            }
-        });
-        anim.start();
-    }
-
-    @UiThread
     public void endTouch(float endVelocity) {
         if (mTouchEndHandled) {
             return;
@@ -339,19 +313,16 @@
 
     @UiThread
     private void resumeLastTask() {
-        TaskKey key = null;
-        if (mLoadPlan != null) {
-            Task task = mLoadPlan.getTaskStack().findTaskWithId(mRunningTaskId);
+        // TODO: We need a better way for this
+        TaskKey key = mRunningTask.key;
+        RecentsTaskLoadPlan loadPlan = RecentsModel.getInstance(mContext).getLastLoadPlan();
+        if (loadPlan != null) {
+            Task task = loadPlan.getTaskStack().findTaskWithId(key.id);
             if (task != null) {
                 key = task.key;
             }
         }
 
-        if (key == null) {
-            // TODO: We need a better way for this
-            key = new TaskKey(mRunningTaskId, 0, null, UserHandle.myUserId(), 0);
-        }
-
         ActivityOptions opts = ActivityOptions.makeCustomAnimation(mContext, 0, 0);
         ActivityManagerWrapper.getInstance().startActivityFromRecentsAsync(key, opts, null, null);
     }
@@ -372,7 +343,8 @@
         if (mInteractionType == TouchInteractionService.INTERACTION_QUICK_SWITCH) {
             for (int i = mRecentsView.getFirstTaskIndex(); i < mRecentsView.getPageCount(); i++) {
                 TaskView taskView = (TaskView) mRecentsView.getPageAt(i);
-                if (taskView.getTask().key.id != mRunningTaskId) {
+                // TODO: Match the keys directly
+                if (taskView.getTask().key.id != mRunningTask.key.id) {
                     mRecentsView.snapToPage(i, QUICK_SWITCH_SNAP_DURATION);
                     taskView.postDelayed(() -> {taskView.launchTask(true);},
                             QUICK_SWITCH_SNAP_DURATION);
diff --git a/quickstep/src/com/android/quickstep/RecentsModel.java b/quickstep/src/com/android/quickstep/RecentsModel.java
index 112f156..22658b2 100644
--- a/quickstep/src/com/android/quickstep/RecentsModel.java
+++ b/quickstep/src/com/android/quickstep/RecentsModel.java
@@ -15,8 +15,10 @@
  */
 package com.android.quickstep;
 
+import android.annotation.TargetApi;
 import android.content.Context;
 import android.content.res.Resources;
+import android.os.Build;
 import android.os.Looper;
 import android.os.UserHandle;
 
@@ -25,7 +27,9 @@
 import com.android.systemui.shared.recents.model.RecentsTaskLoadPlan;
 import com.android.systemui.shared.recents.model.RecentsTaskLoadPlan.PreloadOptions;
 import com.android.systemui.shared.recents.model.RecentsTaskLoader;
+import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.BackgroundExecutor;
+import com.android.systemui.shared.system.TaskStackChangeListener;
 
 import java.util.concurrent.ExecutionException;
 import java.util.function.Consumer;
@@ -33,7 +37,8 @@
 /**
  * Singleton class to load and manage recents model.
  */
-public class RecentsModel {
+@TargetApi(Build.VERSION_CODES.O)
+public class RecentsModel extends TaskStackChangeListener {
 
     // We do not need any synchronization for this variable as its only written on UI thread.
     private static RecentsModel INSTANCE;
@@ -59,6 +64,9 @@
     private final MainThreadExecutor mMainThreadExecutor;
 
     private RecentsTaskLoadPlan mLastLoadPlan;
+    private int mLastLoadPlanId;
+    private int mTaskChangeId;
+
     private RecentsModel(Context context) {
         mContext = context;
 
@@ -69,6 +77,10 @@
         mRecentsTaskLoader.startLoader(mContext);
 
         mMainThreadExecutor = new MainThreadExecutor();
+        ActivityManagerWrapper.getInstance().registerTaskStackListener(this);
+
+        mTaskChangeId = 1;
+        loadTasks(-1, null);
     }
 
     public RecentsTaskLoader getRecentsTaskLoader() {
@@ -80,8 +92,20 @@
      * @param taskId The running task id or -1
      * @param callback The callback to receive the task plan once its complete or null. This is
      *                always called on the UI thread.
+     * @return the request id associated with this call.
      */
-    public void loadTasks(int taskId, Consumer<RecentsTaskLoadPlan> callback) {
+    public int loadTasks(int taskId, Consumer<RecentsTaskLoadPlan> callback) {
+        final int requestId = mTaskChangeId;
+
+        // Fail fast if nothing has changed.
+        if (mLastLoadPlanId == mTaskChangeId) {
+            if (callback != null) {
+                final RecentsTaskLoadPlan plan = mLastLoadPlan;
+                mMainThreadExecutor.execute(() -> callback.accept(plan));
+            }
+            return requestId;
+        }
+
         BackgroundExecutor.get().submit(() -> {
             // Preload the plan
             RecentsTaskLoadPlan loadPlan = new RecentsTaskLoadPlan(mContext);
@@ -91,11 +115,23 @@
             // Set the load plan on UI thread
             mMainThreadExecutor.execute(() -> {
                 mLastLoadPlan = loadPlan;
+                mLastLoadPlanId = requestId;
+
                 if (callback != null) {
                     callback.accept(loadPlan);
                 }
             });
         });
+        return requestId;
+    }
+
+    @Override
+    public void onTaskStackChanged() {
+        mTaskChangeId++;
+    }
+
+    public boolean isLoadPlanValid(int resultId) {
+        return mTaskChangeId == resultId;
     }
 
     public RecentsTaskLoadPlan getLastLoadPlan() {
diff --git a/quickstep/src/com/android/quickstep/RecentsView.java b/quickstep/src/com/android/quickstep/RecentsView.java
index 09003b2..d1bbd23 100644
--- a/quickstep/src/com/android/quickstep/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/RecentsView.java
@@ -16,8 +16,6 @@
 
 package com.android.quickstep;
 
-import static com.android.launcher3.LauncherState.OVERVIEW;
-
 import android.animation.LayoutTransition;
 import android.content.Context;
 import android.graphics.Rect;
@@ -27,10 +25,8 @@
 
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Launcher;
-import com.android.launcher3.LauncherState;
 import com.android.launcher3.PagedView;
 import com.android.launcher3.R;
-import com.android.launcher3.dragndrop.DragLayer;
 import com.android.launcher3.uioverrides.OverviewState;
 import com.android.launcher3.uioverrides.RecentsViewStateController;
 import com.android.systemui.shared.recents.model.RecentsTaskLoadPlan;
@@ -44,6 +40,8 @@
 
 import java.util.ArrayList;
 
+import static com.android.launcher3.LauncherState.NORMAL;
+
 /**
  * A list of recent tasks.
  */
@@ -60,6 +58,9 @@
     private boolean mTaskStackListenerRegistered;
     private LayoutTransition mLayoutTransition;
 
+    /**
+     * TODO: Call reloadIdNeeded in onTaskStackChanged.
+     */
     private TaskStackChangeListener mTaskStackListener = new TaskStackChangeListener() {
         @Override
         public void onTaskSnapshotChanged(int taskId, ThumbnailData snapshot) {
@@ -76,6 +77,11 @@
     private RecentsViewStateController mStateController;
     private int mFirstTaskIndex;
 
+    private final RecentsModel mModel;
+    private int mLoadPlanId = -1;
+
+    private Task mFirstTask;
+
     public RecentsView(Context context) {
         this(context, null);
     }
@@ -93,6 +99,7 @@
 
         mLauncher = Launcher.getLauncher(context);
         mQuickScrubController = new QuickScrubController(mLauncher);
+        mModel = RecentsModel.getInstance(context);
 
         mScrollState.isRtl = mIsRtl;
     }
@@ -115,7 +122,6 @@
         Rect padding =
                 getPadding(Launcher.getLauncher(getContext()).getDeviceProfile(), getContext());
         setPadding(padding.left, padding.top, padding.right, padding.bottom);
-
         mFirstTaskIndex = getPageCount();
     }
 
@@ -154,9 +160,8 @@
         updateTaskStackListenerState();
     }
 
-    public void update(RecentsTaskLoadPlan loadPlan) {
-        final RecentsTaskLoader loader = RecentsModel.getInstance(getContext())
-                .getRecentsTaskLoader();
+    private void applyLoadPlan(RecentsTaskLoadPlan loadPlan) {
+        final RecentsTaskLoader loader = mModel.getRecentsTaskLoader();
         TaskStack stack = loadPlan != null ? loadPlan.getTaskStack() : null;
         if (stack == null) {
             removeAllViews();
@@ -166,10 +171,17 @@
         // Ensure there are as many views as there are tasks in the stack (adding and trimming as
         // necessary)
         final LayoutInflater inflater = LayoutInflater.from(getContext());
-        final ArrayList<Task> tasks = stack.getTasks();
+        final ArrayList<Task> tasks = new ArrayList<>(stack.getTasks());
         setLayoutTransition(null);
-        int requiredChildCount = tasks.size() + mFirstTaskIndex;
 
+        if (mFirstTask != null) {
+            // TODO: Handle this case here once we have a valid implementation for mFirstTask
+            if (tasks.isEmpty() || !keysEquals(tasks.get(tasks.size() - 1), mFirstTask)) {
+                // tasks.add(mFirstTask);
+            }
+        }
+
+        final int requiredChildCount = tasks.size() + mFirstTaskIndex;
         for (int i = getChildCount(); i < requiredChildCount; i++) {
             final TaskView taskView = (TaskView) inflater.inflate(R.layout.task, this, false);
             addView(taskView);
@@ -190,31 +202,17 @@
         }
     }
 
-    public void initToPage(int pageNo) {
-        setCurrentPage(pageNo);
-        if (getPageAt(mCurrentPage) instanceof TaskView) {
-            ((TaskView) getPageAt(mCurrentPage)).setIconScale(0);
-        }
-    }
-
-    public void launchTaskWithId(int taskId) {
-        for (int i = mFirstTaskIndex; i < getChildCount(); i++) {
-            final TaskView taskView = (TaskView) getChildAt(i);
-            if (taskView.getTask().key.id == taskId) {
-                taskView.launchTask(false /* animate */);
-                return;
-            }
-        }
-    }
-
     private void updateTaskStackListenerState() {
         boolean registerStackListener = mOverviewStateEnabled && isAttachedToWindow()
                 && getWindowVisibility() == VISIBLE;
         if (registerStackListener != mTaskStackListenerRegistered) {
             if (registerStackListener) {
-                ActivityManagerWrapper.getInstance().registerTaskStackListener(mTaskStackListener);
+                ActivityManagerWrapper.getInstance()
+                        .registerTaskStackListener(mTaskStackListener);
+                reloadIfNeeded();
             } else {
-                ActivityManagerWrapper.getInstance().unregisterTaskStackListener(mTaskStackListener);
+                ActivityManagerWrapper.getInstance()
+                        .unregisterTaskStackListener(mTaskStackListener);
             }
             mTaskStackListenerRegistered = registerStackListener;
         }
@@ -302,11 +300,74 @@
     public void onTaskDismissed(TaskView taskView) {
         ActivityManagerWrapper.getInstance().removeTask(taskView.getTask().key.id);
         removeView(taskView);
-        if (getChildCount() == mFirstTaskIndex) {
-            mLauncher.getStateManager().goToState(LauncherState.NORMAL);
+        if (getTaskCount() == 0) {
+            mLauncher.getStateManager().goToState(NORMAL);
         }
     }
 
+    public void reset() {
+        mFirstTask = null;
+        setCurrentPage(0);
+    }
+
+    public int getTaskCount() {
+        return getChildCount() - mFirstTaskIndex;
+    }
+
+    /**
+     * Reloads the view if anything in recents changed.
+     */
+    public void reloadIfNeeded() {
+        if (!mModel.isLoadPlanValid(mLoadPlanId)) {
+            int taskId = -1;
+            if (mFirstTask != null) {
+                taskId = mFirstTask.key.id;
+            }
+            mLoadPlanId = mModel.loadTasks(taskId, this::applyLoadPlan);
+        }
+    }
+
+    /**
+     * Ensures that the first task in the view represents {@param task} and reloads the view
+     * if needed. This allows the swipe-up gesture to assume that the first tile always
+     * corresponds to the correct task.
+     * All subsequent calls to reload will keep the task as the first item until {@link #reset()}
+     * is called.
+     * Also scrolls the view to this task
+     */
+    public void showTask(Task task) {
+        boolean needsReload = false;
+        boolean inflateFirstChild = true;
+        if (getTaskCount() > 0) {
+            TaskView tv = (TaskView) getChildAt(mFirstTaskIndex);
+            inflateFirstChild = !keysEquals(tv.getTask(), task);
+        }
+        if (inflateFirstChild) {
+            needsReload = true;
+            setLayoutTransition(null);
+            // Add an empty view for now
+            final TaskView taskView = (TaskView) LayoutInflater.from(getContext())
+                    .inflate(R.layout.task, this, false);
+            addView(taskView, mFirstTaskIndex);
+            taskView.bind(task);
+            setLayoutTransition(mLayoutTransition);
+        }
+        if (!needsReload) {
+            needsReload = !mModel.isLoadPlanValid(mLoadPlanId);
+        }
+        if (needsReload) {
+            mLoadPlanId = mModel.loadTasks(task.key.id, this::applyLoadPlan);
+        }
+        mFirstTask = task;
+        setCurrentPage(mFirstTaskIndex);
+        ((TaskView) getPageAt(mCurrentPage)).setIconScale(0);
+    }
+
+    private static boolean keysEquals(Task t1, Task t2) {
+        // TODO: Match the keys directly
+        return t1.key.id == t2.key.id;
+    }
+
     public QuickScrubController getQuickScrubController() {
         return mQuickScrubController;
     }
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 56f0e89..4cfa1b8 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -161,10 +161,10 @@
                 .setFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
         ResolveInfo info = getPackageManager().resolveActivity(mHomeIntent, 0);
         mLauncher = new ComponentName(getPackageName(), info.activityInfo.name);
-        mHomeIntent.setComponent(mLauncher);
+        // Clear the packageName as system can fail to dedupe it b/64108432
+        mHomeIntent.setComponent(mLauncher).setPackage(null);
 
         mEventQueue = new MotionEventQueue(Choreographer.getInstance(), this::handleMotionEvent);
-        mRecentsModel.loadTasks(-1, null);
         sConnected = true;
     }
 
@@ -307,7 +307,7 @@
         });
 
         // Preload the plan
-        mRecentsModel.loadTasks(mRunningTask.id, handler::setRecentsTaskLoadPlan);
+        mRecentsModel.loadTasks(mRunningTask.id, null);
         mInteractionHandler = handler;
     }
 
diff --git a/res/layout/work_tab_footer.xml b/res/layout/work_tab_footer.xml
index e3416ac..dc0fdd4 100644
--- a/res/layout/work_tab_footer.xml
+++ b/res/layout/work_tab_footer.xml
@@ -73,7 +73,6 @@
         android:lines="1"
         android:minHeight="24dp"
         android:paddingStart="12dp"
-        android:text="@string/managed_by_your_organisation"
         android:textColor="?android:attr/textColorHint"
         android:textSize="13sp"/>
 
diff --git a/res/values/config.xml b/res/values/config.xml
index f48d474..2096200 100644
--- a/res/values/config.xml
+++ b/res/values/config.xml
@@ -101,6 +101,9 @@
     <!-- Package name of the default wallpaper picker. -->
     <string name="wallpaper_picker_package" translatable="false"></string>
 
+    <!-- Whitelisted package to retrieve packagename for badge. Can be empty. -->
+    <string name="shortcutinfocompat_badgepkg_whitelist" translatable="false"></string>
+
     <!-- View ID to use for QSB widget -->
     <item type="id" name="qsb_widget" />
 
diff --git a/res/values/strings.xml b/res/values/strings.xml
index d77065c..ee09946 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -324,13 +324,16 @@
     <!-- Label of tab to indicate work apps -->
     <string name="all_apps_work_tab">Work</string>
 
-    <!-- Label of the work mode toggle -->
+    <!-- This string is in the work profile tab when a user has All Apps open on their phone. This is a label for a toggle to turn the work profile on and off. "Work profile" means a separate profile on a user's phone that's specifically for their work apps and managed by their company. "Work" is used as an adjective.-->
     <string name="work_profile_toggle_label">Work profile</string>
-    <!-- Title in bottom user education view in work tab -->
+    <!-- Title of an overlay in All Apps. This overlay is letting a user know about their work profile, which is managed by their employer. "Work apps" are apps in a user's work profile.-->
     <string name="bottom_work_tab_user_education_title">Find work apps here</string>
-    <!-- Body text in bottom user education view in work tab -->
-    <string name="bottom_work_tab_user_education_body">Each work app has an orange badge, which means it\'s kept secure by your organization. Work apps can be moved to your Home Screen for easier access.</string>
-    <!-- Label in work tab to tell users that work profile is managed by their organisation. -->
-    <string name="managed_by_your_organisation">Managed by your organisation</string>
+    <!-- Text in an overlay in All Apps. This overlay is letting a user know about their work profile, which is managed by their employer.-->
+    <string name="bottom_work_tab_user_education_body">Each work app has an orange badge and is kept secure by your organization. Move apps to your Home screen for easier access.</string>
+    <!-- This string is in the work profile tab when a user has All Apps open on their phone. It describes the label of a toggle, "Work profile," as being managed by the user's employer.
+    "Organization" is used to represent a variety of businesses, non-profits, and educational institutions).-->
+    <string name="work_mode_on_label">Managed by your organization</string>
+    <!-- This string appears under a the label of a toggle in the work profile tab on a user's phone. It describes the status of the toggle, "Work profile," when it's turned off. "Work profile" means a separate profile on a user's phone that's speficially for their work apps and is managed by their company.-->
+    <string name="work_mode_off_label">Notifications and apps are off</string>
 
 </resources>
diff --git a/src/com/android/launcher3/BaseRecyclerView.java b/src/com/android/launcher3/BaseRecyclerView.java
index 76c7845..cc13263 100644
--- a/src/com/android/launcher3/BaseRecyclerView.java
+++ b/src/com/android/launcher3/BaseRecyclerView.java
@@ -91,7 +91,7 @@
         // Move to mScrollbar's coordinate system.
         // We need to take parent into account (view pager's location)
         ViewGroup parent = (ViewGroup) getParent();
-        int left = parent.getLeft() + getLeft() - mScrollbar.getLeft();
+        int left = parent.getLeft() - mScrollbar.getLeft();
         int top = parent.getTop() + getTop() - mScrollbar.getTop() - getScrollBarTop();
         ev.offsetLocation(left, top);
         try {
diff --git a/src/com/android/launcher3/DeviceProfile.java b/src/com/android/launcher3/DeviceProfile.java
index a279633..164efe5 100644
--- a/src/com/android/launcher3/DeviceProfile.java
+++ b/src/com/android/launcher3/DeviceProfile.java
@@ -102,8 +102,6 @@
 
     // All apps
     public int allAppsCellHeightPx;
-    public int allAppsNumCols;
-    public int allAppsNumPredictiveCols;
     public int allAppsIconSizePx;
     public int allAppsIconDrawablePaddingPx;
     public float allAppsIconTextSizePx;
@@ -385,10 +383,6 @@
         return mInsets;
     }
 
-    public void updateAppsViewNumCols() {
-        allAppsNumCols = allAppsNumPredictiveCols = inv.numColumns;
-    }
-
     public Point getCellSize() {
         Point result = new Point();
         // Since we are only concerned with the overall padding, layout direction does
diff --git a/src/com/android/launcher3/Hotseat.java b/src/com/android/launcher3/Hotseat.java
index 1f5aa13..25eacb5 100644
--- a/src/com/android/launcher3/Hotseat.java
+++ b/src/com/android/launcher3/Hotseat.java
@@ -46,7 +46,7 @@
     private CellLayout mContent;
 
     @ViewDebug.ExportedProperty(category = "launcher")
-    private final boolean mHasVerticalHotseat;
+    private boolean mHasVerticalHotseat;
 
     public Hotseat(Context context) {
         this(context, null);
@@ -59,7 +59,6 @@
     public Hotseat(Context context, AttributeSet attrs, int defStyle) {
         super(context, attrs, defStyle);
         mLauncher = Launcher.getLauncher(context);
-        mHasVerticalHotseat = mLauncher.getDeviceProfile().isVerticalBarLayout();
     }
 
     public CellLayout getLayout() {
@@ -91,13 +90,7 @@
     @Override
     protected void onFinishInflate() {
         super.onFinishInflate();
-        DeviceProfile grid = mLauncher.getDeviceProfile();
-        mContent = (CellLayout) findViewById(R.id.layout);
-        if (grid.isVerticalBarLayout()) {
-            mContent.setGridSize(1, grid.inv.numHotseatIcons);
-        } else {
-            mContent.setGridSize(grid.inv.numHotseatIcons, 1);
-        }
+        mContent = findViewById(R.id.layout);
 
         resetLayout();
     }
@@ -165,7 +158,11 @@
     public void setInsets(Rect insets) {
         FrameLayout.LayoutParams lp = (FrameLayout.LayoutParams) getLayoutParams();
         DeviceProfile grid = mLauncher.getDeviceProfile();
+        mHasVerticalHotseat = mLauncher.getDeviceProfile().isVerticalBarLayout();
+
         if (mHasVerticalHotseat) {
+            mContent.setGridSize(1, grid.inv.numHotseatIcons);
+
             lp.height = ViewGroup.LayoutParams.MATCH_PARENT;
             if (insets.left > insets.right) {
                 lp.gravity = Gravity.LEFT;
@@ -180,6 +177,8 @@
                         grid.hotseatBarSidePaddingPx, insets.top, insets.right, insets.bottom);
             }
         } else {
+            mContent.setGridSize(grid.inv.numHotseatIcons, 1);
+
             lp.gravity = Gravity.BOTTOM;
             lp.width = ViewGroup.LayoutParams.MATCH_PARENT;
             lp.height = grid.hotseatBarSizePx + insets.bottom;
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 2a5f453..e3682b4 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -2552,10 +2552,6 @@
         return bounceAnim;
     }
 
-    public boolean useVerticalBarLayout() {
-        return mDeviceProfile.isVerticalBarLayout();
-    }
-
     /**
      * Add the icons for all apps.
      *
diff --git a/src/com/android/launcher3/LauncherState.java b/src/com/android/launcher3/LauncherState.java
index 670f579..472a5a9 100644
--- a/src/com/android/launcher3/LauncherState.java
+++ b/src/com/android/launcher3/LauncherState.java
@@ -27,6 +27,7 @@
 import com.android.launcher3.uioverrides.AllAppsState;
 import com.android.launcher3.states.SpringLoadedState;
 import com.android.launcher3.uioverrides.OverviewState;
+import com.android.launcher3.uioverrides.UiFactory;
 import com.android.launcher3.userevent.nano.LauncherLogProto.ContainerType;
 
 import java.util.Arrays;
@@ -54,6 +55,9 @@
 
     private static final LauncherState[] sAllStates = new LauncherState[4];
 
+    /**
+     * TODO: Create a separate class for NORMAL state.
+     */
     public static final LauncherState NORMAL = new LauncherState(0, ContainerType.WORKSPACE,
             0, FLAG_DISABLE_RESTORE | FLAG_WORKSPACE_ICONS_CAN_BE_DRAGGED);
 
@@ -175,6 +179,15 @@
         return NORMAL;
     }
 
+    /**
+     * Called when the start transition ends and the user settles on this particular state.
+     */
+    public void onStateTransitionEnd(Launcher launcher) {
+        if (this == NORMAL) {
+            UiFactory.resetOverview(launcher);
+        }
+    }
+
     protected static void dispatchWindowStateChanged(Launcher launcher) {
         launcher.getWindow().getDecorView().sendAccessibilityEvent(TYPE_WINDOW_STATE_CHANGED);
     }
diff --git a/src/com/android/launcher3/LauncherStateManager.java b/src/com/android/launcher3/LauncherStateManager.java
index 3a660dc..bcb6252 100644
--- a/src/com/android/launcher3/LauncherStateManager.java
+++ b/src/com/android/launcher3/LauncherStateManager.java
@@ -270,6 +270,7 @@
             mCurrentStableState = state;
         }
 
+        state.onStateTransitionEnd(mLauncher);
         mLauncher.getWorkspace().setClipChildren(!state.disablePageClipping);
         mLauncher.getUserEventDispatcher().resetElapsedContainerMillis();
         mLauncher.finishAutoCancelActionMode();
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 893d820..de3b09a 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -207,7 +207,7 @@
     private boolean mStripScreensOnPageStopMoving = false;
 
     private DragPreviewProvider mOutlineProvider = null;
-    private final boolean mWorkspaceFadeInAdjacentScreens;
+    private boolean mWorkspaceFadeInAdjacentScreens;
 
     final WallpaperOffsetInterpolator mWallpaperOffset;
     private boolean mUnlockWallpaperFromDefaultPageOnLayout;
@@ -292,8 +292,6 @@
 
         mLauncher = Launcher.getLauncher(context);
         mStateTransitionAnimation = new WorkspaceStateTransitionAnimation(mLauncher, this);
-        DeviceProfile grid = mLauncher.getDeviceProfile();
-        mWorkspaceFadeInAdjacentScreens = grid.shouldFadeAdjacentWorkspaceScreens();
         mWallpaperManager = WallpaperManager.getInstance(context);
 
         mWallpaperOffset = new WallpaperOffsetInterpolator(this);
@@ -312,6 +310,9 @@
         mInsets.set(insets);
 
         DeviceProfile grid = mLauncher.getDeviceProfile();
+        mMaxDistanceForFolderCreation = (0.55f * grid.iconSizePx);
+        mWorkspaceFadeInAdjacentScreens = grid.shouldFadeAdjacentWorkspaceScreens();
+
         Rect padding = grid.workspacePadding;
         setPadding(padding.left, padding.top, padding.right, padding.bottom);
 
@@ -324,6 +325,13 @@
             // We assume symmetrical padding in portrait mode.
             setPageSpacing(Math.max(grid.defaultPageSpacingPx, padding.left + 1));
         }
+
+        int paddingLeftRight = grid.cellLayoutPaddingLeftRightPx;
+        int paddingBottom = grid.cellLayoutBottomPaddingPx;
+        for (int i = mWorkspaceScreens.size() - 1; i >= 0; i--) {
+            mWorkspaceScreens.valueAt(i)
+                    .setPadding(paddingLeftRight, 0, paddingLeftRight, paddingBottom);
+        }
     }
 
     /**
@@ -445,12 +453,9 @@
      */
     protected void initWorkspace() {
         mCurrentPage = DEFAULT_PAGE;
-        DeviceProfile grid = mLauncher.getDeviceProfile();
-        setWillNotDraw(false);
         setClipToPadding(false);
 
         setupLayoutTransition();
-        mMaxDistanceForFolderCreation = (0.55f * grid.iconSizePx);
 
         // Set the wallpaper dimensions when Launcher starts up
         setWallpaperDimension();
diff --git a/src/com/android/launcher3/allapps/AllAppsContainerView.java b/src/com/android/launcher3/allapps/AllAppsContainerView.java
index a40f8b3..59bdf4b 100644
--- a/src/com/android/launcher3/allapps/AllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsContainerView.java
@@ -275,8 +275,6 @@
     @Override
     protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
         DeviceProfile grid = mLauncher.getDeviceProfile();
-        // Update the number of items in the grid before we measure the view
-        grid.updateAppsViewNumCols();
 
         if (mNumAppsPerRow != grid.inv.numColumns ||
                 mNumPredictedAppsPerRow != grid.inv.numColumns) {
@@ -364,18 +362,30 @@
     public void updateIconBadges(Set<PackageUserKey> updatedBadges) {
         final PackageUserKey packageUserKey = new PackageUserKey(null, null);
         for (int j = 0; j < mAH.length; j++) {
-            if (mAH[j].recyclerView != null) {
-                final int n = mAH[j].recyclerView.getChildCount();
-                for (int i = 0; i < n; i++) {
-                    View child = mAH[j].recyclerView.getChildAt(i);
-                    if (!(child instanceof BubbleTextView) || !(child.getTag() instanceof ItemInfo)) {
-                        continue;
-                    }
-                    ItemInfo info = (ItemInfo) child.getTag();
-                    if (packageUserKey.updateFromItemInfo(info) && updatedBadges.contains(packageUserKey)) {
-                        ((BubbleTextView) child).applyBadgeState(info, true /* animate */);
-                    }
-                }
+            updateIconBadges(updatedBadges, packageUserKey, mAH[j].recyclerView);
+        }
+        if (mHeader != null) {
+            updateIconBadges(updatedBadges, packageUserKey, mHeader.getPredictionRow());
+        }
+    }
+
+    private void updateIconBadges(Set<PackageUserKey> updatedBadges, PackageUserKey packageUserKey,
+            ViewGroup parent) {
+        if (parent == null) {
+            return;
+        }
+        final int n = parent.getChildCount();
+        for (int i = 0; i < n; i++) {
+            View child = parent.getChildAt(i);
+            if (child instanceof PredictionRowView) {
+                updateIconBadges(updatedBadges, packageUserKey, (PredictionRowView) child);
+            }
+            if (!(child instanceof BubbleTextView) || !(child.getTag() instanceof ItemInfo)) {
+                continue;
+            }
+            ItemInfo info = (ItemInfo) child.getTag();
+            if (packageUserKey.updateFromItemInfo(info) && updatedBadges.contains(packageUserKey)) {
+                ((BubbleTextView) child).applyBadgeState(info, true /* animate */);
             }
         }
     }
diff --git a/src/com/android/launcher3/allapps/AllAppsGridAdapter.java b/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
index 234eb81..769f9ba 100644
--- a/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
+++ b/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
@@ -41,6 +41,7 @@
 import com.android.launcher3.Utilities;
 import com.android.launcher3.allapps.AlphabeticalAppsList.AdapterItem;
 import com.android.launcher3.anim.SpringAnimationHandler;
+import com.android.launcher3.compat.UserManagerCompat;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.discovery.AppDiscoveryAppInfo;
 import com.android.launcher3.discovery.AppDiscoveryItemView;
@@ -377,6 +378,11 @@
             case VIEW_TYPE_WORK_TAB_FOOTER:
                 WorkModeSwitch workModeToggle = holder.itemView.findViewById(R.id.work_mode_toggle);
                 workModeToggle.refresh();
+                TextView managedByLabel = holder.itemView.findViewById(R.id.managed_by_label);
+                boolean anyProfileQuietModeEnabled = UserManagerCompat.getInstance(
+                        managedByLabel.getContext()).isAnyProfileQuietModeEnabled();
+                managedByLabel.setText(anyProfileQuietModeEnabled
+                        ? R.string.work_mode_off_label : R.string.work_mode_on_label);
                 break;
         }
         if (mBindViewCallback != null) {
diff --git a/src/com/android/launcher3/allapps/WorkModeSwitch.java b/src/com/android/launcher3/allapps/WorkModeSwitch.java
index 05d5e53..e7cf092 100644
--- a/src/com/android/launcher3/allapps/WorkModeSwitch.java
+++ b/src/com/android/launcher3/allapps/WorkModeSwitch.java
@@ -55,22 +55,9 @@
     }
 
     public void refresh() {
-        setCheckedInternal(!isAnyProfileQuietModeEnabled());
-        setEnabled(true);
-    }
-
-    private boolean isAnyProfileQuietModeEnabled() {
         UserManagerCompat userManager = UserManagerCompat.getInstance(getContext());
-        List<UserHandle> userProfiles = userManager.getUserProfiles();
-        for (UserHandle userProfile : userProfiles) {
-            if (Process.myUserHandle().equals(userProfile)) {
-                continue;
-            }
-            if (userManager.isQuietModeEnabled(userProfile)) {
-                return true;
-            }
-        }
-        return false;
+        setCheckedInternal(!userManager.isAnyProfileQuietModeEnabled());
+        setEnabled(true);
     }
 
     private void trySetQuietModeEnabledToAllProfilesAsync(boolean enabled) {
diff --git a/src/com/android/launcher3/compat/UserManagerCompat.java b/src/com/android/launcher3/compat/UserManagerCompat.java
index ddb244a..62055dc 100644
--- a/src/com/android/launcher3/compat/UserManagerCompat.java
+++ b/src/com/android/launcher3/compat/UserManagerCompat.java
@@ -64,4 +64,5 @@
 
     public abstract boolean isDemoUser();
     public abstract boolean requestQuietModeEnabled(boolean enableQuietMode, UserHandle user);
+    public abstract boolean isAnyProfileQuietModeEnabled();
 }
diff --git a/src/com/android/launcher3/compat/UserManagerCompatVL.java b/src/com/android/launcher3/compat/UserManagerCompatVL.java
index 5585fe4..e57786d 100644
--- a/src/com/android/launcher3/compat/UserManagerCompatVL.java
+++ b/src/com/android/launcher3/compat/UserManagerCompatVL.java
@@ -88,6 +88,11 @@
     }
 
     @Override
+    public boolean isAnyProfileQuietModeEnabled() {
+        return false;
+    }
+
+    @Override
     public void enableAndResetCache() {
         synchronized (this) {
             mUsers = new LongArrayMap<>();
diff --git a/src/com/android/launcher3/compat/UserManagerCompatVN.java b/src/com/android/launcher3/compat/UserManagerCompatVN.java
index 50a0217..3733565 100644
--- a/src/com/android/launcher3/compat/UserManagerCompatVN.java
+++ b/src/com/android/launcher3/compat/UserManagerCompatVN.java
@@ -19,8 +19,11 @@
 import android.annotation.TargetApi;
 import android.content.Context;
 import android.os.Build;
+import android.os.Process;
 import android.os.UserHandle;
 
+import java.util.List;
+
 @TargetApi(Build.VERSION_CODES.N)
 public class UserManagerCompatVN extends UserManagerCompatVM {
 
@@ -37,5 +40,19 @@
     public boolean isUserUnlocked(UserHandle user) {
         return mUserManager.isUserUnlocked(user);
     }
+
+    @Override
+    public boolean isAnyProfileQuietModeEnabled() {
+        List<UserHandle> userProfiles = getUserProfiles();
+        for (UserHandle userProfile : userProfiles) {
+            if (Process.myUserHandle().equals(userProfile)) {
+                continue;
+            }
+            if (isQuietModeEnabled(userProfile)) {
+                return true;
+            }
+        }
+        return false;
+    }
 }
 
diff --git a/src/com/android/launcher3/dragndrop/DragView.java b/src/com/android/launcher3/dragndrop/DragView.java
index 9732261..a59b899 100644
--- a/src/com/android/launcher3/dragndrop/DragView.java
+++ b/src/com/android/launcher3/dragndrop/DragView.java
@@ -369,9 +369,9 @@
                 return new FixedSizeEmptyDrawable(iconSize);
             }
             ShortcutInfoCompat si = (ShortcutInfoCompat) obj;
-            Bitmap badge = LauncherIcons.getShortcutInfoBadge(si, appState.getIconCache())
-                    .iconBitmap;
-
+            LauncherIcons li = LauncherIcons.obtain(appState.getContext());
+            Bitmap badge = li.getShortcutInfoBadge(si, appState.getIconCache()).iconBitmap;
+            li.recycle();
             float badgeSize = mLauncher.getResources().getDimension(R.dimen.profile_badge_size);
             float insetFraction = (iconSize - badgeSize) / iconSize;
             return new InsetDrawable(new FastBitmapDrawable(badge),
diff --git a/src/com/android/launcher3/graphics/LauncherIcons.java b/src/com/android/launcher3/graphics/LauncherIcons.java
index 0c9f4d9..34fc921 100644
--- a/src/com/android/launcher3/graphics/LauncherIcons.java
+++ b/src/com/android/launcher3/graphics/LauncherIcons.java
@@ -356,10 +356,11 @@
         return result;
     }
 
-    public static ItemInfoWithIcon getShortcutInfoBadge(
-            ShortcutInfoCompat shortcutInfo, IconCache cache) {
+    public ItemInfoWithIcon getShortcutInfoBadge(ShortcutInfoCompat shortcutInfo, IconCache cache) {
         ComponentName cn = shortcutInfo.getActivity();
-        if (cn != null) {
+        String badgePkg = shortcutInfo.getBadgePackage(mContext);
+        boolean hasBadgePkgSet = !badgePkg.equals(shortcutInfo.getPackage());
+        if (cn != null && !hasBadgePkgSet) {
             // Get the app info for the source activity.
             AppInfo appInfo = new AppInfo();
             appInfo.user = shortcutInfo.getUserHandle();
@@ -370,7 +371,7 @@
             cache.getTitleAndIcon(appInfo, false);
             return appInfo;
         } else {
-            PackageItemInfo pkgInfo = new PackageItemInfo(shortcutInfo.getPackage());
+            PackageItemInfo pkgInfo = new PackageItemInfo(badgePkg);
             cache.getTitleAndIconForApp(pkgInfo, false);
             return pkgInfo;
         }
diff --git a/src/com/android/launcher3/shortcuts/ShortcutInfoCompat.java b/src/com/android/launcher3/shortcuts/ShortcutInfoCompat.java
index 9c91c87..325777d 100644
--- a/src/com/android/launcher3/shortcuts/ShortcutInfoCompat.java
+++ b/src/com/android/launcher3/shortcuts/ShortcutInfoCompat.java
@@ -18,11 +18,14 @@
 
 import android.annotation.TargetApi;
 import android.content.ComponentName;
+import android.content.Context;
 import android.content.Intent;
 import android.content.pm.ShortcutInfo;
 import android.os.Build;
 import android.os.UserHandle;
 
+import com.android.launcher3.R;
+
 /**
  * Wrapper class for {@link android.content.pm.ShortcutInfo}, representing deep shortcuts into apps.
  *
@@ -31,8 +34,8 @@
 @TargetApi(Build.VERSION_CODES.N)
 public class ShortcutInfoCompat {
     private static final String INTENT_CATEGORY = "com.android.launcher3.DEEP_SHORTCUT";
+    private static final String EXTRA_BADGEPKG = "badge_package";
     public static final String EXTRA_SHORTCUT_ID = "shortcut_id";
-
     private ShortcutInfo mShortcutInfo;
 
     public ShortcutInfoCompat(ShortcutInfo shortcutInfo) {
@@ -57,6 +60,15 @@
         return mShortcutInfo.getPackage();
     }
 
+    public String getBadgePackage(Context context) {
+        String whitelistedPkg = context.getString(R.string.shortcutinfocompat_badgepkg_whitelist);
+        if (whitelistedPkg.equals(getPackage())
+                && mShortcutInfo.getExtras().containsKey(EXTRA_BADGEPKG)) {
+            return mShortcutInfo.getExtras().getString(EXTRA_BADGEPKG);
+        }
+        return getPackage();
+    }
+
     public String getId() {
         return mShortcutInfo.getId();
     }
diff --git a/src_ui_overrides/com/android/launcher3/uioverrides/UiFactory.java b/src_ui_overrides/com/android/launcher3/uioverrides/UiFactory.java
index fc81e80..2ea10c2 100644
--- a/src_ui_overrides/com/android/launcher3/uioverrides/UiFactory.java
+++ b/src_ui_overrides/com/android/launcher3/uioverrides/UiFactory.java
@@ -56,4 +56,6 @@
         renderer.render(new Canvas(result));
         return result;
     }
+
+    public static void resetOverview(Launcher launcher) { }
 }