From bb2f79cc003df213d34eda642f43e252b9a737ca Mon Sep 17 00:00:00 2001
From: Gabriel Souza <devsrsouza@gmail.com>
Date: Mon, 7 Oct 2024 10:30:47 -0300
Subject: [PATCH] Adding onDispose callback api to Experimental
 LifecycleEffectOnce

---
 .../basicNavigation/BasicNavigationScreen.kt  |  3 +++
 .../core/lifecycle/LifecycleEffectStore.kt    | 16 +++++++++++-----
 .../voyager/core/lifecycle/ScreenLifecycle.kt | 19 ++++++++++++++++---
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/samples/android/src/main/java/cafe/adriel/voyager/sample/basicNavigation/BasicNavigationScreen.kt b/samples/android/src/main/java/cafe/adriel/voyager/sample/basicNavigation/BasicNavigationScreen.kt
index 722b2132..50e4673a 100644
--- a/samples/android/src/main/java/cafe/adriel/voyager/sample/basicNavigation/BasicNavigationScreen.kt
+++ b/samples/android/src/main/java/cafe/adriel/voyager/sample/basicNavigation/BasicNavigationScreen.kt
@@ -39,6 +39,9 @@ data class BasicNavigationScreen(
         )
         LifecycleEffectOnce {
             Log.d("Navigator", "On screen first appear #$index")
+            onDispose {
+                Log.d("Navigator", "On screen dispose")
+            }
         }
 
         val navigator = LocalNavigator.currentOrThrow
diff --git a/voyager-core/src/commonMain/kotlin/cafe/adriel/voyager/core/lifecycle/LifecycleEffectStore.kt b/voyager-core/src/commonMain/kotlin/cafe/adriel/voyager/core/lifecycle/LifecycleEffectStore.kt
index def1a68e..0435a89a 100644
--- a/voyager-core/src/commonMain/kotlin/cafe/adriel/voyager/core/lifecycle/LifecycleEffectStore.kt
+++ b/voyager-core/src/commonMain/kotlin/cafe/adriel/voyager/core/lifecycle/LifecycleEffectStore.kt
@@ -6,17 +6,23 @@ import cafe.adriel.voyager.core.screen.Screen
 import cafe.adriel.voyager.core.screen.ScreenKey
 
 internal object LifecycleEffectStore : ScreenDisposable {
-    private val executedLifecycles = ThreadSafeMap<ScreenKey, ThreadSafeSet<String>>()
+    private val executedLifecycles = ThreadSafeMap<ScreenKey, ThreadSafeSet<LifecycleEffectOnceScope>>()
 
-    fun store(screen: Screen, effectKey: String) {
+    fun store(screen: Screen, effectKey: String): LifecycleEffectOnceScope {
         val set = executedLifecycles.getOrPut(screen.key) { ThreadSafeSet() }
-        set.add(effectKey)
+        val scope = LifecycleEffectOnceScope(uniqueKey = effectKey, set.size + 1)
+        set.add(scope)
+
+        return scope
     }
 
     fun hasExecuted(screen: Screen, effectKey: String): Boolean =
-        executedLifecycles.get(screen.key)?.contains(effectKey) == true
+        executedLifecycles.get(screen.key)?.any { it.uniqueKey == effectKey } == true
 
     override fun onDispose(screen: Screen) {
-        executedLifecycles.remove(screen.key)
+        val scopes = executedLifecycles.remove(screen.key)
+        scopes?.sortedBy { it.registerOrderIndex }?.reversed()?.forEach { scope ->
+            scope.onDisposed?.invoke()
+        }
     }
 }
diff --git a/voyager-core/src/commonMain/kotlin/cafe/adriel/voyager/core/lifecycle/ScreenLifecycle.kt b/voyager-core/src/commonMain/kotlin/cafe/adriel/voyager/core/lifecycle/ScreenLifecycle.kt
index 6c29ab10..dc4580cf 100644
--- a/voyager-core/src/commonMain/kotlin/cafe/adriel/voyager/core/lifecycle/ScreenLifecycle.kt
+++ b/voyager-core/src/commonMain/kotlin/cafe/adriel/voyager/core/lifecycle/ScreenLifecycle.kt
@@ -23,9 +23,22 @@ public fun Screen.LifecycleEffect(
     }
 }
 
+@ExperimentalVoyagerApi
+public data class LifecycleEffectOnceScope(
+    val uniqueKey: String,
+    val registerOrderIndex: Int,
+) {
+    internal var onDisposed: (() -> Unit)? = null
+
+    @ExperimentalVoyagerApi
+    public fun onDispose(onDisposed: () -> Unit) {
+        this.onDisposed = onDisposed
+    }
+}
+
 @ExperimentalVoyagerApi
 @Composable
-public fun Screen.LifecycleEffectOnce(onFirstAppear: () -> Unit) {
+public fun Screen.LifecycleEffectOnce(onFirstAppear: LifecycleEffectOnceScope.() -> Unit) {
     val uniqueCompositionKey = rememberSaveable { randomUuid() }
 
     val lifecycleEffectStore = remember {
@@ -34,8 +47,8 @@ public fun Screen.LifecycleEffectOnce(onFirstAppear: () -> Unit) {
 
     LaunchedEffect(Unit) {
         if (lifecycleEffectStore.hasExecuted(this@LifecycleEffectOnce, uniqueCompositionKey).not()) {
-            lifecycleEffectStore.store(this@LifecycleEffectOnce, uniqueCompositionKey)
-            onFirstAppear()
+            val scope = lifecycleEffectStore.store(this@LifecycleEffectOnce, uniqueCompositionKey)
+            onFirstAppear(scope)
         }
     }
 }