From a6690ff19f8030897d50c5a422891cde76304ed2 Mon Sep 17 00:00:00 2001
From: Thiago dos Santos <thiagodossantos@BRSAOMN044385.local>
Date: Sun, 24 Sep 2023 14:16:09 -0300
Subject: [PATCH] feat: androidx view model scoped support too

---
 .../androidViewModel/AndroidDetailsScreen.kt  |   5 +-
 .../hiltIntegration/HiltDetailsScreen.kt      |  10 +-
 .../hiltIntegration/HiltDetailsViewModel.kt   |   6 +
 .../hiltIntegration/HiltListViewModel.kt      |   5 +
 .../androidx/AndroidScreenLifecycleOwner.kt   | 188 ++----------------
 .../androidx/VoyagerAndroidLifecycleOwner.kt  | 163 +++++++++++++++
 .../cafe/adriel/voyager/ext/ContextExt.kt     |  18 ++
 voyager-hilt/build.gradle.kts                 |   1 +
 .../cafe/adriel/voyager/hilt/ScreenModel.kt   |  12 +-
 .../cafe/adriel/voyager/hilt/ViewModel.kt     |  89 +++++++--
 .../voyager/hilt/internal/ContextExt.kt       |  28 ---
 voyager-navigator/build.gradle.kts            |   1 +
 .../navigator/viewmodel/NavigatorViewModel.kt |  56 ++++++
 .../viewmodel/NavigatorViewModelStore.kt      |  19 ++
 .../navigator/model/NavigatorScreenModel.kt   |  71 +++++++
 .../model/NavigatorScreenModelStore.kt        |  33 +++
 16 files changed, 474 insertions(+), 231 deletions(-)
 create mode 100644 voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/androidx/VoyagerAndroidLifecycleOwner.kt
 create mode 100644 voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/ext/ContextExt.kt
 delete mode 100644 voyager-hilt/src/main/java/cafe/adriel/voyager/hilt/internal/ContextExt.kt
 create mode 100644 voyager-navigator/src/androidMain/kotlin/cafe/adriel/voyager/navigator/viewmodel/NavigatorViewModel.kt
 create mode 100644 voyager-navigator/src/androidMain/kotlin/cafe/adriel/voyager/navigator/viewmodel/NavigatorViewModelStore.kt
 create mode 100644 voyager-navigator/src/commonMain/kotlin/cafe/adriel/voyager/navigator/model/NavigatorScreenModel.kt
 create mode 100644 voyager-navigator/src/commonMain/kotlin/cafe/adriel/voyager/navigator/model/NavigatorScreenModelStore.kt

diff --git a/samples/android/src/main/java/cafe/adriel/voyager/sample/androidViewModel/AndroidDetailsScreen.kt b/samples/android/src/main/java/cafe/adriel/voyager/sample/androidViewModel/AndroidDetailsScreen.kt
index 4812f492..80108a8e 100644
--- a/samples/android/src/main/java/cafe/adriel/voyager/sample/androidViewModel/AndroidDetailsScreen.kt
+++ b/samples/android/src/main/java/cafe/adriel/voyager/sample/androidViewModel/AndroidDetailsScreen.kt
@@ -1,12 +1,11 @@
 package cafe.adriel.voyager.sample.androidViewModel
 
 import androidx.compose.runtime.Composable
+import androidx.lifecycle.viewmodel.compose.viewModel
 import cafe.adriel.voyager.androidx.AndroidScreen
 import cafe.adriel.voyager.navigator.LocalNavigator
 import cafe.adriel.voyager.navigator.currentOrThrow
 import cafe.adriel.voyager.sample.DetailsContent
-import org.koin.androidx.compose.getViewModel
-import org.koin.core.parameter.parametersOf
 
 data class AndroidDetailsScreen(
     val index: Int
@@ -15,7 +14,7 @@ data class AndroidDetailsScreen(
     @Composable
     override fun Content() {
         val navigator = LocalNavigator.currentOrThrow
-        val viewModel = getViewModel<AndroidDetailsViewModel> { parametersOf(index) }
+        val viewModel = viewModel { AndroidDetailsViewModel(index) }
 
         DetailsContent(viewModel, "Item #${viewModel.index}", navigator::pop)
     }
diff --git a/samples/android/src/main/java/cafe/adriel/voyager/sample/hiltIntegration/HiltDetailsScreen.kt b/samples/android/src/main/java/cafe/adriel/voyager/sample/hiltIntegration/HiltDetailsScreen.kt
index cfe32a8c..6e13e052 100644
--- a/samples/android/src/main/java/cafe/adriel/voyager/sample/hiltIntegration/HiltDetailsScreen.kt
+++ b/samples/android/src/main/java/cafe/adriel/voyager/sample/hiltIntegration/HiltDetailsScreen.kt
@@ -2,7 +2,7 @@ package cafe.adriel.voyager.sample.hiltIntegration
 
 import androidx.compose.runtime.Composable
 import cafe.adriel.voyager.androidx.AndroidScreen
-import cafe.adriel.voyager.hilt.getScreenModel
+import cafe.adriel.voyager.hilt.getViewModel
 import cafe.adriel.voyager.navigator.LocalNavigator
 import cafe.adriel.voyager.navigator.currentOrThrow
 import cafe.adriel.voyager.sample.DetailsContent
@@ -17,15 +17,15 @@ data class HiltDetailsScreen(
 
         // Uncomment version below if you want keep using ViewModel instead of to convert it to ScreenModel
         // ViewModelProvider.Factory is not required. Until now Hilt has no support to Assisted Injection by default
-        /*val viewModel: HiltDetailsViewModel = getViewModel(
+        val viewModel: HiltDetailsViewModel = getViewModel(
             viewModelProviderFactory = HiltDetailsViewModel.provideFactory(index)
-        )*/
+        )
 
         // This version include more boilerplate because we are simulating support
         // to Assisted Injection using ScreenModel. See [HiltListScreen] for a simple version
-        val viewModel = getScreenModel<HiltDetailsScreenModel, HiltDetailsScreenModel.Factory> { factory ->
+        /*val viewModel = getScreenModel<HiltDetailsScreenModel, HiltDetailsScreenModel.Factory> { factory ->
             factory.create(index)
-        }
+        }*/
 
         DetailsContent(viewModel, "Item #${viewModel.index}", navigator::pop)
     }
diff --git a/samples/android/src/main/java/cafe/adriel/voyager/sample/hiltIntegration/HiltDetailsViewModel.kt b/samples/android/src/main/java/cafe/adriel/voyager/sample/hiltIntegration/HiltDetailsViewModel.kt
index c7d08f8a..f5b5c35c 100644
--- a/samples/android/src/main/java/cafe/adriel/voyager/sample/hiltIntegration/HiltDetailsViewModel.kt
+++ b/samples/android/src/main/java/cafe/adriel/voyager/sample/hiltIntegration/HiltDetailsViewModel.kt
@@ -1,5 +1,6 @@
 package cafe.adriel.voyager.sample.hiltIntegration
 
+import android.util.Log
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 
@@ -8,6 +9,11 @@ import androidx.lifecycle.ViewModelProvider
 class HiltDetailsViewModel(
     val index: Int
 ) : ViewModel() {
+
+    override fun onCleared() {
+        Log.d(">> TAG <<", "HiltDetailsViewModel cleared with index: $index")
+    }
+
     companion object {
         fun provideFactory(
             index: Int
diff --git a/samples/android/src/main/java/cafe/adriel/voyager/sample/hiltIntegration/HiltListViewModel.kt b/samples/android/src/main/java/cafe/adriel/voyager/sample/hiltIntegration/HiltListViewModel.kt
index ebeb971d..63b2fde5 100644
--- a/samples/android/src/main/java/cafe/adriel/voyager/sample/hiltIntegration/HiltListViewModel.kt
+++ b/samples/android/src/main/java/cafe/adriel/voyager/sample/hiltIntegration/HiltListViewModel.kt
@@ -1,5 +1,6 @@
 package cafe.adriel.voyager.sample.hiltIntegration
 
+import android.util.Log
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.ViewModel
 import cafe.adriel.voyager.sample.sampleItems
@@ -19,4 +20,8 @@ class HiltListViewModel @Inject constructor(
 
     val items: List<String>
         get() = handle["items"] ?: error("Items not found")
+
+    override fun onCleared() {
+        Log.d(">> TAG <<", "HiltListViewModel cleared")
+    }
 }
diff --git a/voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/androidx/AndroidScreenLifecycleOwner.kt b/voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/androidx/AndroidScreenLifecycleOwner.kt
index 8db4fb42..64df49e5 100644
--- a/voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/androidx/AndroidScreenLifecycleOwner.kt
+++ b/voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/androidx/AndroidScreenLifecycleOwner.kt
@@ -1,123 +1,32 @@
 package cafe.adriel.voyager.androidx
 
-import android.app.Activity
-import android.app.Application
-import android.content.Context
-import android.content.ContextWrapper
-import android.os.Bundle
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
-import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.ProvidedValue
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
-import androidx.compose.runtime.saveable.rememberSaveable
-import androidx.compose.runtime.setValue
 import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.platform.LocalLifecycleOwner
 import androidx.compose.ui.platform.LocalSavedStateRegistryOwner
-import androidx.lifecycle.DefaultLifecycleObserver
-import androidx.lifecycle.HasDefaultViewModelProviderFactory
-import androidx.lifecycle.Lifecycle
-import androidx.lifecycle.LifecycleOwner
-import androidx.lifecycle.LifecycleRegistry
-import androidx.lifecycle.SAVED_STATE_REGISTRY_OWNER_KEY
-import androidx.lifecycle.SavedStateViewModelFactory
-import androidx.lifecycle.VIEW_MODEL_STORE_OWNER_KEY
-import androidx.lifecycle.ViewModelProvider
-import androidx.lifecycle.ViewModelProvider.AndroidViewModelFactory
-import androidx.lifecycle.ViewModelStore
-import androidx.lifecycle.ViewModelStoreOwner
-import androidx.lifecycle.enableSavedStateHandles
-import androidx.lifecycle.viewmodel.CreationExtras
-import androidx.lifecycle.viewmodel.MutableCreationExtras
 import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
-import androidx.savedstate.SavedStateRegistry
-import androidx.savedstate.SavedStateRegistryController
-import androidx.savedstate.SavedStateRegistryOwner
 import cafe.adriel.voyager.core.lifecycle.ScreenLifecycleOwner
 import cafe.adriel.voyager.core.lifecycle.ScreenLifecycleStore
 import cafe.adriel.voyager.core.screen.Screen
-import java.util.concurrent.atomic.AtomicReference
 
-public class AndroidScreenLifecycleOwner private constructor() :
-    ScreenLifecycleOwner,
-    LifecycleOwner,
-    ViewModelStoreOwner,
-    SavedStateRegistryOwner,
-    HasDefaultViewModelProviderFactory {
+public class AndroidScreenLifecycleOwner : ScreenLifecycleOwner {
 
-    override val lifecycle: LifecycleRegistry = LifecycleRegistry(this)
-
-    override val viewModelStore: ViewModelStore = ViewModelStore()
-
-    private val atomicContext = AtomicReference<Context>()
-
-    private val controller = SavedStateRegistryController.create(this)
-
-    private var deactivateLifecycleListener: (() -> Unit)? = null
-
-    private var isCreated: Boolean by mutableStateOf(false)
-
-    override val savedStateRegistry: SavedStateRegistry
-        get() = controller.savedStateRegistry
-
-    override val defaultViewModelProviderFactory: ViewModelProvider.Factory
-        get() = SavedStateViewModelFactory(
-            application = atomicContext.get()?.applicationContext?.getApplication(),
-            owner = this
-        )
-
-    override val defaultViewModelCreationExtras: CreationExtras
-        get() = MutableCreationExtras().apply {
-            val application = atomicContext.get()?.applicationContext?.getApplication()
-            if (application != null) {
-                set(AndroidViewModelFactory.APPLICATION_KEY, application)
-            }
-            set(SAVED_STATE_REGISTRY_OWNER_KEY, this@AndroidScreenLifecycleOwner)
-            set(VIEW_MODEL_STORE_OWNER_KEY, this@AndroidScreenLifecycleOwner)
-
-            /* TODO if (getArguments() != null) {
-                extras.set<Bundle>(DEFAULT_ARGS_KEY, getArguments())
-            }*/
-        }
-
-    init {
-        controller.performAttach()
-        enableSavedStateHandles()
-    }
-
-    private fun onCreate(savedState: Bundle?) {
-        check(!isCreated) { "onCreate already called" }
-        isCreated = true
-        controller.performRestore(savedState)
-        initEvents.forEach {
-            lifecycle.handleLifecycleEvent(it)
-        }
-    }
-
-    private fun onStart() {
-        startEvents.forEach {
-            lifecycle.handleLifecycleEvent(it)
-        }
-    }
-
-    private fun onStop() {
-        deactivateLifecycleListener?.invoke()
-        deactivateLifecycleListener = null
-        stopEvents.forEach {
-            lifecycle.handleLifecycleEvent(it)
-        }
-    }
+    private var voyagerAndroidLifecycleOwner: VoyagerAndroidLifecycleOwner? = null
 
     @Composable
     override fun ProvideBeforeScreenContent(
         provideSaveableState: @Composable (suffixKey: String, content: @Composable () -> Unit) -> Unit,
         content: @Composable () -> Unit
     ) {
+        if (voyagerAndroidLifecycleOwner == null) {
+            voyagerAndroidLifecycleOwner = VoyagerAndroidLifecycleOwner(LocalContext.current)
+        }
+
         provideSaveableState("lifecycle") {
-            LifecycleDisposableEffect()
+            voyagerAndroidLifecycleOwner!!.LifecycleDisposableEffect()
 
             val hooks = getHooks()
 
@@ -128,93 +37,26 @@ public class AndroidScreenLifecycleOwner private constructor() :
     }
 
     override fun onDispose(screen: Screen) {
-        val context = atomicContext.getAndSet(null) ?: return
-        if (context is Activity && context.isChangingConfigurations) return
-        viewModelStore.clear()
-        disposeEvents.forEach {
-            lifecycle.handleLifecycleEvent(it)
-        }
-    }
-
-    private fun performSave(outState: Bundle) {
-        controller.performSave(outState)
+        voyagerAndroidLifecycleOwner?.dispose()
+        voyagerAndroidLifecycleOwner = null
     }
 
     @Composable
     private fun getHooks(): List<ProvidedValue<*>> {
-        atomicContext.compareAndSet(null, LocalContext.current)
+        val androidLifecycleOwner = checkNotNull(voyagerAndroidLifecycleOwner) {
+            "ProvideBeforeScreenContent was not called before getHooks"
+        }
 
         return remember(this) {
             listOf(
-                LocalLifecycleOwner provides this,
-                LocalViewModelStoreOwner provides this,
-                LocalSavedStateRegistryOwner provides this
+                LocalLifecycleOwner provides androidLifecycleOwner,
+                LocalViewModelStoreOwner provides androidLifecycleOwner,
+                LocalSavedStateRegistryOwner provides androidLifecycleOwner
             )
         }
     }
 
-    private fun registerLifecycleListener(outState: Bundle) {
-        val activity = atomicContext.get()?.getActivity()
-        if (activity != null && activity is LifecycleOwner) {
-            val observer = object : DefaultLifecycleObserver {
-                override fun onStop(owner: LifecycleOwner) {
-                    performSave(outState)
-                }
-            }
-            val lifecycle = activity.lifecycle
-            lifecycle.addObserver(observer)
-            deactivateLifecycleListener = { lifecycle.removeObserver(observer) }
-        }
-    }
-
-    @Composable
-    private fun LifecycleDisposableEffect() {
-        val savedState = rememberSaveable { Bundle() }
-        if (!isCreated) {
-            onCreate(savedState) // do this in the UI thread to force it to be called before anything else
-        }
-
-        DisposableEffect(this) {
-            registerLifecycleListener(savedState)
-            onStart()
-            onDispose {
-                performSave(savedState)
-                onStop()
-            }
-        }
-    }
-
-    private tailrec fun Context.getActivity(): Activity? = when (this) {
-        is Activity -> this
-        is ContextWrapper -> baseContext.getActivity()
-        else -> null
-    }
-
-    private tailrec fun Context.getApplication(): Application? = when (this) {
-        is Application -> this
-        is ContextWrapper -> baseContext.getApplication()
-        else -> null
-    }
-
     public companion object {
-
-        private val initEvents = arrayOf(
-            Lifecycle.Event.ON_CREATE
-        )
-
-        private val startEvents = arrayOf(
-            Lifecycle.Event.ON_START,
-            Lifecycle.Event.ON_RESUME
-        )
-
-        private val stopEvents = arrayOf(
-            Lifecycle.Event.ON_PAUSE,
-            Lifecycle.Event.ON_STOP
-        )
-
-        private val disposeEvents = arrayOf(
-            Lifecycle.Event.ON_DESTROY
-        )
         public fun get(screen: Screen): ScreenLifecycleOwner {
             return ScreenLifecycleStore.register(screen) { AndroidScreenLifecycleOwner() }
         }
diff --git a/voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/androidx/VoyagerAndroidLifecycleOwner.kt b/voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/androidx/VoyagerAndroidLifecycleOwner.kt
new file mode 100644
index 00000000..42766be0
--- /dev/null
+++ b/voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/androidx/VoyagerAndroidLifecycleOwner.kt
@@ -0,0 +1,163 @@
+package cafe.adriel.voyager.androidx
+
+import android.content.Context
+import android.os.Bundle
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisposableEffect
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.saveable.rememberSaveable
+import androidx.compose.runtime.setValue
+import androidx.lifecycle.DefaultLifecycleObserver
+import androidx.lifecycle.HasDefaultViewModelProviderFactory
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.LifecycleRegistry
+import androidx.lifecycle.SAVED_STATE_REGISTRY_OWNER_KEY
+import androidx.lifecycle.SavedStateViewModelFactory
+import androidx.lifecycle.VIEW_MODEL_STORE_OWNER_KEY
+import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.ViewModelProvider.AndroidViewModelFactory
+import androidx.lifecycle.ViewModelStore
+import androidx.lifecycle.ViewModelStoreOwner
+import androidx.lifecycle.enableSavedStateHandles
+import androidx.lifecycle.viewmodel.CreationExtras
+import androidx.lifecycle.viewmodel.MutableCreationExtras
+import androidx.savedstate.SavedStateRegistry
+import androidx.savedstate.SavedStateRegistryController
+import androidx.savedstate.SavedStateRegistryOwner
+import cafe.adriel.voyager.ext.application
+import cafe.adriel.voyager.ext.componentActivity
+
+public class VoyagerAndroidLifecycleOwner(
+    private val context: Context
+) : LifecycleOwner,
+    ViewModelStoreOwner,
+    SavedStateRegistryOwner,
+    HasDefaultViewModelProviderFactory {
+
+    override val lifecycle: LifecycleRegistry = LifecycleRegistry(this)
+
+    override val viewModelStore: ViewModelStore = ViewModelStore()
+
+    private val controller = SavedStateRegistryController.create(this)
+
+    private var deactivateLifecycleListener: (() -> Unit)? = null
+
+    private var isCreated: Boolean by mutableStateOf(false)
+
+    override val savedStateRegistry: SavedStateRegistry
+        get() = controller.savedStateRegistry
+
+    override val defaultViewModelProviderFactory: ViewModelProvider.Factory
+        get() = SavedStateViewModelFactory(
+            application = context.application,
+            owner = this
+        )
+
+    override val defaultViewModelCreationExtras: CreationExtras
+        get() = MutableCreationExtras().apply {
+            val application = context.application
+            if (application != null) {
+                set(AndroidViewModelFactory.APPLICATION_KEY, application)
+            }
+            set(SAVED_STATE_REGISTRY_OWNER_KEY, this@VoyagerAndroidLifecycleOwner)
+            set(VIEW_MODEL_STORE_OWNER_KEY, this@VoyagerAndroidLifecycleOwner)
+
+            /* TODO if (getArguments() != null) {
+                extras.set<Bundle>(DEFAULT_ARGS_KEY, getArguments())
+            }*/
+        }
+
+    init {
+        controller.performAttach()
+        enableSavedStateHandles()
+    }
+
+    @Composable
+    public fun LifecycleDisposableEffect() {
+        val savedState = rememberSaveable { Bundle() }
+        if (!isCreated) {
+            isCreated = true
+            onCreate(savedState) // do this in the UI thread to force it to be called before anything else
+        }
+
+        DisposableEffect(this) {
+            registerLifecycleListener(savedState)
+            onStart()
+            onDispose {
+                performSave(savedState)
+                onStop()
+            }
+        }
+    }
+
+    public fun dispose() {
+        val activity = context.componentActivity
+        if (activity?.isChangingConfigurations == true) return
+        viewModelStore.clear()
+        disposeEvents.forEach {
+            lifecycle.handleLifecycleEvent(it)
+        }
+    }
+
+    private fun onCreate(savedState: Bundle?) {
+        controller.performRestore(savedState)
+        initEvents.forEach {
+            lifecycle.handleLifecycleEvent(it)
+        }
+    }
+
+    private fun onStart() {
+        startEvents.forEach {
+            lifecycle.handleLifecycleEvent(it)
+        }
+    }
+
+    private fun onStop() {
+        deactivateLifecycleListener?.invoke()
+        deactivateLifecycleListener = null
+        stopEvents.forEach {
+            lifecycle.handleLifecycleEvent(it)
+        }
+    }
+
+    private fun performSave(outState: Bundle) {
+        controller.performSave(outState)
+    }
+
+    private fun registerLifecycleListener(outState: Bundle) {
+        val activity = context.componentActivity
+        if (activity is LifecycleOwner) {
+            val observer = object : DefaultLifecycleObserver {
+                override fun onStop(owner: LifecycleOwner) {
+                    performSave(outState)
+                }
+            }
+            val lifecycle = activity.lifecycle
+            lifecycle.addObserver(observer)
+            deactivateLifecycleListener = { lifecycle.removeObserver(observer) }
+        }
+    }
+
+    private companion object {
+
+        private val initEvents = arrayOf(
+            Lifecycle.Event.ON_CREATE
+        )
+
+        private val startEvents = arrayOf(
+            Lifecycle.Event.ON_START,
+            Lifecycle.Event.ON_RESUME
+        )
+
+        private val stopEvents = arrayOf(
+            Lifecycle.Event.ON_PAUSE,
+            Lifecycle.Event.ON_STOP
+        )
+
+        private val disposeEvents = arrayOf(
+            Lifecycle.Event.ON_DESTROY
+        )
+    }
+}
diff --git a/voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/ext/ContextExt.kt b/voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/ext/ContextExt.kt
new file mode 100644
index 00000000..f109d529
--- /dev/null
+++ b/voyager-core/src/androidMain/kotlin/cafe/adriel/voyager/ext/ContextExt.kt
@@ -0,0 +1,18 @@
+package cafe.adriel.voyager.ext
+
+import android.app.Application
+import android.content.Context
+import android.content.ContextWrapper
+import androidx.activity.ComponentActivity
+
+public inline fun <reified T> Context.findOwner(
+    noinline nextFunction: (Context) -> Context? = { (it as? ContextWrapper)?.baseContext }
+): T? = generateSequence(seed = this, nextFunction = nextFunction).mapNotNull { context ->
+    context as? T
+}.firstOrNull()
+
+public val Context.application: Application?
+    get() = findOwner<Application> { it.applicationContext }
+
+public val Context.componentActivity: ComponentActivity?
+    get() = findOwner<ComponentActivity>()
diff --git a/voyager-hilt/build.gradle.kts b/voyager-hilt/build.gradle.kts
index 61d99376..cbab21b5 100644
--- a/voyager-hilt/build.gradle.kts
+++ b/voyager-hilt/build.gradle.kts
@@ -25,6 +25,7 @@ dependencies {
     implementation(libs.compose.ui)
     implementation(libs.lifecycle.savedState)
     implementation(libs.lifecycle.viewModelKtx)
+    implementation(libs.lifecycle.viewModelCompose)
     implementation(libs.hilt.android)
     kapt(libs.hilt.compiler)
 
diff --git a/voyager-hilt/src/main/java/cafe/adriel/voyager/hilt/ScreenModel.kt b/voyager-hilt/src/main/java/cafe/adriel/voyager/hilt/ScreenModel.kt
index e08461e5..df94e8da 100644
--- a/voyager-hilt/src/main/java/cafe/adriel/voyager/hilt/ScreenModel.kt
+++ b/voyager-hilt/src/main/java/cafe/adriel/voyager/hilt/ScreenModel.kt
@@ -5,7 +5,7 @@ import androidx.compose.ui.platform.LocalContext
 import cafe.adriel.voyager.core.model.ScreenModel
 import cafe.adriel.voyager.core.model.rememberScreenModel
 import cafe.adriel.voyager.core.screen.Screen
-import cafe.adriel.voyager.hilt.internal.componentActivity
+import cafe.adriel.voyager.ext.componentActivity
 import dagger.hilt.android.EntryPointAccessors
 
 /**
@@ -19,8 +19,11 @@ public inline fun <reified T : ScreenModel> Screen.getScreenModel(
 ): T {
     val context = LocalContext.current
     return rememberScreenModel(tag) {
+        val activity = checkNotNull(context.componentActivity) {
+            "No Activity found in the context: $context"
+        }
         val screenModels = EntryPointAccessors
-            .fromActivity(context.componentActivity, ScreenModelEntryPoint::class.java)
+            .fromActivity(activity, ScreenModelEntryPoint::class.java)
             .screenModels()
         val model = screenModels[T::class.java]?.get()
             ?: error(
@@ -45,8 +48,11 @@ public inline fun <reified T : ScreenModel, reified F : ScreenModelFactory> Scre
 ): T {
     val context = LocalContext.current
     return rememberScreenModel(tag) {
+        val activity = checkNotNull(context.componentActivity) {
+            "No Activity found in the context: $context"
+        }
         val screenFactories = EntryPointAccessors
-            .fromActivity(context.componentActivity, ScreenModelEntryPoint::class.java)
+            .fromActivity(activity, ScreenModelEntryPoint::class.java)
             .screenModelFactories()
         val screenFactory = screenFactories[F::class.java]?.get()
             ?: error(
diff --git a/voyager-hilt/src/main/java/cafe/adriel/voyager/hilt/ViewModel.kt b/voyager-hilt/src/main/java/cafe/adriel/voyager/hilt/ViewModel.kt
index f35ff51a..c9fd0f14 100644
--- a/voyager-hilt/src/main/java/cafe/adriel/voyager/hilt/ViewModel.kt
+++ b/voyager-hilt/src/main/java/cafe/adriel/voyager/hilt/ViewModel.kt
@@ -3,41 +3,92 @@ package cafe.adriel.voyager.hilt
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.remember
 import androidx.compose.ui.platform.LocalContext
+import androidx.lifecycle.HasDefaultViewModelProviderFactory
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
-import androidx.lifecycle.ViewModelStore
-import cafe.adriel.voyager.androidx.AndroidScreenLifecycleOwner
-import cafe.adriel.voyager.core.lifecycle.ScreenLifecycleProvider
+import androidx.lifecycle.ViewModelStoreOwner
+import androidx.lifecycle.viewmodel.CreationExtras
+import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
 import cafe.adriel.voyager.core.screen.Screen
-import cafe.adriel.voyager.hilt.internal.componentActivity
+import cafe.adriel.voyager.ext.componentActivity
 
 /**
- * A function to provide a [dagger.hilt.android.lifecycle.HiltViewModel] managed by voyager ViewModelLifecycleOwner
- * instead of using Activity ViewModelLifecycleOwner.
- * There is compatibility with Activity ViewModelLifecycleOwner too but it must be avoided because your ViewModels
- * will be cleared when activity is totally destroyed only.
+ * A function to provide a [dagger.hilt.android.lifecycle.HiltViewModel] managed by voyager [Screen]
+ * [cafe.adriel.voyager.androidx.VoyagerAndroidLifecycleOwner] instead of using Activity or Fragment LifecycleOwner.
  *
- * @param viewModelProviderFactory A custom factory commonly used with Assisted Injection
- * @return A new instance of [ViewModel] or the existent instance in the [ViewModelStore]
+ * @param key The key to use to identify the [ViewModel].
+ * @param viewModelProviderFactory The [ViewModelProvider.Factory] that should be used to create the [ViewModel]
+ * @param viewModelStoreOwner The owner of the [ViewModel] that controls the scope and lifetime
+ * of the returned [ViewModel]. Defaults to using [LocalViewModelStoreOwner].
+ * or null if you would like to use the default factory from the [LocalViewModelStoreOwner]
+ * @param extras The default extras used to create the [ViewModel].
+ *
+ * @return A [ViewModel] that is an instance of the given [T] type.
  */
 @Composable
+public inline fun <reified T : ViewModel> Screen.hiltViewModel(
+    key: Any? = T::class,
+    viewModelProviderFactory: ViewModelProvider.Factory? = null,
+    viewModelStoreOwner: ViewModelStoreOwner = checkNotNull(LocalViewModelStoreOwner.current) {
+        "No ViewModelStoreOwner was provided via LocalViewModelStoreOwner"
+    },
+    extras: CreationExtras = if (viewModelStoreOwner is HasDefaultViewModelProviderFactory) {
+        viewModelStoreOwner.defaultViewModelCreationExtras
+    } else {
+        CreationExtras.Empty
+    }
+): T = getViewModel(
+    key = key,
+    viewModelProviderFactory = viewModelProviderFactory,
+    viewModelStoreOwner = viewModelStoreOwner,
+    extras = extras
+)
+
+/**
+ * A function to provide a [dagger.hilt.android.lifecycle.HiltViewModel] managed by voyager [Screen]
+ * [cafe.adriel.voyager.androidx.VoyagerAndroidLifecycleOwner] instead of using Activity or Fragment LifecycleOwner.
+ *
+ * @param key The key to use to identify the [ViewModel].
+ * @param viewModelProviderFactory The [ViewModelProvider.Factory] that should be used to create the [ViewModel]
+ * @param viewModelStoreOwner The owner of the [ViewModel] that controls the scope and lifetime
+ * of the returned [ViewModel]. Defaults to using [LocalViewModelStoreOwner].
+ * or null if you would like to use the default factory from the [LocalViewModelStoreOwner]
+ * @param extras The default extras used to create the [ViewModel].
+ *
+ * @return A [ViewModel] that is an instance of the given [T] type.
+ */
+@Suppress("UnusedReceiverParameter")
+@Composable
 public inline fun <reified T : ViewModel> Screen.getViewModel(
-    viewModelProviderFactory: ViewModelProvider.Factory? = null
+    key: Any? = T::class,
+    viewModelProviderFactory: ViewModelProvider.Factory? = null,
+    viewModelStoreOwner: ViewModelStoreOwner = checkNotNull(LocalViewModelStoreOwner.current) {
+        "No ViewModelStoreOwner was provided via LocalViewModelStoreOwner"
+    },
+    extras: CreationExtras = if (viewModelStoreOwner is HasDefaultViewModelProviderFactory) {
+        viewModelStoreOwner.defaultViewModelCreationExtras
+    } else {
+        CreationExtras.Empty
+    }
 ): T {
     val context = LocalContext.current
-    return remember(key1 = T::class) {
-        val activity = context.componentActivity
-        val lifecycleOwner = (this as? ScreenLifecycleProvider)
-            ?.getLifecycleOwner() as? AndroidScreenLifecycleOwner
-            ?: activity
+    return remember(key1 = key) {
+        val activity = checkNotNull(context.componentActivity) {
+            "No Activity found in the context: $context"
+        }
+        val delegateFactory = when {
+            viewModelProviderFactory != null -> viewModelProviderFactory
+            viewModelStoreOwner is HasDefaultViewModelProviderFactory -> viewModelStoreOwner.defaultViewModelProviderFactory
+            else -> error("A custom or default ViewModelProvider.Factory is required")
+        }
         val factory = VoyagerHiltViewModelFactories.getVoyagerFactory(
             activity = activity,
-            delegateFactory = viewModelProviderFactory ?: lifecycleOwner.defaultViewModelProviderFactory
+            delegateFactory = delegateFactory
         )
         val provider = ViewModelProvider(
-            store = lifecycleOwner.viewModelStore,
+            store = viewModelStoreOwner.viewModelStore,
             factory = factory,
-            defaultCreationExtras = lifecycleOwner.defaultViewModelCreationExtras
+            defaultCreationExtras = extras
         )
         provider[T::class.java]
     }
diff --git a/voyager-hilt/src/main/java/cafe/adriel/voyager/hilt/internal/ContextExt.kt b/voyager-hilt/src/main/java/cafe/adriel/voyager/hilt/internal/ContextExt.kt
deleted file mode 100644
index 88659969..00000000
--- a/voyager-hilt/src/main/java/cafe/adriel/voyager/hilt/internal/ContextExt.kt
+++ /dev/null
@@ -1,28 +0,0 @@
-package cafe.adriel.voyager.hilt.internal
-
-import android.content.Context
-import android.content.ContextWrapper
-import androidx.activity.ComponentActivity
-import androidx.lifecycle.ViewModelProvider
-
-// Unfortunately findOwner function is internal in activity-compose
-// TODO: Maybe move to androidx module because we'll need this function when implement onCloseRequest support
-internal inline fun <reified T> findOwner(context: Context): T? {
-    var innerContext = context
-    while (innerContext is ContextWrapper) {
-        if (innerContext is T) {
-            return innerContext
-        }
-        innerContext = innerContext.baseContext
-    }
-    return null
-}
-
-@PublishedApi
-internal val Context.componentActivity: ComponentActivity
-    get() = findOwner<ComponentActivity>(this)
-        ?: error("Context must be a androidx.activity.ComponentActivity. Current is $this")
-
-@PublishedApi
-internal val Context.defaultViewModelProviderFactory: ViewModelProvider.Factory
-    get() = componentActivity.defaultViewModelProviderFactory
diff --git a/voyager-navigator/build.gradle.kts b/voyager-navigator/build.gradle.kts
index b717c6eb..1d518eb0 100644
--- a/voyager-navigator/build.gradle.kts
+++ b/voyager-navigator/build.gradle.kts
@@ -31,6 +31,7 @@ kotlin {
         val androidMain by getting {
             dependencies {
                 implementation(libs.compose.activity)
+                implementation(libs.lifecycle.viewModelCompose)
             }
         }
     }
diff --git a/voyager-navigator/src/androidMain/kotlin/cafe/adriel/voyager/navigator/viewmodel/NavigatorViewModel.kt b/voyager-navigator/src/androidMain/kotlin/cafe/adriel/voyager/navigator/viewmodel/NavigatorViewModel.kt
new file mode 100644
index 00000000..e3fcb460
--- /dev/null
+++ b/voyager-navigator/src/androidMain/kotlin/cafe/adriel/voyager/navigator/viewmodel/NavigatorViewModel.kt
@@ -0,0 +1,56 @@
+package cafe.adriel.voyager.navigator.viewmodel
+
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.platform.LocalContext
+import androidx.lifecycle.HasDefaultViewModelProviderFactory
+import androidx.lifecycle.ViewModel
+import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.ViewModelStoreOwner
+import androidx.lifecycle.viewmodel.CreationExtras
+import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
+import androidx.lifecycle.viewmodel.compose.viewModel
+import cafe.adriel.voyager.core.screen.Screen
+import cafe.adriel.voyager.navigator.LocalNavigator
+import cafe.adriel.voyager.navigator.currentOrThrow
+import cafe.adriel.voyager.navigator.lifecycle.NavigatorLifecycleStore
+
+@Suppress("UnusedReceiverParameter")
+public val Screen.navigatorViewModelStoreOwner: ViewModelStoreOwner
+    @Composable get() {
+        val context = LocalContext.current
+        val navigator = LocalNavigator.currentOrThrow
+        val manager = NavigatorLifecycleStore.register(
+            navigator = navigator,
+            factory = { NavigatorViewModelStore(context) }
+        )
+        return manager.viewModelStoreOwner
+    }
+
+/**
+ * Remember a [ViewModel] that will be scoped to current [LocalNavigator]
+ *
+ * @param key The key to use to identify the [ViewModel].
+ * @param factory The [ViewModelProvider.Factory] that should be used to create the [ViewModel]
+ * or null if you would like to use the default factory from the [LocalViewModelStoreOwner]
+ * @param extras The default extras used to create the [ViewModel].
+ *
+ * @return A [ViewModel] that is an instance of the given [T] type.
+ */
+@Composable
+public inline fun <reified T : ViewModel> Screen.rememberNavigatorViewModel(
+    key: String? = null,
+    factory: ViewModelProvider.Factory? = null,
+    extras: CreationExtras? = null
+): T {
+    val storeOwner = navigatorViewModelStoreOwner
+    return viewModel(
+        key = key,
+        factory = factory,
+        viewModelStoreOwner = storeOwner,
+        extras = when {
+            extras != null -> extras
+            storeOwner is HasDefaultViewModelProviderFactory -> storeOwner.defaultViewModelCreationExtras
+            else -> CreationExtras.Empty
+        }
+    )
+}
diff --git a/voyager-navigator/src/androidMain/kotlin/cafe/adriel/voyager/navigator/viewmodel/NavigatorViewModelStore.kt b/voyager-navigator/src/androidMain/kotlin/cafe/adriel/voyager/navigator/viewmodel/NavigatorViewModelStore.kt
new file mode 100644
index 00000000..40dd8773
--- /dev/null
+++ b/voyager-navigator/src/androidMain/kotlin/cafe/adriel/voyager/navigator/viewmodel/NavigatorViewModelStore.kt
@@ -0,0 +1,19 @@
+package cafe.adriel.voyager.navigator.viewmodel
+
+import android.content.Context
+import androidx.compose.runtime.Composable
+import androidx.lifecycle.ViewModelStoreOwner
+import cafe.adriel.voyager.androidx.VoyagerAndroidLifecycleOwner
+import cafe.adriel.voyager.navigator.Navigator
+import cafe.adriel.voyager.navigator.lifecycle.NavigatorDisposable
+
+public class NavigatorViewModelStore(context: Context) : NavigatorDisposable {
+    private val voyagerAndroidLifecycleOwner = VoyagerAndroidLifecycleOwner(context)
+
+    public val viewModelStoreOwner: ViewModelStoreOwner
+        @Composable get() = voyagerAndroidLifecycleOwner.apply { LifecycleDisposableEffect() }
+
+    override fun onDispose(navigator: Navigator) {
+        voyagerAndroidLifecycleOwner.dispose()
+    }
+}
diff --git a/voyager-navigator/src/commonMain/kotlin/cafe/adriel/voyager/navigator/model/NavigatorScreenModel.kt b/voyager-navigator/src/commonMain/kotlin/cafe/adriel/voyager/navigator/model/NavigatorScreenModel.kt
new file mode 100644
index 00000000..3fbd77fc
--- /dev/null
+++ b/voyager-navigator/src/commonMain/kotlin/cafe/adriel/voyager/navigator/model/NavigatorScreenModel.kt
@@ -0,0 +1,71 @@
+package cafe.adriel.voyager.navigator.model
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisallowComposableCalls
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.setValue
+import cafe.adriel.voyager.core.model.ScreenModel
+import cafe.adriel.voyager.core.platform.multiplatformName
+import cafe.adriel.voyager.core.screen.Screen
+import cafe.adriel.voyager.navigator.LocalNavigator
+import cafe.adriel.voyager.navigator.currentOrThrow
+import cafe.adriel.voyager.navigator.lifecycle.NavigatorLifecycleStore
+
+/**
+ * Lookup for a [ScreenModel] that was remembered in the current [LocalNavigator] or in your parents.
+ * If an instance was not found, an [IllegalStateException] will be thrown
+ *
+ * @param tag A custom tag used to "tag" the [ScreenModel]
+ * @throws IllegalStateException if now [ScreenModel] was found in [LocalNavigator] or in your parents
+ *
+ * @return The [ScreenModel] instance found
+ */
+@Suppress("UnusedReceiverParameter")
+@Throws(IllegalStateException::class)
+@Composable
+public inline fun <reified T : ScreenModel> Screen.navigatorScreenModel(
+    tag: String? = null
+): T {
+    val navigator = LocalNavigator.currentOrThrow
+    var screenModel: T? by remember(tag) { mutableStateOf(null) }
+    if (screenModel == null) {
+        var currentNavigator = navigator
+        do {
+            val key = "${currentNavigator.key}:${T::class.multiplatformName}:${tag ?: "default"}"
+            val manager = NavigatorLifecycleStore.register(
+                navigator = currentNavigator,
+                factory = { NavigatorScreenModelStore() }
+            )
+            screenModel = manager.get(key)
+            currentNavigator = currentNavigator.parent ?: break
+        } while (screenModel == null)
+    }
+    return checkNotNull(screenModel) {
+        "${T::class} was not found in $navigator and it parents"
+    }
+}
+
+/**
+ * Remember a [ScreenModel] that will be scoped to current [LocalNavigator]
+ *
+ * @param tag A custom tag used to "tag" the [ScreenModel]
+ * @param factory A function to create a new instance if one is not remembered yet
+ *
+ * @return The [ScreenModel] instance
+ */
+@Suppress("UnusedReceiverParameter")
+@Composable
+public inline fun <reified T : ScreenModel> Screen.rememberNavigatorScreenModel(
+    tag: String? = null,
+    noinline factory: @DisallowComposableCalls () -> T
+): T {
+    val navigator = LocalNavigator.currentOrThrow
+    val manager = NavigatorLifecycleStore.register(
+        navigator = navigator,
+        factory = { NavigatorScreenModelStore() }
+    )
+    val key = "${navigator.key}:${T::class.multiplatformName}:${tag ?: "default"}"
+    return manager.getOrPut(key, factory)
+}
diff --git a/voyager-navigator/src/commonMain/kotlin/cafe/adriel/voyager/navigator/model/NavigatorScreenModelStore.kt b/voyager-navigator/src/commonMain/kotlin/cafe/adriel/voyager/navigator/model/NavigatorScreenModelStore.kt
new file mode 100644
index 00000000..f8bf1ed1
--- /dev/null
+++ b/voyager-navigator/src/commonMain/kotlin/cafe/adriel/voyager/navigator/model/NavigatorScreenModelStore.kt
@@ -0,0 +1,33 @@
+package cafe.adriel.voyager.navigator.model
+
+import cafe.adriel.voyager.core.concurrent.ThreadSafeMap
+import cafe.adriel.voyager.core.model.ScreenModel
+import cafe.adriel.voyager.navigator.Navigator
+import cafe.adriel.voyager.navigator.lifecycle.NavigatorDisposable
+
+public class NavigatorScreenModelStore : NavigatorDisposable {
+    private val screenModels: MutableMap<String, ScreenModel> = ThreadSafeMap()
+
+    override fun onDispose(navigator: Navigator) {
+        screenModels.forEach { entry ->
+            entry.value.onDispose()
+        }
+        screenModels.clear()
+    }
+
+    @Suppress("UNCHECKED_CAST")
+    public fun <T : ScreenModel> get(key: String): T? = screenModels[key] as? T
+
+    @Suppress("UNCHECKED_CAST")
+    public fun <T : ScreenModel> getOrPut(
+        key: String,
+        factory: () -> T
+    ): T {
+        var model = screenModels[key] as? T
+        if (model == null) {
+            model = factory()
+            screenModels[key] = model
+        }
+        return model
+    }
+}