Move viewModel() and LocalViewModelStoreOwner from compose:ui:ui to a lifecycle:lifecycle-viewmodel-compose

Relnote: viewModel() composable and LocalViewModelStoreOwner were moved to androidx.lifecycle.viewmodel.compose. You will now need to add a separate dependency androidx.lifecycle:lifecycle-viewmodel.compose:1.0.0-alpha01 in order to use it.
Test: moved existing tests
Change-Id: I7a374b76168a6387e585337c131a988bddcb912b
diff --git a/compose/integration-tests/docs-snippets/build.gradle b/compose/integration-tests/docs-snippets/build.gradle
index 68bb102..8259db5 100644
--- a/compose/integration-tests/docs-snippets/build.gradle
+++ b/compose/integration-tests/docs-snippets/build.gradle
@@ -40,6 +40,7 @@
     implementation project(":compose:ui:ui-viewbinding")
     implementation project(":navigation:navigation-compose")
     implementation project(":activity:activity-compose")
+    implementation project(":lifecycle:lifecycle-viewmodel-compose")
 
 
     implementation(KOTLIN_STDLIB)
diff --git a/compose/integration-tests/docs-snippets/src/main/java/androidx/compose/integration/docs/interoperability/Interoperability.kt b/compose/integration-tests/docs-snippets/src/main/java/androidx/compose/integration/docs/interoperability/Interoperability.kt
index 509abfd..ca8e4eb 100644
--- a/compose/integration-tests/docs-snippets/src/main/java/androidx/compose/integration/docs/interoperability/Interoperability.kt
+++ b/compose/integration-tests/docs-snippets/src/main/java/androidx/compose/integration/docs/interoperability/Interoperability.kt
@@ -84,9 +84,9 @@
 import androidx.compose.ui.unit.dp
 import androidx.compose.ui.viewinterop.AndroidView
 import androidx.compose.ui.viewinterop.AndroidViewBinding
-import androidx.compose.ui.viewinterop.viewModel
 import androidx.lifecycle.MutableLiveData
 import androidx.lifecycle.ViewModel
+import androidx.lifecycle.viewmodel.compose.viewModel
 
 /**
  * This file lets DevRel track changes to snippets present in
@@ -226,7 +226,8 @@
     }
 }
 
-@Composable private fun RowScope.InteropSnippet7() {
+@Composable
+private fun RowScope.InteropSnippet7() {
     Text(
         text = stringResource(R.string.ok),
         modifier = Modifier.padding(dimensionResource(R.dimen.padding_small))
@@ -270,6 +271,7 @@
         }
     }
 }
+
 /* ktlint-enable indent */
 private object InteropSnippet10 {
     class ExampleViewModel : ViewModel() { /*...*/ }
@@ -719,21 +721,26 @@
     object layout {
         const val fragment_example = 1
     }
+
     object id {
         const val compose_view = 2
         const val compose_view_x = 3
     }
+
     object string {
         const val ok = 4
         const val plane_description = 5
         const val login = 6
     }
+
     object dimen {
         const val padding_small = 7
     }
+
     object drawable {
         const val ic_plane = 8
     }
+
     object color {
         const val Blue700 = 9
     }
@@ -745,12 +752,15 @@
     val coordinator = Coord()
     lateinit var myView: View
 }
+
 private class DataExample(val title: String = "")
+
 private val data = DataExample()
 private fun startActivity(): Nothing = TODO()
 class ExampleViewModel : ViewModel() {
     val exampleLiveData = MutableLiveData(" ")
 }
+
 private fun ShowData(dataExample: State<String?>): Nothing = TODO()
 private class ExampleImageLoader {
     fun load(url: String): DummyInto = TODO()
@@ -759,21 +769,44 @@
     open class Listener {
         open fun onSuccess(bitmap: Bitmap): Unit = TODO()
     }
+
     companion object {
         fun get() = ExampleImageLoader()
     }
 }
+
 private class DummyInto {
-    fun into(listener: ExampleImageLoader.Listener) { }
+    fun into(listener: ExampleImageLoader.Listener) {}
 }
-private fun ExampleComposable() { }
-@Composable private fun MdcTheme(content: @Composable () -> Unit) { }
-@Composable private fun AppCompatTheme(content: @Composable () -> Unit) { }
-@Composable private fun BlueTheme(content: @Composable () -> Unit) { }
-@Composable private fun PinkTheme(content: @Composable () -> Unit) { }
-@Composable private fun YourAppTheme(content: @Composable () -> Unit) { }
-@Composable private fun ProvideWindowInsets(content: @Composable () -> Unit) { }
-@Composable private fun Icon() { }
+
+private fun ExampleComposable() {}
+@Composable
+private fun MdcTheme(content: @Composable () -> Unit) {
+}
+
+@Composable
+private fun AppCompatTheme(content: @Composable () -> Unit) {
+}
+
+@Composable
+private fun BlueTheme(content: @Composable () -> Unit) {
+}
+
+@Composable
+private fun PinkTheme(content: @Composable () -> Unit) {
+}
+
+@Composable
+private fun YourAppTheme(content: @Composable () -> Unit) {
+}
+
+@Composable
+private fun ProvideWindowInsets(content: @Composable () -> Unit) {
+}
+
+@Composable
+private fun Icon() {
+}
 
 private open class Fragment {
 
@@ -785,15 +818,17 @@
     ): View {
         TODO("not implemented")
     }
+
     fun requireContext(): Context = TODO()
 }
 
 private class AppCompatActivity {
     val window: Any = Any()
 }
+
 private class WindowCompat {
     companion object {
-        fun setDecorFitsSystemWindows(window: Any, bool: Boolean) { }
+        fun setDecorFitsSystemWindows(window: Any, bool: Boolean) {}
     }
 }
 
diff --git a/compose/integration-tests/docs-snippets/src/main/java/androidx/compose/integration/docs/state/State.kt b/compose/integration-tests/docs-snippets/src/main/java/androidx/compose/integration/docs/state/State.kt
index f4cafe1..954c142 100644
--- a/compose/integration-tests/docs-snippets/src/main/java/androidx/compose/integration/docs/state/State.kt
+++ b/compose/integration-tests/docs-snippets/src/main/java/androidx/compose/integration/docs/state/State.kt
@@ -49,10 +49,10 @@
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.unit.dp
-import androidx.compose.ui.viewinterop.viewModel
 import androidx.lifecycle.LiveData
 import androidx.lifecycle.MutableLiveData
 import androidx.lifecycle.ViewModel
+import androidx.lifecycle.viewmodel.compose.viewModel
 
 /**
  * This file lets DevRel track changes to snippets present in
diff --git a/compose/ui/ui/api/current.txt b/compose/ui/ui/api/current.txt
index c36481e..10251fc 100644
--- a/compose/ui/ui/api/current.txt
+++ b/compose/ui/ui/api/current.txt
@@ -2083,7 +2083,7 @@
     method public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.lifecycle.LifecycleOwner> getLocalLifecycleOwner();
     method public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.savedstate.SavedStateRegistryOwner> getLocalSavedStateRegistryOwner();
     method public static androidx.compose.runtime.ProvidableCompositionLocal<android.view.View> getLocalView();
-    method public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.lifecycle.ViewModelStoreOwner> getLocalViewModelStoreOwner();
+    method @Deprecated public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.lifecycle.ViewModelStoreOwner> getLocalViewModelStoreOwner();
   }
 
   public final class AndroidClipboardManagerKt {
@@ -2724,7 +2724,7 @@
   }
 
   public final class ViewModelKt {
-    method @androidx.compose.runtime.Composable public static <T extends androidx.lifecycle.ViewModel> T viewModel(Class<T> modelClass, optional String? key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
+    method @Deprecated @androidx.compose.runtime.Composable public static <T extends androidx.lifecycle.ViewModel> T viewModel(Class<T> modelClass, optional String? key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
     method @androidx.compose.runtime.Composable public static inline <reified T extends androidx.lifecycle.ViewModel> T! viewModel(optional String key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
   }
 
diff --git a/compose/ui/ui/api/public_plus_experimental_current.txt b/compose/ui/ui/api/public_plus_experimental_current.txt
index c36481e..10251fc 100644
--- a/compose/ui/ui/api/public_plus_experimental_current.txt
+++ b/compose/ui/ui/api/public_plus_experimental_current.txt
@@ -2083,7 +2083,7 @@
     method public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.lifecycle.LifecycleOwner> getLocalLifecycleOwner();
     method public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.savedstate.SavedStateRegistryOwner> getLocalSavedStateRegistryOwner();
     method public static androidx.compose.runtime.ProvidableCompositionLocal<android.view.View> getLocalView();
-    method public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.lifecycle.ViewModelStoreOwner> getLocalViewModelStoreOwner();
+    method @Deprecated public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.lifecycle.ViewModelStoreOwner> getLocalViewModelStoreOwner();
   }
 
   public final class AndroidClipboardManagerKt {
@@ -2724,7 +2724,7 @@
   }
 
   public final class ViewModelKt {
-    method @androidx.compose.runtime.Composable public static <T extends androidx.lifecycle.ViewModel> T viewModel(Class<T> modelClass, optional String? key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
+    method @Deprecated @androidx.compose.runtime.Composable public static <T extends androidx.lifecycle.ViewModel> T viewModel(Class<T> modelClass, optional String? key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
     method @androidx.compose.runtime.Composable public static inline <reified T extends androidx.lifecycle.ViewModel> T! viewModel(optional String key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
   }
 
diff --git a/compose/ui/ui/api/restricted_current.txt b/compose/ui/ui/api/restricted_current.txt
index 0152d28..d88f197 100644
--- a/compose/ui/ui/api/restricted_current.txt
+++ b/compose/ui/ui/api/restricted_current.txt
@@ -2160,7 +2160,7 @@
     method public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.lifecycle.LifecycleOwner> getLocalLifecycleOwner();
     method public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.savedstate.SavedStateRegistryOwner> getLocalSavedStateRegistryOwner();
     method public static androidx.compose.runtime.ProvidableCompositionLocal<android.view.View> getLocalView();
-    method public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.lifecycle.ViewModelStoreOwner> getLocalViewModelStoreOwner();
+    method @Deprecated public static androidx.compose.runtime.ProvidableCompositionLocal<androidx.lifecycle.ViewModelStoreOwner> getLocalViewModelStoreOwner();
   }
 
   public final class AndroidClipboardManagerKt {
@@ -2803,7 +2803,7 @@
   }
 
   public final class ViewModelKt {
-    method @androidx.compose.runtime.Composable public static <T extends androidx.lifecycle.ViewModel> T viewModel(Class<T> modelClass, optional String? key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
+    method @Deprecated @androidx.compose.runtime.Composable public static <T extends androidx.lifecycle.ViewModel> T viewModel(Class<T> modelClass, optional String? key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
     method @androidx.compose.runtime.Composable public static inline <reified T extends androidx.lifecycle.ViewModel> T! viewModel(optional String key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
   }
 
diff --git a/compose/ui/ui/integration-tests/ui-demos/src/main/java/androidx/compose/ui/demos/UiDemos.kt b/compose/ui/ui/integration-tests/ui-demos/src/main/java/androidx/compose/ui/demos/UiDemos.kt
index 460dcdf..28c6c4a 100644
--- a/compose/ui/ui/integration-tests/ui-demos/src/main/java/androidx/compose/ui/demos/UiDemos.kt
+++ b/compose/ui/ui/integration-tests/ui-demos/src/main/java/androidx/compose/ui/demos/UiDemos.kt
@@ -124,7 +124,6 @@
         ComposableDemo("Popup") { PopupDemo() },
         GraphicsDemos,
         GestureDemos,
-        ComposableDemo("Views interoperability") { ViewInteropDemo() },
-        ComposableDemo("ViewModel") { ViewModelDemo() }
+        ComposableDemo("Views interoperability") { ViewInteropDemo() }
     )
 )
diff --git a/compose/ui/ui/integration-tests/ui-demos/src/main/java/androidx/compose/ui/demos/ViewModelDemo.kt b/compose/ui/ui/integration-tests/ui-demos/src/main/java/androidx/compose/ui/demos/ViewModelDemo.kt
deleted file mode 100644
index 4414ae5..0000000
--- a/compose/ui/ui/integration-tests/ui-demos/src/main/java/androidx/compose/ui/demos/ViewModelDemo.kt
+++ /dev/null
@@ -1,57 +0,0 @@
-/*
- * Copyright 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package androidx.compose.ui.demos
-
-import androidx.compose.foundation.layout.Arrangement
-import androidx.compose.foundation.layout.Column
-import androidx.compose.material.Button
-import androidx.compose.material.Text
-import androidx.compose.runtime.Composable
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.livedata.observeAsState
-import androidx.compose.ui.Alignment
-import androidx.compose.ui.viewinterop.viewModel
-import androidx.lifecycle.LiveData
-import androidx.lifecycle.SavedStateHandle
-import androidx.lifecycle.ViewModel
-
-@Composable
-fun ViewModelDemo() {
-    Column(
-        horizontalAlignment = Alignment.CenterHorizontally,
-        verticalArrangement = Arrangement.SpaceEvenly
-    ) {
-        val countViewModel = viewModel<CountViewModel>()
-        val count by countViewModel.count.observeAsState()
-        Text("Count is $count")
-        Button(onClick = { countViewModel.increaseCount() }) {
-            Text("Increase")
-        }
-    }
-}
-
-class CountViewModel(private val savedStateHandle: SavedStateHandle) : ViewModel() {
-
-    private val CountKey = "${javaClass.simpleName}-count"
-
-    val count: LiveData<Int> =
-        savedStateHandle.getLiveData(CountKey, 0)
-
-    fun increaseCount() {
-        savedStateHandle[CountKey] = (count.value ?: 0) + 1
-    }
-}
diff --git a/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/AndroidAmbients.kt b/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/AndroidAmbients.kt
index 6f0bbfc..a6323477 100644
--- a/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/AndroidAmbients.kt
+++ b/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/AndroidAmbients.kt
@@ -132,6 +132,7 @@
 /**
  * The CompositionLocal containing the current [ViewModelStoreOwner].
  */
+@Suppress("DEPRECATION")
 @Deprecated(
     "Renamed to LocalViewModelStoreOwner",
     replaceWith = ReplaceWith(
@@ -144,6 +145,12 @@
 /**
  * The CompositionLocal containing the current [ViewModelStoreOwner].
  */
+@Suppress("DeprecatedCallableAddReplaceWith")
+@Deprecated(
+    "It was moved to androidx.lifecycle.viewmodel.compose package. You should add a " +
+        "dependency on androidx.lifecycle:lifecycle-viewmodel.compose:1.0.0-alpha01 in order to " +
+        "use it"
+)
 val LocalViewModelStoreOwner = staticCompositionLocalOf<ViewModelStoreOwner>()
 
 @Composable
@@ -190,6 +197,7 @@
         LocalSavedStateRegistryOwner provides viewTreeOwners.savedStateRegistryOwner,
         LocalSaveableStateRegistry provides saveableStateRegistry,
         LocalView provides owner.view,
+        @Suppress("DEPRECATION")
         LocalViewModelStoreOwner provides viewTreeOwners.viewModelStoreOwner
     ) {
         ProvideCommonCompositionLocals(
diff --git a/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/viewinterop/ViewModel.kt b/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/viewinterop/ViewModel.kt
index 5e8c29794..cb6f4ca 100644
--- a/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/viewinterop/ViewModel.kt
+++ b/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/viewinterop/ViewModel.kt
@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+@file:Suppress("DeprecatedCallableAddReplaceWith", "DEPRECATION")
+
 package androidx.compose.ui.viewinterop
 
 import androidx.compose.runtime.Composable
@@ -34,6 +36,11 @@
  * @return A [ViewModel] that is an instance of the given [T] type.
  */
 @Composable
+@Deprecated(
+    "It was moved to androidx.lifecycle.viewmodel.compose package. You should add a " +
+        "dependency on androidx.lifecycle:lifecycle-viewmodel.compose:1.0.0-alpha01 in order to " +
+        "use it"
+)
 inline fun <reified T : ViewModel> viewModel(
     key: String? = null,
     factory: ViewModelProvider.Factory? = null
@@ -53,6 +60,11 @@
  * @return A [ViewModel] that is an instance of the given [T] type.
  */
 @Composable
+@Deprecated(
+    "It was moved to androidx.lifecycle.viewmodel.compose package. You should add a " +
+        "dependency on androidx.lifecycle:lifecycle-viewmodel.compose:1.0.0-alpha01 in order to " +
+        "use it"
+)
 fun <T : ViewModel> viewModel(
     modelClass: Class<T>,
     key: String? = null,
diff --git a/lifecycle/lifecycle-viewmodel-compose/api/current.txt b/lifecycle/lifecycle-viewmodel-compose/api/current.txt
index e6f50d0..72357ec 100644
--- a/lifecycle/lifecycle-viewmodel-compose/api/current.txt
+++ b/lifecycle/lifecycle-viewmodel-compose/api/current.txt
@@ -1 +1,17 @@
 // Signature format: 4.0
+package androidx.lifecycle.viewmodel.compose {
+
+  public final class LocalViewModelStoreOwner {
+    method public androidx.compose.runtime.ProvidableCompositionLocal<androidx.lifecycle.ViewModelStoreOwner> asProvidableCompositionLocal();
+    method @androidx.compose.runtime.Composable public androidx.lifecycle.ViewModelStoreOwner getCurrent();
+    property @androidx.compose.runtime.Composable public final androidx.lifecycle.ViewModelStoreOwner current;
+    field public static final androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner INSTANCE;
+  }
+
+  public final class ViewModelKt {
+    method @androidx.compose.runtime.Composable public static <T extends androidx.lifecycle.ViewModel> T viewModel(Class<T> modelClass, optional String? key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
+    method @androidx.compose.runtime.Composable public static inline <reified T extends androidx.lifecycle.ViewModel> T! viewModel(optional String key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
+  }
+
+}
+
diff --git a/lifecycle/lifecycle-viewmodel-compose/api/public_plus_experimental_current.txt b/lifecycle/lifecycle-viewmodel-compose/api/public_plus_experimental_current.txt
index e6f50d0..72357ec 100644
--- a/lifecycle/lifecycle-viewmodel-compose/api/public_plus_experimental_current.txt
+++ b/lifecycle/lifecycle-viewmodel-compose/api/public_plus_experimental_current.txt
@@ -1 +1,17 @@
 // Signature format: 4.0
+package androidx.lifecycle.viewmodel.compose {
+
+  public final class LocalViewModelStoreOwner {
+    method public androidx.compose.runtime.ProvidableCompositionLocal<androidx.lifecycle.ViewModelStoreOwner> asProvidableCompositionLocal();
+    method @androidx.compose.runtime.Composable public androidx.lifecycle.ViewModelStoreOwner getCurrent();
+    property @androidx.compose.runtime.Composable public final androidx.lifecycle.ViewModelStoreOwner current;
+    field public static final androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner INSTANCE;
+  }
+
+  public final class ViewModelKt {
+    method @androidx.compose.runtime.Composable public static <T extends androidx.lifecycle.ViewModel> T viewModel(Class<T> modelClass, optional String? key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
+    method @androidx.compose.runtime.Composable public static inline <reified T extends androidx.lifecycle.ViewModel> T! viewModel(optional String key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
+  }
+
+}
+
diff --git a/lifecycle/lifecycle-viewmodel-compose/api/restricted_current.txt b/lifecycle/lifecycle-viewmodel-compose/api/restricted_current.txt
index e6f50d0..72357ec 100644
--- a/lifecycle/lifecycle-viewmodel-compose/api/restricted_current.txt
+++ b/lifecycle/lifecycle-viewmodel-compose/api/restricted_current.txt
@@ -1 +1,17 @@
 // Signature format: 4.0
+package androidx.lifecycle.viewmodel.compose {
+
+  public final class LocalViewModelStoreOwner {
+    method public androidx.compose.runtime.ProvidableCompositionLocal<androidx.lifecycle.ViewModelStoreOwner> asProvidableCompositionLocal();
+    method @androidx.compose.runtime.Composable public androidx.lifecycle.ViewModelStoreOwner getCurrent();
+    property @androidx.compose.runtime.Composable public final androidx.lifecycle.ViewModelStoreOwner current;
+    field public static final androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner INSTANCE;
+  }
+
+  public final class ViewModelKt {
+    method @androidx.compose.runtime.Composable public static <T extends androidx.lifecycle.ViewModel> T viewModel(Class<T> modelClass, optional String? key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
+    method @androidx.compose.runtime.Composable public static inline <reified T extends androidx.lifecycle.ViewModel> T! viewModel(optional String key, optional androidx.lifecycle.ViewModelProvider.Factory? factory);
+  }
+
+}
+
diff --git a/lifecycle/lifecycle-viewmodel-compose/build.gradle b/lifecycle/lifecycle-viewmodel-compose/build.gradle
index a23a7ff..a89b2a0 100644
--- a/lifecycle/lifecycle-viewmodel-compose/build.gradle
+++ b/lifecycle/lifecycle-viewmodel-compose/build.gradle
@@ -33,8 +33,19 @@
     kotlinPlugin projectOrArtifact(":compose:compiler:compiler")
 
     api "androidx.lifecycle:lifecycle-viewmodel-ktx:2.3.0"
+    api projectOrArtifact(":compose:runtime:runtime")
+    api projectOrArtifact(":compose:ui:ui")
 
     implementation(KOTLIN_STDLIB)
+
+    androidTestImplementation projectOrArtifact(":compose:ui:ui-test-junit4")
+    androidTestImplementation(ANDROIDX_TEST_RULES)
+    androidTestImplementation(ANDROIDX_TEST_RUNNER)
+    androidTestImplementation(JUNIT)
+    androidTestImplementation(TRUTH)
+    androidTestImplementation "androidx.fragment:fragment:1.2.4"
+    androidTestImplementation "androidx.appcompat:appcompat:1.1.0"
+    androidTestImplementation projectOrArtifact(":activity:activity-compose")
 }
 
 androidx {
diff --git a/lifecycle/lifecycle-viewmodel-compose/src/androidTest/AndroidManifest.xml b/lifecycle/lifecycle-viewmodel-compose/src/androidTest/AndroidManifest.xml
index 33022bf..6de27ab 100644
--- a/lifecycle/lifecycle-viewmodel-compose/src/androidTest/AndroidManifest.xml
+++ b/lifecycle/lifecycle-viewmodel-compose/src/androidTest/AndroidManifest.xml
@@ -1,5 +1,4 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
+<?xml version="1.0" encoding="utf-8"?><!--
   Copyright 2021 The Android Open Source Project
 
   Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,4 +13,13 @@
   See the License for the specific language governing permissions and
   limitations under the License.
   -->
-<manifest package="androidx.lifecycle.viewmodel.compose" />
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+    package="androidx.lifecycle.viewmodel.compose">
+
+    <application>
+        <activity
+            android:name="androidx.appcompat.app.AppCompatActivity"
+            android:theme="@style/Theme.AppCompat" />
+        <activity android:name="androidx.fragment.app.FragmentActivity" />
+    </application>
+</manifest>
diff --git a/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/TestViewModelClasses.kt b/lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/TestViewModelClasses.kt
similarity index 68%
rename from compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/TestViewModelClasses.kt
rename to lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/TestViewModelClasses.kt
index 1c56197..e8f0678 100644
--- a/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/TestViewModelClasses.kt
+++ b/lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/TestViewModelClasses.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020 The Android Open Source Project
+ * Copyright 2021 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,28 +14,29 @@
  * limitations under the License.
  */
 
-package androidx.compose.ui.viewinterop
+package androidx.lifecycle.viewmodel.compose
 
 import android.app.Application
 import androidx.lifecycle.AndroidViewModel
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.ViewModel
 
-val viewModelClasses = arrayOf(
+public val viewModelClasses: Array<Class<out ViewModel>> = arrayOf(
     SimpleViewModel::class.java,
     ApplicationViewModel::class.java,
     SavedStateHandleViewModel::class.java,
     SavedStateHandleAndApplicationViewModel::class.java
 )
 
-class SimpleViewModel : ViewModel()
+public class SimpleViewModel : ViewModel()
 
-class SavedStateHandleViewModel(@Suppress("UNUSED_PARAMETER") savedStateHandle: SavedStateHandle) :
-    ViewModel()
+public class SavedStateHandleViewModel(
+    @Suppress("UNUSED_PARAMETER") savedStateHandle: SavedStateHandle
+) : ViewModel()
 
-class ApplicationViewModel(application: Application) : AndroidViewModel(application)
+public class ApplicationViewModel(application: Application) : AndroidViewModel(application)
 
-class SavedStateHandleAndApplicationViewModel(
+public class SavedStateHandleAndApplicationViewModel(
     application: Application,
     @Suppress("UNUSED_PARAMETER") savedStateHandle: SavedStateHandle
 ) : AndroidViewModel(application)
\ No newline at end of file
diff --git a/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/ViewModelInAppCompatActivityTest.kt b/lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/ViewModelInAppCompatActivityTest.kt
similarity index 82%
rename from compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/ViewModelInAppCompatActivityTest.kt
rename to lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/ViewModelInAppCompatActivityTest.kt
index 6e5fa84..7681326 100644
--- a/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/ViewModelInAppCompatActivityTest.kt
+++ b/lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/ViewModelInAppCompatActivityTest.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020 The Android Open Source Project
+ * Copyright 2021 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package androidx.compose.ui.viewinterop
+package androidx.lifecycle.viewmodel.compose
 
 import androidx.activity.compose.setContent
 import androidx.appcompat.app.AppCompatActivity
@@ -34,21 +34,20 @@
 
 @MediumTest
 @RunWith(AndroidJUnit4::class)
-class ViewModelInAppCompatActivityTest {
+public class ViewModelInAppCompatActivityTest {
     @Suppress("DEPRECATION")
     @get:Rule
-    val activityTestRule = androidx.test.rule.ActivityTestRule<AppCompatActivity>(
-        AppCompatActivity::class.java
-    )
+    public val activityTestRule: androidx.test.rule.ActivityTestRule<AppCompatActivity> =
+        androidx.test.rule.ActivityTestRule(AppCompatActivity::class.java)
     private lateinit var activity: AppCompatActivity
 
     @Before
-    fun setup() {
+    public fun setup() {
         activity = activityTestRule.activity
     }
 
     @Test
-    fun lifecycleOwnerIsAvailable() {
+    public fun lifecycleOwnerIsAvailable() {
         val latch = CountDownLatch(1)
         var owner: LifecycleOwner? = null
 
@@ -64,7 +63,7 @@
     }
 
     @Test
-    fun lifecycleOwnerIsAvailableWhenComposedIntoViewGroup() {
+    public fun lifecycleOwnerIsAvailableWhenComposedIntoViewGroup() {
         val latch = CountDownLatch(1)
         var owner: LifecycleOwner? = null
 
diff --git a/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/ViewModelInComponentActivityTest.kt b/lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/ViewModelInComponentActivityTest.kt
similarity index 82%
rename from compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/ViewModelInComponentActivityTest.kt
rename to lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/ViewModelInComponentActivityTest.kt
index 10b603a..fbc90b5 100644
--- a/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/ViewModelInComponentActivityTest.kt
+++ b/lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/ViewModelInComponentActivityTest.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020 The Android Open Source Project
+ * Copyright 2021 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package androidx.compose.ui.viewinterop
+package androidx.lifecycle.viewmodel.compose
 
 import androidx.activity.ComponentActivity
 import androidx.activity.compose.setContent
@@ -34,21 +34,20 @@
 
 @MediumTest
 @RunWith(AndroidJUnit4::class)
-class ViewModelInComponentActivityTest {
+public class ViewModelInComponentActivityTest {
     @Suppress("DEPRECATION")
     @get:Rule
-    val activityTestRule = androidx.test.rule.ActivityTestRule<ComponentActivity>(
-        ComponentActivity::class.java
-    )
+    public val activityTestRule: androidx.test.rule.ActivityTestRule<ComponentActivity> =
+        androidx.test.rule.ActivityTestRule(ComponentActivity::class.java)
     private lateinit var activity: ComponentActivity
 
     @Before
-    fun setup() {
+    public fun setup() {
         activity = activityTestRule.activity
     }
 
     @Test
-    fun lifecycleOwnerIsAvailable() {
+    public fun lifecycleOwnerIsAvailable() {
         val latch = CountDownLatch(1)
         var owner: LifecycleOwner? = null
 
@@ -64,7 +63,7 @@
     }
 
     @Test
-    fun lifecycleOwnerIsAvailableWhenComposedIntoViewGroup() {
+    public fun lifecycleOwnerIsAvailableWhenComposedIntoViewGroup() {
         val latch = CountDownLatch(1)
         var owner: LifecycleOwner? = null
 
diff --git a/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/ViewModelInFragmentTest.kt b/lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/ViewModelInFragmentTest.kt
similarity index 73%
rename from compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/ViewModelInFragmentTest.kt
rename to lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/ViewModelInFragmentTest.kt
index 64a7db0..a6ba005 100644
--- a/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/ViewModelInFragmentTest.kt
+++ b/lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/ViewModelInFragmentTest.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020 The Android Open Source Project
+ * Copyright 2021 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package androidx.compose.ui.viewinterop
+package androidx.lifecycle.viewmodel.compose
 
 import android.os.Bundle
 import android.view.LayoutInflater
@@ -23,7 +23,6 @@
 import androidx.fragment.app.Fragment
 import androidx.fragment.app.FragmentActivity
 import androidx.fragment.app.FragmentContainerView
-import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.ViewModel
 import androidx.test.filters.MediumTest
 import org.junit.Assert.assertTrue
@@ -37,28 +36,27 @@
 
 @MediumTest
 @RunWith(Parameterized::class)
-class ViewModelInFragmentTest(private val viewModelClass: Class<out ViewModel>) {
+public class ViewModelInFragmentTest(private val viewModelClass: Class<out ViewModel>) {
 
-    companion object {
+    public companion object {
         @JvmStatic
         @Parameterized.Parameters(name = "{0}")
-        fun initParameters() = viewModelClasses
+        public fun initParameters(): Array<Class<out ViewModel>> = viewModelClasses
     }
 
     @Suppress("DEPRECATION")
     @get:Rule
-    val activityTestRule = androidx.test.rule.ActivityTestRule<FragmentActivity>(
-        FragmentActivity::class.java
-    )
+    public val activityTestRule: androidx.test.rule.ActivityTestRule<FragmentActivity> =
+        androidx.test.rule.ActivityTestRule(FragmentActivity::class.java)
     private lateinit var activity: FragmentActivity
 
     @Before
-    fun setup() {
+    public fun setup() {
         activity = activityTestRule.activity
     }
 
     @Test
-    fun viewModelCreatedInFragment() {
+    public fun viewModelCreatedInFragment() {
         val fragment = TestFragment(viewModelClass)
 
         activityTestRule.runOnUiThread {
@@ -74,16 +72,15 @@
     }
 }
 
-class TestFragment(private val viewModelClass: Class<out ViewModel>) : Fragment() {
+public class TestFragment(private val viewModelClass: Class<out ViewModel>) : Fragment() {
 
-    var owner: LifecycleOwner? = null
-    val latch = CountDownLatch(1)
+    public val latch: CountDownLatch = CountDownLatch(1)
 
     override fun onCreateView(
         inflater: LayoutInflater,
         container: ViewGroup?,
         savedInstanceState: Bundle?
-    ) = ComposeView(requireContext()).apply {
+    ): ComposeView = ComposeView(requireContext()).apply {
         setContent {
             viewModel(viewModelClass)
             latch.countDown()
diff --git a/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/ViewModelTest.kt b/lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/ViewModelTest.kt
similarity index 73%
rename from compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/ViewModelTest.kt
rename to lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/ViewModelTest.kt
index 4347258..ce0e504 100644
--- a/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/viewinterop/ViewModelTest.kt
+++ b/lifecycle/lifecycle-viewmodel-compose/src/androidTest/java/androidx/lifecycle/viewmodel/compose/ViewModelTest.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020 The Android Open Source Project
+ * Copyright 2021 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,10 +14,10 @@
  * limitations under the License.
  */
 
-package androidx.compose.ui.viewinterop
+package androidx.lifecycle.viewmodel.compose
 
 import androidx.compose.runtime.Providers
-import androidx.compose.ui.platform.LocalViewModelStoreOwner
+import androidx.compose.ui.test.junit4.ComposeContentTestRule
 import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.lifecycle.HasDefaultViewModelProviderFactory
 import androidx.lifecycle.ViewModel
@@ -33,16 +33,16 @@
 
 @MediumTest
 @RunWith(AndroidJUnit4::class)
-class ViewModelTest {
+public class ViewModelTest {
 
     @get:Rule
-    val rule = createComposeRule()
+    public val rule: ComposeContentTestRule = createComposeRule()
 
     @Test
-    fun viewModelCreatedViaDefaultFactory() {
+    public fun viewModelCreatedViaDefaultFactory() {
         val owner = FakeViewModelStoreOwner()
         rule.setContent {
-            Providers(LocalViewModelStoreOwner provides owner) {
+            Providers(LocalViewModelStoreOwner.asProvidableCompositionLocal() provides owner) {
                 viewModel<TestViewModel>()
             }
         }
@@ -51,11 +51,11 @@
     }
 
     @Test
-    fun viewModelCreatedViaDefaultFactoryWithKey() {
+    public fun viewModelCreatedViaDefaultFactoryWithKey() {
         val owner = FakeViewModelStoreOwner()
         var createdInComposition: Any? = null
         rule.setContent {
-            Providers(LocalViewModelStoreOwner provides owner) {
+            Providers(LocalViewModelStoreOwner.asProvidableCompositionLocal() provides owner) {
                 createdInComposition = viewModel<TestViewModel>()
             }
         }
@@ -66,11 +66,11 @@
     }
 
     @Test
-    fun createdViewModelIsEqualsToCreatedManually() {
+    public fun createdViewModelIsEqualsToCreatedManually() {
         val owner = FakeViewModelStoreOwner()
         var createdInComposition: Any? = null
         rule.setContent {
-            Providers(LocalViewModelStoreOwner provides owner) {
+            Providers(LocalViewModelStoreOwner.asProvidableCompositionLocal() provides owner) {
                 createdInComposition = viewModel<TestViewModel>()
             }
         }
@@ -81,12 +81,13 @@
     }
 
     @Test
-    fun createdViewModelIsEqualsToCreatedManuallyWithKey() {
+    public fun createdViewModelIsEqualsToCreatedManuallyWithKey() {
         val owner = FakeViewModelStoreOwner()
         var createdInComposition: Any? = null
         rule.setContent {
-            Providers(LocalViewModelStoreOwner provides owner) {
-                createdInComposition = viewModel<TestViewModel>(key = "test")
+            Providers(LocalViewModelStoreOwner.asProvidableCompositionLocal() provides owner) {
+                createdInComposition =
+                    viewModel<TestViewModel>(key = "test")
             }
         }
 
@@ -96,11 +97,11 @@
     }
 
     @Test
-    fun customFactoryIsUsedWhenProvided() {
+    public fun customFactoryIsUsedWhenProvided() {
         val owner = FakeViewModelStoreOwner()
         val customFactory = FakeViewModelProviderFactory()
         rule.setContent {
-            Providers(LocalViewModelStoreOwner provides owner) {
+            Providers(LocalViewModelStoreOwner.asProvidableCompositionLocal() provides owner) {
                 viewModel<TestViewModel>(factory = customFactory)
             }
         }
@@ -109,11 +110,11 @@
     }
 
     @Test
-    fun defaultFactoryIsNotUsedWhenCustomProvided() {
+    public fun defaultFactoryIsNotUsedWhenCustomProvided() {
         val owner = FakeViewModelStoreOwner()
         val customFactory = FakeViewModelProviderFactory()
         rule.setContent {
-            Providers(LocalViewModelStoreOwner provides owner) {
+            Providers(LocalViewModelStoreOwner.asProvidableCompositionLocal() provides owner) {
                 viewModel<TestViewModel>(factory = customFactory)
             }
         }
@@ -122,12 +123,12 @@
     }
 
     @Test
-    fun createdWithCustomFactoryViewModelIsEqualsToCreatedManually() {
+    public fun createdWithCustomFactoryViewModelIsEqualsToCreatedManually() {
         val owner = FakeViewModelStoreOwner()
         var createdInComposition: Any? = null
         val customFactory = FakeViewModelProviderFactory()
         rule.setContent {
-            Providers(LocalViewModelStoreOwner provides owner) {
+            Providers(LocalViewModelStoreOwner.asProvidableCompositionLocal() provides owner) {
                 createdInComposition = viewModel<TestViewModel>()
             }
         }
@@ -138,13 +139,14 @@
     }
 
     @Test
-    fun createdWithCustomFactoryViewModelIsEqualsToCreatedManuallyWithKey() {
+    public fun createdWithCustomFactoryViewModelIsEqualsToCreatedManuallyWithKey() {
         val owner = FakeViewModelStoreOwner()
         var createdInComposition: Any? = null
         val customFactory = FakeViewModelProviderFactory()
         rule.setContent {
-            Providers(LocalViewModelStoreOwner provides owner) {
-                createdInComposition = viewModel<TestViewModel>(key = "test")
+            Providers(LocalViewModelStoreOwner.asProvidableCompositionLocal() provides owner) {
+                createdInComposition =
+                    viewModel<TestViewModel>(key = "test")
             }
         }
 
@@ -173,4 +175,4 @@
 
     override fun getViewModelStore(): ViewModelStore = store
     override fun getDefaultViewModelProviderFactory(): ViewModelProvider.Factory = factory
-}
\ No newline at end of file
+}
diff --git a/lifecycle/lifecycle-viewmodel-compose/src/main/java/androidx/lifecycle/viewmodel/compose/LocalViewModelStoreOwner.kt b/lifecycle/lifecycle-viewmodel-compose/src/main/java/androidx/lifecycle/viewmodel/compose/LocalViewModelStoreOwner.kt
new file mode 100644
index 0000000..c3ea604
--- /dev/null
+++ b/lifecycle/lifecycle-viewmodel-compose/src/main/java/androidx/lifecycle/viewmodel/compose/LocalViewModelStoreOwner.kt
@@ -0,0 +1,48 @@
+/*
+ * Copyright 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.lifecycle.viewmodel.compose
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.ProvidableCompositionLocal
+import androidx.compose.runtime.compositionLocalOf
+import androidx.compose.ui.platform.LocalView
+import androidx.lifecycle.ViewModelStoreOwner
+import androidx.lifecycle.ViewTreeViewModelStoreOwner
+
+/**
+ * The CompositionLocal containing the current [ViewModelStoreOwner].
+ */
+public object LocalViewModelStoreOwner {
+    private val LocalViewModelStoreOwner =
+        compositionLocalOf<ViewModelStoreOwner?> { null }
+
+    /**
+     * Returns current composition local value for the owner.
+     */
+    public val current: ViewModelStoreOwner
+        @Composable
+        get() = LocalViewModelStoreOwner.current
+            ?: ViewTreeViewModelStoreOwner.get(LocalView.current)
+            ?: error("No ViewModelStoreOwner provided")
+
+    /**
+     * Returns a [ProvidableCompositionLocal] which you can use to override the value for the
+     * subtree.
+     */
+    public fun asProvidableCompositionLocal():
+        ProvidableCompositionLocal<ViewModelStoreOwner?> = LocalViewModelStoreOwner
+}
diff --git a/lifecycle/lifecycle-viewmodel-compose/src/main/java/androidx/lifecycle/viewmodel/compose/ViewModel.kt b/lifecycle/lifecycle-viewmodel-compose/src/main/java/androidx/lifecycle/viewmodel/compose/ViewModel.kt
new file mode 100644
index 0000000..3166d7f
--- /dev/null
+++ b/lifecycle/lifecycle-viewmodel-compose/src/main/java/androidx/lifecycle/viewmodel/compose/ViewModel.kt
@@ -0,0 +1,78 @@
+/*
+ * Copyright 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.lifecycle.viewmodel.compose
+
+import androidx.compose.runtime.Composable
+import androidx.lifecycle.ViewModel
+import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.ViewModelStoreOwner
+
+/**
+ * Returns an existing [ViewModel] or creates a new one in the scope (usually, a fragment or
+ * an activity)
+ *
+ * The created [ViewModel] is associated with the given scope and will be retained
+ * as long as the scope is alive (e.g. if it is an activity, until it is
+ * finished or process is killed).
+ *
+ * @param key The key to use to identify the [ViewModel].
+ * @return A [ViewModel] that is an instance of the given [T] type.
+ */
+@Suppress("MissingJvmstatic")
+@Composable
+public inline fun <reified T : ViewModel> viewModel(
+    key: String? = null,
+    factory: ViewModelProvider.Factory? = null
+): T = viewModel(T::class.java, key, factory)
+
+/**
+ * Returns an existing [ViewModel] or creates a new one in the scope (usually, a fragment or
+ * an activity)
+ *
+ * The created [ViewModel] is associated with the given scope and will be retained
+ * as long as the scope is alive (e.g. if it is an activity, until it is
+ * finished or process is killed).
+ *
+ * @param modelClass The class of the [ViewModel] to create an instance of it if it is not
+ * present.
+ * @param key The key to use to identify the [ViewModel].
+ * @return A [ViewModel] that is an instance of the given [T] type.
+ */
+@Suppress("MissingJvmstatic")
+@Composable
+public fun <T : ViewModel> viewModel(
+    modelClass: Class<T>,
+    key: String? = null,
+    factory: ViewModelProvider.Factory? = null
+): T = LocalViewModelStoreOwner.current.get(modelClass, key, factory)
+
+private fun <T : ViewModel> ViewModelStoreOwner.get(
+    javaClass: Class<T>,
+    key: String? = null,
+    factory: ViewModelProvider.Factory? = null
+): T {
+    val provider = if (factory != null) {
+        ViewModelProvider(this, factory)
+    } else {
+        ViewModelProvider(this)
+    }
+    return if (key != null) {
+        provider.get(key, javaClass)
+    } else {
+        provider.get(javaClass)
+    }
+}
diff --git a/navigation/navigation-compose/build.gradle b/navigation/navigation-compose/build.gradle
index 9de1beb..5a4779d 100644
--- a/navigation/navigation-compose/build.gradle
+++ b/navigation/navigation-compose/build.gradle
@@ -35,6 +35,7 @@
     api projectOrArtifact(":compose:runtime:runtime")
     api projectOrArtifact(":compose:runtime:runtime-saveable")
     api projectOrArtifact(":compose:ui:ui")
+    api projectOrArtifact(":lifecycle:lifecycle-viewmodel-compose")
     api "androidx.navigation:navigation-runtime-ktx:2.3.1"
 
     androidTestImplementation projectOrArtifact(":compose:material:material")
diff --git a/navigation/navigation-compose/src/main/java/androidx/navigation/compose/NavHost.kt b/navigation/navigation-compose/src/main/java/androidx/navigation/compose/NavHost.kt
index 166dbec..d999432 100644
--- a/navigation/navigation-compose/src/main/java/androidx/navigation/compose/NavHost.kt
+++ b/navigation/navigation-compose/src/main/java/androidx/navigation/compose/NavHost.kt
@@ -26,10 +26,10 @@
 import androidx.compose.runtime.saveable.rememberSaveableStateHolder
 import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.platform.LocalLifecycleOwner
-import androidx.compose.ui.platform.LocalViewModelStoreOwner
-import androidx.compose.ui.viewinterop.viewModel
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.ViewModel
+import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
+import androidx.lifecycle.viewmodel.compose.viewModel
 import androidx.navigation.NavGraph
 import androidx.navigation.NavGraphBuilder
 import androidx.navigation.NavHostController
@@ -124,7 +124,11 @@
             // while in the scope of the composable, we provide the navBackStackEntry as the
             // ViewModelStoreOwner and LifecycleOwner
             Providers(
-                LocalViewModelStoreOwner provides currentNavBackStackEntry,
+                LocalViewModelStoreOwner.asProvidableCompositionLocal()
+                    provides currentNavBackStackEntry,
+                @Suppress("DEPRECATION") // To be removed when we remove the one from compose:ui
+                androidx.compose.ui.platform.LocalViewModelStoreOwner provides
+                    currentNavBackStackEntry,
                 LocalLifecycleOwner provides currentNavBackStackEntry
             ) {
                 saveableStateHolder.SaveableStateProvider {