diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/AllAppsTipView.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/AllAppsTipView.java
index d93aea44a7..8477b103a2 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/AllAppsTipView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/AllAppsTipView.java
@@ -26,10 +26,10 @@ import android.os.UserManager;
 import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.allapps.FloatingHeaderView;
+import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.views.ArrowTipView;
 import com.android.systemui.shared.system.LauncherEventUtil;
 
@@ -71,17 +71,16 @@ public class AllAppsTipView {
 
     public static void scheduleShowIfNeeded(Launcher launcher) {
         if (!hasSeenAllAppsTip(launcher)) {
-            launcher.getStateManager().addStateListener(
-                    new LauncherStateManager.StateListener() {
-                        @Override
-                        public void onStateTransitionComplete(LauncherState finalState) {
-                            if (finalState == ALL_APPS) {
-                                if (showAllAppsTipIfNecessary(launcher)) {
-                                    launcher.getStateManager().removeStateListener(this);
-                                }
-                            }
+            launcher.getStateManager().addStateListener(new StateListener<LauncherState>() {
+                @Override
+                public void onStateTransitionComplete(LauncherState finalState) {
+                    if (finalState == ALL_APPS) {
+                        if (showAllAppsTipIfNecessary(launcher)) {
+                            launcher.getStateManager().removeStateListener(this);
                         }
-                    });
+                    }
+                }
+            });
         }
     }
 }
diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/AppsDividerView.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/AppsDividerView.java
index 81a6070664..914d9e9774 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/AppsDividerView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/AppsDividerView.java
@@ -38,18 +38,18 @@ import androidx.core.content.ContextCompat;
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager;
 import com.android.launcher3.R;
 import com.android.launcher3.allapps.FloatingHeaderRow;
 import com.android.launcher3.allapps.FloatingHeaderView;
 import com.android.launcher3.anim.PropertySetter;
+import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.util.Themes;
 
 /**
  * A view which shows a horizontal divider
  */
 @TargetApi(Build.VERSION_CODES.O)
-public class AppsDividerView extends View implements LauncherStateManager.StateListener,
+public class AppsDividerView extends View implements StateListener<LauncherState>,
         FloatingHeaderRow {
 
     private static final String ALL_APPS_VISITED_COUNT = "launcher.all_apps_visited_count";
diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionUiStateManager.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionUiStateManager.java
index e68627a8f6..ab3c71ae86 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionUiStateManager.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/appprediction/PredictionUiStateManager.java
@@ -32,7 +32,6 @@ import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager.StateListener;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.allapps.AllAppsContainerView;
 import com.android.launcher3.allapps.AllAppsStore.OnUpdateListener;
@@ -41,6 +40,7 @@ import com.android.launcher3.icons.IconCache.ItemInfoUpdateReceiver;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.ItemInfoWithIcon;
 import com.android.launcher3.shortcuts.ShortcutKey;
+import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.userevent.nano.LauncherLogProto;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.MainThreadInitializedObject;
@@ -63,8 +63,8 @@ import java.util.stream.IntStream;
  * 4) Maintains the current active client id (for the predictions) and all updates are performed on
  * that client id.
  */
-public class PredictionUiStateManager implements StateListener, ItemInfoUpdateReceiver,
-        OnIDPChangeListener, OnUpdateListener {
+public class PredictionUiStateManager implements StateListener<LauncherState>,
+        ItemInfoUpdateReceiver, OnIDPChangeListener, OnUpdateListener {
 
     public static final String LAST_PREDICTION_ENABLED_STATE = "last_prediction_enabled_state";
 
diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 30a9416d29..6cfc846432 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -39,7 +39,6 @@ import com.android.launcher3.BaseQuickstepLauncher;
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager.AtomicAnimationFactory;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.allapps.DiscoveryBounce;
 import com.android.launcher3.anim.AnimatorPlaybackController;
@@ -50,6 +49,7 @@ import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.popup.SystemShortcut;
+import com.android.launcher3.statemanager.StateManager.AtomicAnimationFactory;
 import com.android.launcher3.testing.TestProtocol;
 import com.android.launcher3.uioverrides.states.QuickstepAtomicAnimationFactory;
 import com.android.launcher3.uioverrides.touchcontrollers.FlingAndHoldTouchController;
diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/states/QuickstepAtomicAnimationFactory.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/states/QuickstepAtomicAnimationFactory.java
index 6bc69f949a..11593a16c1 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/states/QuickstepAtomicAnimationFactory.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/states/QuickstepAtomicAnimationFactory.java
@@ -62,12 +62,12 @@ import com.android.launcher3.CellLayout;
 import com.android.launcher3.Hotseat;
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.LauncherState.ScaleAndTranslation;
-import com.android.launcher3.LauncherStateManager;
-import com.android.launcher3.LauncherStateManager.AtomicAnimationFactory;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.allapps.AllAppsContainerView;
 import com.android.launcher3.allapps.AllAppsTransitionController;
 import com.android.launcher3.anim.SpringAnimationBuilder;
+import com.android.launcher3.statemanager.StateManager;
+import com.android.launcher3.statemanager.StateManager.AtomicAnimationFactory;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.uioverrides.QuickstepLauncher;
 import com.android.quickstep.SysUINavigationMode;
@@ -76,7 +76,7 @@ import com.android.quickstep.views.RecentsView;
 /**
  * Animation factory for quickstep specific transitions
  */
-public class QuickstepAtomicAnimationFactory extends AtomicAnimationFactory {
+public class QuickstepAtomicAnimationFactory extends AtomicAnimationFactory<LauncherState> {
 
     // Scale recents takes before animating in
     private static final float RECENTS_PREPARE_SCALE = 1.33f;
@@ -153,7 +153,7 @@ public class QuickstepAtomicAnimationFactory extends AtomicAnimationFactory {
                     config.setInterpolator(ANIM_HOTSEAT_TRANSLATE, OVERSHOOT_1_2);
                 }
 
-                LauncherStateManager stateManager = mLauncher.getStateManager();
+                StateManager<LauncherState> stateManager = mLauncher.getStateManager();
                 return stateManager.createAtomicAnimation(
                         stateManager.getCurrentStableState(), OVERVIEW, config);
             }
diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/touchcontrollers/NoButtonNavbarToOverviewTouchController.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/touchcontrollers/NoButtonNavbarToOverviewTouchController.java
index f288b7c070..966e25bff5 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/touchcontrollers/NoButtonNavbarToOverviewTouchController.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/touchcontrollers/NoButtonNavbarToOverviewTouchController.java
@@ -32,9 +32,9 @@ import android.view.MotionEvent;
 
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimationSuccessListener;
+import com.android.launcher3.statemanager.StateManager;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.testing.TestProtocol;
 import com.android.launcher3.userevent.nano.LauncherLogProto.Action.Touch;
@@ -173,7 +173,7 @@ public class NoButtonNavbarToOverviewTouchController extends FlingAndHoldTouchCo
     protected void goToOverviewOnDragEnd(float velocity) {
         float velocityDp = dpiFromPx(velocity);
         boolean isFling = Math.abs(velocityDp) > 1;
-        LauncherStateManager stateManager = mLauncher.getStateManager();
+        StateManager<LauncherState> stateManager = mLauncher.getStateManager();
         boolean goToHomeInsteadOfOverview = isFling;
         if (goToHomeInsteadOfOverview) {
             if (velocity > 0) {
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java
index 6d3077e0e9..9d7efc4ddb 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java
@@ -41,11 +41,11 @@ import android.widget.FrameLayout;
 import com.android.launcher3.BaseQuickstepLauncher;
 import com.android.launcher3.Hotseat;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager.StateListener;
 import com.android.launcher3.anim.Interpolators;
 import com.android.launcher3.appprediction.PredictionUiStateManager;
 import com.android.launcher3.appprediction.PredictionUiStateManager.Client;
 import com.android.launcher3.statehandlers.DepthController;
+import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.TraceHelper;
 import com.android.launcher3.views.ScrimView;
@@ -59,7 +59,7 @@ import com.android.systemui.plugins.RecentsExtraCard;
  */
 @TargetApi(Build.VERSION_CODES.O)
 public class LauncherRecentsView extends RecentsView<BaseQuickstepLauncher>
-        implements StateListener {
+        implements StateListener<LauncherState> {
 
     private final TransformParams mTransformParams = new TransformParams();
 
diff --git a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
index 9d64d09cc8..629a74b87a 100644
--- a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
@@ -30,7 +30,6 @@ import android.content.SharedPreferences;
 import android.os.Bundle;
 import android.os.CancellationSignal;
 
-import com.android.launcher3.LauncherStateManager.StateHandler;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.model.WellbeingModel;
 import com.android.launcher3.popup.SystemShortcut;
@@ -38,6 +37,7 @@ import com.android.launcher3.proxy.ProxyActivityStarter;
 import com.android.launcher3.proxy.StartActivityParams;
 import com.android.launcher3.statehandlers.BackButtonAlphaHandler;
 import com.android.launcher3.statehandlers.DepthController;
+import com.android.launcher3.statemanager.StateManager.StateHandler;
 import com.android.launcher3.uioverrides.RecentsViewStateController;
 import com.android.launcher3.util.OnboardingPrefs;
 import com.android.launcher3.util.UiThreadHelper;
@@ -153,6 +153,7 @@ public abstract class BaseQuickstepLauncher extends Launcher
 
     @Override
     protected void onDeferredResumed() {
+        super.onDeferredResumed();
         if (mPendingActivityRequestCode != -1 && isInState(NORMAL)) {
             // Remove any active ProxyActivityStarter task and send RESULT_CANCELED to Launcher.
             onActivityResult(mPendingActivityRequestCode, RESULT_CANCELED, null);
@@ -194,7 +195,7 @@ public abstract class BaseQuickstepLauncher extends Launcher
     }
 
     @Override
-    protected StateHandler[] createStateHandlers() {
+    protected StateHandler<LauncherState>[] createStateHandlers() {
         return new StateHandler[] {
                 getAllAppsController(),
                 getWorkspace(),
@@ -208,9 +209,8 @@ public abstract class BaseQuickstepLauncher extends Launcher
     }
 
     @Override
-    protected OnboardingPrefs createOnboardingPrefs(SharedPreferences sharedPrefs,
-            LauncherStateManager stateManager) {
-        return new QuickstepOnboardingPrefs(this, sharedPrefs, stateManager);
+    protected OnboardingPrefs createOnboardingPrefs(SharedPreferences sharedPrefs) {
+        return new QuickstepOnboardingPrefs(this, sharedPrefs);
     }
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/statehandlers/BackButtonAlphaHandler.java b/quickstep/src/com/android/launcher3/statehandlers/BackButtonAlphaHandler.java
index 075a4833fa..13501a452b 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/BackButtonAlphaHandler.java
+++ b/quickstep/src/com/android/launcher3/statehandlers/BackButtonAlphaHandler.java
@@ -16,14 +16,13 @@
 
 package com.android.launcher3.statehandlers;
 
-import static com.android.launcher3.LauncherState.FLAG_HIDE_BACK_BUTTON;
 import static com.android.launcher3.anim.Interpolators.LINEAR;
 import static com.android.quickstep.AnimatedFloat.VALUE;
 
 import com.android.launcher3.BaseQuickstepLauncher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager;
 import com.android.launcher3.anim.PendingAnimation;
+import com.android.launcher3.statemanager.StateManager.StateHandler;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.util.UiThreadHelper;
 import com.android.quickstep.AnimatedFloat;
@@ -33,7 +32,7 @@ import com.android.quickstep.SystemUiProxy;
 /**
  * State handler for animating back button alpha
  */
-public class BackButtonAlphaHandler implements LauncherStateManager.StateHandler {
+public class BackButtonAlphaHandler implements StateHandler<LauncherState> {
 
     private final BaseQuickstepLauncher mLauncher;
     private final AnimatedFloat mBackAlpha = new AnimatedFloat(this::updateBackAlpha);
diff --git a/quickstep/src/com/android/launcher3/statehandlers/DepthController.java b/quickstep/src/com/android/launcher3/statehandlers/DepthController.java
index 8c778c0653..8292a92af6 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/DepthController.java
+++ b/quickstep/src/com/android/launcher3/statehandlers/DepthController.java
@@ -26,10 +26,10 @@ import android.view.ViewTreeObserver;
 
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.PendingAnimation;
+import com.android.launcher3.statemanager.StateManager.StateHandler;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
 import com.android.systemui.shared.system.SurfaceControlCompat;
@@ -39,7 +39,7 @@ import com.android.systemui.shared.system.WallpaperManagerCompat;
 /**
  * Controls blur and wallpaper zoom, for the Launcher surface only.
  */
-public class DepthController implements LauncherStateManager.StateHandler {
+public class DepthController implements StateHandler<LauncherState> {
 
     public static final FloatProperty<DepthController> DEPTH =
             new FloatProperty<DepthController>("depth") {
diff --git a/quickstep/src/com/android/launcher3/uioverrides/BaseRecentsViewStateController.java b/quickstep/src/com/android/launcher3/uioverrides/BaseRecentsViewStateController.java
index 47fff5ee32..ec3a4905f5 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/BaseRecentsViewStateController.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/BaseRecentsViewStateController.java
@@ -36,9 +36,9 @@ import androidx.annotation.NonNull;
 
 import com.android.launcher3.BaseQuickstepLauncher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager.StateHandler;
 import com.android.launcher3.anim.PendingAnimation;
 import com.android.launcher3.graphics.OverviewScrim;
+import com.android.launcher3.statemanager.StateManager.StateHandler;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.quickstep.views.RecentsView;
 
@@ -49,7 +49,7 @@ import com.android.quickstep.views.RecentsView;
  * @param <T> the recents view
  */
 public abstract class BaseRecentsViewStateController<T extends RecentsView>
-        implements StateHandler {
+        implements StateHandler<LauncherState> {
     protected final T mRecentsView;
     protected final BaseQuickstepLauncher mLauncher;
 
diff --git a/quickstep/src/com/android/quickstep/util/QuickstepOnboardingPrefs.java b/quickstep/src/com/android/quickstep/util/QuickstepOnboardingPrefs.java
index aa6d56a568..2d8bba2896 100644
--- a/quickstep/src/com/android/quickstep/util/QuickstepOnboardingPrefs.java
+++ b/quickstep/src/com/android/quickstep/util/QuickstepOnboardingPrefs.java
@@ -25,8 +25,8 @@ import android.content.SharedPreferences;
 
 import com.android.launcher3.BaseQuickstepLauncher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager;
-import com.android.launcher3.LauncherStateManager.StateListener;
+import com.android.launcher3.statemanager.StateManager;
+import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.util.OnboardingPrefs;
 import com.android.quickstep.SysUINavigationMode;
 
@@ -35,23 +35,23 @@ import com.android.quickstep.SysUINavigationMode;
  */
 public class QuickstepOnboardingPrefs extends OnboardingPrefs<BaseQuickstepLauncher> {
 
-    public QuickstepOnboardingPrefs(BaseQuickstepLauncher launcher, SharedPreferences sharedPrefs,
-            LauncherStateManager stateManager) {
-        super(launcher, sharedPrefs, stateManager);
+    public QuickstepOnboardingPrefs(BaseQuickstepLauncher launcher, SharedPreferences sharedPrefs) {
+        super(launcher, sharedPrefs);
 
+        StateManager<LauncherState> stateManager = launcher.getStateManager();
         if (!getBoolean(HOME_BOUNCE_SEEN)) {
-            mStateManager.addStateListener(new StateListener() {
+            stateManager.addStateListener(new StateListener<LauncherState>() {
                 @Override
                 public void onStateTransitionComplete(LauncherState finalState) {
                     boolean swipeUpEnabled = SysUINavigationMode.INSTANCE
                             .get(mLauncher).getMode().hasGestures;
-                    LauncherState prevState = mStateManager.getLastState();
+                    LauncherState prevState = stateManager.getLastState();
 
                     if (((swipeUpEnabled && finalState == OVERVIEW) || (!swipeUpEnabled
                             && finalState == ALL_APPS && prevState == NORMAL) ||
                             hasReachedMaxCount(HOME_BOUNCE_COUNT))) {
                         mSharedPrefs.edit().putBoolean(HOME_BOUNCE_SEEN, true).apply();
-                        mStateManager.removeStateListener(this);
+                        stateManager.removeStateListener(this);
                     }
                 }
             });
@@ -65,27 +65,27 @@ public class QuickstepOnboardingPrefs extends OnboardingPrefs<BaseQuickstepLaunc
             mSharedPrefs.edit().putBoolean(SHELF_BOUNCE_SEEN, shelfBounceSeen).apply();
         }
         if (!shelfBounceSeen) {
-            mStateManager.addStateListener(new StateListener() {
+            stateManager.addStateListener(new StateListener<LauncherState>() {
                 @Override
                 public void onStateTransitionComplete(LauncherState finalState) {
-                    LauncherState prevState = mStateManager.getLastState();
+                    LauncherState prevState = stateManager.getLastState();
 
                     if ((finalState == ALL_APPS && prevState == OVERVIEW) ||
                             hasReachedMaxCount(SHELF_BOUNCE_COUNT)) {
                         mSharedPrefs.edit().putBoolean(SHELF_BOUNCE_SEEN, true).apply();
-                        mStateManager.removeStateListener(this);
+                        stateManager.removeStateListener(this);
                     }
                 }
             });
         }
 
         if (!hasReachedMaxCount(ALL_APPS_COUNT)) {
-            mStateManager.addStateListener(new StateListener() {
+            stateManager.addStateListener(new StateListener<LauncherState>() {
                 @Override
                 public void onStateTransitionComplete(LauncherState finalState) {
                     if (finalState == ALL_APPS) {
                         if (incrementEventCount(ALL_APPS_COUNT)) {
-                            mStateManager.removeStateListener(this);
+                            stateManager.removeStateListener(this);
                             mLauncher.getScrimView().updateDragHandleVisibility();
                         }
                     }
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 873b066eec..59476dd5a1 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -67,7 +67,6 @@ import android.database.sqlite.SQLiteDatabase;
 import android.os.Build;
 import android.os.Bundle;
 import android.os.CancellationSignal;
-import android.os.Handler;
 import android.os.Parcelable;
 import android.os.Process;
 import android.os.StrictMode;
@@ -87,13 +86,12 @@ import android.view.accessibility.AccessibilityEvent;
 import android.view.animation.OvershootInterpolator;
 import android.widget.Toast;
 
+import androidx.annotation.CallSuper;
 import androidx.annotation.Nullable;
 import androidx.annotation.StringRes;
 import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.DropTarget.DragObject;
-import com.android.launcher3.LauncherStateManager.AtomicAnimationFactory;
-import com.android.launcher3.LauncherStateManager.StateHandler;
 import com.android.launcher3.accessibility.LauncherAccessibilityDelegate;
 import com.android.launcher3.allapps.AllAppsContainerView;
 import com.android.launcher3.allapps.AllAppsStore;
@@ -131,6 +129,10 @@ import com.android.launcher3.popup.PopupContainerWithArrow;
 import com.android.launcher3.popup.PopupDataProvider;
 import com.android.launcher3.popup.SystemShortcut;
 import com.android.launcher3.qsb.QsbContainerView;
+import com.android.launcher3.statemanager.StateManager;
+import com.android.launcher3.statemanager.StateManager.StateHandler;
+import com.android.launcher3.statemanager.StateManager.StateListener;
+import com.android.launcher3.statemanager.StatefulActivity;
 import com.android.launcher3.states.RotationHelper;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.TestProtocol;
@@ -194,7 +196,7 @@ import java.util.stream.Stream;
 /**
  * Default launcher application.
  */
-public class Launcher extends BaseDraggingActivity implements LauncherExterns,
+public class Launcher extends StatefulActivity<LauncherState> implements LauncherExterns,
         Callbacks, InvariantDeviceProfile.OnIDPChangeListener, PluginListener<OverlayPlugin> {
     public static final String TAG = "Launcher";
 
@@ -241,7 +243,7 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
     public static final String ON_RESUME_EVT = "Launcher.onResume";
     public static final String ON_NEW_INTENT_EVT = "Launcher.onNewIntent";
 
-    private LauncherStateManager mStateManager;
+    private StateManager<LauncherState> mStateManager;
 
     private static final int ON_ACTIVITY_RESULT_ANIMATION_DELAY = 500;
 
@@ -325,10 +327,6 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
 
     private RotationHelper mRotationHelper;
 
-    final Handler mHandler = new Handler();
-    private final Runnable mHandleDeferredResume = this::handleDeferredResume;
-    private boolean mDeferredResumePending;
-
     private float mCurrentAssistantVisibility = 0f;
 
     protected LauncherOverlayManager mOverlayManager;
@@ -375,9 +373,9 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
 
         mDragController = new DragController(this);
         mAllAppsController = new AllAppsTransitionController(this);
-        mStateManager = new LauncherStateManager(this);
+        mStateManager = new StateManager<>(this, NORMAL);
 
-        mOnboardingPrefs = createOnboardingPrefs(mSharedPrefs, mStateManager);
+        mOnboardingPrefs = createOnboardingPrefs(mSharedPrefs);
 
         mAppWidgetManager = new WidgetManagerHelper(this);
         mAppWidgetHost = new LauncherAppWidgetHost(this,
@@ -440,7 +438,7 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
 
         mRotationHelper.initialize();
 
-        mStateManager.addStateListener(new LauncherStateManager.StateListener() {
+        mStateManager.addStateListener(new StateListener<LauncherState>() {
 
             @Override
             public void onStateTransitionComplete(LauncherState finalState) {
@@ -467,9 +465,8 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
         return new LauncherOverlayManager() { };
     }
 
-    protected OnboardingPrefs createOnboardingPrefs(SharedPreferences sharedPrefs,
-            LauncherStateManager stateManager) {
-        return new OnboardingPrefs<>(this, sharedPrefs, stateManager);
+    protected OnboardingPrefs createOnboardingPrefs(SharedPreferences sharedPrefs) {
+        return new OnboardingPrefs<>(this, sharedPrefs);
     }
 
     public OnboardingPrefs getOnboardingPrefs() {
@@ -523,13 +520,9 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
     }
 
     @Override
-    public void reapplyUi() {
-        reapplyUi(true /* cancelCurrentAnimation */);
-    }
-
     public void reapplyUi(boolean cancelCurrentAnimation) {
         getRootView().dispatchInsets();
-        getStateManager().reapplyState(cancelCurrentAnimation);
+        super.reapplyUi(cancelCurrentAnimation);
     }
 
     @Override
@@ -583,7 +576,8 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
         return mFocusHandler;
     }
 
-    public LauncherStateManager getStateManager() {
+    @Override
+    public StateManager<LauncherState> getStateManager() {
         return mStateManager;
     }
 
@@ -890,11 +884,7 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
 
     @Override
     protected void onStop() {
-        final boolean wasActive = isUserActive();
-        final LauncherState origState = getStateManager().getState();
-        final int origDragLayerChildCount = mDragLayer.getChildCount();
         super.onStop();
-
         if (mDeferOverlayCallbacks) {
             checkIfOverlayStillDeferred();
         } else {
@@ -902,28 +892,8 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
         }
 
         logStopAndResume(Action.Command.STOP);
-
         mAppWidgetHost.setListenIfResumed(false);
-
         NotificationListener.removeNotificationsChangedListener();
-        getStateManager().moveToRestState();
-
-        // Workaround for b/78520668, explicitly trim memory once UI is hidden
-        onTrimMemory(TRIM_MEMORY_UI_HIDDEN);
-
-        if (wasActive) {
-            // The expected condition is that this activity is stopped because the device goes to
-            // sleep and the UI may have noticeable changes.
-            mDragLayer.post(() -> {
-                if ((!getStateManager().isInStableState(origState)
-                        // The drag layer may be animating (e.g. dismissing QSB).
-                        || mDragLayer.getAlpha() < 1
-                        // Maybe an ArrowPopup is closed.
-                        || mDragLayer.getChildCount() != origDragLayerChildCount)) {
-                    onUiChangedWhileSleeping();
-                }
-            });
-        }
     }
 
     @Override
@@ -939,35 +909,27 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
         TraceHelper.INSTANCE.endSection(traceToken);
     }
 
-    private void handleDeferredResume() {
-        if (hasBeenResumed() && !mStateManager.getState().hasFlag(FLAG_NON_INTERACTIVE)) {
-            logStopAndResume(Action.Command.RESUME);
-            getUserEventDispatcher().startSession();
+    @Override
+    @CallSuper
+    protected void onDeferredResumed() {
+        logStopAndResume(Action.Command.RESUME);
+        getUserEventDispatcher().startSession();
 
-            AppLaunchTracker.INSTANCE.get(this).onReturnedToHome();
+        AppLaunchTracker.INSTANCE.get(this).onReturnedToHome();
 
-            // Process any items that were added while Launcher was away.
-            InstallShortcutReceiver.disableAndFlushInstallQueue(
-                    InstallShortcutReceiver.FLAG_ACTIVITY_PAUSED, this);
+        // Process any items that were added while Launcher was away.
+        InstallShortcutReceiver.disableAndFlushInstallQueue(
+                InstallShortcutReceiver.FLAG_ACTIVITY_PAUSED, this);
 
-            // Refresh shortcuts if the permission changed.
-            mModel.refreshShortcutsIfRequired();
+        // Refresh shortcuts if the permission changed.
+        mModel.refreshShortcutsIfRequired();
 
-            // Set the notification listener and fetch updated notifications when we resume
-            NotificationListener.setNotificationsChangedListener(mPopupDataProvider);
+        // Set the notification listener and fetch updated notifications when we resume
+        NotificationListener.setNotificationsChangedListener(mPopupDataProvider);
 
-            DiscoveryBounce.showForHomeIfNeeded(this);
-
-            onDeferredResumed();
-            addActivityFlags(ACTIVITY_STATE_DEFERRED_RESUMED);
-
-            mDeferredResumePending = false;
-        } else {
-            mDeferredResumePending = true;
-        }
+        DiscoveryBounce.showForHomeIfNeeded(this);
     }
 
-    protected void onDeferredResumed() { }
 
     private void logStopAndResume(int command) {
         int containerType = mStateManager.getState().containerType;
@@ -1016,10 +978,9 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
         return mOverlayManager;
     }
 
+    @Override
     public void onStateSetStart(LauncherState state) {
-        if (mDeferredResumePending) {
-            handleDeferredResume();
-        }
+        super.onStateSetStart(state);
         if (mDeferOverlayCallbacks) {
             scheduleDeferredCheck();
         }
@@ -1042,7 +1003,9 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
         mWorkspace.getPageIndicator().setShouldAutoHide(!state.hasFlag(FLAG_MULTI_PAGE));
     }
 
+    @Override
     public void onStateSetEnd(LauncherState state) {
+        super.onStateSetStart(state);
         getAppWidgetHost().setResumed(state == LauncherState.NORMAL);
         getWorkspace().setClipChildren(!state.hasFlag(FLAG_MULTI_PAGE));
 
@@ -1068,9 +1031,6 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
                 TraceHelper.FLAG_UI_EVENT);
         super.onResume();
 
-        mHandler.removeCallbacks(mHandleDeferredResume);
-        Utilities.postAsyncCallback(mHandler, mHandleDeferredResume);
-
         if (!mOnResumeCallbacks.isEmpty()) {
             final ArrayList<OnResumeCallback> resumeCallbacks = new ArrayList<>(mOnResumeCallbacks);
             mOnResumeCallbacks.clear();
@@ -1113,10 +1073,6 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
         }
     }
 
-    public boolean isInState(LauncherState state) {
-        return mStateManager.getState() == state;
-    }
-
     /**
      * Restores the previous state, if it exists.
      *
@@ -1355,8 +1311,6 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
         }
     };
 
-    protected void onUiChangedWhileSleeping() { }
-
     private void updateNotificationDots(Predicate<PackageUserKey> updatedDots) {
         mWorkspace.updateNotificationDots(updatedDots);
         mAppsView.getAppsStore().updateNotificationDots(updatedDots);
@@ -2721,17 +2675,10 @@ public class Launcher extends BaseDraggingActivity implements LauncherExterns,
         return super.onKeyUp(keyCode, event);
     }
 
-    protected StateHandler[] createStateHandlers() {
+    protected StateHandler<LauncherState>[] createStateHandlers() {
         return new StateHandler[] { getAllAppsController(), getWorkspace() };
     }
 
-    /**
-     * Creates a factory for atomic state animations
-     */
-    public AtomicAnimationFactory createAtomicAnimationFactory() {
-        return new AtomicAnimationFactory(0);
-    }
-
     public TouchController[] createTouchControllers() {
         return new TouchController[] {getDragController(), new AllAppsSwipeController(this)};
     }
diff --git a/src/com/android/launcher3/LauncherState.java b/src/com/android/launcher3/LauncherState.java
index e133d318f1..db2a6cd967 100644
--- a/src/com/android/launcher3/LauncherState.java
+++ b/src/com/android/launcher3/LauncherState.java
@@ -29,6 +29,8 @@ import static com.android.launcher3.testing.TestProtocol.SPRING_LOADED_STATE_ORD
 import android.content.Context;
 import android.view.animation.Interpolator;
 
+import com.android.launcher3.statemanager.BaseState;
+import com.android.launcher3.statemanager.StateManager;
 import com.android.launcher3.states.HintState;
 import com.android.launcher3.states.SpringLoadedState;
 import com.android.launcher3.uioverrides.states.AllAppsState;
@@ -40,7 +42,7 @@ import java.util.Arrays;
 /**
  * Base state for various states used for the Launcher
  */
-public abstract class LauncherState {
+public abstract class LauncherState implements BaseState<LauncherState> {
 
     /**
      * Set of elements indicating various workspace elements which change visibility across states
@@ -60,25 +62,22 @@ public abstract class LauncherState {
             HOTSEAT_SEARCH_BOX | ALL_APPS_HEADER | ALL_APPS_HEADER_EXTRA | ALL_APPS_CONTENT;
 
     // Flag indicating workspace has multiple pages visible.
-    public static final int FLAG_MULTI_PAGE = 1 << 0;
+    public static final int FLAG_MULTI_PAGE = BaseState.getFlag(0);
     // Flag indicating that workspace and its contents are not accessible
-    public static final int FLAG_WORKSPACE_INACCESSIBLE = 1 << 1;
+    public static final int FLAG_WORKSPACE_INACCESSIBLE = BaseState.getFlag(1);
 
-    public static final int FLAG_DISABLE_RESTORE = 1 << 2;
     // Flag indicating the state allows workspace icons to be dragged.
-    public static final int FLAG_WORKSPACE_ICONS_CAN_BE_DRAGGED = 1 << 3;
+    public static final int FLAG_WORKSPACE_ICONS_CAN_BE_DRAGGED = BaseState.getFlag(2);
     // Flag to indicate that workspace should draw page background
-    public static final int FLAG_WORKSPACE_HAS_BACKGROUNDS = 1 << 4;
-    // Flag to indicate that Launcher is non-interactive in this state
-    public static final int FLAG_NON_INTERACTIVE = 1 << 5;
+    public static final int FLAG_WORKSPACE_HAS_BACKGROUNDS = BaseState.getFlag(3);
     // True if the back button should be hidden when in this state (assuming no floating views are
     // open, launcher has window focus, etc).
-    public static final int FLAG_HIDE_BACK_BUTTON = 1 << 6;
+    public static final int FLAG_HIDE_BACK_BUTTON = BaseState.getFlag(4);
     // Flag to indicate if the state would have scrim over sysui region: statu sbar and nav bar
-    public static final int FLAG_HAS_SYS_UI_SCRIM = 1 << 7;
+    public static final int FLAG_HAS_SYS_UI_SCRIM = BaseState.getFlag(5);
     // Flag to inticate that all popups should be closed when this state is enabled.
-    public static final int FLAG_CLOSE_POPUPS = 1 << 8;
-    public static final int FLAG_OVERVIEW_UI = 1 << 9;
+    public static final int FLAG_CLOSE_POPUPS = BaseState.getFlag(6);
+    public static final int FLAG_OVERVIEW_UI = BaseState.getFlag(7);
 
 
     public static final float NO_OFFSET = 0;
@@ -151,26 +150,15 @@ public abstract class LauncherState {
     /**
      * Returns if the state has the provided flag
      */
+    @Override
     public final boolean hasFlag(int mask) {
         return (mFlags & mask) != 0;
     }
 
-    /**
-     * @return true if the state can be persisted across activity restarts.
-     */
-    public final boolean shouldDisableRestore() {
-        return hasFlag(FLAG_DISABLE_RESTORE);
-    }
-
     public static LauncherState[] values() {
         return Arrays.copyOf(sAllStates, sAllStates.length);
     }
 
-    /**
-     * @return How long the animation to this state should take (or from this state to NORMAL).
-     */
-    public abstract int getTransitionDuration(Context context);
-
     public ScaleAndTranslation getWorkspaceScaleAndTranslation(Launcher launcher) {
         return new ScaleAndTranslation(NO_SCALE, NO_OFFSET, NO_OFFSET);
     }
@@ -264,14 +252,20 @@ public abstract class LauncherState {
         };
     }
 
+    @Override
     public LauncherState getHistoryForState(LauncherState previousState) {
         // No history is supported
         return NORMAL;
     }
 
+    @Override
+    public String toString() {
+        return "Ordinal-" + ordinal;
+    }
+
     public void onBackPressed(Launcher launcher) {
         if (this != NORMAL) {
-            LauncherStateManager lsm = launcher.getStateManager();
+            StateManager<LauncherState> lsm = launcher.getStateManager();
             LauncherState lastState = lsm.getLastState();
             lsm.goToState(lastState);
         }
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 8dab8187e8..286b522636 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -89,6 +89,7 @@ import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.pageindicators.WorkspacePageIndicator;
 import com.android.launcher3.popup.PopupContainerWithArrow;
+import com.android.launcher3.statemanager.StateManager.StateHandler;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.touch.WorkspaceTouchListener;
 import com.android.launcher3.userevent.nano.LauncherLogProto.Action;
@@ -119,7 +120,7 @@ import java.util.function.Predicate;
  */
 public class Workspace extends PagedView<WorkspacePageIndicator>
         implements DropTarget, DragSource, View.OnTouchListener,
-        DragController.DragListener, Insettable, LauncherStateManager.StateHandler,
+        DragController.DragListener, Insettable, StateHandler<LauncherState>,
         WorkspaceLayoutManager {
 
     /** The value that {@link #mTransitionProgress} must be greater than for
diff --git a/src/com/android/launcher3/WorkspaceStateTransitionAnimation.java b/src/com/android/launcher3/WorkspaceStateTransitionAnimation.java
index 29cf803c3a..06a73dbc28 100644
--- a/src/com/android/launcher3/WorkspaceStateTransitionAnimation.java
+++ b/src/com/android/launcher3/WorkspaceStateTransitionAnimation.java
@@ -66,7 +66,7 @@ public class WorkspaceStateTransitionAnimation {
     }
 
     /**
-     * @see com.android.launcher3.LauncherStateManager.StateHandler#setStateWithAnimation
+     * @see com.android.launcher3.statemanager.StateManager.StateHandler#setStateWithAnimation
      */
     public void setStateWithAnimation(
             LauncherState toState, StateAnimationConfig config, PendingAnimation animation) {
diff --git a/src/com/android/launcher3/allapps/AllAppsTransitionController.java b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
index 21dd141a85..f0570368ac 100644
--- a/src/com/android/launcher3/allapps/AllAppsTransitionController.java
+++ b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
@@ -28,11 +28,11 @@ import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.DeviceProfile.OnDeviceProfileChangeListener;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager.StateHandler;
 import com.android.launcher3.R;
 import com.android.launcher3.anim.AnimationSuccessListener;
 import com.android.launcher3.anim.PendingAnimation;
 import com.android.launcher3.anim.PropertySetter;
+import com.android.launcher3.statemanager.StateManager.StateHandler;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.Themes;
@@ -50,8 +50,8 @@ import com.android.systemui.plugins.PluginListener;
  * If release velocity < THRES1, snap according to either top or bottom depending on whether it's
  * closer to top or closer to the page indicator.
  */
-public class AllAppsTransitionController implements StateHandler, OnDeviceProfileChangeListener,
-        PluginListener<AllAppsSearchPlugin> {
+public class AllAppsTransitionController implements StateHandler<LauncherState>,
+        OnDeviceProfileChangeListener, PluginListener<AllAppsSearchPlugin> {
 
     public static final FloatProperty<AllAppsTransitionController> ALL_APPS_PROGRESS =
             new FloatProperty<AllAppsTransitionController>("allAppsProgress") {
diff --git a/src/com/android/launcher3/allapps/DiscoveryBounce.java b/src/com/android/launcher3/allapps/DiscoveryBounce.java
index 064868208d..539794230a 100644
--- a/src/com/android/launcher3/allapps/DiscoveryBounce.java
+++ b/src/com/android/launcher3/allapps/DiscoveryBounce.java
@@ -31,9 +31,9 @@ import android.view.MotionEvent;
 import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager.StateListener;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.util.OnboardingPrefs;
 
 /**
@@ -46,7 +46,7 @@ public class DiscoveryBounce extends AbstractFloatingView {
     private final Launcher mLauncher;
     private final Animator mDiscoBounceAnimation;
 
-    private final StateListener mStateListener = new StateListener() {
+    private final StateListener<LauncherState> mStateListener = new StateListener<LauncherState>() {
         @Override
         public void onStateTransitionStart(LauncherState toState) {
             handleClose(false);
diff --git a/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java b/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
index f6766c4023..80b6a5ab96 100644
--- a/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/LauncherAllAppsContainerView.java
@@ -22,7 +22,7 @@ import android.view.MotionEvent;
 
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager;
+import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.views.WorkEduView;
 
 /**
@@ -32,7 +32,7 @@ public class LauncherAllAppsContainerView extends AllAppsContainerView {
 
     private final Launcher mLauncher;
 
-    private LauncherStateManager.StateListener mWorkTabListener;
+    private StateListener<LauncherState> mWorkTabListener;
 
     public LauncherAllAppsContainerView(Context context) {
         this(context, null);
diff --git a/src/com/android/launcher3/dragndrop/DragView.java b/src/com/android/launcher3/dragndrop/DragView.java
index 1e23bb6a71..de0fa1a1f7 100644
--- a/src/com/android/launcher3/dragndrop/DragView.java
+++ b/src/com/android/launcher3/dragndrop/DragView.java
@@ -49,18 +49,18 @@ import com.android.launcher3.FirstFrameAnimatorHelper;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.Interpolators;
 import com.android.launcher3.icons.LauncherIcons;
 import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.util.Thunk;
 
 import java.util.Arrays;
 
-public class DragView extends View implements LauncherStateManager.StateListener {
+public class DragView extends View implements StateListener<LauncherState> {
     private static final ColorMatrix sTempMatrix1 = new ColorMatrix();
     private static final ColorMatrix sTempMatrix2 = new ColorMatrix();
 
diff --git a/src/com/android/launcher3/statemanager/BaseState.java b/src/com/android/launcher3/statemanager/BaseState.java
new file mode 100644
index 0000000000..daec1d8508
--- /dev/null
+++ b/src/com/android/launcher3/statemanager/BaseState.java
@@ -0,0 +1,55 @@
+/*
+ * 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.statemanager;
+
+import android.content.Context;
+
+/**
+ * Interface representing a state of a StatefulActivity
+ */
+public interface BaseState<T extends BaseState> {
+
+    // Flag to indicate that Launcher is non-interactive in this state
+    int FLAG_NON_INTERACTIVE = 1 << 0;
+    int FLAG_DISABLE_RESTORE = 1 << 1;
+
+    static int getFlag(int index) {
+        // reserve few spots to base flags
+        return 1 << (index + 2);
+    }
+
+    /**
+     * @return How long the animation to this state should take (or from this state to NORMAL).
+     */
+    int getTransitionDuration(Context context);
+
+    /**
+     * Returns the state to go back to from this state
+     */
+    T getHistoryForState(T previousState);
+
+    /**
+     * @return true if the state can be persisted across activity restarts.
+     */
+    default boolean shouldDisableRestore() {
+        return hasFlag(FLAG_DISABLE_RESTORE);
+    }
+
+    /**
+     * Returns if the state has the provided flag
+     */
+    boolean hasFlag(int flagMask);
+}
diff --git a/src/com/android/launcher3/LauncherStateManager.java b/src/com/android/launcher3/statemanager/StateManager.java
similarity index 77%
rename from src/com/android/launcher3/LauncherStateManager.java
rename to src/com/android/launcher3/statemanager/StateManager.java
index f6de48e4e6..44471660af 100644
--- a/src/com/android/launcher3/LauncherStateManager.java
+++ b/src/com/android/launcher3/statemanager/StateManager.java
@@ -14,9 +14,8 @@
  * limitations under the License.
  */
 
-package com.android.launcher3;
+package com.android.launcher3.statemanager;
 
-import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.states.StateAnimationConfig.ANIM_ALL_COMPONENTS;
 
 import android.animation.Animator;
@@ -27,6 +26,7 @@ import android.os.Handler;
 import android.os.Looper;
 import android.util.Log;
 
+import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimationSuccessListener;
 import com.android.launcher3.anim.AnimatorPlaybackController;
 import com.android.launcher3.anim.PendingAnimation;
@@ -38,83 +38,48 @@ import java.io.PrintWriter;
 import java.util.ArrayList;
 
 /**
- * TODO: figure out what kind of tests we can write for this
- *
- * Things to test when changing the following class.
- *   - Home from workspace
- *          - from center screen
- *          - from other screens
- *   - Home from all apps
- *          - from center screen
- *          - from other screens
- *   - Back from all apps
- *          - from center screen
- *          - from other screens
- *   - Launch app from workspace and quit
- *          - with back
- *          - with home
- *   - Launch app from all apps and quit
- *          - with back
- *          - with home
- *   - Go to a screen that's not the default, then all
- *     apps, and launch and app, and go back
- *          - with back
- *          -with home
- *   - On workspace, long press power and go back
- *          - with back
- *          - with home
- *   - On all apps, long press power and go back
- *          - with back
- *          - with home
- *   - On workspace, power off
- *   - On all apps, power off
- *   - Launch an app and turn off the screen while in that app
- *          - Go back with home key
- *          - Go back with back key  TODO: make this not go to workspace
- *          - From all apps
- *          - From workspace
- *   - Enter and exit car mode (becase it causes an extra configuration changed)
- *          - From all apps
- *          - From the center workspace
- *          - From another workspace
+ * Class to manage transitions between different states for a StatefulActivity based on different
+ * states
  */
-public class LauncherStateManager {
+public class StateManager<STATE_TYPE extends BaseState<STATE_TYPE>> {
 
     public static final String TAG = "StateManager";
 
     private final AnimationState mConfig = new AnimationState();
     private final Handler mUiHandler;
-    private final Launcher mLauncher;
-    private final ArrayList<StateListener> mListeners = new ArrayList<>();
+    private final StatefulActivity<STATE_TYPE> mActivity;
+    private final ArrayList<StateListener<STATE_TYPE>> mListeners = new ArrayList<>();
+    private final STATE_TYPE mBaseState;
 
     // Animators which are run on properties also controlled by state animations.
     private final AtomicAnimationFactory mAtomicAnimationFactory;
 
-    private StateHandler[] mStateHandlers;
-    private LauncherState mState = NORMAL;
+    private StateHandler<STATE_TYPE>[] mStateHandlers;
+    private STATE_TYPE mState;
 
-    private LauncherState mLastStableState = NORMAL;
-    private LauncherState mCurrentStableState = NORMAL;
+    private STATE_TYPE mLastStableState;
+    private STATE_TYPE mCurrentStableState;
 
-    private LauncherState mRestState;
+    private STATE_TYPE mRestState;
 
-    public LauncherStateManager(Launcher l) {
+    public StateManager(StatefulActivity<STATE_TYPE> l, STATE_TYPE baseState) {
         mUiHandler = new Handler(Looper.getMainLooper());
-        mLauncher = l;
-
+        mActivity = l;
+        mBaseState = baseState;
+        mState = mLastStableState = mCurrentStableState = baseState;
         mAtomicAnimationFactory = l.createAtomicAnimationFactory();
     }
 
-    public LauncherState getState() {
+    public STATE_TYPE getState() {
         return mState;
     }
 
-    public LauncherState getCurrentStableState() {
+    public STATE_TYPE getCurrentStableState() {
         return mCurrentStableState;
     }
 
     public void dump(String prefix, PrintWriter writer) {
-        writer.println(prefix + "LauncherState:");
+        writer.println(prefix + "StateManager:");
         writer.println(prefix + "\tmLastStableState:" + mLastStableState);
         writer.println(prefix + "\tmCurrentStableState:" + mCurrentStableState);
         writer.println(prefix + "\tmState:" + mState);
@@ -124,7 +89,7 @@ public class LauncherStateManager {
 
     public StateHandler[] getStateHandlers() {
         if (mStateHandlers == null) {
-            mStateHandlers = mLauncher.createStateHandlers();
+            mStateHandlers = mActivity.createStateHandlers();
         }
         return mStateHandlers;
     }
@@ -141,29 +106,29 @@ public class LauncherStateManager {
      * Returns true if the state changes should be animated.
      */
     public boolean shouldAnimateStateChange() {
-        return !mLauncher.isForceInvisible() && mLauncher.isStarted();
+        return !mActivity.isForceInvisible() && mActivity.isStarted();
     }
 
     /**
      * @return {@code true} if the state matches the current state and there is no active
      *         transition to different state.
      */
-    public boolean isInStableState(LauncherState state) {
+    public boolean isInStableState(STATE_TYPE state) {
         return mState == state && mCurrentStableState == state
                 && (mConfig.targetState == null || mConfig.targetState == state);
     }
 
     /**
-     * @see #goToState(LauncherState, boolean, Runnable)
+     * @see #goToState(STATE_TYPE, boolean, Runnable)
      */
-    public void goToState(LauncherState state) {
+    public void goToState(STATE_TYPE state) {
         goToState(state, shouldAnimateStateChange());
     }
 
     /**
-     * @see #goToState(LauncherState, boolean, Runnable)
+     * @see #goToState(STATE_TYPE, boolean, Runnable)
      */
-    public void goToState(LauncherState state, boolean animated) {
+    public void goToState(STATE_TYPE state, boolean animated) {
         goToState(state, animated, 0, null);
     }
 
@@ -174,21 +139,21 @@ public class LauncherStateManager {
      *                true otherwise
      * @paras onCompleteRunnable any action to perform at the end of the transition, of null.
      */
-    public void goToState(LauncherState state, boolean animated, Runnable onCompleteRunnable) {
+    public void goToState(STATE_TYPE state, boolean animated, Runnable onCompleteRunnable) {
         goToState(state, animated, 0, onCompleteRunnable);
     }
 
     /**
      * Changes the Launcher state to the provided state after the given delay.
      */
-    public void goToState(LauncherState state, long delay, Runnable onCompleteRunnable) {
+    public void goToState(STATE_TYPE state, long delay, Runnable onCompleteRunnable) {
         goToState(state, true, delay, onCompleteRunnable);
     }
 
     /**
      * Changes the Launcher state to the provided state after the given delay.
      */
-    public void goToState(LauncherState state, long delay) {
+    public void goToState(STATE_TYPE state, long delay) {
         goToState(state, true, delay, null);
     }
 
@@ -212,10 +177,10 @@ public class LauncherStateManager {
         }
     }
 
-    private void goToState(LauncherState state, boolean animated, long delay,
+    private void goToState(STATE_TYPE state, boolean animated, long delay,
             final Runnable onCompleteRunnable) {
-        animated &= Utilities.areAnimationsEnabled(mLauncher);
-        if (mLauncher.isInState(state)) {
+        animated &= Utilities.areAnimationsEnabled(mActivity);
+        if (mActivity.isInState(state)) {
             if (mConfig.currentAnimation == null) {
                 // Run any queued runnable
                 if (onCompleteRunnable != null) {
@@ -233,7 +198,7 @@ public class LauncherStateManager {
         }
 
         // Cancel the current animation. This will reset mState to mCurrentStableState, so store it.
-        LauncherState fromState = mState;
+        STATE_TYPE fromState = mState;
         mConfig.reset();
 
         if (!animated) {
@@ -266,13 +231,13 @@ public class LauncherStateManager {
         }
     }
 
-    private void goToStateAnimated(LauncherState state, LauncherState fromState,
+    private void goToStateAnimated(STATE_TYPE state, STATE_TYPE fromState,
             Runnable onCompleteRunnable) {
-        // Since state NORMAL can be reached from multiple states, just assume that the
+        // Since state mBaseState can be reached from multiple states, just assume that the
         // transition plays in reverse and use the same duration as previous state.
-        mConfig.duration = state == NORMAL
-                ? fromState.getTransitionDuration(mLauncher)
-                : state.getTransitionDuration(mLauncher);
+        mConfig.duration = state == mBaseState
+                ? fromState.getTransitionDuration(mActivity)
+                : state.getTransitionDuration(mActivity);
         prepareForAtomicAnimation(fromState, state, mConfig);
         AnimatorSet animation = createAnimationToNewWorkspaceInternal(state).getAnim();
         if (onCompleteRunnable != null) {
@@ -286,7 +251,7 @@ public class LauncherStateManager {
      * - Setting interpolators for various animations included in the state transition.
      * - Setting some start values (e.g. scale) for views that are hidden but about to be shown.
      */
-    public void prepareForAtomicAnimation(LauncherState fromState, LauncherState toState,
+    public void prepareForAtomicAnimation(STATE_TYPE fromState, STATE_TYPE toState,
             StateAnimationConfig config) {
         mAtomicAnimationFactory.prepareForAtomicAnimation(fromState, toState, config);
     }
@@ -295,11 +260,11 @@ public class LauncherStateManager {
      * Creates an animation representing atomic transitions between the provided states
      */
     public AnimatorSet createAtomicAnimation(
-            LauncherState fromState, LauncherState toState, StateAnimationConfig config) {
+            STATE_TYPE fromState, STATE_TYPE toState, StateAnimationConfig config) {
         PendingAnimation builder = new PendingAnimation(config.duration);
         prepareForAtomicAnimation(fromState, toState, config);
 
-        for (StateHandler handler : mLauncher.getStateManager().getStateHandlers()) {
+        for (StateHandler handler : mActivity.getStateManager().getStateHandlers()) {
             handler.setStateWithAnimation(toState, config, builder);
         }
         return builder.getAnim();
@@ -309,23 +274,23 @@ public class LauncherStateManager {
      * Creates a {@link AnimatorPlaybackController} that can be used for a controlled
      * state transition.
      * @param state the final state for the transition.
-     * @param duration intended duration for normal playback. Use higher duration for better
+     * @param duration intended duration for state playback. Use higher duration for better
      *                accuracy.
      */
     public AnimatorPlaybackController createAnimationToNewWorkspace(
-            LauncherState state, long duration) {
+            STATE_TYPE state, long duration) {
         return createAnimationToNewWorkspace(state, duration, ANIM_ALL_COMPONENTS);
     }
 
     public AnimatorPlaybackController createAnimationToNewWorkspace(
-            LauncherState state, long duration, @AnimationFlags int animComponents) {
+            STATE_TYPE state, long duration, @AnimationFlags int animComponents) {
         StateAnimationConfig config = new StateAnimationConfig();
         config.duration = duration;
         config.animFlags = animComponents;
         return createAnimationToNewWorkspace(state, config);
     }
 
-    public AnimatorPlaybackController createAnimationToNewWorkspace(LauncherState state,
+    public AnimatorPlaybackController createAnimationToNewWorkspace(STATE_TYPE state,
             StateAnimationConfig config) {
         config.userControlled = true;
         mConfig.reset();
@@ -335,10 +300,10 @@ public class LauncherStateManager {
         return mConfig.playbackController;
     }
 
-    private PendingAnimation createAnimationToNewWorkspaceInternal(final LauncherState state) {
+    private PendingAnimation createAnimationToNewWorkspaceInternal(final STATE_TYPE state) {
         if (TestProtocol.sDebugTracing) {
             Log.d(TestProtocol.OVERIEW_NOT_ALLAPPS, "createAnimationToNewWorkspaceInternal: "
-                    + state.ordinal);
+                    + state);
         }
         PendingAnimation builder = new PendingAnimation(mConfig.duration);
         for (StateHandler handler : getStateHandlers()) {
@@ -355,7 +320,7 @@ public class LauncherStateManager {
             @Override
             public void onAnimationSuccess(Animator animator) {
                 if (TestProtocol.sDebugTracing) {
-                    Log.d(TestProtocol.OVERIEW_NOT_ALLAPPS, "onAnimationSuccess: " + state.ordinal);
+                    Log.d(TestProtocol.OVERIEW_NOT_ALLAPPS, "onAnimationSuccess: " + state);
                 }
                 onStateTransitionEnd(state);
             }
@@ -364,24 +329,24 @@ public class LauncherStateManager {
         return builder;
     }
 
-    private void onStateTransitionStart(LauncherState state) {
+    private void onStateTransitionStart(STATE_TYPE state) {
         mState = state;
-        mLauncher.onStateSetStart(mState);
+        mActivity.onStateSetStart(mState);
 
         for (int i = mListeners.size() - 1; i >= 0; i--) {
             mListeners.get(i).onStateTransitionStart(state);
         }
     }
 
-    private void onStateTransitionEnd(LauncherState state) {
+    private void onStateTransitionEnd(STATE_TYPE state) {
         // Only change the stable states after the transitions have finished
         if (state != mCurrentStableState) {
             mLastStableState = state.getHistoryForState(mCurrentStableState);
             mCurrentStableState = state;
         }
 
-        mLauncher.onStateSetEnd(state);
-        if (state == NORMAL) {
+        mActivity.onStateSetEnd(state);
+        if (state == mBaseState) {
             setRestState(null);
         }
 
@@ -390,7 +355,7 @@ public class LauncherStateManager {
         }
     }
 
-    public LauncherState getLastState() {
+    public STATE_TYPE getLastState() {
         return mLastStableState;
     }
 
@@ -402,15 +367,15 @@ public class LauncherStateManager {
         if (mState.shouldDisableRestore()) {
             goToState(getRestState());
             // Reset history
-            mLastStableState = NORMAL;
+            mLastStableState = mBaseState;
         }
     }
 
-    public LauncherState getRestState() {
-        return mRestState == null ? NORMAL : mRestState;
+    public STATE_TYPE getRestState() {
+        return mRestState == null ? mBaseState : mRestState;
     }
 
-    public void setRestState(LauncherState restState) {
+    public void setRestState(STATE_TYPE restState) {
         mRestState = restState;
     }
 
@@ -505,13 +470,14 @@ public class LauncherStateManager {
         }
     }
 
-    private static class AnimationState extends StateAnimationConfig implements AnimatorListener {
+    private static class AnimationState<STATE_TYPE> extends StateAnimationConfig
+            implements AnimatorListener {
 
         private static final StateAnimationConfig DEFAULT = new StateAnimationConfig();
 
         public AnimatorPlaybackController playbackController;
         public AnimatorSet currentAnimation;
-        public LauncherState targetState;
+        public STATE_TYPE targetState;
 
         // Id to keep track of config changes, to tie an animation with the corresponding request
         public int changeId = 0;
@@ -546,7 +512,7 @@ public class LauncherStateManager {
             }
         }
 
-        public void setAnimation(AnimatorSet animation, LauncherState targetState) {
+        public void setAnimation(AnimatorSet animation, STATE_TYPE targetState) {
             currentAnimation = animation;
             this.targetState = targetState;
             currentAnimation.addListener(this);
@@ -562,31 +528,31 @@ public class LauncherStateManager {
         public void onAnimationRepeat(Animator animator) { }
     }
 
-    public interface StateHandler {
+    public interface StateHandler<STATE_TYPE> {
 
         /**
          * Updates the UI to {@param state} without any animations
          */
-        void setState(LauncherState state);
+        void setState(STATE_TYPE state);
 
         /**
          * Sets the UI to {@param state} by animating any changes.
          */
         void setStateWithAnimation(
-                LauncherState toState, StateAnimationConfig config, PendingAnimation animation);
+                STATE_TYPE toState, StateAnimationConfig config, PendingAnimation animation);
     }
 
-    public interface StateListener {
+    public interface StateListener<STATE_TYPE> {
 
-        default void onStateTransitionStart(LauncherState toState) { }
+        default void onStateTransitionStart(STATE_TYPE toState) { }
 
-        default void onStateTransitionComplete(LauncherState finalState) { }
+        default void onStateTransitionComplete(STATE_TYPE finalState) { }
     }
 
     /**
      * Factory class to configure and create atomic animations.
      */
-    public static class AtomicAnimationFactory {
+    public static class AtomicAnimationFactory<STATE_TYPE> {
 
         private final Animator[] mStateElementAnimators;
 
@@ -622,6 +588,6 @@ public class LauncherStateManager {
          * - Setting some start values (e.g. scale) for views that are hidden but about to be shown.
          */
         public void prepareForAtomicAnimation(
-                LauncherState fromState, LauncherState toState, StateAnimationConfig config) { }
+                STATE_TYPE fromState, STATE_TYPE toState, StateAnimationConfig config) { }
     }
 }
diff --git a/src/com/android/launcher3/statemanager/StatefulActivity.java b/src/com/android/launcher3/statemanager/StatefulActivity.java
new file mode 100644
index 0000000000..0a1607cf04
--- /dev/null
+++ b/src/com/android/launcher3/statemanager/StatefulActivity.java
@@ -0,0 +1,149 @@
+/*
+ * 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.statemanager;
+
+import static com.android.launcher3.LauncherState.FLAG_NON_INTERACTIVE;
+
+import android.os.Handler;
+
+import androidx.annotation.CallSuper;
+
+import com.android.launcher3.BaseDraggingActivity;
+import com.android.launcher3.Utilities;
+import com.android.launcher3.statemanager.StateManager.AtomicAnimationFactory;
+import com.android.launcher3.statemanager.StateManager.StateHandler;
+import com.android.launcher3.views.BaseDragLayer;
+
+/**
+ * Abstract activity with state management
+ * @param <STATE_TYPE> Type of state object
+ */
+public abstract class StatefulActivity<STATE_TYPE extends BaseState<STATE_TYPE>>
+        extends BaseDraggingActivity {
+
+    public final Handler mHandler = new Handler();
+    private final Runnable mHandleDeferredResume = this::handleDeferredResume;
+    private boolean mDeferredResumePending;
+
+    /**
+     * Create handlers to control the property changes for this activity
+     */
+    protected abstract StateHandler<STATE_TYPE>[] createStateHandlers();
+
+    /**
+     * Returns true if the activity is in the provided state
+     */
+    public boolean isInState(STATE_TYPE state) {
+        return getStateManager().getState() == state;
+    }
+
+    /**
+     * Returns the state manager for this activity
+     */
+    public abstract StateManager<STATE_TYPE> getStateManager();
+
+    /**
+     * Called when transition to the state starts
+     */
+    @CallSuper
+    public void onStateSetStart(STATE_TYPE state) {
+        if (mDeferredResumePending) {
+            handleDeferredResume();
+        }
+    }
+
+    /**
+     * Called when transition to state ends
+     */
+    public void onStateSetEnd(STATE_TYPE state) { }
+
+    /**
+     * Creates a factory for atomic state animations
+     */
+    public AtomicAnimationFactory<STATE_TYPE> createAtomicAnimationFactory() {
+        return new AtomicAnimationFactory(0);
+    }
+
+    @Override
+    public void reapplyUi() {
+        reapplyUi(true /* cancelCurrentAnimation */);
+    }
+
+    /**
+     * Re-applies if any state transition is not running, optionally cancelling
+     * the transition if requested.
+     */
+    public void reapplyUi(boolean cancelCurrentAnimation) {
+        getStateManager().reapplyState(cancelCurrentAnimation);
+    }
+
+    @Override
+    protected void onStop() {
+        BaseDragLayer dragLayer = getDragLayer();
+        final boolean wasActive = isUserActive();
+        final STATE_TYPE origState = getStateManager().getState();
+        final int origDragLayerChildCount = dragLayer.getChildCount();
+        super.onStop();
+
+        getStateManager().moveToRestState();
+
+        // Workaround for b/78520668, explicitly trim memory once UI is hidden
+        onTrimMemory(TRIM_MEMORY_UI_HIDDEN);
+
+        if (wasActive) {
+            // The expected condition is that this activity is stopped because the device goes to
+            // sleep and the UI may have noticeable changes.
+            dragLayer.post(() -> {
+                if ((!getStateManager().isInStableState(origState)
+                        // The drag layer may be animating (e.g. dismissing QSB).
+                        || dragLayer.getAlpha() < 1
+                        // Maybe an ArrowPopup is closed.
+                        || dragLayer.getChildCount() != origDragLayerChildCount)) {
+                    onUiChangedWhileSleeping();
+                }
+            });
+        }
+    }
+
+    /**
+     * Called if the Activity UI changed while the activity was not visible
+     */
+    protected void onUiChangedWhileSleeping() { }
+
+    private void handleDeferredResume() {
+        if (hasBeenResumed() && !getStateManager().getState().hasFlag(FLAG_NON_INTERACTIVE)) {
+            onDeferredResumed();
+            addActivityFlags(ACTIVITY_STATE_DEFERRED_RESUMED);
+
+            mDeferredResumePending = false;
+        } else {
+            mDeferredResumePending = true;
+        }
+    }
+
+    /**
+     * Called want the activity has stayed resumed for 1 frame.
+     */
+    protected void onDeferredResumed() { }
+
+    @Override
+    protected void onResume() {
+        super.onResume();
+
+        mHandler.removeCallbacks(mHandleDeferredResume);
+        Utilities.postAsyncCallback(mHandler, mHandleDeferredResume);
+    }
+}
diff --git a/src/com/android/launcher3/util/OnboardingPrefs.java b/src/com/android/launcher3/util/OnboardingPrefs.java
index baa1eeefcd..162028991d 100644
--- a/src/com/android/launcher3/util/OnboardingPrefs.java
+++ b/src/com/android/launcher3/util/OnboardingPrefs.java
@@ -21,7 +21,6 @@ import android.util.ArrayMap;
 import androidx.annotation.StringDef;
 
 import com.android.launcher3.Launcher;
-import com.android.launcher3.LauncherStateManager;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -71,13 +70,10 @@ public class OnboardingPrefs<T extends Launcher> {
 
     protected final T mLauncher;
     protected final SharedPreferences mSharedPrefs;
-    protected final LauncherStateManager mStateManager;
 
-    public OnboardingPrefs(T launcher, SharedPreferences sharedPrefs,
-            LauncherStateManager stateManager) {
+    public OnboardingPrefs(T launcher, SharedPreferences sharedPrefs) {
         mLauncher = launcher;
         mSharedPrefs = sharedPrefs;
-        mStateManager = stateManager;
     }
 
     /** @return The number of times we have seen the given event. */
diff --git a/src/com/android/launcher3/views/ScrimView.java b/src/com/android/launcher3/views/ScrimView.java
index da874cf3e1..a2c7d14d52 100644
--- a/src/com/android/launcher3/views/ScrimView.java
+++ b/src/com/android/launcher3/views/ScrimView.java
@@ -62,10 +62,10 @@ import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Insettable;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager;
-import com.android.launcher3.LauncherStateManager.StateListener;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.statemanager.StateManager;
+import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.uioverrides.WallpaperColorInfo;
 import com.android.launcher3.uioverrides.WallpaperColorInfo.OnChangeListener;
 import com.android.launcher3.userevent.nano.LauncherLogProto.Action;
@@ -116,7 +116,8 @@ public class ScrimView<T extends Launcher> extends View implements Insettable, O
     private final AccessibilityManager mAM;
     protected final int mEndScrim;
 
-    private final StateListener mAccessibilityLauncherStateListener = new StateListener() {
+    private final StateListener<LauncherState> mAccessibilityLauncherStateListener =
+            new StateListener<LauncherState>() {
         @Override
         public void onStateTransitionComplete(LauncherState finalState) {
             setImportantForAccessibility(finalState == ALL_APPS
@@ -383,7 +384,7 @@ public class ScrimView<T extends Launcher> extends View implements Insettable, O
 
     @Override
     public void onAccessibilityStateChanged(boolean enabled) {
-        LauncherStateManager stateManager = mLauncher.getStateManager();
+        StateManager<LauncherState> stateManager = mLauncher.getStateManager();
         stateManager.removeStateListener(mAccessibilityLauncherStateListener);
 
         if (enabled) {
diff --git a/src/com/android/launcher3/views/WorkEduView.java b/src/com/android/launcher3/views/WorkEduView.java
index 859b9d0e6c..d35a38fba8 100644
--- a/src/com/android/launcher3/views/WorkEduView.java
+++ b/src/com/android/launcher3/views/WorkEduView.java
@@ -32,19 +32,19 @@ import androidx.annotation.Nullable;
 import com.android.launcher3.Insettable;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager;
-import com.android.launcher3.LauncherStateManager.StateListener;
 import com.android.launcher3.R;
 import com.android.launcher3.allapps.AllAppsContainerView;
 import com.android.launcher3.allapps.AllAppsPagedView;
 import com.android.launcher3.anim.AnimationSuccessListener;
 import com.android.launcher3.anim.Interpolators;
+import com.android.launcher3.statemanager.StateManager.StateListener;
 import com.android.launcher3.userevent.nano.LauncherLogProto;
 
 /**
  * On boarding flow for users right after setting up work profile
  */
-public class WorkEduView extends AbstractSlideInView implements Insettable, StateListener {
+public class WorkEduView extends AbstractSlideInView
+        implements Insettable, StateListener<LauncherState> {
 
     private static final int DEFAULT_CLOSE_DURATION = 200;
     public static final String KEY_WORK_EDU_STEP = "showed_work_profile_edu";
@@ -185,8 +185,8 @@ public class WorkEduView extends AbstractSlideInView implements Insettable, Stat
     /**
      * Checks if user has not seen onboarding UI yet and shows it when user navigates to all apps
      */
-    public static LauncherStateManager.StateListener showEduFlowIfNeeded(Launcher launcher,
-            @Nullable LauncherStateManager.StateListener oldListener) {
+    public static StateListener<LauncherState> showEduFlowIfNeeded(Launcher launcher,
+            @Nullable StateListener<LauncherState> oldListener) {
         if (oldListener != null) {
             launcher.getStateManager().removeStateListener(oldListener);
         }
@@ -195,7 +195,7 @@ public class WorkEduView extends AbstractSlideInView implements Insettable, Stat
             return null;
         }
 
-        LauncherStateManager.StateListener listener = new LauncherStateManager.StateListener() {
+        StateListener<LauncherState> listener = new StateListener<LauncherState>() {
             @Override
             public void onStateTransitionComplete(LauncherState finalState) {
                 if (finalState != LauncherState.ALL_APPS) return;
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 192e69fed0..e7da0621c8 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -49,11 +49,11 @@ import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.LauncherStateManager;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.common.WidgetUtils;
 import com.android.launcher3.model.AppLaunchTracker;
 import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.statemanager.StateManager;
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.tapl.LauncherInstrumentation.ContainerType;
 import com.android.launcher3.tapl.TestHelpers;
@@ -526,7 +526,7 @@ public abstract class AbstractLauncherUiTest {
     private static void checkLauncherIntegrity(
             Launcher launcher, ContainerType expectedContainerType) {
         if (launcher != null) {
-            final LauncherStateManager stateManager = launcher.getStateManager();
+            final StateManager<LauncherState> stateManager = launcher.getStateManager();
             final LauncherState stableState = stateManager.getCurrentStableState();
 
             assertTrue("Stable state != state: " + stableState.getClass().getSimpleName() + ", "