Fixed Dialog state and wrong Composition Local providing

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

@@ -7,7 +7,6 @@ import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.neverEqualPolicy
 import androidx.compose.runtime.saveable.SaveableStateHolder
-import androidx.compose.runtime.saveable.rememberSaveableStateHolder
 import androidx.compose.runtime.setValue
 import androidx.compose.runtime.staticCompositionLocalOf
 import com.github.terrakok.modo.model.ScreenModelStore
@@ -61,15 +60,13 @@ internal class ComposeRenderer<State : NavigationState>(
         screen: Screen,
         content: RendererContent<State> = defaultRendererContent
     ) {
-        // use single state holder for whole hierarchy, store it on top of it
-        val stateHolder: SaveableStateHolder = LocalSaveableStateHolder.current ?: rememberSaveableStateHolder()
+        val stateHolder: SaveableStateHolder = LocalSaveableStateHolder.currentOrThrow
         DisposableEffect(key1 = state) {
             onDispose {
                 clearScreens(stateHolder)
             }
         }
         CompositionLocalProvider(
-            LocalSaveableStateHolder providesDefault stateHolder,
             LocalContainerScreen provides containerScreen
         ) {
             ComposeRendererScope(lastState, state, screen).content()

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

@@ -21,21 +21,21 @@ object Modo {
      * @param savedState - container with modo state and graph
      * @param rootScreenProvider invokes when [savedState] is null and [inMemoryScreen] is null and we need to provide root screen.
      */
-    fun <T : Screen> init(savedState: Bundle?, inMemoryScreen: T?, rootScreenProvider: () -> T): T {
-        val modoGraph = savedState?.getParcelable<T>(MODO_GRAPH)
+    fun <T : Screen> init(savedState: Bundle?, inMemoryScreen: RootScreen<T>?, rootScreenProvider: () -> T): RootScreen<T> {
+        val modoGraph = savedState?.getParcelable<RootScreen<T>>(MODO_GRAPH)
         return if (modoGraph != null) {
             restoreScreenCounter(savedState.getInt(MODO_SCREEN_COUNTER_KEY))
             modoGraph
         } else {
-            inMemoryScreen ?: rootScreenProvider()
+            inMemoryScreen ?: RootScreen(rootScreenProvider())
         }
     }
 
     /**
      * Must be called to clear all data from [ScreenModelStore], related with removed screens.
      */
-    fun onRootScreenFinished(rootScreen: ContainerScreen<*>?) {
-        rootScreen?.let(::clearScreenModel)
+    fun <T: ContainerScreen<*>> onRootScreenFinished(rootScreen: RootScreen<T>?) {
+        rootScreen?.screen?.let(::clearScreenModel)
     }
 
     private fun clearScreenModel(screen: Screen) {

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

@@ -0,0 +1,28 @@
+package com.github.terrakok.modo
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.saveable.SaveableStateHolder
+import androidx.compose.runtime.saveable.rememberSaveableStateHolder
+import kotlinx.parcelize.Parcelize
+
+/**
+ * Screen for single source of providing [LocalSaveableStateHolder]. Should be used with [Modo.init].
+ */
+@Parcelize
+class RootScreen<T : Screen>(
+    val screen: T,
+    override val screenKey: ScreenKey = generateScreenKey()
+) : Screen {
+
+    @Composable
+    override fun Content() {
+        val stateHolder: SaveableStateHolder = LocalSaveableStateHolder.current ?: rememberSaveableStateHolder()
+        CompositionLocalProvider(
+            LocalSaveableStateHolder providesDefault stateHolder
+        ) {
+            screen.Content()
+        }
+    }
+
+}

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

@@ -59,7 +59,7 @@ abstract class StackScreen(
                 onDismissRequest = { back() },
                 properties = remember { dialog.provideDialogProperties() }
             ) {
-                Content(screen = dialog)
+                Content(dialog, content)
             }
         }
     }

sample/src/main/java/com/github/terrakok/androidcomposeapp/AppActivity.kt

@@ -8,11 +8,12 @@ import androidx.compose.material.Surface
 import com.github.terrakok.androidcomposeapp.screens.SampleScreen
 import com.github.terrakok.androidcomposeapp.screens.containers.SampleStack
 import com.github.terrakok.modo.Modo
+import com.github.terrakok.modo.RootScreen
 import com.github.terrakok.modo.stack.StackScreen
 
 class AppActivity : AppCompatActivity() {
 
-    private var rootScreen: StackScreen? = null
+    private var rootScreen: RootScreen<StackScreen>? = null
 
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)

sample/src/main/java/com/github/terrakok/androidcomposeapp/screens/SampleScreen.kt

@@ -23,7 +23,7 @@ import com.github.terrakok.androidcomposeapp.ListScreen
 import com.github.terrakok.androidcomposeapp.ModelSampleScreen
 import com.github.terrakok.androidcomposeapp.randomBackground
 import com.github.terrakok.androidcomposeapp.screens.containers.SampleMultiScreen
-import com.github.terrakok.androidcomposeapp.screens.containers.SampleStackScreen
+import com.github.terrakok.androidcomposeapp.screens.containers.SampleContainerScreen
 import com.github.terrakok.modo.LocalContainerScreen
 import com.github.terrakok.modo.NavigationContainer
 import com.github.terrakok.modo.Screen
@@ -131,7 +131,7 @@ private fun rememberButtons(
                 navigator.backTo(it)
             }
         },
-        "Container" to { navigator.forward(SampleStackScreen(i + 1)) },
+        "Container" to { navigator.forward(SampleContainerScreen(i + 1)) },
         "Multiscreen" to { navigator.forward(SampleMultiScreen()) },
         "List/Details" to { navigator.forward(ListScreen()) },
         // Just experiments

sample/src/main/java/com/github/terrakok/androidcomposeapp/screens/containers/ContainerScreen.kt → sample/src/main/java/com/github/terrakok/androidcomposeapp/screens/containers/SampleContainerScreen.kt

@@ -25,7 +25,7 @@ import com.github.terrakok.modo.stack.back
 import kotlinx.parcelize.Parcelize
 
 @Parcelize
-class SampleStackScreen(
+class SampleContainerScreen(
     private val i: Int,
     private val navModel: StackNavModel
 ) : StackScreen(navModel) {
@@ -68,5 +68,5 @@ class SampleStackScreen(
 @Preview
 @Composable
 private fun PreviewContainerScreen() {
-    SampleStackScreen(1).Content()
+    SampleContainerScreen(1).Content()
 }