Jetpack Compose LazyRow custom sticky header

While creating a UI on Jetpack Compose, there was a need for a sticky header for LazyRow. But the current implementation is embedded in the list as a row element.

I would like the sticky header to be above the elements like here: Example

UPD: I started to try to solve the problem myself but I am having the following problems:

When my OffsetX is in Leaving state then it breaks when adding itemSpacing for LazyRow

Also with a long StickyHeader, I can't figure out how to properly set the offset to take its length into account.

@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

Demonstration of my implementation: video

I've worked on a Compose Multiplatform library that provides solution for this:

https://github.com/gregkorossy/lazy-sticky-headers

Preview showing both horizontal and vertical sticky headers:

P.S.: It's usually not cool to provide only a link to a solution, but in this case it would be difficult to copy the source code over here.