Compose + Anvil

AnvilViewModelFactory.kt

import android.content.Context
import android.os.Bundle
import androidx.activity.ComponentActivity
import androidx.compose.runtime.Composable
import androidx.compose.ui.platform.LocalContext
import androidx.lifecycle.AbstractSavedStateViewModelFactory
import androidx.lifecycle.SavedStateHandle
import androidx.lifecycle.ViewModel
import androidx.lifecycle.ViewModelStoreOwner
import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
import androidx.lifecycle.viewmodel.compose.viewModel
import androidx.navigation.NavBackStackEntry
import androidx.savedstate.SavedStateRegistryOwner
import com.welbits.transit.common.di.scopedBindings
import logcat.logcat
import javax.inject.Provider

typealias ViewModelMap = Map<Class<out ViewModel>, @JvmSuppressWildcards Provider<ViewModel>>

@Composable
inline fun <reified VM : ViewModel> anvilActivityViewModel(): VM =
    anvilViewModel<VM>(LocalContext.current as ViewModelStoreOwner)

@Composable
inline fun <reified VM : ViewModel> anvilViewModel(
    viewModelStoreOwner: ViewModelStoreOwner = checkNotNull(LocalViewModelStoreOwner.current) {
        "No ViewModelStoreOwner was provided via LocalViewModelStoreOwner"
    },
): VM {
    require(viewModelStoreOwner is SavedStateRegistryOwner)
    val factory = AnvilViewModelFactory(
        context = LocalContext.current,
        owner = viewModelStoreOwner,
        defaultArgs = when (viewModelStoreOwner) {
            is NavBackStackEntry -> viewModelStoreOwner.arguments
            is ComponentActivity -> viewModelStoreOwner.intent.extras
            else -> null
        },
    )
    return viewModel(viewModelStoreOwner, factory = factory)
}

@PublishedApi
internal class AnvilViewModelFactory(
    private val context: Context,
    owner: SavedStateRegistryOwner,
    defaultArgs: Bundle? = null,
) : AbstractSavedStateViewModelFactory(owner, defaultArgs) {

    override fun <T : ViewModel> create(
        key: String,
        modelClass: Class<T>,
        handle: SavedStateHandle,
    ): T {

        val subcomponent = context
            .scopedBindings<ViewModelComponent.ParentBindings>()
            .viewModelComponentFactory
            .create(handle)

        val provider = subcomponent.viewModelMap[modelClass]
            ?: error("${modelClass.simpleName} not in ViewModelMap: ${subcomponent.viewModelMap.map { it.key.simpleName }}")

        logcat { "${modelClass.simpleName} found. SavedStateHandle.keys() = ${handle.keys()}" }

        @Suppress("UNCHECKED_CAST")
        return provider.get() as T
    }
}

ViewModelComponent.kt

import androidx.lifecycle.SavedStateHandle
import com.squareup.anvil.annotations.ContributesSubcomponent
import com.squareup.anvil.annotations.ContributesTo
import com.welbits.transit.common.di.scope.ActivityScope
import com.welbits.transit.common.di.scope.SingleIn
import com.welbits.transit.common.di.scope.ViewModelScope
import dagger.BindsInstance

@SingleIn(ViewModelScope::class)
@ContributesSubcomponent(
    scope = ViewModelScope::class,
    parentScope = ActivityScope::class,
)
interface ViewModelComponent {
    @ContributesSubcomponent.Factory
    interface Factory {
        fun create(@BindsInstance savedStateHandle: SavedStateHandle): ViewModelComponent
    }

    @ContributesTo(ActivityScope::class)
    interface ParentBindings {
        val viewModelComponentFactory: Factory
    }

    val viewModelMap: ViewModelMap
}

ViewModelKey.kt

import androidx.lifecycle.ViewModel
import dagger.MapKey
import kotlin.reflect.KClass

/**
 * A [MapKey] annotation for maps with [KClass] of [ViewModel] keys.
 *
 * Note this was designed to be used only with [CompositeViewModelFactory].
 */
@Target(
    AnnotationTarget.CLASS,
    AnnotationTarget.FUNCTION,
)
@Retention(value = AnnotationRetention.RUNTIME)
@MapKey
annotation class ViewModelKey(val value: KClass<out ViewModel>)