diff --git a/libraries/stdlib/src/generated/_Filtering.kt b/libraries/stdlib/src/generated/_Filtering.kt index 9dcba9e8c0bc7..b09784587ef19 100644 --- a/libraries/stdlib/src/generated/_Filtering.kt +++ b/libraries/stdlib/src/generated/_Filtering.kt @@ -1001,9 +1001,11 @@ public fun Collection.take(n: Int) : List { public fun Iterable.take(n: Int) : List { var count = 0 val list = ArrayList(n) - for (item in this) - if (count++ >= n) - list.add(item) + for (item in this) { + if (count++ == n) + break + list.add(item) + } return list } diff --git a/libraries/stdlib/test/collections/CollectionTest.kt b/libraries/stdlib/test/collections/CollectionTest.kt index 8f431cb91a5e4..5fe212678903c 100644 --- a/libraries/stdlib/test/collections/CollectionTest.kt +++ b/libraries/stdlib/test/collections/CollectionTest.kt @@ -1,9 +1,7 @@ package test.collections -import kotlin.test.* - import java.util.* - +import kotlin.test.* import org.junit.Test as test class CollectionTest { @@ -408,4 +406,15 @@ class CollectionTest { expect(3.0.toFloat()) { arrayListOf(1.0.toFloat(), 2.0.toFloat()).sum() } } + test fun takeReturnsFirstNElements() { + expect(listOf(1, 2, 3, 4, 5)) { (1..10) take 5 } + expect(listOf(1, 2, 3, 4, 5)) { (1..10).toList().take(5) } + expect(listOf(1, 2)) { (1..10) take 2 } + expect(listOf(1, 2)) { (1..10).toList().take(2) } + expect(listOf()) { (0L..5L) take 0 } + expect(listOf()) { listOf(1L) take 0 } + expect(listOf(1)) { (1..1) take 10 } + expect(listOf(1)) { listOf(1) take 10 } + expect(setOf(1, 2)) { sortedSetOf(1, 2, 3, 4, 5).take(2).toSet() } + } } diff --git a/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt b/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt index 32f9f897703c3..b20b297579a15 100644 --- a/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt +++ b/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt @@ -51,9 +51,11 @@ fun filtering(): List { """ var count = 0 val list = ArrayList(n) - for (item in this) + for (item in this) { if (count++ >= n) - list.add(item) + break + list.add(item) + } return list """ }