Remember pageloader in ViewModel (#1780)

* Remember pageloader in ViewModel

* Fix hang on

* Refine

* Move ***InVM functions to ViewModelState

* Add rememberInVM overload takes no key

* Support asyncInVM & launchInVM with key

* Move UNCHECKED_CAST code

---------

Co-authored-by: Dude so hot <djohn@fbi.gov>

app/src/main/kotlin/com/hippo/ehviewer/ui/reader/ReaderScreen.kt

@@ -48,7 +48,7 @@ import androidx.compose.ui.input.pointer.pointerInput
 import androidx.compose.ui.res.colorResource
 import androidx.compose.ui.res.stringResource
 import androidx.core.view.WindowInsetsControllerCompat
-import arrow.core.Either
+import arrow.core.Either.Companion.catch
 import arrow.core.raise.ensure
 import com.google.accompanist.systemuicontroller.rememberSystemUiController
 import com.hippo.ehviewer.R
@@ -69,6 +69,8 @@ import com.hippo.ehviewer.ui.theme.EhTheme
 import com.hippo.ehviewer.ui.tools.Await
 import com.hippo.ehviewer.ui.tools.DialogState
 import com.hippo.ehviewer.ui.tools.EmptyWindowInsets
+import com.hippo.ehviewer.ui.tools.asyncInVM
+import com.hippo.ehviewer.ui.tools.launchInVM
 import com.hippo.ehviewer.util.displayString
 import com.hippo.files.toOkioPath
 import com.ramcosta.composedestinations.annotation.Destination
@@ -82,6 +84,7 @@ import eu.kanade.tachiyomi.ui.reader.setting.ReadingModeType
 import eu.kanade.tachiyomi.util.lang.launchIO
 import eu.kanade.tachiyomi.util.lang.withIOContext
 import kotlin.coroutines.resume
+import kotlinx.coroutines.awaitCancellation
 import kotlinx.coroutines.flow.merge
 import kotlinx.coroutines.flow.onCompletion
 import kotlinx.coroutines.flow.sample
@@ -108,11 +111,14 @@ private fun Background(
 @Destination<RootGraph>
 @Composable
 fun AnimatedVisibilityScope.ReaderScreen(args: ReaderScreenArgs, navigator: DestinationsNavigator) = composing(navigator) {
-    Await({ preparePageLoader(args) }) { pageLoader ->
-        DisposableEffect(pageLoader) {
-            pageLoader.start()
-            onDispose {
-                pageLoader.stop()
+    val pageLoader by asyncInVM(args) { preparePageLoader(args) }
+    Await({ pageLoader.await() }) { loader ->
+        launchInVM(loader) {
+            loader.start()
+            try {
+                awaitCancellation()
+            } finally {
+                loader.stop()
             }
         }
         val bgColor by collectBackgroundColorAsState()
@@ -130,11 +136,7 @@ fun AnimatedVisibilityScope.ReaderScreen(args: ReaderScreenArgs, navigator: Dest
             }
         }
         Await(
-            {
-                Either.catch {
-                    readingFailed.takeUnless { pageLoader.awaitReady() }
-                }.fold({ it.displayString() }, { it })
-            },
+            block = { catch { readingFailed.takeUnless { loader.awaitReady() } }.fold({ it.displayString() }, { it }) },
             placeholder = {
                 Background(bgColor) {
                     CircularProgressIndicator()
@@ -143,8 +145,8 @@ fun AnimatedVisibilityScope.ReaderScreen(args: ReaderScreenArgs, navigator: Dest
         ) { error ->
             if (error == null) {
                 val info = (args as? ReaderScreenArgs.Gallery)?.info
-                key(pageLoader) {
-                    ReaderScreen(pageLoader, info, navigator)
+                key(loader) {
+                    ReaderScreen(loader, info, navigator)
                 }
             } else {
                 Background(bgColor) {

app/src/main/kotlin/com/hippo/ehviewer/ui/tools/DateStoreState.kt

@@ -1,27 +1,16 @@
 package com.hippo.ehviewer.ui.tools
 
-import androidx.collection.mutableIntObjectMapOf
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisallowComposableCalls
-import androidx.compose.runtime.DisposableEffect
-import androidx.compose.runtime.currentCompositeKeyHash
-import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
-import androidx.compose.runtime.rememberUpdatedState
 import androidx.datastore.dataStoreFile
 import androidx.datastore.preferences.core.PreferenceDataStoreFactory
 import androidx.datastore.preferences.core.byteArrayPreferencesKey
 import androidx.datastore.preferences.core.edit
 import androidx.datastore.preferences.core.emptyPreferences
-import androidx.lifecycle.ViewModel
 import androidx.lifecycle.compose.LifecycleResumeEffect
-import androidx.lifecycle.viewModelScope
-import androidx.lifecycle.viewmodel.compose.viewModel
-import kotlin.coroutines.CoroutineContext
-import kotlin.coroutines.EmptyCoroutineContext
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.flow.SharingStarted
@@ -36,29 +25,6 @@ val dataStoreScope = CoroutineScope(Dispatchers.IO + SupervisorJob())
 val dataStore = PreferenceDataStoreFactory.create { appCtx.dataStoreFile("Remembered.preferences_pb") }
 val dataStateFlow = dataStore.data.stateIn(dataStoreScope, SharingStarted.Eagerly, emptyPreferences())
 
-class StateMapViewModel : ViewModel() {
-    val statesMap = mutableIntObjectMapOf<ArrayDeque<Any>>()
-}
-
-@Composable
-inline fun <reified T : Any> rememberInVM(
-    userKey: Any? = null,
-    crossinline init: @DisallowComposableCalls ViewModel.() -> T,
-) = with(viewModel<StateMapViewModel>()) {
-    val key = currentCompositeKeyHash
-    remember(userKey) {
-        val states = statesMap.getOrPut(key, ::ArrayDeque)
-        states.removeLastOrNull() as T? ?: init()
-    }.also { value ->
-        val valueState by rememberUpdatedState(value)
-        DisposableEffect(key) {
-            onDispose {
-                statesMap[key]?.addFirst(valueState)
-            }
-        }
-    }
-}
-
 // Find out how make this work with any generic `Serializable` type
 @Composable
 inline fun <reified T> rememberMutableStateInDataStore(
@@ -82,15 +48,3 @@ inline fun <reified T> rememberMutableStateInDataStore(
         }
     }
 }
-
-@Composable
-fun launchInVM(
-    context: CoroutineContext = EmptyCoroutineContext,
-    start: CoroutineStart = CoroutineStart.DEFAULT,
-    block: suspend CoroutineScope.() -> Unit,
-) = rememberInVM {
-    viewModelScope.launch(context, start, block)
-}
-
-@Composable
-fun <T> rememberUpdatedStateInVM(newValue: T) = rememberInVM { mutableStateOf(newValue) }.apply { value = newValue }

app/src/main/kotlin/com/hippo/ehviewer/ui/tools/ViewModelState.kt

@@ -0,0 +1,129 @@
+package com.hippo.ehviewer.ui.tools
+
+import androidx.collection.mutableIntObjectMapOf
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisallowComposableCalls
+import androidx.compose.runtime.DisposableEffect
+import androidx.compose.runtime.MutableState
+import androidx.compose.runtime.State
+import androidx.compose.runtime.currentCompositeKeyHash
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.rememberUpdatedState
+import androidx.compose.runtime.snapshotFlow
+import androidx.lifecycle.ViewModel
+import androidx.lifecycle.viewModelScope
+import androidx.lifecycle.viewmodel.compose.viewModel
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.CoroutineStart
+import kotlinx.coroutines.async
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.launch
+
+@Suppress("UNCHECKED_CAST")
+inline fun <T : Any> nonNullState(block: MutableState<T>.() -> Unit) = (mutableStateOf<T?>(null) as MutableState<T>).apply(block) as State<T>
+
+class StateMapViewModel : ViewModel() {
+    val statesMap = mutableIntObjectMapOf<ArrayDeque<Any>>()
+}
+
+@Composable
+inline fun <reified T : Any> rememberInVM(
+    crossinline init: @DisallowComposableCalls ViewModel.() -> T,
+) = with(viewModel<StateMapViewModel>()) {
+    val compositeKey = currentCompositeKeyHash
+    remember {
+        val states = statesMap.getOrPut(compositeKey, ::ArrayDeque)
+        states.removeLastOrNull() as T? ?: init()
+    }.also { value ->
+        val valueState by rememberUpdatedState(value)
+        DisposableEffect(compositeKey) {
+            onDispose {
+                statesMap[compositeKey]?.addFirst(valueState)
+            }
+        }
+    }
+}
+
+@Composable
+inline fun <reified T : Any> rememberInVM(
+    key: Any?,
+    crossinline init: @DisallowComposableCalls ViewModel.() -> T,
+) = with(viewModel<StateMapViewModel>()) {
+    val compositeKey = currentCompositeKeyHash
+    remember(key) {
+        val states = statesMap.getOrPut(compositeKey, ::ArrayDeque)
+        states.removeLastOrNull() as T? ?: init()
+    }.also { value ->
+        val valueState by rememberUpdatedState(value)
+        DisposableEffect(compositeKey) {
+            onDispose {
+                statesMap[compositeKey]?.addFirst(valueState)
+            }
+        }
+    }
+}
+
+@Composable
+fun launchInVM(
+    context: CoroutineContext = EmptyCoroutineContext,
+    start: CoroutineStart = CoroutineStart.DEFAULT,
+    block: suspend CoroutineScope.() -> Unit,
+) = rememberInVM {
+    viewModelScope.launch(context, start, block)
+}
+
+@Composable
+fun launchInVM(
+    key: Any?,
+    context: CoroutineContext = EmptyCoroutineContext,
+    start: CoroutineStart = CoroutineStart.DEFAULT,
+    block: suspend CoroutineScope.() -> Unit,
+) = rememberUpdatedStateInVM(key).let { key ->
+    val f by rememberUpdatedStateInVM(block)
+    rememberInVM {
+        nonNullState {
+            viewModelScope.launch(start = CoroutineStart.UNDISPATCHED) {
+                snapshotFlow { key.value }.mapLatest {
+                    coroutineScope { value = launch(context, start, f) }
+                }.collect()
+            }
+        }
+    }
+}
+
+@Composable
+fun <R> asyncInVM(
+    context: CoroutineContext = EmptyCoroutineContext,
+    start: CoroutineStart = CoroutineStart.DEFAULT,
+    block: suspend CoroutineScope.() -> R,
+) = rememberInVM {
+    viewModelScope.async(context, start, block)
+}
+
+@Composable
+fun <R> asyncInVM(
+    key: Any?,
+    context: CoroutineContext = EmptyCoroutineContext,
+    start: CoroutineStart = CoroutineStart.DEFAULT,
+    block: suspend CoroutineScope.() -> R,
+) = rememberUpdatedStateInVM(key).let { key ->
+    val f by rememberUpdatedStateInVM(block)
+    rememberInVM {
+        nonNullState {
+            viewModelScope.launch(start = CoroutineStart.UNDISPATCHED) {
+                snapshotFlow { key.value }.mapLatest {
+                    coroutineScope { value = async(context, start, f) }
+                }.collect()
+            }
+        }
+    }
+}
+
+@Composable
+fun <T> rememberUpdatedStateInVM(newValue: T) = rememberInVM { mutableStateOf(newValue) }.apply { value = newValue }