[JetpackCompose] SavedStateHandleでUiStateを保持しながら、StateFlowでライフサイクルを意識する

前提

lifecycle-runtime-compose:2.6.0からcollectAsStateWithLifecycleが追加されました。
これを使用することで、ライフサイクルを意識した方法でFlowを収集することができます。デフォルトではonStartからonStopの間にFlowが更新されたとき、再Composeが走ります。
例えば私は下記のように使っています。

AuthorViewModel.kt

@HiltViewModel
class AuthorViewModel @Inject constructor() : ViewModel() {
    val _uiState: MutableStateFlow<UiState> = MutableStateFlow(UiState())
    val uiState: StateFlow<UiState> = _uiState.asStateFlow()
}

AuthorRoute.kt

@OptIn(ExperimentalLifecycleComposeApi::class)
@Composable
fun AuthorRoute(
    viewModel: AuthorViewModel = hiltViewModel()
) {
    val uiState: AuthorScreenUiState by viewModel.uiState.collectAsStateWithLifecycle()
    // uiStateの内容を表示
}

またlifecycle-viewmodel-compose:2.5.0で、試験運用版ではありますが、SavedStateHandle#saveable APIが追加されました。
このAPIを使用すると、SavedStateHandleとComposeのSaverの相互運用が可能になります。
例えば下記のように使用することで、Activityが再生成されてもUiStateが復元されます。

AuthorViewModel.kt

@HiltViewModel
class AuthorViewModel @Inject constructor(
    savedStateHandle: SavedStateHandle,
) : ViewModel() {
    var uiState: UiState by savedStateHandle.saveable { mutableStateOf(UiState()) }
        private set
}

SavedStateHandleに値を保存しているのにも関わらず、値を格納/取り出すためのキーを指定する必要がないのは、プロパティ名をキーとしているからです。

前置きが長くなりましたが、これらを活用して本題に取り掛かりましょう。

環境

app/build.gradle

def lifecycle = "2.6.2"
implementation "androidx.lifecycle:lifecycle-viewmodel-compose:${lifecycle}"
implementation "androidx.lifecycle:lifecycle-runtime-compose:${lifecycle}"

MutableStateFlowを保存する拡張関数を作る

下記の拡張関数を定義します。

AuthorViewModel.kt

@OptIn(SavedStateHandleSaveableApi::class)
fun <T : Any> SavedStateHandle.saveableMutableStateFlow(init: () -> MutableStateFlow<T>): PropertyDelegateProvider<Any?, ReadOnlyProperty<Any?, MutableStateFlow<T>>> {
    return saveable(
        saver = Saver(
            save = {
                it.value
            },
            restore = {
                MutableStateFlow(it)
            }),
        init = init
    )
}

PropertyDelegateProviderはプロパティ名を取得するために必要です。
カスタムSaverを作ることで、MutableStateFlowの再生成に対応しています。

使う

使用する側は特に意識することなく、saveable APIと同じように使います。

AuthorViewModel.kt

@HiltViewModel
class AuthorViewModel @Inject constructor(
    savedStateHandle: SavedStateHandle,
) : ViewModel() {
    private val _uiState: MutableStateFlow<UiState> by savedStateHandle.saveableMutableStateFlow {
        MutableStateFlow(UiState())
    }
    val uiState: StateFlow<UiState> = _uiState.asStateFlow()
}

AuthorRoute.kt

@OptIn(ExperimentalLifecycleComposeApi::class)
@Composable
fun AuthorRoute(
    viewModel: AuthorViewModel = hiltViewModel()
) {
    val uiState: AuthorScreenUiState by viewModel.uiState.collectAsStateWithLifecycle()
    // uiStateの内容を表示
}

これでSavedStateHandleでUiStateを保持しながら、StateFlowでライフサイクルを意識することができました。

さいごに

saveable APIは試験運用版のため、この方法は今後修正が入る可能性があります。
拡張関数saveableMutableStateFlow()の内部を書き直す程度の修正であればいいのですが、呼び出し側にも修正が必要な可能性があり、乱用はリスクが伴いますのでご注意ください。
早くsaveable APIが安定板になりますように...。