Jetpack Compose LazyRow自定义悬浮粘性头部实现问题求助

我在Jetpack Compose中构建UI时，需要给LazyRow实现一个悬浮粘性头部，但目前的实现是把头部作为列表内的一个行元素嵌入的。我想要的效果是粘性头部始终显示在元素上方，就像标准的横向滚动粘性头部那样。

后来我尝试自己动手解决这个问题，但碰到了两个棘手的问题：

• 当OffsetX处于Leaving状态时，如果给LazyRow添加了itemSpacing，整个布局逻辑就会出现异常
• 面对较长的粘性头部，我不知道该如何正确计算并设置偏移量，以适配它的实际长度

下面是我目前写的实现代码：

@Composable
fun <K, V> LazyRowWithStickyHeader(
    items: Map<K, List<V>>,
    modifier: Modifier = Modifier,
    state: LazyListState = rememberLazyListState(),
    contentPadding: PaddingValues = PaddingValues(0.dp),
    reverseLayout: Boolean = false,
    horizontalArrangement: Arrangement.Horizontal = if (!reverseLayout) Arrangement.Start else Arrangement.End,
    verticalAlignment: Alignment.Vertical = Alignment.Top,
    flingBehavior: FlingBehavior = ScrollableDefaults.flingBehavior(),
    userScrollEnabled: Boolean = true,
    stickyHeader: StickyHeaderScope.(K) -> Unit,
    itemContent: @Composable LazyItemScope.(V) -> Unit
) {
    val itemsWithKeys = remember(items) {
        items.flatMap { entry -> entry.value.map { entry.key to it } }
    }
    val textMeasurer = rememberTextMeasurer()
    var itemWidth by remember { mutableIntStateOf(0) }
    var stickyHeaderHeight by remember { mutableStateOf(0.dp) }

    Box(
        modifier = Modifier.drawWithCache {
            onDrawBehind {
                var previousKey: K? = null
                val startPadding = state.layoutInfo.beforeContentPadding
                if (itemWidth == 0) {
                    itemWidth = state.layoutInfo.visibleItemsInfo.firstOrNull()?.size ?: 0
                }
                state.layoutInfo.visibleItemsInfo.forEachIndexed { index, itemInfo ->
                    val currentKey = itemsWithKeys.getOrNull(itemInfo.index)?.first
                    val nextItemKey = itemsWithKeys.getOrNull(itemInfo.index + 1)?.first
                    if (currentKey == null || currentKey == previousKey) {
                        return@forEachIndexed
                    }

                    StickyHeaderScopeImpl(
                        drawScope = this,
                        textMeasurer = textMeasurer,
                        offsetProvider = { size ->
                            stickyHeaderHeight = size.height.toDp()
                            val offsetX = when {
                                //Stickying
                                currentKey == nextItemKey && index == 0 -> {
                                    startPadding
                                }
                                //Coming
                                currentKey == nextItemKey -> {
                                    (itemInfo.offset + startPadding).coerceAtLeast(startPadding)
                                }
                                //Leaving
                                else -> {
                                    itemInfo.offset + startPadding
                                }
                            }
                            Offset(x = offsetX.toFloat(), y = 0f)
                        }
                    ).stickyHeader(currentKey)
                    previousKey = currentKey
                }
            }
        }
    ) {
        LazyRow(
            modifier = modifier.padding(top = stickyHeaderHeight),
            state = state,
            contentPadding = contentPadding,
            reverseLayout = reverseLayout,
            horizontalArrangement = horizontalArrangement,
            verticalAlignment = verticalAlignment,
            flingBehavior = flingBehavior,
            userScrollEnabled = userScrollEnabled
        ) {
            items(
                items = itemsWithKeys
            ) { (_, value) ->
                itemContent(value)
            }
        }
    }
}

fun StickyHeaderScope.drawStickyHeader(
    text: String,
    style: TextStyle,
    color: Color = style.color,
    overflow: TextOverflow = TextOverflow.Clip,
    softWrap: Boolean = true,
    maxLines: Int = Int.MAX_VALUE,
) {
    val textLayout = textMeasurer.measure(
        text = AnnotatedString(text),
        style = style,
        overflow = overflow,
        softWrap = softWrap,
        maxLines = maxLines
    )
    with(drawScope) {
        drawText(
            textLayoutResult = textLayout,
            color = color,
            topLeft = offsetProvider(textLayout.size)
        )
    }
}

interface StickyHeaderScope {
    val drawScope: DrawScope
    val textMeasurer: TextMeasurer
    val offsetProvider: (IntSize) -> Offset
}

private class StickyHeaderScopeImpl(
    override val drawScope: DrawScope,
    override val textMeasurer: TextMeasurer,
    override val offsetProvider: (IntSize) -> Offset
) : StickyHeaderScope

备注：内容来源于stack exchange，提问作者xephosbot