【Compose】コンテンツを削除して高さが変わる時に高さの変化をアニメーションさせる

Columnの中の途中のコンテンツを消す時に、いきなり高さが0になるのではなく、アニメーションさせて消します。
これは、 AnimatedVisibility を使用すればいいのですが、中のコンテンツを表示させるためのものがnullableであり、nullになったら消したい場合は、消す時に高さがいきなり0になってしまいます。

うまくいかない例

class UiState {
    val text: String = ""
}
val uiState: UiState? = null
AnimatedVisibility(visible = uiState != null) {
    if (uiState != null) {
        Text(text = uiState.text)
    }
}

表示自体はいきなり消えてしまいますが、高さは次第に小さくなります。

Code

@Composable
public fun HeightSmoothChangeAnimation(
    modifier: Modifier = Modifier,
    contentAlignment: Alignment.Vertical,
    animationSpec: AnimationSpec<Int> = remember {
        spring(visibilityThreshold = null)
    },
    content: @Composable () -> Unit,
) {
    var initialized by remember { mutableStateOf(false) }
    val animateHeight = remember {
        Animatable(
            initialValue = 0,
            typeConverter = Int.VectorConverter
        )
    }
    BoxWithConstraints(
        modifier = modifier,
    ) {
        val boxConstraints = this
        Box(
            modifier = Modifier
                .height(
                    height = with(LocalDensity.current) {
                        animateHeight.value.toDp()
                    }
                )
                .clipToBounds(),
        ) {
            val coroutineScope = rememberCoroutineScope()
            Layout(
                modifier = Modifier
                    .layout { measurable, constraints ->
                        val placeable = measurable.measure(constraints.copy(maxHeight = Constraints.Infinity))
                        val height = kotlin.math.min(constraints.maxHeight, placeable.height)

                        val y = when (contentAlignment) {
                            Alignment.Top -> 0
                            Alignment.Bottom -> height - animateHeight.targetValue
                            else -> (height - animateHeight.targetValue) / 2
                        }

                        layout(placeable.width, height) {
                            placeable.place(0, y)
                        }
                    },
                content = {
                    content()
                }
            ) { measurables, _ ->
                val newConstraints = Constraints(
                    minWidth = boxConstraints.minWidth.toPx().toInt(),
                    maxWidth = boxConstraints.maxWidth.toPx().toInt(),
                    minHeight = boxConstraints.minHeight.toPx().toInt(),
                    maxHeight = boxConstraints.maxHeight.toPx().toInt(),
                )
                val placeable = measurables.map { it.measure(newConstraints) }
                val maxHeight = placeable.maxOfOrNull { it.height } ?: 0
                val maxWidth = placeable.maxOfOrNull { it.width } ?: 0
                coroutineScope.launch {
                    if (initialized) {
                        animateHeight.animateTo(maxHeight, animationSpec)
                    } else {
                        animateHeight.snapTo(maxHeight)
                    }
                    initialized = true
                }

                layout(maxWidth, maxHeight) {
                    placeable.forEach { it.place(0, 0) }
                }
            }
        }
    }
}

Preview

@Preview(showBackground = true)
@Composable
private fun Preview() {
    Column(Modifier.fillMaxSize()) {
        var visibleContent by remember {
            mutableStateOf(true)
        }
        var alignment by remember {
            mutableStateOf(Alignment.Top)
        }
        Row(modifier = Modifier.fillMaxWidth()) {
            Button(onClick = { visibleContent = !visibleContent }) {
                Text(text = "Toggle")
            }
            Spacer(modifier = Modifier.width(16.dp))
            Button(
                onClick = {
                    alignment = when (alignment) {
                        Alignment.Top -> Alignment.CenterVertically
                        Alignment.CenterVertically -> Alignment.Bottom
                        else -> Alignment.Top
                    }
                },
            ) {
                val text by remember {
                    derivedStateOf {
                        when (alignment) {
                            Alignment.Top -> "TOP"
                            Alignment.CenterVertically -> "CenterVertically"
                            else -> "Bottom"
                        }
                    }
                }
                Text(text = text)
            }
        }
        Spacer(
            modifier = Modifier
                .fillMaxWidth()
                .height(height = 24.dp)
                .background(Color.Cyan)
        )
        HeightSmoothChangeAnimation(
            contentAlignment = alignment,
            animationSpec = remember {
                spring(
                    stiffness = 30f,
                    visibilityThreshold = null,
                )
            },
        ) {
            Column(Modifier.fillMaxWidth()) {
                if (visibleContent) {
                    Text(text = "Top")
                    Spacer(modifier = Modifier.height(18.dp))
                    Text(text = "Center")
                    Spacer(modifier = Modifier.height(18.dp))
                    Text(text = "Bottom")
                }
            }
        }
        Spacer(
            modifier = Modifier
                .background(Color.Yellow)
                .fillMaxWidth()
                .height(height = 24.dp)
        )
    }
}