Making ActivityContext extend SavedStateRegistryOwner

Bug: 390572144
Test: Presubmit
Flag: EXEMPT refactor

Change-Id: I8272e95a8d2da95b3c93ec616fdf877b89db5b26
diff --git a/quickstep/src/com/android/launcher3/taskbar/BaseTaskbarContext.java b/quickstep/src/com/android/launcher3/taskbar/BaseTaskbarContext.java
index b33fd38..09a8670 100644
--- a/quickstep/src/com/android/launcher3/taskbar/BaseTaskbarContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/BaseTaskbarContext.java
@@ -19,25 +19,19 @@
 import android.content.Intent;
 import android.content.pm.ShortcutInfo;
 import android.os.UserHandle;
-import android.view.ContextThemeWrapper;
 import android.view.LayoutInflater;
 
-import com.android.launcher3.DeviceProfile.OnDeviceProfileChangeListener;
 import com.android.launcher3.popup.SystemShortcut;
+import com.android.launcher3.util.BaseContext;
 import com.android.launcher3.util.Themes;
-import com.android.launcher3.views.ActivityContext;
 import com.android.quickstep.SystemUiProxy;
 
-import java.util.ArrayList;
-import java.util.List;
-
 // TODO(b/218912746): Share more behavior to avoid all apps context depending directly on taskbar.
 /** Base for common behavior between taskbar window contexts. */
-public abstract class BaseTaskbarContext extends ContextThemeWrapper implements ActivityContext,
-        SystemShortcut.BubbleActivityStarter {
+public abstract class BaseTaskbarContext extends BaseContext
+        implements SystemShortcut.BubbleActivityStarter {
 
     protected final LayoutInflater mLayoutInflater;
-    private final List<OnDeviceProfileChangeListener> mDPChangeListeners = new ArrayList<>();
 
     public BaseTaskbarContext(Context windowContext) {
         super(windowContext, Themes.getActivityThemeRes(windowContext));
@@ -50,11 +44,6 @@
     }
 
     @Override
-    public final List<OnDeviceProfileChangeListener> getOnDeviceProfileChangeListeners() {
-        return mDPChangeListeners;
-    }
-
-    @Override
     public void showShortcutBubble(ShortcutInfo info) {
         if (info == null) return;
         SystemUiProxy.INSTANCE.get(this).showShortcutBubble(info);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 4ac05ff..fc8ea87 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -71,7 +71,6 @@
 import android.os.Trace;
 import android.provider.Settings;
 import android.util.Log;
-import android.view.Display;
 import android.view.Gravity;
 import android.view.Surface;
 import android.view.View;
@@ -159,7 +158,6 @@
 import com.android.launcher3.util.SplitConfigurationOptions.SplitSelectSource;
 import com.android.launcher3.util.TraceHelper;
 import com.android.launcher3.util.VibratorWrapper;
-import com.android.launcher3.util.ViewCache;
 import com.android.launcher3.views.ActivityContext;
 import com.android.quickstep.NavHandle;
 import com.android.quickstep.RecentsModel;
@@ -225,7 +223,6 @@
 
     private NavigationMode mNavMode;
     private boolean mImeDrawsImeNavBar;
-    private final ViewCache mViewCache = new ViewCache();
 
     private final boolean mIsSafeModeEnabled;
     private final boolean mIsUserSetupComplete;
@@ -285,7 +282,6 @@
         mIsNavBarForceVisible = mIsNavBarKidsMode;
 
         // Get display and corners first, as views might use them in constructor.
-        Display display = windowContext.getDisplay();
         Context c = getApplicationContext();
         mWindowManager = c.getSystemService(WindowManager.class);
 
@@ -387,6 +383,7 @@
                         DesktopVisibilityController.INSTANCE.get(this)));
 
         mLauncherPrefs = LauncherPrefs.get(this);
+        onViewCreated();
     }
 
     /** Updates {@link DeviceProfile} instances for any Taskbar windows. */
@@ -801,11 +798,6 @@
     }
 
     @Override
-    public ViewCache getViewCache() {
-        return mViewCache;
-    }
-
-    @Override
     public View.OnClickListener getItemOnClickListener() {
         return this::onTaskbarIconClicked;
     }
@@ -994,6 +986,7 @@
      * Called when this instance of taskbar is no longer needed
      */
     public void onDestroy() {
+        onViewDestroyed();
         mIsDestroyed = true;
         mTaskbarFeatureEvaluator.onDestroy();
         setUIController(TaskbarUIController.DEFAULT);
diff --git a/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayContext.java b/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayContext.java
index 636d89b..55bb0f9 100644
--- a/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayContext.java
@@ -64,6 +64,7 @@
         mStashedTaskbarHeight = controllers.taskbarStashController.getStashedHeight();
 
         mUiController = controllers.uiController;
+        onViewCreated();
     }
 
     public @Nullable TaskbarSearchSessionController getSearchSessionController() {
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowContext.kt b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowContext.kt
index 9a408ad..d70d7eb 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowContext.kt
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowContext.kt
@@ -18,7 +18,6 @@
 
 import android.content.Context
 import android.graphics.PixelFormat
-import android.view.ContextThemeWrapper
 import android.view.Display
 import android.view.ViewGroup
 import android.view.WindowManager
@@ -26,10 +25,8 @@
 import android.view.WindowManager.LayoutParams.PRIVATE_FLAG_CONSUME_IME_INSETS
 import com.android.launcher3.DeviceProfile
 import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.util.BaseContext
 import com.android.launcher3.util.Themes
-import com.android.launcher3.views.ActivityContext
-import com.android.launcher3.views.BaseDragLayer
-import com.android.quickstep.fallback.RecentsDragLayer
 
 /**
  * Window context for the Overview overlays.
@@ -37,18 +34,14 @@
  * <p>
  * Overlays have their own window and need a window context.
  */
-open class RecentsWindowContext(windowContext: Context, wallpaperColorHints: Int) :
-    ContextThemeWrapper(
-        windowContext,
-        Themes.getActivityThemeRes(windowContext, wallpaperColorHints),
-    ),
-    ActivityContext {
+abstract class RecentsWindowContext(windowContext: Context, wallpaperColorHints: Int) :
+    BaseContext(
+        base = windowContext,
+        themeResId = Themes.getActivityThemeRes(windowContext, wallpaperColorHints),
+        destroyOnDetach = false,
+    ) {
 
     private var deviceProfile: DeviceProfile? = null
-    private var dragLayer: RecentsDragLayer<RecentsWindowManager> = RecentsDragLayer(this, null)
-    private val deviceProfileChangeListeners:
-        MutableList<DeviceProfile.OnDeviceProfileChangeListener> =
-        ArrayList()
 
     private val windowTitle: String = "RecentsWindow"
 
@@ -58,10 +51,6 @@
             windowTitle,
         )
 
-    override fun getDragLayer(): BaseDragLayer<RecentsWindowManager> {
-        return dragLayer
-    }
-
     fun initDeviceProfile() {
         deviceProfile =
             if (displayId == Display.DEFAULT_DISPLAY)
@@ -76,11 +65,6 @@
         return deviceProfile!!
     }
 
-    override fun getOnDeviceProfileChangeListeners():
-        List<DeviceProfile.OnDeviceProfileChangeListener> {
-        return deviceProfileChangeListeners
-    }
-
     /**
      * Creates LayoutParams for adding a view directly to WindowManager as a new window.
      *
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
index 6a13927..1a3a2e3 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
@@ -50,6 +50,7 @@
 import com.android.launcher3.testing.shared.TestProtocol.SEQUENCE_MAIN
 import com.android.launcher3.util.ContextTracker
 import com.android.launcher3.util.DisplayController
+import com.android.launcher3.util.Executors
 import com.android.launcher3.util.RunnableList
 import com.android.launcher3.util.SystemUiController
 import com.android.launcher3.views.BaseDragLayer
@@ -213,6 +214,7 @@
 
     override fun destroy() {
         super.destroy()
+        Executors.MAIN_EXECUTOR.execute { onViewDestroyed() }
         cleanupRecentsWindow()
         TaskStackChangeListeners.getInstance().unregisterTaskStackListener(taskStackChangeListener)
         callbacks?.removeListener(recentsAnimationListener)
@@ -264,6 +266,7 @@
 
         this.callbacks = callbacks
         callbacks?.addListener(recentsAnimationListener)
+        onViewCreated()
     }
 
     override fun startHome() {
diff --git a/src/com/android/launcher3/BaseActivity.java b/src/com/android/launcher3/BaseActivity.java
index 84c8040..3e6b4dd 100644
--- a/src/com/android/launcher3/BaseActivity.java
+++ b/src/com/android/launcher3/BaseActivity.java
@@ -37,6 +37,10 @@
 import androidx.annotation.IntDef;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.lifecycle.Lifecycle;
+import androidx.lifecycle.LifecycleRegistry;
+import androidx.savedstate.SavedStateRegistry;
+import androidx.savedstate.SavedStateRegistryController;
 
 import com.android.launcher3.DeviceProfile.OnDeviceProfileChangeListener;
 import com.android.launcher3.logging.StatsLogManager;
@@ -47,6 +51,7 @@
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.DisplayController.DisplayInfoChangeListener;
 import com.android.launcher3.util.DisplayController.Info;
+import com.android.launcher3.util.LifecycleHelper;
 import com.android.launcher3.util.RunnableList;
 import com.android.launcher3.util.SystemUiController;
 import com.android.launcher3.util.ViewCache;
@@ -97,11 +102,14 @@
     private final ArrayList<MultiWindowModeChangedListener> mMultiWindowModeChangedListeners =
             new ArrayList<>();
 
+    private final SavedStateRegistryController mSavedStateRegistryController =
+            SavedStateRegistryController.create(this);
+    private final LifecycleRegistry mLifecycleRegistry = new LifecycleRegistry(this);
+
     protected DeviceProfile mDeviceProfile;
     protected SystemUiController mSystemUiController;
     private StatsLogManager mStatsLogManager;
 
-
     public static final int ACTIVITY_STATE_STARTED = 1 << 0;
     public static final int ACTIVITY_STATE_RESUMED = 1 << 1;
 
@@ -178,6 +186,12 @@
 
     private ActionMode mCurrentActionMode;
 
+    public BaseActivity() {
+        mSavedStateRegistryController.performAttach();
+        registerActivityLifecycleCallbacks(
+                new LifecycleHelper(this, mSavedStateRegistryController, mLifecycleRegistry));
+    }
+
     @Override
     public ViewCache getViewCache() {
         return mViewCache;
@@ -478,6 +492,18 @@
 
     protected void reapplyUi() {}
 
+    @NonNull
+    @Override
+    public SavedStateRegistry getSavedStateRegistry() {
+        return mSavedStateRegistryController.getSavedStateRegistry();
+    }
+
+    @NonNull
+    @Override
+    public Lifecycle getLifecycle() {
+        return mLifecycleRegistry;
+    }
+
     public static <T extends BaseActivity> T fromContext(Context context) {
         if (context instanceof BaseActivity) {
             return (T) context;
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 5644051..289f175 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -252,7 +252,6 @@
 import com.android.launcher3.util.TouchController;
 import com.android.launcher3.util.TraceHelper;
 import com.android.launcher3.views.ActivityContext;
-import com.android.launcher3.views.ComposeInitializer;
 import com.android.launcher3.views.FloatingIconView;
 import com.android.launcher3.views.FloatingSurfaceView;
 import com.android.launcher3.views.OptionsPopupView;
@@ -578,7 +577,6 @@
         setDefaultKeyMode(DEFAULT_KEYS_SEARCH_LOCAL);
 
         setContentView(getRootView());
-        ComposeInitializer.initCompose(this);
 
         if (mOnInitialBindListener != null) {
             getRootView().getViewTreeObserver().addOnPreDrawListener(mOnInitialBindListener);
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 3836f7d..b80238c 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -38,7 +38,6 @@
 import android.appwidget.AppWidgetHostView;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
-import android.content.ContextWrapper;
 import android.content.res.Configuration;
 import android.content.res.TypedArray;
 import android.graphics.PointF;
@@ -64,7 +63,6 @@
 import com.android.launcher3.BubbleTextView;
 import com.android.launcher3.CellLayout;
 import com.android.launcher3.DeviceProfile;
-import com.android.launcher3.DeviceProfile.OnDeviceProfileChangeListener;
 import com.android.launcher3.Hotseat;
 import com.android.launcher3.InsettableFrameLayout;
 import com.android.launcher3.InvariantDeviceProfile;
@@ -90,11 +88,13 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.util.BaseContext;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.SandboxContext;
+import com.android.launcher3.util.Themes;
 import com.android.launcher3.util.WindowBounds;
 import com.android.launcher3.util.window.WindowManagerProxy;
 import com.android.launcher3.views.ActivityContext;
@@ -109,7 +109,6 @@
 import dagger.BindsInstance;
 import dagger.Component;
 
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -127,7 +126,7 @@
  *   3) Place appropriate elements like icons and first-page qsb
  *   4) Measure and draw the view on a canvas
  */
-public class LauncherPreviewRenderer extends ContextWrapper
+public class LauncherPreviewRenderer extends BaseContext
         implements ActivityContext, WorkspaceLayoutManager, LayoutInflater.Factory2 {
 
     /**
@@ -156,7 +155,6 @@
         }
     }
 
-    private final List<OnDeviceProfileChangeListener> mDpChangeListeners = new ArrayList<>();
     private final Handler mUiHandler;
     private final Context mContext;
     private final InvariantDeviceProfile mIdp;
@@ -184,7 +182,7 @@
             WallpaperColors wallpaperColorsOverride,
             @Nullable final SparseArray<Size> launcherWidgetSpanInfo) {
 
-        super(context);
+        super(context, Themes.getActivityThemeRes(context));
         mUiHandler = new Handler(Looper.getMainLooper());
         mContext = context;
         mIdp = idp;
@@ -263,6 +261,13 @@
                     : null;
         }
         mAppWidgetHost = new LauncherPreviewAppWidgetHost(context);
+
+        onViewCreated();
+    }
+
+    @Override
+    public InsettableFrameLayout getRootView() {
+        return mRootView;
     }
 
     /**
@@ -349,11 +354,6 @@
     }
 
     @Override
-    public List<OnDeviceProfileChangeListener> getOnDeviceProfileChangeListeners() {
-        return mDpChangeListeners;
-    }
-
-    @Override
     public Hotseat getHotseat() {
         return mHotseat;
     }
diff --git a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
index fd8b0e7..c4fed71 100644
--- a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
+++ b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
@@ -94,27 +94,6 @@
         mModel = LauncherAppState.getInstance(this).getModel();
         mDragController = new SecondaryDragController(this);
         mSecondaryDisplayPredictions = SecondaryDisplayPredictions.newInstance(this);
-        if (getWindow().getDecorView().isAttachedToWindow()) {
-            initUi();
-        }
-    }
-
-    @Override
-    public void onAttachedToWindow() {
-        super.onAttachedToWindow();
-        initUi();
-    }
-
-    @Override
-    public void onDetachedFromWindow() {
-        super.onDetachedFromWindow();
-        this.getDragController().removeDragListener(this);
-    }
-
-    private void initUi() {
-        if (mDragLayer != null) {
-            return;
-        }
 
         mDeviceProfile = InvariantDeviceProfile.INSTANCE.get(this)
                 .createDeviceProfileForSecondaryDisplay(this);
diff --git a/src/com/android/launcher3/util/BaseContext.kt b/src/com/android/launcher3/util/BaseContext.kt
new file mode 100644
index 0000000..819470b
--- /dev/null
+++ b/src/com/android/launcher3/util/BaseContext.kt
@@ -0,0 +1,131 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util
+
+import android.content.Context
+import android.view.ContextThemeWrapper
+import android.view.View
+import android.view.View.OnAttachStateChangeListener
+import android.view.ViewTreeObserver.OnWindowFocusChangeListener
+import android.view.ViewTreeObserver.OnWindowVisibilityChangeListener
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.Lifecycle.State.CREATED
+import androidx.lifecycle.Lifecycle.State.DESTROYED
+import androidx.lifecycle.Lifecycle.State.RESUMED
+import androidx.lifecycle.Lifecycle.State.STARTED
+import androidx.lifecycle.LifecycleRegistry
+import androidx.lifecycle.setViewTreeLifecycleOwner
+import androidx.savedstate.SavedStateRegistry
+import androidx.savedstate.SavedStateRegistryController
+import androidx.savedstate.setViewTreeSavedStateRegistryOwner
+import com.android.launcher3.DeviceProfile.OnDeviceProfileChangeListener
+import com.android.launcher3.Utilities
+import com.android.launcher3.views.ActivityContext
+
+/**
+ * A context wrapper with lifecycle tracking based on the window events on the rootView of the
+ * [ActivityContext]
+ */
+abstract class BaseContext
+@JvmOverloads
+constructor(base: Context, themeResId: Int, private val destroyOnDetach: Boolean = true) :
+    ContextThemeWrapper(base, themeResId), ActivityContext {
+
+    private val listeners = mutableListOf<OnDeviceProfileChangeListener>()
+
+    private val savedStateRegistryController = SavedStateRegistryController.create(this)
+    private val lifecycleRegistry = LifecycleRegistry(this)
+
+    override val savedStateRegistry: SavedStateRegistry
+        get() = savedStateRegistryController.savedStateRegistry
+
+    override val lifecycle: Lifecycle
+        get() = lifecycleRegistry
+
+    private val viewCache = ViewCache()
+
+    init {
+        Executors.MAIN_EXECUTOR.execute {
+            savedStateRegistryController.performAttach()
+            savedStateRegistryController.performRestore(null)
+        }
+    }
+
+    override fun getOnDeviceProfileChangeListeners() = listeners
+
+    private val finishActions = RunnableList()
+
+    /** Called when the root view is created for this context */
+    fun onViewCreated() {
+        val view = rootView
+        val attachListener =
+            object : OnAttachStateChangeListener {
+
+                override fun onViewAttachedToWindow(view: View) {
+                    view.rootView.setViewTreeLifecycleOwner(this@BaseContext)
+                    view.rootView.setViewTreeSavedStateRegistryOwner(this@BaseContext)
+                    lifecycleRegistry.handleLifecycleEvent(Lifecycle.Event.ON_CREATE)
+
+                    val treeObserver = view.viewTreeObserver
+
+                    val focusListener = OnWindowFocusChangeListener { updateState() }
+                    treeObserver.addOnWindowFocusChangeListener(focusListener)
+                    finishActions.add {
+                        treeObserver.removeOnWindowFocusChangeListener(focusListener)
+                    }
+
+                    if (Utilities.ATLEAST_V) {
+                        val visibilityListener = OnWindowVisibilityChangeListener { updateState() }
+                        treeObserver.addOnWindowVisibilityChangeListener(visibilityListener)
+                        finishActions.add {
+                            treeObserver.removeOnWindowVisibilityChangeListener(visibilityListener)
+                        }
+                    }
+                }
+
+                override fun onViewDetachedFromWindow(view: View) {
+                    if (destroyOnDetach) onViewDestroyed()
+                }
+            }
+        view.addOnAttachStateChangeListener(attachListener)
+        finishActions.add { view.removeOnAttachStateChangeListener(attachListener) }
+
+        if (view.isAttachedToWindow) attachListener.onViewAttachedToWindow(view)
+        updateState()
+    }
+
+    override fun getViewCache() = viewCache
+
+    private fun updateState() {
+        if (lifecycleRegistry.currentState.isAtLeast(CREATED)) {
+            lifecycleRegistry.currentState =
+                if (rootView.windowVisibility != View.VISIBLE) CREATED
+                else (if (!rootView.hasWindowFocus()) STARTED else RESUMED)
+        }
+    }
+
+    fun onViewDestroyed() {
+        if (
+            !lifecycleRegistry.currentState.isAtLeast(CREATED) &&
+                lifecycleRegistry.currentState != DESTROYED
+        ) {
+            lifecycleRegistry.currentState = CREATED
+        }
+        lifecycleRegistry.handleLifecycleEvent(Lifecycle.Event.ON_DESTROY)
+        finishActions.executeAllAndDestroy()
+    }
+}
diff --git a/src/com/android/launcher3/util/LifecycleHelper.kt b/src/com/android/launcher3/util/LifecycleHelper.kt
new file mode 100644
index 0000000..803adae
--- /dev/null
+++ b/src/com/android/launcher3/util/LifecycleHelper.kt
@@ -0,0 +1,69 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util
+
+import android.app.Activity
+import android.os.Bundle
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleRegistry
+import androidx.lifecycle.setViewTreeLifecycleOwner
+import androidx.savedstate.SavedStateRegistryController
+import androidx.savedstate.SavedStateRegistryOwner
+import androidx.savedstate.setViewTreeSavedStateRegistryOwner
+
+/** Utility class for triggering various lifecycle events based on activity callbacks */
+class LifecycleHelper(
+    private val owner: SavedStateRegistryOwner,
+    private val savedStateRegistryController: SavedStateRegistryController,
+    private val lifecycleRegistry: LifecycleRegistry,
+) : ActivityLifecycleCallbacksAdapter {
+
+    override fun onActivityPreCreated(activity: Activity, savedInstanceState: Bundle?) {
+        savedStateRegistryController.performRestore(savedInstanceState)
+    }
+
+    override fun onActivityPostCreated(activity: Activity, savedInstanceState: Bundle?) {
+        activity.window.decorView.setViewTreeLifecycleOwner(owner)
+        activity.window.decorView.setViewTreeSavedStateRegistryOwner(owner)
+        lifecycleRegistry.handleLifecycleEvent(Lifecycle.Event.ON_CREATE)
+    }
+
+    override fun onActivityPostStarted(activity: Activity) {
+        lifecycleRegistry.handleLifecycleEvent(Lifecycle.Event.ON_START)
+    }
+
+    override fun onActivityPostResumed(activity: Activity) {
+        lifecycleRegistry.handleLifecycleEvent(Lifecycle.Event.ON_RESUME)
+    }
+
+    override fun onActivityPrePaused(activity: Activity) {
+        lifecycleRegistry.handleLifecycleEvent(Lifecycle.Event.ON_PAUSE)
+    }
+
+    override fun onActivityPreStopped(activity: Activity) {
+        lifecycleRegistry.handleLifecycleEvent(Lifecycle.Event.ON_STOP)
+    }
+
+    override fun onActivityPreDestroyed(activity: Activity) {
+        lifecycleRegistry.handleLifecycleEvent(Lifecycle.Event.ON_DESTROY)
+    }
+
+    override fun onActivitySaveInstanceState(activity: Activity, bundle: Bundle) {
+        lifecycleRegistry.currentState = Lifecycle.State.CREATED
+        savedStateRegistryController.performSave(bundle)
+    }
+}
diff --git a/src/com/android/launcher3/views/ActivityContext.java b/src/com/android/launcher3/views/ActivityContext.java
index 49adaec..bcb9295 100644
--- a/src/com/android/launcher3/views/ActivityContext.java
+++ b/src/com/android/launcher3/views/ActivityContext.java
@@ -56,6 +56,7 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.core.view.WindowInsetsCompat;
+import androidx.savedstate.SavedStateRegistryOwner;
 
 import com.android.launcher3.BubbleTextView;
 import com.android.launcher3.DeviceProfile;
@@ -94,7 +95,7 @@
  * An interface to be used along with a context for various activities in Launcher. This allows a
  * generic class to depend on Context subclass instead of an Activity.
  */
-public interface ActivityContext {
+public interface ActivityContext extends SavedStateRegistryOwner {
 
     String TAG = "ActivityContext";
 
@@ -229,9 +230,7 @@
         getOnDeviceProfileChangeListeners().remove(listener);
     }
 
-    default ViewCache getViewCache() {
-        return new ViewCache();
-    }
+    ViewCache getViewCache();
 
     /**
      * Controller for supporting item drag-and-drop
diff --git a/src/com/android/launcher3/views/ComposeInitializer.java b/src/com/android/launcher3/views/ComposeInitializer.java
deleted file mode 100644
index 0929885..0000000
--- a/src/com/android/launcher3/views/ComposeInitializer.java
+++ /dev/null
@@ -1,229 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3.views;
-
-import android.os.Build;
-import android.view.View;
-import android.view.ViewParent;
-import android.view.ViewTreeObserver;
-
-import androidx.annotation.NonNull;
-import androidx.annotation.RequiresApi;
-import androidx.lifecycle.Lifecycle;
-import androidx.lifecycle.LifecycleOwner;
-import androidx.lifecycle.LifecycleRegistry;
-import androidx.lifecycle.ViewTreeLifecycleOwner;
-import androidx.savedstate.SavedStateRegistry;
-import androidx.savedstate.SavedStateRegistryController;
-import androidx.savedstate.SavedStateRegistryOwner;
-import androidx.savedstate.ViewTreeSavedStateRegistryOwner;
-
-import com.android.launcher3.Utilities;
-
-/**
- * An initializer to use Compose for classes implementing {@code ActivityContext}. This allows
- * adding ComposeView to ViewTree outside a {@link androidx.activity.ComponentActivity}.
- */
-public final class ComposeInitializer {
-    /**
-     * Performs the initialization to use Compose in the ViewTree of {@code target}.
-     */
-    public static void initCompose(ActivityContext target) {
-        getContentChild(target).addOnAttachStateChangeListener(
-                new View.OnAttachStateChangeListener() {
-
-                    @Override
-                    public void onViewAttachedToWindow(View v) {
-                        ComposeInitializer.onAttachedToWindow(v);
-                    }
-
-                    @Override
-                    public void onViewDetachedFromWindow(View v) {
-                        ComposeInitializer.onDetachedFromWindow(v);
-                    }
-                });
-    }
-
-    /**
-     * Find the "content child" for {@code target}.
-     *
-     * @see "WindowRecomposer.android.kt: [View.contentChild]"
-     */
-    private static View getContentChild(ActivityContext target) {
-        View self = target.getDragLayer();
-        ViewParent parent = self.getParent();
-        while (parent instanceof View parentView) {
-            if (parentView.getId() == android.R.id.content) return self;
-            self = parentView;
-            parent = self.getParent();
-        }
-        return self;
-    }
-
-    /**
-     * Function to be called on your window root view's [View.onAttachedToWindow] function.
-     */
-    private static void onAttachedToWindow(View root) {
-        if (ViewTreeLifecycleOwner.get(root) != null) {
-            throw new IllegalStateException(
-                    "View " + root + " already has a LifecycleOwner");
-        }
-
-        ViewParent parent = root.getParent();
-        if (parent instanceof View && ((View) parent).getId() != android.R.id.content) {
-            throw new IllegalStateException(
-                    "ComposeInitializer.onContentChildAttachedToWindow(View) must be called on "
-                            + "the content child. Outside of activities and dialogs, this is "
-                            + "usually the top-most View of a window.");
-        }
-
-        // The lifecycle owner, which is STARTED when [root] is visible and RESUMED when [root]
-        // is both visible and focused.
-        ViewLifecycleOwner lifecycleOwner = new ViewLifecycleOwner(root);
-
-        // We must call [ViewLifecycleOwner.onCreate] after creating the
-        // [SavedStateRegistryOwner] because `onCreate` might move the lifecycle state to STARTED
-        // which will make [SavedStateRegistryController.performRestore] throw.
-        lifecycleOwner.onCreate();
-
-        // Set the owners on the root. They will be reused by any ComposeView inside the root
-        // hierarchy.
-        ViewTreeLifecycleOwner.set(root, lifecycleOwner);
-        ViewTreeSavedStateRegistryOwner.set(root, lifecycleOwner);
-    }
-
-    /**
-     * Function to be called on your window root view's [View.onDetachedFromWindow] function.
-     */
-    private static void onDetachedFromWindow(View root) {
-        final LifecycleOwner lifecycleOwner = ViewTreeLifecycleOwner.get(root);
-        if (lifecycleOwner != null) {
-            ((ViewLifecycleOwner) lifecycleOwner).onDestroy();
-        }
-        ViewTreeLifecycleOwner.set(root, null);
-        ViewTreeSavedStateRegistryOwner.set(root, null);
-    }
-
-    /**
-     * A [LifecycleOwner] for a [View] that updates lifecycle state based on window state.
-     *
-     * Also a trivial implementation of [SavedStateRegistryOwner] that does not do any save or
-     * restore. This works for processes similar to the SystemUI process, which is always running
-     * and top-level windows using this initialization are created once, when the process is
-     * started.
-     *
-     * The implementation requires the caller to call [onCreate] and [onDestroy] when the view is
-     * attached to or detached from a view hierarchy. After [onCreate] and before [onDestroy] is
-     * called, the implementation monitors window state in the following way
-     * * If the window is not visible, we are in the [Lifecycle.State.CREATED] state
-     * * If the window is visible but not focused, we are in the [Lifecycle.State.STARTED] state
-     * * If the window is visible and focused, we are in the [Lifecycle.State.RESUMED] state
-     *
-     * Or in table format:
-     * ```
-     * ┌───────────────┬───────────────────┬──────────────┬─────────────────┐
-     * │ View attached │ Window Visibility │ Window Focus │ Lifecycle State │
-     * ├───────────────┼───────────────────┴──────────────┼─────────────────┤
-     * │ Not attached  │                 Any              │       N/A       │
-     * ├───────────────┼───────────────────┬──────────────┼─────────────────┤
-     * │               │    Not visible    │     Any      │     CREATED     │
-     * │               ├───────────────────┼──────────────┼─────────────────┤
-     * │   Attached    │                   │   No focus   │     STARTED     │
-     * │               │      Visible      ├──────────────┼─────────────────┤
-     * │               │                   │  Has focus   │     RESUMED     │
-     * └───────────────┴───────────────────┴──────────────┴─────────────────┘
-     * ```
-     */
-    private static class ViewLifecycleOwner implements SavedStateRegistryOwner {
-        private final ViewTreeObserver.OnWindowFocusChangeListener mWindowFocusListener =
-                hasFocus -> updateState();
-        private final LifecycleRegistry mLifecycleRegistry = new LifecycleRegistry(this);
-
-        private final SavedStateRegistryController mSavedStateRegistryController =
-                SavedStateRegistryController.create(this);
-
-        private final View mView;
-        private final Api34Impl mApi34Impl;
-
-        ViewLifecycleOwner(View view) {
-            mView = view;
-            if (Utilities.ATLEAST_U) {
-                mApi34Impl = new Api34Impl();
-            } else {
-                mApi34Impl = null;
-            }
-
-            mSavedStateRegistryController.performRestore(null);
-        }
-
-        @NonNull
-        @Override
-        public Lifecycle getLifecycle() {
-            return mLifecycleRegistry;
-        }
-
-        @NonNull
-        @Override
-        public SavedStateRegistry getSavedStateRegistry() {
-            return mSavedStateRegistryController.getSavedStateRegistry();
-        }
-
-        void onCreate() {
-            mLifecycleRegistry.setCurrentState(Lifecycle.State.CREATED);
-            if (Utilities.ATLEAST_U) {
-                mApi34Impl.addOnWindowVisibilityChangeListener();
-            }
-            mView.getViewTreeObserver().addOnWindowFocusChangeListener(
-                    mWindowFocusListener);
-            updateState();
-        }
-
-        void onDestroy() {
-            if (Utilities.ATLEAST_U) {
-                mApi34Impl.removeOnWindowVisibilityChangeListener();
-            }
-            mView.getViewTreeObserver().removeOnWindowFocusChangeListener(
-                    mWindowFocusListener);
-            mLifecycleRegistry.setCurrentState(Lifecycle.State.DESTROYED);
-        }
-
-        private void updateState() {
-            Lifecycle.State state =
-                    mView.getWindowVisibility() != View.VISIBLE ? Lifecycle.State.CREATED
-                            : (!mView.hasWindowFocus() ? Lifecycle.State.STARTED
-                                    : Lifecycle.State.RESUMED);
-            mLifecycleRegistry.setCurrentState(state);
-        }
-
-        @RequiresApi(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
-        private class Api34Impl {
-            private final ViewTreeObserver.OnWindowVisibilityChangeListener
-                    mWindowVisibilityListener =
-                    visibility -> updateState();
-
-            void addOnWindowVisibilityChangeListener() {
-                mView.getViewTreeObserver().addOnWindowVisibilityChangeListener(
-                        mWindowVisibilityListener);
-            }
-
-            void removeOnWindowVisibilityChangeListener() {
-                mView.getViewTreeObserver().removeOnWindowVisibilityChangeListener(
-                        mWindowVisibilityListener);
-            }
-        }
-    }
-}
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/ActivityContextWrapper.java b/tests/multivalentTests/src/com/android/launcher3/util/ActivityContextWrapper.java
index 4217d22..f4dc88d 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/ActivityContextWrapper.java
+++ b/tests/multivalentTests/src/com/android/launcher3/util/ActivityContextWrapper.java
@@ -18,23 +18,15 @@
 import android.R;
 import android.content.Context;
 import android.content.ContextWrapper;
-import android.view.ContextThemeWrapper;
 
 import com.android.launcher3.DeviceProfile;
-import com.android.launcher3.DeviceProfile.OnDeviceProfileChangeListener;
 import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
 
-import java.util.ArrayList;
-import java.util.List;
-
 /**
  * {@link ContextWrapper} with internal Launcher interface for testing
  */
-public class ActivityContextWrapper extends ContextThemeWrapper implements ActivityContext {
-
-    private final List<OnDeviceProfileChangeListener> mDpChangeListeners = new ArrayList<>();
+public class ActivityContextWrapper extends BaseContext {
 
     private final DeviceProfile mProfile;
     private final MyDragLayer mMyDragLayer;
@@ -47,20 +39,15 @@
         super(base, theme);
         mProfile = InvariantDeviceProfile.INSTANCE.get(base).getDeviceProfile(base).copy(base);
         mMyDragLayer = new MyDragLayer(this);
+        Executors.MAIN_EXECUTOR.execute(this::onViewCreated);
     }
 
-
     @Override
     public BaseDragLayer getDragLayer() {
         return mMyDragLayer;
     }
 
     @Override
-    public List<OnDeviceProfileChangeListener> getOnDeviceProfileChangeListeners() {
-        return mDpChangeListeners;
-    }
-
-    @Override
     public DeviceProfile getDeviceProfile() {
         return mProfile;
     }