How can I yield in a recursively generated sequence

user_923847 0
Viewed 181

I have a function generating values. It does so by recursively calling itself, and whenever the base case is reached, a value is emitted. This is just a toy example, the real one is much more complicated, and thus harder to transform to something non-recursive.

I would like to do this as a sequence. But doing this won't compile on the line I yield, with the message

Suspension functions can be called only within coroutine body

Here's the attempt:

fun test1() = sequence {
    fun a(i: Int) {
        if (i > 4) {
            yield(i) // <- fails to compile
        } else {
            a(i + 1)
            a(i + 2)
            a(i + 3)
        }
    }
    a(1)
}

If I try to add the suspend keyword to function a, I instead get the error

Restricted suspending functions can only invoke member or extension suspending functions on their restricted coroutine scope

I can make it work, by instead returning all the values and doing a yieldAll. The problem with this, however, is that it needs to do the whole computation and hold all the results in memory before the sequence can be used. For my case, this won't work as I'd run out of memory, or may even have an almost infinite amount of results where I only want to take a few.

So this works, but is not optimal

fun test2() = sequence {
    fun a(i: Int): List<Int> {
        if (i > 4) {
            return listOf(i)
        } else {
            return listOf(
                a(i + 1),
                a(i + 2),
                a(i + 3)
            ).flatten()
        }
    }
    yieldAll(a(1))
}

Any ideas on how to combine sequences with a recursive function, allowing me to yield values without pre-computing and allocating memory for all of them?

kotlin
2 Answers

How about making a return a lazy Sequence?

fun a(i: Int): Sequence<Int> =
    sequence {
        if (i > 4) {
            // println("yielding $i")
            yield(i)
        } else {
            yieldAll(a(i + 1)) // this yieldAll is lazy
            yieldAll(a(i + 2))
            yieldAll(a(i + 3))
        }
    }

fun test1() = a(1)

If you uncomment the println, and loop over the sequence, you will see that it is indeed lazy:

for (i in test1()) {
    println("printing $i")
}

Output:

yielding 5
printing 5
yielding 6
printing 6
yielding 7
printing 7
yielding 5
printing 5
yielding 6
printing 6
yielding 5
printing 5
yielding 6
printing 6
yielding 7
printing 7
...
user_51335850

You can make the inner function a suspend function that is an extension of SequenceScope. The Sequence builder uses a special kind of coroutine that's limited to only calling suspend function extensions of SequenceScope. Presumably this is because Sequences are intended to be used only as synchronous, non-suspending iterables. Since yield() is one of these SequenceScope suspend functions, you have to make your function also a suspend extension function.

fun test1() = sequence {
    suspend fun SequenceScope<Int>.a(i: Int) {
        if (i > 4) {
            yield(i)
        } else {
            a(i + 1)
            a(i + 2)
            a(i + 3)
        }
    }
    a(1)
}
user_5067960
Related