How to collect items from a Flow until a particular condition is met?

I have a Flow<List<Int?>> and I want to collect this flow but only until I get a null Int. Then the flow should get cancelled. For example,

val flow = flowOf(
    listOf(1, 2),
    listOf(3, null),
    listOf(4)
)
flow.collectUntilNull {
    println(it)
}

The output that I want is:

[1, 2]
[3, null]

I know there is a function Flow.takeWhile but it doesn't emit the value where predicate returns false. In my case I want that one also.

public fun <T> Flow<T>.takeWhile(predicate: suspend (T) -> Boolean): Flow<T> = flow {
    return@flow collectWhile { value ->
        if (predicate(value)) {
            emit(value)
            true
        } else {
            // I want a "emit(value)" here
            false
        }
    }
}

Since collectWhile is internal I can't use this code. Although I guess I can copy paste that collectWhile implementation in my code. But is there another way to do this?

transformWhile is a more flexible/generalized takeWhile.

flow
    .transformWhile {
        emit(it)
        myCondition(it)
    }
    .collect {
        println(it)
    }