Add integration for window-based overlays

modo-compose/src/main/java/com/github/terrakok/modo/ComposeRender.kt

@@ -18,6 +18,7 @@ import androidx.lifecycle.Lifecycle.Event.ON_RESUME
 import androidx.lifecycle.Lifecycle.Event.ON_START
 import androidx.lifecycle.Lifecycle.Event.ON_STOP
 import com.github.terrakok.modo.android.ModoScreenAndroidAdapter
+import com.github.terrakok.modo.android.overlaySaveableStateKey
 import com.github.terrakok.modo.animation.ScreenTransition
 import com.github.terrakok.modo.animation.cleanupProtectedScreens
 import com.github.terrakok.modo.animation.preDisposeProtectedScreens
@@ -253,6 +254,7 @@ internal class ComposeRenderer<State : NavigationState>(
         }
         ScreenModelStore.remove(this)
         stateHolder.removeState(saveableStateKey)
+        stateHolder.removeState(overlaySaveableStateKey)
 
         ModoDevOptions.onScreenDisposeListener?.invoke(this)
         // clear nested screens using recursion

modo-compose/src/main/java/com/github/terrakok/modo/android/ProvideOverlayIntegration.kt

@@ -0,0 +1,83 @@
+package com.github.terrakok.modo.android
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.saveable.LocalSaveableStateRegistry
+import androidx.compose.runtime.saveable.rememberSaveable
+import androidx.lifecycle.compose.LocalLifecycleOwner
+import androidx.savedstate.compose.LocalSavedStateRegistryOwner
+import com.github.terrakok.modo.LocalSaveableStateHolder
+import com.github.terrakok.modo.Screen
+import com.github.terrakok.modo.util.currentOrThrow
+
+/**
+ * Integrates window-based overlays (Dialog, ModalBottomSheet, Popup, etc.) with Modo screen lifecycle and state management.
+ *
+ * Wrap your window-based overlay content with this function if:
+ * - Your code inside uses composition locals: [LocalLifecycleOwner], [LocalSavedStateRegistryOwner],
+ *   [LocalSaveableStateRegistry].
+ * - You need [rememberSaveable] state to survive navigation inside overlay (e.g., overlay hidden on forward navigation,
+ *   restored with saved state on back navigation).
+ *
+ * ## Why you need this
+ *
+ * Window-based overlays (Dialog, ModalBottomSheet, Popup) create a separate `AndroidComposeView` with its own
+ * composition tree. This new composition tree receives default composition locals from the Activity/Window level,
+ * not from your Screen's composition.
+ *
+ * **Problem 1: Composition locals point to wrong owners**
+ *
+ * The overlay inherits Activity-level `LocalLifecycleOwner`, `LocalViewModelStoreOwner`, and `LocalSavedStateRegistryOwner`
+ * instead of Screen-level ones. This means:
+ * - ViewModels get scoped to Activity instead of Screen (survive when they shouldn't)
+ * - Lifecycle observers observe Activity lifecycle instead of Screen lifecycle
+ * - SavedState gets tied to Activity instead of Screen's navigation state
+ *
+ * **Problem 2: State restoration breaks during navigation**
+ *
+ * `rememberSaveable` inside the overlay uses the Activity's SavedStateRegistry, which only survives configuration changes.
+ * When you navigate away from a Screen (moving it to backstack), Modo preserves the Screen's state, but the overlay's
+ * `rememberSaveable` state is disconnected from this system. When you navigate back, a **new** `AndroidComposeView` is
+ * created from scratch with empty state - the overlay's previous state is lost because it was never saved in Modo's navigation state.
+ *
+ * ## Usage
+ *
+ * ```kotlin
+ * MyScreen(...): Screen {
+ *     @Composable
+ *     override fun Content(modifier: Modifier) {
+ *         ModalBottomSheet(
+ *             onDismissRequest = { navigation.back() }
+ *         ) {
+ *             ProvideOverlayIntegration {
+ *                 MyBottomSheetContent()
+ *             }
+ *         }
+ *     }
+ * }
+ * ```
+ *
+ * @param content The overlay content with restored screen composition locals and proper state management
+ */
+@Composable
+fun Screen.ProvideOverlayIntegration(
+    content: @Composable () -> Unit
+) {
+    val androidAdapter = remember(this) {
+        ModoScreenAndroidAdapter.getOrNull(this) ?: error(
+            "ModoScreenAndroidAdapter is not found. Ensure Modo is properly initialized."
+        )
+    }
+    val saveableStateHolder = LocalSaveableStateHolder.currentOrThrow
+    // Create a dialog-scoped state container within Modo's state management.
+    // It is not enough just propagate LocalSaveableStateRegistry and LocalSavedStateRegistryOwner because of the logic of saving rememberSaveable.
+    // So we need to use a new SaveableStateProvider with a new key (the same key as before will lead to the exception).
+    saveableStateHolder.SaveableStateProvider(overlaySaveableStateKey) {
+        // Provide screen-level composition locals (lifecycle, ViewModelStore, SavedStateRegistry)
+        androidAdapter.ProvideCompositionLocals {
+            content()
+        }
+    }
+}
+
+internal val Screen.overlaySaveableStateKey: String get() = screenKey.value + ".dialog"

sample/src/main/java/com/github/terrakok/modo/sample/Animations.kt

@@ -35,8 +35,7 @@ fun ComposeRendererScope<StackState>.SlideTransition(
                     scaleIn(initialScale = 2f, animationSpec = animationSpec) + fadeIn(animationSpec) togetherWith
                         fadeOut(animationSpec)
                 }
-                oldState?.stack?.last() is DialogScreen ||
-                    oldState?.stack?.last() !is DialogScreen && newState?.stack?.last() is DialogScreen -> {
+                screen is DialogScreen -> {
                     val animationSpec = tween<Float>(durationMillis = SampleAppConfig.animationDurationMs)
                     fadeIn(animationSpec) togetherWith fadeOut(animationSpec)
                 }

sample/src/main/java/com/github/terrakok/modo/sample/screens/base/ButtonsScreenContent.kt

@@ -12,8 +12,12 @@ import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.layout.size
 import androidx.compose.foundation.layout.systemBars
 import androidx.compose.foundation.layout.windowInsetsPadding
+import androidx.compose.material.Icon
+import androidx.compose.material.IconButton
 import androidx.compose.material.MaterialTheme
 import androidx.compose.material.Text
+import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.filled.ArrowDropDown
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.IntState
@@ -23,6 +27,7 @@ import androidx.compose.runtime.mutableIntStateOf
 import androidx.compose.runtime.saveable.rememberSaveable
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.vector.rememberVectorPainter
 import androidx.compose.ui.platform.LocalInspectionMode
 import androidx.compose.ui.text.style.TextAlign
 import androidx.compose.ui.tooling.preview.Preview
@@ -34,6 +39,7 @@ import com.github.terrakok.modo.Screen
 import com.github.terrakok.modo.ScreenKey
 import com.github.terrakok.modo.sample.SampleAppConfig
 import com.github.terrakok.modo.sample.components.BackButton
+import com.github.terrakok.modo.sample.logs.logcat
 import com.github.terrakok.modo.sample.randomBackground
 import com.github.terrakok.modo.sample.screens.ButtonsState
 import com.github.terrakok.modo.sample.screens.GroupedButtonsList
@@ -43,7 +49,6 @@ import com.github.terrakok.modo.stack.LocalStackNavigation
 import com.github.terrakok.modo.stack.back
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.isActive
-import logcat.logcat
 
 internal const val COUNTER_DELAY_MS = 100L
 
@@ -53,10 +58,26 @@ internal fun Screen.ButtonsScreenContent(
     screenName: String,
     state: GroupedButtonsState,
     modifier: Modifier = Modifier,
+    windowInsets: WindowInsets = WindowInsets.systemBars,
+    logLifecycle: Boolean = true,
+    enableCounter: Boolean = true,
+    @Suppress("ComposableLambdaParameterNaming")
+    topRightButtonSlot: @Composable () -> Unit = {},
 ) {
-    LogLifecycle()
+    if (logLifecycle) {
+        LogLifecycle()
+    }
     val counter by rememberCounterState()
-    ButtonsScreenContent(screenIndex, screenName, counter, screenKey, state, modifier)
+    ButtonsScreenContent(
+        screenIndex = screenIndex,
+        screenName = screenName,
+        counter = if (enableCounter) counter else 0,
+        screenKey = screenKey,
+        state = state,
+        topRightButtonSlot = topRightButtonSlot,
+        windowInsets = windowInsets,
+        modifier = modifier
+    )
 }
 
 @Composable
@@ -81,12 +102,17 @@ internal fun ButtonsScreenContent(
     screenKey: ScreenKey,
     state: GroupedButtonsState,
     modifier: Modifier = Modifier,
+    windowInsets: WindowInsets = WindowInsets.systemBars,
+    @Suppress("ComposableLambdaParameterNaming")
+    topRightButtonSlot: @Composable () -> Unit = {}
 ) {
     SampleScreenContent(
         screenIndex = screenIndex,
         screenName = screenName,
         counter = counter,
         screenKey = screenKey,
+        topRightButtonSlot = topRightButtonSlot,
+        windowInsets = windowInsets,
         modifier = modifier,
     ) {
         GroupedButtonsList(
@@ -102,27 +128,36 @@ internal fun Screen.SampleScreenContent(
     screenName: String,
     screenKey: ScreenKey,
     modifier: Modifier = Modifier,
+    windowInsets: WindowInsets = WindowInsets.systemBars,
     content: @Composable ColumnScope.() -> Unit
 ) {
     LogLifecycle()
     val counter by rememberCounterState()
-    SampleScreenContent(screenIndex, screenName, counter, screenKey, modifier, content)
+    SampleScreenContent(
+        screenIndex = screenIndex,
+        screenName = screenName,
+        counter = counter,
+        screenKey = screenKey,
+        modifier = modifier,
+        windowInsets = windowInsets,
+        content = content
+    )
 }
 
 @OptIn(ExperimentalModoApi::class)
 @Composable
-fun Screen.LogLifecycle(prefix: String = this::class.simpleName.orEmpty()) {
+fun Screen.LogLifecycle(prefix: String = "") {
     val lifecycleOwner = LocalLifecycleOwner.current
 
     // You will not be able to observe updates of lifecycleOwner when this content is not in the composition
     DisposableEffect(lifecycleOwner) {
-        logcat(tag = "LifecycleDebug") { "$prefix $screenKey DisposableEffect".trim() }
+        logcat("LogLifecycle") { "$prefix DisposableEffect $lifecycleOwner".trim() }
         val observer = LifecycleEventObserver { _, event ->
-            logcat(tag = "LifecycleDebug") { "$prefix $screenKey DisposableEffect $event".trim() }
+            logcat("LogLifecycle") { "$prefix DisposableEffect $event $lifecycleOwner".trim() }
         }
         lifecycleOwner.lifecycle.addObserver(observer)
         onDispose {
-            logcat(tag = "LifecycleDebug") { "$prefix $screenKey DisposableEffect onDispose" }
+            logcat("LogLifecycle") { "$prefix DisposableEffect.onDispose $lifecycleOwner".trim() }
             lifecycleOwner.lifecycle.removeObserver(observer)
         }
     }
@@ -140,12 +175,14 @@ internal fun SampleScreenContent(
     counter: Int,
     screenKey: ScreenKey,
     modifier: Modifier = Modifier,
+    windowInsets: WindowInsets = WindowInsets.systemBars,
+    topRightButtonSlot: @Composable () -> Unit = {},
     content: @Composable ColumnScope.() -> Unit
 ) {
     Box(
         modifier = modifier
             .randomBackground()
-            .windowInsetsPadding(WindowInsets.systemBars),
+            .windowInsetsPadding(windowInsets),
     ) {
         Column(Modifier.padding(8.dp)) {
             Row(
@@ -156,8 +193,10 @@ internal fun SampleScreenContent(
                     onClick = { stackNavigation?.back() },
                 )
                 Text(
-                    text = counter.toString()
+                    text = counter.toString(),
+                    modifier = Modifier.weight(1f)
                 )
+                topRightButtonSlot()
             }
             Text(
                 text = "$screenName $screenIndex",
@@ -194,6 +233,14 @@ private fun ButtonsPreview() {
                 ModoButtonSpec("Button with a very long text") {},
             )
         ),
+        topRightButtonSlot = {
+            IconButton(onClick = {}) {
+                Icon(
+                    painter = rememberVectorPainter(image = Icons.Filled.ArrowDropDown),
+                    contentDescription = null
+                )
+            }
+        },
         modifier = Modifier.fillMaxSize()
     )
 }