takeWhile which includes the actual value matching the predicate (takeWhileInclusive)

Question

Often I see the problem that I want to collect values from a list until a value matches, but where I also need the matching value itself. While takeWhile is nearly perfect for that matter, it actually doesn't allow keeping the last (or basically the matching) entry.

A simple example: show the class hierarchy of an object up to the first class that implements a specific interface

generateSequence(obj::class.java, Class<*>::getSuperclass)
        .takeWhile { interestedType !in it.interfaces }
        .joinToString(" > ")
        .run(::println)

for obj=arrayListOf(1) and interestedType=Collection::class.java I want to see something like:

class java.util.ArrayList > class java.util.AbstractList > class java.util.AbstractCollection

and I hoped it would be as easy as:

generateSequence(obj::class.java, Class<*>::getSuperclass)
        .takeWhileInclusive { interestedType !in it.interfaces }
        .joinToString(" > ")
        .run(::println)

But such a function doesn't exist (yet?). But maybe there is some other function that really comes close to that? Or maybe with at most 2 consecutive function calls it is already easily implementable and I just don't see it?

What I am not looking for: how I can solve that particular issue regarding which class in the hierarchy implements an interface. That's just a simple example. What I am also not looking for: how I can implement that with Iterator or a basic while-/for-loop... (except: if it is easily readable and doesn't take more than 3 lines then... maybe ;-)).

What I found: Is this implementation of takeWhileInclusive safe? which also links its own implementation (and its inspiration) for takeWhileInclusive. However I don't really like that it's using the var to register whether it found a match... I also am a bit unsure when I read the comments ("assume sequential order") whether this implementation really makes sense/really is safe.

Answer 1

I did not find a suitable existing function yet and I also do not really like the linked solution, so I played around a bit and ended up with the following extension function:

fun <T> Sequence<T>.takeWhileInclusive(predicate: (T) -> Boolean) = sequence {
    with(iterator()) {
        while (hasNext()) {
            val next = next()
            yield(next)
            if (!predicate(next)) break
        }
    }
}

It makes use of sequence, lazily yielding values when needed. At least I can omit the intermediate var and I assume that it might be more beneficial that way...

Answer 2

Sadly several of the functions used for takeWhile are internal. But if we copy them into an extension file we can build takeWhileInclusive also for flows:

import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.FlowCollector
import kotlinx.coroutines.flow.flow
import java.util.concurrent.CancellationException

fun <T> Flow<T>.takeWhileInclusive(predicate: suspend (T) -> Boolean): Flow<T> = flow {
    // This return is needed to work around a bug in JS BE: KT-39227
    return@flow collectWhile { value ->
        emit(value)
        predicate(value)
    }
}

private suspend inline fun <T> Flow<T>.collectWhile(crossinline predicate: suspend (value: T) -> Boolean) {
    val collector = object : FlowCollector<T> {
        override suspend fun emit(value: T) {
            // Note: we are checking predicate first, then throw. If the predicate does suspend (calls emit, for example)
            // the the resulting code is never tail-suspending and produces a state-machine
            if (!predicate(value)) {
                throw AbortFlowException(this)
            }
        }
    }
    try {
        collect(collector)
    } catch (e: AbortFlowException) {
        e.checkOwnership(collector)
    }
}

private class AbortFlowException constructor(
    @JvmField @Transient val owner: FlowCollector<*>
) : CancellationException("Flow was aborted, no more elements needed") {

    override fun fillInStackTrace(): Throwable {
        // Prevent Android <= 6.0 bug, #1866
        stackTrace = emptyArray()
        return this
    }
}

private fun AbortFlowException.checkOwnership(owner: FlowCollector<*>) {
    if (this.owner !== owner) throw this
}

I don't like the copy paste of kotlin internal code, but haven't found a better way to achieve it.

Answer 3

Nowaday, you can use transformWhile() this way:

flow
  .transformWhile {
    emit(it)
    condition(it)
  }