Jetpack Compose Lazy Column insertion and deletion animations with multiple item types

What's the best way to animate insertion and deletion animations in lazy column or row with multiple item types similar to how it's done using DiffUtil?

https://issuetracker.google.com/issues/150812265

Modifier.animateItemPlacement() was created for this reason, but to do it with multiple item types is less straight forward.

Animation demo: https://youtube.com/shorts/FBwMV1HoAoQ?feature=share

Ps (for demo)

1. RewardItem click used to remove reward items and CartHeader click adds them back
2. CartItem remove button click to remove item from cart item and modify button click to add it back

Sealed Class:

sealed class CartListItems(open val id: String = "") {
    class RewardHeaderItem(override val id: String, val title: String) : CartListItems()
    class RewardListItem(override val id: String, val rewards: List<RewardItem>) : CartListItems()
    class CartHeaderItem(override val id: String, val title: String) : CartListItems()
    class CartListItem(override val id: String, val cartItem: CartItem) : CartListItems()
}

Inside ViewModel:

val cartListItems: StateFlow<List<CartListItems>> =
        combine(
            rewardItems,
            cartItems
        ) { rewardItems, cartItems ->
            buildCartList(rewardItems, cartItems)
        }.stateIn(
            scope = viewModelScope,
            started = Eagerly,
            initialValue = emptyList()
        )


private fun buildCartList(rewardItems: List<RewardItem>, cartItems: List<CartItem>): List<CartListItems> {
        val items = ArrayList<CartListItems>()

        if (rewardItems.isNotEmpty()) {
            items.add(
                CartListItems.RewardHeaderItem("rewards-header", "Your Rewards")
            )
            items.add(
                CartListItems.RewardListItem("rewards-list", rewardItems)
            )
        }
        if (cartItems.isNotEmpty()) {
            items.add(
                CartListItems.CartHeaderItem("cart-header", "Your Cart")
            )
            items.addAll(
                cartItems.map { CartListItems.CartListItem("cart-item:${it.id}", it) }
            )
        }

        return items
    }

List Composable:

@Composable
private fun CartList(
    cartViewModel: CartViewModel = viewModel()
) {
    val listItems by cartViewModel.cartListItems.collectAsState()
    
    LazyColumn(
        modifier = Modifier.fillMaxSize(),
        horizontalAlignment = Alignment.CenterHorizontally,
        contentPadding = PaddingValues(vertical = 16.dp),
        verticalArrangement = Arrangement.spacedBy(16.dp)
    ) {

        items(listItems, key = { it.id }) { listItem ->
            when (listItem) {
                is CartListItems.RewardHeaderItem -> {
                    Box(modifier = Modifier.animateItemPlacement()) {
                        RewardsHeader()
                    }
                }
                is CartListItems.RewardListItem -> {
                    Box(modifier = Modifier.animateItemPlacement()) {
                        RewardsList(listItem.rewards)
                    }
                }
                is CartListItems.CartHeaderItem -> {
                    Box(modifier = Modifier.animateItemPlacement()) {
                        CartHeader()
                    }
                }
                is CartListItems.CartListItem -> {
                    Box(modifier = Modifier.animateItemPlacement()) {
                        CartItem(listItem.cartItem)
                    }
                }
            }
        }
    }
}