Fix implementation of Iterable<T>.take

#KT-4780 Fixed
JetBrains · Mar 28, 2014 · e3fffe2 · e3fffe2
1 parent 829cd95
commit e3fffe2
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 8 deletions.
diff --git a/libraries/stdlib/src/generated/_Filtering.kt b/libraries/stdlib/src/generated/_Filtering.kt
@@ -1001,9 +1001,11 @@ public fun <T> Collection<T>.take(n: Int) : List<T> {
 public fun <T> Iterable<T>.take(n: Int) : List<T> {
     var count = 0
     val list = ArrayList<T>(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
@@ -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<Float>(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<Long>()) { (0L..5L) take 0 }
+        expect(listOf<Long>()) { 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
@@ -51,9 +51,11 @@ fun filtering(): List<GenericFunction> {
             """
             var count = 0
             val list = ArrayList<T>(n)
-            for (item in this)
+            for (item in this) {
                 if (count++ >= n)
-                    list.add(item)
+                    break
+                list.add(item)
+            }
             return list
             """
         }