Fix screen model store clear

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

@@ -30,7 +30,7 @@ inline fun <reified T : ScreenModel> Screen.rememberScreenModel(
     tag: String? = null,
     crossinline factory: @DisallowComposableCalls () -> T
 ): T =
-    remember(ScreenModelStore.getKey<T>(this, tag)) {
+    remember(ScreenModelStore.getScreenModelKey<T>(this, tag)) {
         ScreenModelStore.getOrPut(this, tag, factory)
     }
 

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

@@ -4,6 +4,8 @@ import androidx.compose.runtime.DisallowComposableCalls
 import com.github.terrakok.modo.ModoDevOptions
 import com.github.terrakok.modo.Screen
 import com.github.terrakok.modo.ScreenKey
+import com.github.terrakok.modo.model.ScreenModelStore.getDependencyKey
+import com.github.terrakok.modo.model.ScreenModelStore.getScreenModelKey
 import kotlinx.coroutines.flow.MutableStateFlow
 import org.jetbrains.annotations.VisibleForTesting
 import java.util.concurrent.ConcurrentHashMap
@@ -30,6 +32,16 @@ internal data class DependencyWithRemoveOrder(
     val dependency: Dependency
 )
 
+/**
+ * Internal store for screen models and dependencies.
+ *
+ * Thread Safety:
+ * While the store uses thread-safe collections, it is primarily designed for use from the UI thread.
+ * Race conditions may occur only when multiple threads attempt to initialize a dependency or
+ * screen model for the same key simultaneously for the first time.
+ * Since these models are typically created and accessed within the Compose runtime,
+ * they generally benefit from the sequential execution guaranteed by Compose effects.
+ */
 object ScreenModelStore {
 
     @PublishedApi
@@ -111,9 +123,6 @@ object ScreenModelStore {
         ?.dependency
         ?.dependencyInstance as? T
 
-    fun getDependencyKey(screen: Screen, name: String, tag: String? = null) =
-        "${screen.screenKey.value}:$name${if (tag != null) ":$tag" else ""}"
-
     fun remove(screen: Screen) {
         if (removedScreenKeys[screen.screenKey] != null) {
             ModoDevOptions.onIllegalScreenModelStoreAccess.validationFailed(
@@ -142,15 +151,24 @@ object ScreenModelStore {
         screenDependenciesInternal(screen).sortedBy { it.value.removePriority }.map { it.value.dependency.dependencyInstance }
 
     internal fun screenDependenciesInternal(screen: Screen): Sequence<Map.Entry<DependencyKey, DependencyWithRemoveOrder>> =
-        dependencies.asSequence().filter { it.key.startsWith(screen.screenKey.value) }
+        dependencies.asSequence().filter { isKeyOfScreen(it.key, screen) }
 
     /**
-     * Generates a key based on input parameters.
+     * Generates a screen model key based on screen, type, and tag.
+     * Format: "screenKey:ScreenModelClassName:tag"
      */
     @PublishedApi
-    internal inline fun <reified T : ScreenModel> getKey(screen: Screen, tag: String?): ScreenModelKey =
+    internal inline fun <reified T : ScreenModel> getScreenModelKey(screen: Screen, tag: String?): ScreenModelKey =
         "${screen.screenKey.value}:${T::class.qualifiedName}:${tag ?: "default"}"
 
+    /**
+     * Generates a dependency key based on screen, name, and optional tag.
+     * Format: "screenKey:name" or "screenKey:name:tag"
+     */
+    @PublishedApi
+    internal fun getDependencyKey(screen: Screen, name: String, tag: String? = null) =
+        "${screen.screenKey.value}:$name${if (tag != null) ":$tag" else ""}"
+
     @PublishedApi
     internal fun getDependencyKey(screenModel: ScreenModel, name: String): DependencyKey =
         screenModels
@@ -167,12 +185,21 @@ object ScreenModelStore {
         tag: String?,
         factory: @DisallowComposableCalls () -> T
     ): T {
-        val key = getKey<T>(screen, tag)
+        val key = getScreenModelKey<T>(screen, tag)
         lastScreenModelKey.value = key
         assertGetOrPutScreenModelsCorrect(screen, screenModels[key])
         return screenModels.getOrPut(key, factory) as T
     }
 
+    @VisibleForTesting
+    internal inline fun <reified T : ScreenModel> getOrNull(
+        screen: Screen,
+        tag: String?,
+    ): T? {
+        val key = getScreenModelKey<T>(screen, tag)
+        return screenModels.get(key) as T?
+    }
+
     @PublishedApi
     internal fun assertGetOrPutScreenModelsCorrect(screen: Screen, valueInMap: Any?) {
         assertGetOrPutCorrect(
@@ -217,7 +244,19 @@ object ScreenModelStore {
 
     private fun <T> Map<String, T>.onEach(screen: Screen, block: (String) -> Unit) =
         asSequence()
-            .filter { it.key.startsWith(screen.screenKey.value) }
+            .filter { isKeyOfScreen(it.key, screen) }
             .map { it.key }
             .forEach(block)
+
+    /**
+     * Checks if the given key (dependency or screen model) belongs to the specified screen.
+     *
+     * Works with keys created by:
+     * - [getDependencyKey]: "screenKey:name" or "screenKey:name:tag"
+     * - [getScreenModelKey]: "screenKey:ClassName:tag"
+     *
+     * All keys follow the format "screenKey:...", so we check for the prefix.
+     */
+    private fun isKeyOfScreen(key: String, screen: Screen): Boolean =
+        key.startsWith("${screen.screenKey.value}:")
 }

modo-compose/src/test/java/com/github/terrakok/modo/MockScreenModel.kt

@@ -2,4 +2,8 @@ package com.github.terrakok.modo
 
 import com.github.terrakok.modo.model.ScreenModel
 
-class MockScreenModel(val id: String = "") : ScreenModel
+class MockScreenModel(val id: String = "", val onDispose: () -> Unit = {}) : ScreenModel {
+    override fun onDispose() {
+        onDispose.invoke()
+    }
+}