Clearing AccountViewModels more precisely

amethyst/src/main/java/com/vitorpamplona/amethyst/ui/screen/AccountSessionManager.kt

@@ -23,6 +23,7 @@ package com.vitorpamplona.amethyst.ui.screen
 import androidx.compose.runtime.Stable
 import com.vitorpamplona.amethyst.AccountInfo
 import com.vitorpamplona.amethyst.LocalPreferences
+import com.vitorpamplona.amethyst.model.Account
 import com.vitorpamplona.amethyst.model.AccountSettings
 import com.vitorpamplona.amethyst.model.DefaultChannels
 import com.vitorpamplona.amethyst.model.DefaultDMRelayList
@@ -75,6 +76,18 @@ import java.util.regex.Pattern
 
 val EMAIL_PATTERN: Pattern = Pattern.compile(".+@.+\\.[a-z]+")
 
+sealed class AccountState {
+    object Loading : AccountState()
+
+    object LoggedOff : AccountState()
+
+    @Stable
+    class LoggedIn(
+        val account: Account,
+        var route: Route? = null,
+    ) : AccountState()
+}
+
 @Stable
 class AccountSessionManager(
     val accountsCache: AccountCacheState,
@@ -175,15 +188,6 @@ class AccountSessionManager(
         }
     }
 
-    private fun prepareLogoutOrSwitch() =
-        when (val state = _accountContent.value) {
-            is AccountState.LoggedIn -> {
-                state.currentViewModelStore.viewModelStore.clear()
-            }
-
-            else -> {}
-        }
-
     fun login(
         key: String,
         password: String,
@@ -256,9 +260,6 @@ class AccountSessionManager(
         onError: (String) -> Unit,
     ) {
         try {
-            if (_accountContent.value is AccountState.LoggedIn) {
-                prepareLogoutOrSwitch()
-            }
             loginAndStartUI(key, transientAccount, loginWithExternalSigner, packageName)
         } catch (e: Exception) {
             if (e is CancellationException) throw e
@@ -269,10 +270,6 @@ class AccountSessionManager(
 
     fun newKey(name: String? = null) {
         scope.launch(Dispatchers.IO) {
-            if (_accountContent.value is AccountState.LoggedIn) {
-                prepareLogoutOrSwitch()
-            }
-
             _accountContent.update { AccountState.Loading }
 
             val accountSettings = createNewAccount(name)
@@ -342,7 +339,6 @@ class AccountSessionManager(
         accountInfo: AccountInfo,
         route: Route? = null,
     ) {
-        prepareLogoutOrSwitch()
         localPreferences.switchToAccount(accountInfo)
         loginWithDefaultAccount(route)
     }
@@ -364,7 +360,6 @@ class AccountSessionManager(
         scope.launch(Dispatchers.IO) {
             if (accountInfo.npub == currentAccountNPub()) {
                 // log off and relogin with the 0 account
-                prepareLogoutOrSwitch()
                 localPreferences.deleteAccount(accountInfo)
                 accountsCache.removeAccount(accountInfo.npub.bechToBytes().toHexKey())
                 loginWithDefaultAccount()

amethyst/src/main/java/com/vitorpamplona/amethyst/ui/screen/AccountState.kt

@@ -20,49 +20,110 @@
  */
 package com.vitorpamplona.amethyst.ui.screen
 
+import androidx.activity.ComponentActivity
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
-import androidx.compose.runtime.Stable
+import androidx.compose.runtime.RememberObserver
+import androidx.compose.runtime.remember
+import androidx.compose.ui.platform.LocalContext
+import androidx.lifecycle.DefaultLifecycleObserver
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelStore
 import androidx.lifecycle.ViewModelStoreOwner
 import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
-import com.vitorpamplona.amethyst.model.Account
-import com.vitorpamplona.amethyst.ui.navigation.routes.Route
-
-sealed class AccountState {
-    object Loading : AccountState()
-
-    object LoggedOff : AccountState()
-
-    @Stable
-    class LoggedIn(
-        val account: Account,
-        var route: Route? = null,
-    ) : AccountState() {
-        val currentViewModelStore = AccountCentricViewModelStore()
-    }
-}
+import androidx.lifecycle.viewmodel.compose.viewModel
+import com.vitorpamplona.amethyst.ui.components.getActivity
 
+/**
+ * Creates a new scope for the given ViewModel type.
+ */
 @Composable
 fun SetAccountCentricViewModelStore(
     state: AccountState.LoggedIn,
     content: @Composable () -> Unit,
 ) {
+    val activity = LocalContext.current.getActivity()
+    val vmStore: StoreOwnerRegistry = viewModel(viewModelStoreOwner = activity)
+    vmStore.checkAttached(activity)
+
+    val owner = vmStore.getOwner(state)
+
+    val observer = remember { CompositionObserver(vmStore, state) }
+
     CompositionLocalProvider(
-        LocalViewModelStoreOwner provides state.currentViewModelStore,
-    ) {
-        content()
+        LocalViewModelStoreOwner provides owner,
+        content = content,
+    )
+}
+
+/**
+ * This class is responsible for notifying the [StoreOwnerRegistry] when a composable is detached so
+ * that the viewmodel can be cleared.
+ */
+class CompositionObserver(
+    private val vmStore: StoreOwnerRegistry,
+    private val key: Any,
+) : RememberObserver {
+    override fun onRemembered() {}
+
+    override fun onForgotten() = vmStore.composableDetached(key)
+
+    override fun onAbandoned() = vmStore.composableDetached(key)
+}
+
+/**
+ * Registry for [ViewModelStoreOwner]s that are scoped to a particular composition.
+ * This ViewModel is registered with the Activity's lifecycle and will clear the viewmodels.
+ */
+class StoreOwnerRegistry : ViewModel() {
+    private var isActivityRegistered: Boolean = false
+    private var isChangingConfigurations: Boolean = false
+    private val map = mutableMapOf<Any, ViewModelStoreOwner>()
+
+    override fun onCleared() {
+        map.values.forEach { it.viewModelStore.clear() }
+        super.onCleared()
     }
 
-    // moved this clearing activity to the viewmodel account
-    // because the new composable might run before the onDispose.
-    // DisposableEffect(key1 = state) {
-    //    onDispose {
-    //        state.currentViewModelStore.viewModelStore.clear()
-    //    }
-    // }
+    fun getOwner(key: Any): ViewModelStoreOwner = map[key] ?: ScopedViewModelStoreOwner().also { map[key] = it }
+
+    fun composableDetached(key: Any) {
+        // TODO: This prevents the viewmodel from being cleared when the Composable is detached due
+        //  to a configuration change. We need to make sure that the viewmodel is cleared when the
+        //  Composition is recreated without the Composable. E.g. by observing the Composition
+        if (isChangingConfigurations) return
+        map.remove(key)?.also { owner -> owner.viewModelStore.clear() }
+    }
+
+    fun checkAttached(activity: ComponentActivity) {
+        if (!isActivityRegistered) {
+            isActivityRegistered = true
+            activity.lifecycle.addObserver(
+                object : DefaultLifecycleObserver {
+                    override fun onStart(owner: LifecycleOwner) {
+                        isChangingConfigurations = false
+                    }
+
+                    override fun onStop(owner: LifecycleOwner) {
+                        if (activity.isChangingConfigurations) {
+                            isChangingConfigurations = true
+                        }
+                    }
+
+                    override fun onDestroy(owner: LifecycleOwner) {
+                        isActivityRegistered = false
+                        owner.lifecycle.removeObserver(this)
+                    }
+                },
+            )
+        }
+    }
 }
 
-class AccountCentricViewModelStore : ViewModelStoreOwner {
-    override val viewModelStore = ViewModelStore()
+/**
+ * Simple ViewModelStoreOwner that can be used to create a new scope.
+ */
+class ScopedViewModelStoreOwner : ViewModelStoreOwner {
+    override val viewModelStore: ViewModelStore = ViewModelStore()
 }