Безопасное обновления состояния в ViewModel

Это краткая заметка о подходе, который я выработал для себя, чтобы обновлять состояние экрана при использовании MVI-like паттерна в ViewModel.

Сразу оговорюсь, что классический "полноценный" MVI подразумевает использование редьюсеров, которые решают часть проблем, описанных в этой заметке, но сугубо на мой вкус, этот подход заставляет писать много бойлерплейтного кода.

Вводные

Предположим, у нас есть базовая ViewModel такого вида (опущу работу с эффектами и прочее, т.к. для примера это излишне):

abstract class BaseViewModel<State : Any>(initialState: State) : ViewModel() {

    private val _state = MutableStateFlow(initialState)

    protected var state: State
        get() = _state.value
        set(value) {
            _state.value = value
        }

    fun states() = _state.asStateFlow()

}

И моделька состояния экрана:

data class ProfileScreenState(
    val user: User,
    val notificationsNumber: Int
) {
    data class User(
        val name: String,
        val balance: Balance
    ) {
        data class Balance(
            val total: Double,
            val accounts: Map<Int, Double>
        )
    }
}

Для обновления состояния внутри ViewModel мы пишем что-то такое:

state = state.copy(
    user = state.user.copy(balance = getBalance()),
    notificationsNumber = fetchNotificationsNumber()
)

Проблемы

При таком подходе есть две проблемы.

1. Если State содержит много вложенных классов, то можно быстро утонуть в конструкциях вида:

state = state.copy(
	user = state.user.copy(
		balance = state.user.balance.copy(total = 100.0)
	)
)

2. Работа с асинхронным кодом при обновлении состояния может привести к невалидному состоянию. В момент вызова state.copy происходит захват состояния, и если дальше мы где-то вызываем корутину, то к моменту окончания её выполнения захваченное состояние уже может быть неактуальным.

suspend fun fetchUser(): User { /* ... */ }

suspend fun fetchNotificationsNumber(): Int { /* ... */ }

fun updateUser() {
    viewModelScope.launch {
        state = state.copy(
            user = fetchUser()
        )
    }
}

fun updateNotifications() {
    viewModelScope.launch {
        state = state.copy(
            notificationsNumber = fetchNotificationsNumber()
        )
    }
}

Если функции updateUser и updateNotifications были вызваны одновременно, то с большой вероятность победит какая-то одна из них. То есть мы или получим актуального юзера, или актуальное количество уведомлений, но не то и другое вместе.

Второй случай решается вызовом асинхронных операций до обновления состояния:

val notificationsNumber = fetchNotificationsNumber()
state = state.copy(notificationsNumber = notificationsNumber)

Но об этом надо либо помнить, либо озаботиться линтерами.

Простое решение

Обе проблемы решаются банально, но не без своих недостатков.

Первую можно решить избавившись от вложенных классов и сделав State максимально "плоским":

data class ProfileScreenState(
    val userName: String,
    val totalBalance: Double,
    val accountBalances: Map<Int, Double>,
    val notificationsNumber: Int
)

Минус очевиден: на сложных компонентах такое состояние может вырасти до десятков полей и ориентироваться в нём будет не так просто.

Вторая проблема тоже не rocket science, просто переносим обновление состояния в лямбду внутри которой запрещаем вызов корутин:

// BaseViewModel

fun updateState(block: (State) -> State) {
    _state.update(block)
}

// Использование

updateState { state ->
    state.copy(
        notificationsNumber = notificationsNumber
    )
}

Для меня минус в получившемся синтаксисе. Это уже не красивое присваивание в свойство класса, а вызов функции. Минус спорный, но это мои загоны :)

Моё решение

Я решил две этих проблемы двумя плагинами (да не переусложняю я!).
Во-первых Arrow, который реализует концепцию оптики из функционального программирования.
Во-вторых плагина от JetBrains, позволяющего переопределять оператор присваивания (=).

Сразу покажу, как теперь выглядит обновление состояния в моём коде:

val freshUser = fetchUser()
val freshNotificationsNumber = fetchNotificationsNumber()
state = {
    State.user.totalBalance set freshUser.totalBalance
    State.notificationsNumber set freshNotificationsNumber
}

Как видно, это решает первую проблему, потому что больше нет необходимости писать множественные вложенные copy. А также решает вторую, потому что внутри фигурных скобок не получится вызвать корутину, и мы вынуждены получить все данные перед обновлением состояния. При этом сохранился синтаксис с присваиванием значения в свойство класса.

Реализация

Нам понадобится аннотация, которую нужно навешивать на классы состояния, для переопределения операции присваивания:

@Suppress("ClassName")
@Retention(AnnotationRetention.SOURCE)
@Target(AnnotationTarget.CLASS)
annotation class assigned

Название аннотации можно выбрать на свой вкус.

Далее подготовим базовую ViewModel. Добавляем зависимость на Arrow Optics:

implementation("io.arrow-kt:arrow-optics:2.0.0")

И изменияем ViewModel:

abstract class BaseViewModel<State : Any>(initialState: State) : ViewModel() {

    private val _state = MutableStateFlow(initialState)

    protected val state: State
        get() = _state.value

    fun states() = _state.asStateFlow()

    protected fun State.assign(@BuilderInference block: Copy<State>.() -> Unit) {
        _state.update { it.copy(block) }
    }

}

Так как аннотации должны быть у конечного класса состояния, я использую gradle-скрипт, который подключает все необходимые зависимости в фиче-модуль:

// buildSrc/src/main/kotlin/optics-setup.gradle.kts

import org.jetbrains.kotlin.gradle.tasks.KotlinCompilationTask

plugins {
    id("kotlin-multiplatform")
    id("com.android.library")
    id("com.google.devtools.ksp")
    id("org.jetbrains.kotlin.plugin.assignment")
}

kotlin {
    sourceSets {
        commonMain {
            kotlin.srcDir("build/generated/ksp/metadata/commonMain/kotlin")
            dependencies {
                implementation("io.arrow-kt:arrow-optics:2.0.0")
            }
        }
    }
}

dependencies {
    add("kspCommonMainMetadata", "io.arrow-kt:arrow-optics-ksp-plugin:2.0.0")
}

assignment {
    annotation("org.app.example.assigned")
}

project.tasks.withType(KotlinCompilationTask::class.java).configureEach {
    if (name != "kspCommonMainKotlinMetadata") {
        dependsOn("kspCommonMainKotlinMetadata")
    }
}

Подключаем его в фиче-модулях:

plugins {
	// ...
    id("optics-setup")
}

И теперь всё что нам остается, это повесить на классы стейта нужные аннотации и написать пустые companion object:

@optics
@assigned
data class ProfileScreenState(
    val user: User,
    val notificationsNumber: Int
) {
    companion object

    @optics
    data class User(
        val name: String,
        val balance: Balance
	) {
        companion object

        @optics
        data class Balance(
            val total: Double,
            val accounts: Map<Int, Double>
		) {
            companion object
		}
	}
}

Вывод

Я ни в коем случае не навязываю использование такого подхода. Из его минусов могу выделить:

• Зависимость от двух плагинов. Но Arrow существует и развивается столько, сколько я знаю Kotlin, а второй плагин — от JetBrains;

• Небольшой бойлерплейт в виде двух аннотаций, которые надо не забывать объявлять для классов состояния;

• Специфический синтаксис.

Но если вам близко моё видение, надеюсь, эта заметка будет полезной и упростит ваш опыт разработки.