Implementing takeWhileInclusive in Kotlin

written June 2, 2018 in collections, extensions, kotlin, takewhile

Implementing takeWhileInclusive extension function in Kotlin.

TL,DR (aka “just show me the code” ): https://gist.github.com/jivimberg/ff5aad3f5c6315deb420fd508a145c61

You probably know about takeWhile operation that returns a List containing the first elements satisfying the given predicate.

fun main(args: Array<String>) { //sampleStart val someNumbers = listOf(1, 5, 3, 22, 4, 8, 14, 23, 49, 77, 2, 49) println(someNumbers.takeWhile { it % 7 != 0 }) //sampleEnd }

I was in need of an inclusive version of the takeWhile. In other words I needed a function that returned the first elements satisfying the given predicate, plus the first element that didn’t satisfy it.

So in the provided example takeWhile returns [1, 5, 3, 22, 4, 8] whereas takeWhileInclusive would return [1, 5, 3, 22, 4, 8, 14] .

A quick search showed me I was not alone. matklad already had a simple implementation working for Sequence:

fun <T> Sequence<T>.takeWhileInclusive(pred: (T) -> Boolean): Sequence<T> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = pred(it) result } }

Using extension functions we are able to add a new function to the standard library type Sequence. If this is your first encounter with an extension function I’d encourage you to read more about them here and then play with them here.

I found this implementation of takeWhileInclusive quite elegant. It uses the original takeWhile with the given predicated, but keeps a shouldContinue variable to delay the predicate evaluation by one step. In other words the evaluation of the predicate passed to takeWhile on element i will actually be the result of applying the predicate function on i - 1. Which if you think about it is exactly what we need. Let’s give it a try:

fun <T> Sequence<T>.takeWhileInclusive( predicate: (T) -> Boolean ): Sequence<T> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } fun main(args: Array<String>) { //sampleStart val someNumbers = listOf(1, 5, 3, 22, 4, 8, 14, 23, 49, 77, 2, 49).asSequence() println(someNumbers.takeWhileInclusive { it % 7 != 0 }.toList()) //sampleEnd }

This was good! Only this comment posted on the original gist made me doubt about the safety of this approach:

“I love this. If Sequence were parallel, though, wouldn’t there be worries about using out-of-closure state?”

That send me down the rabbit hole and I spent some days reading about parallelStreams and coroutines until I convinced my self the approach was ok.¹

With that out of the way I decided to port this solution to be supported in all the places where the takeWhile exists (including String and CharArray). So that we don’t have to convert to a from a Sequence just to use this function.

Here’s the end result, ready to be dropped on your project:

// kotlin.collections inline fun <T> Array<out T>.takeWhileInclusive( predicate: (T) -> Boolean ): List<T> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } inline fun ByteArray.takeWhileInclusive( predicate: (Byte) -> Boolean ): List<Byte> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } inline fun ShortArray.takeWhileInclusive( predicate: (Short) -> Boolean ): List<Short> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } inline fun IntArray.takeWhileInclusive( predicate: (Int) -> Boolean ): List<Int> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } inline fun LongArray.takeWhileInclusive( predicate: (Long) -> Boolean ): List<Long> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } inline fun FloatArray.takeWhileInclusive( predicate: (Float) -> Boolean ): List<Float> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } inline fun DoubleArray.takeWhileInclusive( predicate: (Double) -> Boolean ): List<Double> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } inline fun BooleanArray.takeWhileInclusive( predicate: (Boolean) -> Boolean ): List<Boolean> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } inline fun CharArray.takeWhileInclusive( predicate: (Char) -> Boolean ): List<Char> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } inline fun <T> Iterable<T>.takeWhileInclusive( predicate: (T) -> Boolean ): List<T> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } // kotlin.sequences fun <T> Sequence<T>.takeWhileInclusive( predicate: (T) -> Boolean ): Sequence<T> { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } // kotlin.text inline fun CharSequence.takeWhileInclusive( predicate: (Char) -> Boolean ): CharSequence { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } } inline fun String.takeWhileInclusive( predicate: (Char) -> Boolean ): String { var shouldContinue = true return takeWhile { val result = shouldContinue shouldContinue = predicate(it) result } }

Hope you find this useful!

If you want to read more about my discoveries, check this question at StackOverflow ↩

Implementing takeWhileInclusive in Kotlin

Comments