[Compose Sample] Mavericks and Navigation cleanup (#1763)

Use the official mavericks-compose artifact and pass in NavController rather than using a composition local.
diff --git a/sample-compose/build.gradle b/sample-compose/build.gradle
index e5b3540..6c57c64 100644
--- a/sample-compose/build.gradle
+++ b/sample-compose/build.gradle
@@ -1,7 +1,6 @@
 plugins {
   id 'com.android.application'
   id 'kotlin-android'
-  id 'kotlin-android-extensions'
   id 'kotlin-kapt'
 }
 
@@ -37,7 +36,6 @@
         "-Xuse-experimental=androidx.compose.animation.ExperimentalAnimationApi",
         "-Xopt-in=androidx.compose.material.ExperimentalMaterialApi",
         "-Xopt-in=kotlin.RequiresOptIn",
-        "-Xopt-in=com.airbnb.mvrx.InternalMavericksApi"
     ]
   }
   buildFeatures {
@@ -53,10 +51,6 @@
   correctErrorTypes = true
 }
 
-androidExtensions {
-  features = ["parcelize"]
-}
-
 dependencies {
   implementation project(':lottie-compose')
   implementation 'androidx.core:core-ktx:1.3.2'
@@ -73,7 +67,7 @@
   implementation 'androidx.lifecycle:lifecycle-runtime-ktx:2.3.0'
   implementation 'androidx.lifecycle:lifecycle-viewmodel-ktx:2.3.0'
 
-  implementation "androidx.navigation:navigation-ui-ktx:2.3.3"
+  implementation "androidx.navigation:navigation-ui-ktx:2.3.4"
 
   implementation 'org.jetbrains.kotlinx:kotlinx-coroutines-core:1.4.2'
   implementation 'org.jetbrains.kotlinx:kotlinx-coroutines-android:1.4.2'
@@ -84,9 +78,10 @@
   implementation 'com.squareup.retrofit2:retrofit:2.9.0'
   implementation 'com.squareup.retrofit2:converter-moshi:2.9.0'
   implementation "dev.chrisbanes.accompanist:accompanist-coil:0.6.0"
-  implementation 'com.airbnb.android:mavericks:2.0.0'
+  implementation 'com.airbnb.android:mavericks:2.1.0'
+  implementation 'com.airbnb.android:mavericks-compose:2.1.0-alpha01'
 
-  testImplementation 'junit:junit:4.13.1'
+  testImplementation 'junit:junit:4.13.2'
   androidTestImplementation 'androidx.test.ext:junit:1.1.2'
   androidTestImplementation 'androidx.test.espresso:espresso-core:3.3.0'
 }
\ No newline at end of file
diff --git a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/ComposeActivity.kt b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/ComposeActivity.kt
index 48ee9d4..8b01ac5 100644
--- a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/ComposeActivity.kt
+++ b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/ComposeActivity.kt
@@ -9,7 +9,6 @@
 import androidx.compose.foundation.layout.padding
 import androidx.compose.material.*
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.CompositionLocalProvider
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
@@ -25,7 +24,6 @@
 import com.airbnb.lottie.sample.compose.ui.LottieTheme
 import com.airbnb.lottie.sample.compose.ui.Teal
 import com.airbnb.lottie.sample.compose.ui.toColorSafe
-import com.airbnb.lottie.sample.compose.utils.LocalNavController
 import com.airbnb.lottie.sample.compose.utils.getBase64String
 
 class ComposeActivity : AppCompatActivity() {
@@ -41,65 +39,61 @@
     private fun LottieScaffold() {
         val navController = rememberNavController()
 
-        CompositionLocalProvider(
-            LocalNavController provides navController,
-        ) {
-            LottieTheme {
-                Scaffold(
-                    bottomBar = {
-                        BottomNavigation(
-                            backgroundColor = Color(0xFFF7F7F7),
-                            elevation = 8.dp,
-                            contentColor = Teal,
-                        ) {
-                            val navBackStackEntry by navController.currentBackStackEntryAsState()
-                            val currentRoute = navBackStackEntry?.arguments?.getString(KEY_ROUTE)
+        LottieTheme {
+            Scaffold(
+                bottomBar = {
+                    BottomNavigation(
+                        backgroundColor = Color(0xFFF7F7F7),
+                        elevation = 8.dp,
+                        contentColor = Teal,
+                    ) {
+                        val navBackStackEntry by navController.currentBackStackEntryAsState()
+                        val currentRoute = navBackStackEntry?.arguments?.getString(KEY_ROUTE)
 
-                            BottomNavItemData.values().forEach { item ->
-                                BottomNavigationItem(
-                                    icon = {
-                                        Icon(
-                                            painter = painterResource(item.iconRes),
-                                            contentDescription = null
-                                        )
-                                    },
-                                    label = { Text(stringResource(item.labelRes)) },
-                                    selected = currentRoute == item.route.route,
-                                    onClick = {
-                                        if (currentRoute != item.route.route) {
-                                            navController.navigate(item.route.route)
-                                        }
-                                    },
-                                )
-                            }
+                        BottomNavItemData.values().forEach { item ->
+                            BottomNavigationItem(
+                                icon = {
+                                    Icon(
+                                        painter = painterResource(item.iconRes),
+                                        contentDescription = null
+                                    )
+                                },
+                                label = { Text(stringResource(item.labelRes)) },
+                                selected = currentRoute == item.route.route,
+                                onClick = {
+                                    if (currentRoute != item.route.route) {
+                                        navController.navigate(item.route.route)
+                                    }
+                                },
+                            )
                         }
                     }
-                ) { innerPadding ->
-                    Box(
-                        modifier = Modifier.padding(innerPadding)
-                    ) {
-                        NavHost(navController, startDestination = Route.Showcase.route) {
-                            composable(Route.Showcase.route, arguments = Route.Showcase.args) { ShowcasePage() }
-                            composable(Route.Preview.route, arguments = Route.Preview.args) { PreviewPage() }
-                            composable(Route.LottieFiles.route, arguments = Route.LottieFiles.args) { LottieFilesPage() }
-                            composable(Route.Learn.route, arguments = Route.Learn.args) { ShowcasePage() }
-                            composable(
-                                Route.Player.fullRoute,
-                                arguments = Route.Player.args
-                            ) { entry ->
-                                val arguments = entry.arguments ?: error("No arguments provided to ${Route.Player}")
-                                val spec = when {
-                                    arguments.getString("url") != null -> LottieAnimationSpec.Url(arguments.getBase64String("url"))
-                                    arguments.getString("file") != null -> LottieAnimationSpec.File(arguments.getBase64String("file"))
-                                    arguments.getString("asset") != null -> LottieAnimationSpec.Asset(arguments.getBase64String("asset"))
-                                    else -> error("You must specify a url, file, or asset")
-                                }
-                                val backgroundColor = when (arguments.getString("backgroundColor") != null) {
-                                    true -> arguments.getBase64String("backgroundColor").toColorSafe()
-                                    else -> null
-                                }
-                                PlayerPage(spec, backgroundColor)
+                }
+            ) { innerPadding ->
+                Box(
+                    modifier = Modifier.padding(innerPadding)
+                ) {
+                    NavHost(navController, startDestination = Route.Showcase.route) {
+                        composable(Route.Showcase.route, arguments = Route.Showcase.args) { ShowcasePage(navController) }
+                        composable(Route.Preview.route, arguments = Route.Preview.args) { PreviewPage(navController) }
+                        composable(Route.LottieFiles.route, arguments = Route.LottieFiles.args) { LottieFilesPage(navController) }
+                        composable(Route.Learn.route, arguments = Route.Learn.args) { ShowcasePage(navController) }
+                        composable(
+                            Route.Player.fullRoute,
+                            arguments = Route.Player.args
+                        ) { entry ->
+                            val arguments = entry.arguments ?: error("No arguments provided to ${Route.Player}")
+                            val spec = when {
+                                arguments.getString("url") != null -> LottieAnimationSpec.Url(arguments.getBase64String("url"))
+                                arguments.getString("file") != null -> LottieAnimationSpec.File(arguments.getBase64String("file"))
+                                arguments.getString("asset") != null -> LottieAnimationSpec.Asset(arguments.getBase64String("asset"))
+                                else -> error("You must specify a url, file, or asset")
                             }
+                            val backgroundColor = when (arguments.getString("backgroundColor") != null) {
+                                true -> arguments.getBase64String("backgroundColor").toColorSafe()
+                                else -> null
+                            }
+                            PlayerPage(spec, backgroundColor)
                         }
                     }
                 }
diff --git a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/dagger/DaggerMavericksViewModelFactory.kt b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/dagger/DaggerMavericksViewModelFactory.kt
index 4d25a91..4f1c3ae 100644
--- a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/dagger/DaggerMavericksViewModelFactory.kt
+++ b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/dagger/DaggerMavericksViewModelFactory.kt
@@ -1,24 +1,27 @@
 package com.airbnb.lottie.sample.compose.dagger
 
-import androidx.fragment.app.FragmentActivity
+import androidx.activity.ComponentActivity
 import com.airbnb.lottie.sample.compose.LottieComposeApplication
-import com.airbnb.mvrx.*
+import com.airbnb.mvrx.MavericksState
+import com.airbnb.mvrx.MavericksViewModel
+import com.airbnb.mvrx.MavericksViewModelFactory
+import com.airbnb.mvrx.ViewModelContext
 
-inline fun <reified VM : MavericksViewModel<S>, S : MavericksState> daggerMavericksViewModelFactory() = DaggerMavericksViewModelFactory<VM, S>(VM::class.java)
+inline fun <reified VM : MavericksViewModel<S>, S : MavericksState> daggerMavericksViewModelFactory() =
+    DaggerMavericksViewModelFactory<VM, S>(VM::class.java)
 
 
 class DaggerMavericksViewModelFactory<VM : MavericksViewModel<S>, S : MavericksState>(
-        private val viewModelClass: Class<out MavericksViewModel<S>>
+    private val viewModelClass: Class<out MavericksViewModel<S>>
 ) : MavericksViewModelFactory<VM, S> {
 
-    override fun create(viewModelContext: ViewModelContext, state: S): VM? {
+    override fun create(viewModelContext: ViewModelContext, state: S): VM {
         return createViewModel(viewModelContext.activity, state)
     }
 
     private fun <VM : MavericksViewModel<S>, S : MavericksState> createViewModel(
-            fragmentActivity:
-            FragmentActivity,
-            initialState: S
+        fragmentActivity: ComponentActivity,
+        initialState: S
     ): VM {
         val viewModelFactoryMap = (fragmentActivity.application as LottieComposeApplication).component.viewModelFactories()
         val viewModelFactory = viewModelFactoryMap[viewModelClass]
diff --git a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/lottiefiles/LottieFilesPage.kt b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/lottiefiles/LottieFilesPage.kt
index af6942e..64ffcaf 100644
--- a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/lottiefiles/LottieFilesPage.kt
+++ b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/lottiefiles/LottieFilesPage.kt
@@ -16,6 +16,7 @@
 import androidx.compose.ui.platform.LocalDensity
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.unit.dp
+import androidx.navigation.NavController
 import com.airbnb.lottie.sample.compose.R
 import com.airbnb.lottie.sample.compose.composables.Marquee
 
@@ -26,7 +27,7 @@
 }
 
 @Composable
-fun LottieFilesPage() {
+fun LottieFilesPage(navController: NavController) {
     var tab by rememberSaveable { mutableStateOf(LottieFilesTab.Recent) }
 
     Column {
@@ -36,9 +37,9 @@
             onTabSelected = { tab = it },
         )
         when (tab) {
-            LottieFilesTab.Recent -> LottieFilesRecentAndPopularPage(LottieFilesMode.Recent)
-            LottieFilesTab.Popular -> LottieFilesRecentAndPopularPage(LottieFilesMode.Popular)
-            LottieFilesTab.Search -> LottieFilesSearchPage()
+            LottieFilesTab.Recent -> LottieFilesRecentAndPopularPage(navController, LottieFilesMode.Recent)
+            LottieFilesTab.Popular -> LottieFilesRecentAndPopularPage(navController, LottieFilesMode.Popular)
+            LottieFilesTab.Search -> LottieFilesSearchPage(navController)
         }
     }
 
diff --git a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/lottiefiles/LottieFilesRecentsAndPopularPage.kt b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/lottiefiles/LottieFilesRecentsAndPopularPage.kt
index 1348d19..36ed4be 100644
--- a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/lottiefiles/LottieFilesRecentsAndPopularPage.kt
+++ b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/lottiefiles/LottieFilesRecentsAndPopularPage.kt
@@ -13,10 +13,12 @@
 import androidx.compose.material.icons.filled.Repeat
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.SideEffect
+import androidx.compose.runtime.getValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.unit.dp
+import androidx.navigation.NavController
 import androidx.navigation.compose.navigate
 import com.airbnb.lottie.sample.compose.Route
 import com.airbnb.lottie.sample.compose.api.AnimationDataV2
@@ -24,12 +26,11 @@
 import com.airbnb.lottie.sample.compose.composables.AnimationRow
 import com.airbnb.lottie.sample.compose.dagger.AssistedViewModelFactory
 import com.airbnb.lottie.sample.compose.dagger.daggerMavericksViewModelFactory
-import com.airbnb.lottie.sample.compose.utils.collectState
-import com.airbnb.lottie.sample.compose.utils.findNavController
-import com.airbnb.lottie.sample.compose.utils.mavericksViewModel
 import com.airbnb.mvrx.MavericksState
 import com.airbnb.mvrx.MavericksViewModel
 import com.airbnb.mvrx.MavericksViewModelFactory
+import com.airbnb.mvrx.compose.collectAsState
+import com.airbnb.mvrx.compose.mavericksViewModel
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
@@ -104,10 +105,9 @@
 }
 
 @Composable
-fun LottieFilesRecentAndPopularPage(mode: LottieFilesMode) {
+fun LottieFilesRecentAndPopularPage(navController: NavController, mode: LottieFilesMode) {
     val viewModel: LottieFilesRecentAndPopularViewModel = mavericksViewModel()
-    val state = viewModel.collectState()
-    val navController = findNavController()
+    val state by viewModel.collectAsState()
     SideEffect {
         viewModel.setMode(mode)
     }
diff --git a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/lottiefiles/LottieFilesSearchPage.kt b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/lottiefiles/LottieFilesSearchPage.kt
index 4933326..1e1d34e 100644
--- a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/lottiefiles/LottieFilesSearchPage.kt
+++ b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/lottiefiles/LottieFilesSearchPage.kt
@@ -12,12 +12,14 @@
 import androidx.compose.material.icons.filled.Repeat
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.SideEffect
+import androidx.compose.runtime.getValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.tooling.preview.Preview
 import androidx.compose.ui.unit.dp
+import androidx.navigation.NavController
 import androidx.navigation.compose.navigate
 import com.airbnb.lottie.sample.compose.R
 import com.airbnb.lottie.sample.compose.Route
@@ -26,16 +28,14 @@
 import com.airbnb.lottie.sample.compose.composables.AnimationRow
 import com.airbnb.lottie.sample.compose.dagger.AssistedViewModelFactory
 import com.airbnb.lottie.sample.compose.dagger.daggerMavericksViewModelFactory
-import com.airbnb.lottie.sample.compose.utils.collectState
-import com.airbnb.lottie.sample.compose.utils.findNavController
-import com.airbnb.lottie.sample.compose.utils.mavericksViewModel
 import com.airbnb.mvrx.MavericksState
 import com.airbnb.mvrx.MavericksViewModel
 import com.airbnb.mvrx.MavericksViewModelFactory
+import com.airbnb.mvrx.compose.collectAsState
+import com.airbnb.mvrx.compose.mavericksViewModel
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
-import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.launch
 
@@ -111,10 +111,9 @@
 }
 
 @Composable
-fun LottieFilesSearchPage() {
+fun LottieFilesSearchPage(navController: NavController) {
     val viewModel: LottieFilesSearchViewModel = mavericksViewModel()
-    val state = viewModel.collectState()
-    val navController = findNavController()
+    val state by viewModel.collectAsState()
     LottieFilesSearchPage(
         state,
         viewModel::setQuery,
@@ -141,7 +140,9 @@
                 value = state.query,
                 onValueChange = onQueryChanged,
                 label = { Text(stringResource(R.string.query)) },
-                modifier = Modifier.fillMaxWidth().padding(horizontal = 16.dp),
+                modifier = Modifier
+                    .fillMaxWidth()
+                    .padding(horizontal = 16.dp),
                 singleLine = true
             )
             LazyColumn(
diff --git a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/preview/PreviewPage.kt b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/preview/PreviewPage.kt
index e67e03a..9aa1c92 100644
--- a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/preview/PreviewPage.kt
+++ b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/preview/PreviewPage.kt
@@ -18,18 +18,18 @@
 import androidx.compose.ui.unit.dp
 import androidx.compose.ui.unit.sp
 import androidx.compose.ui.window.Dialog
+import androidx.navigation.NavController
 import androidx.navigation.compose.navigate
+import androidx.navigation.compose.rememberNavController
 import com.airbnb.lottie.sample.compose.R
 import com.airbnb.lottie.sample.compose.Route
 import com.airbnb.lottie.sample.compose.composables.Marquee
 import com.airbnb.lottie.sample.compose.ui.LottieTheme
-import com.airbnb.lottie.sample.compose.utils.findNavController
 
 @Composable
-fun PreviewPage() {
+fun PreviewPage(navController: NavController) {
     var showingAssetsDialog by remember { mutableStateOf(false) }
     var showingUrlDialog by remember { mutableStateOf(false) }
-    val navController = findNavController()
 
     Column {
         Marquee(stringResource(R.string.tab_preview))
@@ -172,9 +172,10 @@
 @Preview
 @Composable
 fun PreviewPagePreview() {
+    val navController = rememberNavController()
     LottieTheme {
         Box(modifier = Modifier.background(Color.White)) {
-            PreviewPage()
+            PreviewPage(navController)
         }
     }
 }
diff --git a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/showcase/ShowcasePage.kt b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/showcase/ShowcasePage.kt
index e00de2d..852d950 100644
--- a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/showcase/ShowcasePage.kt
+++ b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/showcase/ShowcasePage.kt
@@ -6,28 +6,28 @@
 import androidx.compose.foundation.lazy.items
 import androidx.compose.material.Divider
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.getValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.tooling.preview.Preview
+import androidx.navigation.NavController
 import androidx.navigation.compose.navigate
+import androidx.navigation.compose.rememberNavController
 import com.airbnb.lottie.sample.compose.Route
 import com.airbnb.lottie.sample.compose.composables.AnimationRow
 import com.airbnb.lottie.sample.compose.composables.Loader
 import com.airbnb.lottie.sample.compose.composables.Marquee
 import com.airbnb.lottie.sample.compose.ui.LottieTheme
-import com.airbnb.lottie.sample.compose.utils.collectState
-import com.airbnb.lottie.sample.compose.utils.findNavController
-import com.airbnb.lottie.sample.compose.utils.mavericksViewModel
 import com.airbnb.mvrx.Loading
 import com.airbnb.mvrx.Uninitialized
+import com.airbnb.mvrx.compose.collectAsState
+import com.airbnb.mvrx.compose.mavericksViewModel
 
 @Composable
-fun ShowcasePage() {
+fun ShowcasePage(navController: NavController) {
     val viewModel: ShowcaseViewModel = mavericksViewModel()
-    val state = viewModel.collectState()
-    val featuredAnimations = state.animations
-    val navController = findNavController()
+    val featuredAnimations by viewModel.collectAsState(ShowcaseState::animations)
     Box(
         modifier = Modifier.fillMaxSize()
     ) {
@@ -55,7 +55,8 @@
 @Preview(showBackground = true)
 @Composable
 fun DefaultPreview() {
+    val navController = rememberNavController()
     LottieTheme {
-        ShowcasePage()
+        ShowcasePage(navController)
     }
 }
\ No newline at end of file
diff --git a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/utils/MavericksExtensions.kt b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/utils/MavericksExtensions.kt
deleted file mode 100644
index 7571440..0000000
--- a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/utils/MavericksExtensions.kt
+++ /dev/null
@@ -1,53 +0,0 @@
-package com.airbnb.lottie.sample.compose.utils
-
-
-import androidx.compose.runtime.Composable
-import androidx.compose.runtime.collectAsState
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.remember
-import androidx.compose.ui.platform.LocalContext
-import androidx.compose.ui.platform.LocalLifecycleOwner
-import androidx.fragment.app.Fragment
-import androidx.fragment.app.FragmentActivity
-import androidx.navigation.NavBackStackEntry
-import com.airbnb.mvrx.*
-
-@Composable
-fun <VM : MavericksViewModel<S>, S : MavericksState> VM.collectState(): S {
-    val state by stateFlow.collectAsState(initial = withState(this) { it })
-    return state
-}
-
-@Composable
-inline fun <reified VM : MavericksViewModel<S>, reified S : MavericksState> mavericksViewModel(): VM {
-    val viewModelClass = VM::class
-    val context = LocalContext.current
-    val viewModelContext = when (val lifecycleOwner = LocalLifecycleOwner.current) {
-        is Fragment -> {
-            val activity = lifecycleOwner.requireActivity()
-            val args = lifecycleOwner.arguments?.get(Mavericks.KEY_ARG)
-            FragmentViewModelContext(activity, args, lifecycleOwner)
-        }
-        is FragmentActivity -> {
-            val args = lifecycleOwner.intent.extras?.get(Mavericks.KEY_ARG)
-            ActivityViewModelContext(lifecycleOwner, args)
-        }
-        is NavBackStackEntry -> {
-            val args = lifecycleOwner.arguments?.get(Mavericks.KEY_ARG)
-            val activity = context as? FragmentActivity ?: error("Local context should be a FragmentActivity but it is a ${context::class.simpleName}!")
-            ActivityViewModelContext(activity, args)
-        }
-        else -> error("Unknown LifecycleOwner ${lifecycleOwner::class.simpleName}. Must be Fragment or Activity for now.")
-    }
-    val activity = context as? FragmentActivity ?: error("Composable is not hosted in a FragmentActivity")
-    return remember(viewModelClass, activity) {
-        val keyFactory = { viewModelClass.java.name }
-        @Suppress("RestrictedApi")
-        MavericksViewModelProvider.get(
-            viewModelClass = viewModelClass.java,
-            stateClass = S::class.java,
-            viewModelContext = viewModelContext,
-            key = keyFactory()
-        )
-    }
-}
\ No newline at end of file
diff --git a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/utils/NavigationExtensions.kt b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/utils/NavigationExtensions.kt
index 976ae33..17a5871 100644
--- a/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/utils/NavigationExtensions.kt
+++ b/sample-compose/src/main/java/com/airbnb/lottie/sample/compose/utils/NavigationExtensions.kt
@@ -2,14 +2,6 @@
 
 import android.os.Bundle
 import android.util.Base64
-import androidx.compose.runtime.Composable
-import androidx.compose.runtime.staticCompositionLocalOf
-import androidx.navigation.NavController
 import java.nio.charset.StandardCharsets
 
-val LocalNavController = staticCompositionLocalOf<NavController> { error("You must specify a NavController.") }
-
-@Composable
-fun findNavController() = LocalNavController.current
-
 fun Bundle.getBase64String(key: String) = String(Base64.decode(getString(key), Base64.DEFAULT), StandardCharsets.UTF_8)