diff --git a/libraries/stdlib/src/generated/_Sequences.kt b/libraries/stdlib/src/generated/_Sequences.kt index 64c60dae6ab..7aed5ab8476 100644 --- a/libraries/stdlib/src/generated/_Sequences.kt +++ b/libraries/stdlib/src/generated/_Sequences.kt @@ -568,7 +568,7 @@ public fun Sequence.toSortedSet(comparator: Comparator): SortedSet< * Returns a single sequence of all elements from results of [transform] function being invoked on each element of original sequence. */ public fun Sequence.flatMap(transform: (T) -> Sequence): Sequence { - return FlatteningSequence(this, transform) + return FlatteningSequence(this, transform, { it.iterator() }) } /** diff --git a/libraries/stdlib/src/kotlin/collections/Sequence.kt b/libraries/stdlib/src/kotlin/collections/Sequence.kt index c1152f82b7e..6e3e560c12f 100644 --- a/libraries/stdlib/src/kotlin/collections/Sequence.kt +++ b/libraries/stdlib/src/kotlin/collections/Sequence.kt @@ -68,9 +68,21 @@ private object EmptySequence : Sequence { /** * Returns a sequence of all elements from all sequences in this sequence. */ -public fun Sequence>.flatten(): Sequence { - return MultiSequence(this) +public fun Sequence>.flatten(): Sequence = flatten { it.iterator() } + +/** + * Returns a sequence of all elements from all iterables in this sequence. + */ +@kotlin.jvm.JvmName("flattenSequenceOfIterable") +public fun Sequence>.flatten(): Sequence = flatten { it.iterator() } + +private fun Sequence.flatten(iterator: (T) -> Iterator): Sequence { + if (this is TransformingSequence<*, *>) { + return (this as TransformingSequence<*, T>).flatten(iterator) + } + return FlatteningSequence(this, { it }, iterator) } + /** * Returns a pair of lists, where * *first* list is built from the first values of each pair from this sequence, @@ -151,6 +163,10 @@ constructor(private val sequence: Sequence, private val transformer: (T) -> R return iterator.hasNext() } } + + internal fun flatten(iterator: (R) -> Iterator): Sequence { + return FlatteningSequence(sequence, transformer, iterator) + } } /** @@ -215,15 +231,17 @@ internal class MergingSequence } } -internal class FlatteningSequence - constructor(private val sequence: Sequence, - private val transformer: (T) -> Sequence - ) : Sequence { - override fun iterator(): Iterator = object : Iterator { +internal class FlatteningSequence + constructor( + private val sequence: Sequence, + private val transformer: (T) -> R, + private val iterator: (R) -> Iterator + ) : Sequence { + override fun iterator(): Iterator = object : Iterator { val iterator = sequence.iterator() - var itemIterator: Iterator? = null + var itemIterator: Iterator? = null - override fun next(): R { + override fun next(): E { if (!ensureItemIterator()) throw NoSuchElementException() return itemIterator!!.next() @@ -242,44 +260,7 @@ internal class FlatteningSequence return false } else { val element = iterator.next() - val nextItemIterator = transformer(element).iterator() - if (nextItemIterator.hasNext()) { - itemIterator = nextItemIterator - return true - } - } - } - return true - } - } -} - -internal class MultiSequence -constructor(private val sequence: Sequence>) : Sequence { - override fun iterator(): Iterator = object : Iterator { - val iterator = sequence.iterator() - var itemIterator: Iterator? = null - - override fun next(): T { - if (!ensureItemIterator()) - throw NoSuchElementException() - return itemIterator!!.next() - } - - override fun hasNext(): Boolean { - return ensureItemIterator() - } - - private fun ensureItemIterator(): Boolean { - if (itemIterator?.hasNext() == false) - itemIterator = null - - while (itemIterator == null) { - if (!iterator.hasNext()) { - return false - } else { - val element = iterator.next() - val nextItemIterator = element.iterator() + val nextItemIterator = iterator(transformer(element)) if (nextItemIterator.hasNext()) { itemIterator = nextItemIterator return true diff --git a/libraries/stdlib/test/collections/SequenceTest.kt b/libraries/stdlib/test/collections/SequenceTest.kt index 35af6513cd7..11aaa9531f8 100644 --- a/libraries/stdlib/test/collections/SequenceTest.kt +++ b/libraries/stdlib/test/collections/SequenceTest.kt @@ -319,8 +319,19 @@ public class SequenceTest { } @test fun flatten() { - val data = sequenceOf(1, 2).map { sequenceOf(0..it) } - assertEquals(listOf(0, 1, 0, 1, 2), data.flatten().toList()) + val expected = listOf(0, 1, 0, 1, 2) + + val seq = sequenceOf((0..1).asSequence(), (0..2).asSequence()).flatten() + assertEquals(expected, seq.toList()) + + val seqMappedSeq = sequenceOf(1, 2).map { (0..it).asSequence() }.flatten() + assertEquals(expected, seqMappedSeq.toList()) + + val seqOfIterable = sequenceOf(0..1, 0..2).flatten() + assertEquals(expected, seqOfIterable.toList()) + + val seqMappedIterable = sequenceOf(1, 2).map { 0..it }.flatten() + assertEquals(expected, seqMappedIterable.toList()) } @test fun distinct() { diff --git a/libraries/tools/kotlin-stdlib-gen/src/templates/Mapping.kt b/libraries/tools/kotlin-stdlib-gen/src/templates/Mapping.kt index f3c4f311470..53e6447fdf4 100644 --- a/libraries/tools/kotlin-stdlib-gen/src/templates/Mapping.kt +++ b/libraries/tools/kotlin-stdlib-gen/src/templates/Mapping.kt @@ -235,7 +235,7 @@ fun mapping(): List { typeParam("R") returns("Sequence") body { - "return FlatteningSequence(this, transform)" + "return FlatteningSequence(this, transform, { it.iterator() })" } }