From 25cfd36e7c201e7bb24fa7222b896cf833c3200e Mon Sep 17 00:00:00 2001
From: Ilya Gorbunov <ilya.a.gorbunov@yandex.ru>
Date: Tue, 24 Mar 2015 03:33:34 +0300
Subject: [PATCH] drop, take: combine implementations for Collections, Lists
 and Iterables. slice: predict length of resulting list.

---
 libraries/stdlib/src/generated/_Filtering.kt  | 129 +++++++-----------
 .../src/templates/Filtering.kt                |  63 +++++----
 2 files changed, 88 insertions(+), 104 deletions(-)

diff --git a/libraries/stdlib/src/generated/_Filtering.kt b/libraries/stdlib/src/generated/_Filtering.kt
index 5e2f230caac..3666e5e9d8d 100644
--- a/libraries/stdlib/src/generated/_Filtering.kt
+++ b/libraries/stdlib/src/generated/_Filtering.kt
@@ -135,48 +135,34 @@ public fun ShortArray.drop(n: Int): List<Short> {
     return list
 }
 
-/**
- * Returns a list containing all elements except first [n] elements
- */
-public fun <T> Collection<T>.drop(n: Int): List<T> {
-    require(n >= 0, { "Requested element count $n is less than zero." })
-    if (n >= size())
-        return emptyList()
-    var count = 0
-    val list = ArrayList<T>(size() - n)
-    for (item in this) {
-        if (count++ >= n) list.add(item)
-    }
-    return list
-}
-
 /**
  * Returns a list containing all elements except first [n] elements
  */
 public fun <T> Iterable<T>.drop(n: Int): List<T> {
     require(n >= 0, { "Requested element count $n is less than zero." })
+    val list: ArrayList<T>
+    if (this is Collection<*>) {
+        val resultSize = size() - n
+        if (resultSize <= 0)
+            return emptyList()
+        list = ArrayList<T>(resultSize)
+        if (this is List<T>) {
+            for (index in n..size() - 1) {
+                list.add(this[index])
+            }
+            return list
+        }
+    }
+    else {
+        list = ArrayList<T>()
+    }
     var count = 0
-    val list = ArrayList<T>()
     for (item in this) {
         if (count++ >= n) list.add(item)
     }
     return list
 }
 
-/**
- * Returns a list containing all elements except first [n] elements
- */
-public fun <T> List<T>.drop(n: Int): List<T> {
-    require(n >= 0, { "Requested element count $n is less than zero." })
-    if (n >= size())
-        return emptyList()
-    val list = ArrayList<T>(size() - n)
-    for (index in n..size() - 1) {
-        list.add(this[index])
-    }
-    return list
-}
-
 /**
  * Returns a sequence containing all elements except first [n] elements
  */
@@ -858,7 +844,7 @@ public inline fun <C : Appendable> String.filterTo(destination: C, predicate: (C
  * Returns a list containing elements at specified positions
  */
 public fun <T> Array<out T>.slice(indices: Iterable<Int>): List<T> {
-    val list = ArrayList<T>()
+    val list = ArrayList<T>(indices.collectionSizeOrDefault(10))
     for (index in indices) {
         list.add(get(index))
     }
@@ -869,7 +855,7 @@ public fun <T> Array<out T>.slice(indices: Iterable<Int>): List<T> {
  * Returns a list containing elements at specified positions
  */
 public fun BooleanArray.slice(indices: Iterable<Int>): List<Boolean> {
-    val list = ArrayList<Boolean>()
+    val list = ArrayList<Boolean>(indices.collectionSizeOrDefault(10))
     for (index in indices) {
         list.add(get(index))
     }
@@ -880,7 +866,7 @@ public fun BooleanArray.slice(indices: Iterable<Int>): List<Boolean> {
  * Returns a list containing elements at specified positions
  */
 public fun ByteArray.slice(indices: Iterable<Int>): List<Byte> {
-    val list = ArrayList<Byte>()
+    val list = ArrayList<Byte>(indices.collectionSizeOrDefault(10))
     for (index in indices) {
         list.add(get(index))
     }
@@ -891,7 +877,7 @@ public fun ByteArray.slice(indices: Iterable<Int>): List<Byte> {
  * Returns a list containing elements at specified positions
  */
 public fun CharArray.slice(indices: Iterable<Int>): List<Char> {
-    val list = ArrayList<Char>()
+    val list = ArrayList<Char>(indices.collectionSizeOrDefault(10))
     for (index in indices) {
         list.add(get(index))
     }
@@ -902,7 +888,7 @@ public fun CharArray.slice(indices: Iterable<Int>): List<Char> {
  * Returns a list containing elements at specified positions
  */
 public fun DoubleArray.slice(indices: Iterable<Int>): List<Double> {
-    val list = ArrayList<Double>()
+    val list = ArrayList<Double>(indices.collectionSizeOrDefault(10))
     for (index in indices) {
         list.add(get(index))
     }
@@ -913,7 +899,7 @@ public fun DoubleArray.slice(indices: Iterable<Int>): List<Double> {
  * Returns a list containing elements at specified positions
  */
 public fun FloatArray.slice(indices: Iterable<Int>): List<Float> {
-    val list = ArrayList<Float>()
+    val list = ArrayList<Float>(indices.collectionSizeOrDefault(10))
     for (index in indices) {
         list.add(get(index))
     }
@@ -924,7 +910,7 @@ public fun FloatArray.slice(indices: Iterable<Int>): List<Float> {
  * Returns a list containing elements at specified positions
  */
 public fun IntArray.slice(indices: Iterable<Int>): List<Int> {
-    val list = ArrayList<Int>()
+    val list = ArrayList<Int>(indices.collectionSizeOrDefault(10))
     for (index in indices) {
         list.add(get(index))
     }
@@ -935,7 +921,7 @@ public fun IntArray.slice(indices: Iterable<Int>): List<Int> {
  * Returns a list containing elements at specified positions
  */
 public fun LongArray.slice(indices: Iterable<Int>): List<Long> {
-    val list = ArrayList<Long>()
+    val list = ArrayList<Long>(indices.collectionSizeOrDefault(10))
     for (index in indices) {
         list.add(get(index))
     }
@@ -946,7 +932,7 @@ public fun LongArray.slice(indices: Iterable<Int>): List<Long> {
  * Returns a list containing elements at specified positions
  */
 public fun ShortArray.slice(indices: Iterable<Int>): List<Short> {
-    val list = ArrayList<Short>()
+    val list = ArrayList<Short>(indices.collectionSizeOrDefault(10))
     for (index in indices) {
         list.add(get(index))
     }
@@ -957,7 +943,7 @@ public fun ShortArray.slice(indices: Iterable<Int>): List<Short> {
  * Returns a list containing elements at specified positions
  */
 public fun <T> List<T>.slice(indices: Iterable<Int>): List<T> {
-    val list = ArrayList<T>()
+    val list = ArrayList<T>(indices.collectionSizeOrDefault(10))
     for (index in indices) {
         list.add(get(index))
     }
@@ -968,7 +954,7 @@ public fun <T> List<T>.slice(indices: Iterable<Int>): List<T> {
  * Returns a string containing characters at specified positions
  */
 public fun String.slice(indices: Iterable<Int>): String {
-    val result = StringBuilder()
+    val result = StringBuilder(indices.collectionSizeOrDefault(10))
     for (i in indices) {
         result.append(get(i))
     }
@@ -981,7 +967,7 @@ public fun String.slice(indices: Iterable<Int>): String {
 public fun <T> Array<out T>.take(n: Int): List<T> {
     require(n >= 0, "Requested element count $n is less than zero.")
     var count = 0
-    val realN = if (n > size()) size() else n
+    val realN = Math.min(n, size())
     val list = ArrayList<T>(realN)
     for (item in this) {
         if (count++ == realN)
@@ -997,7 +983,7 @@ public fun <T> Array<out T>.take(n: Int): List<T> {
 public fun BooleanArray.take(n: Int): List<Boolean> {
     require(n >= 0, "Requested element count $n is less than zero.")
     var count = 0
-    val realN = if (n > size()) size() else n
+    val realN = Math.min(n, size())
     val list = ArrayList<Boolean>(realN)
     for (item in this) {
         if (count++ == realN)
@@ -1013,7 +999,7 @@ public fun BooleanArray.take(n: Int): List<Boolean> {
 public fun ByteArray.take(n: Int): List<Byte> {
     require(n >= 0, "Requested element count $n is less than zero.")
     var count = 0
-    val realN = if (n > size()) size() else n
+    val realN = Math.min(n, size())
     val list = ArrayList<Byte>(realN)
     for (item in this) {
         if (count++ == realN)
@@ -1029,7 +1015,7 @@ public fun ByteArray.take(n: Int): List<Byte> {
 public fun CharArray.take(n: Int): List<Char> {
     require(n >= 0, "Requested element count $n is less than zero.")
     var count = 0
-    val realN = if (n > size()) size() else n
+    val realN = Math.min(n, size())
     val list = ArrayList<Char>(realN)
     for (item in this) {
         if (count++ == realN)
@@ -1045,7 +1031,7 @@ public fun CharArray.take(n: Int): List<Char> {
 public fun DoubleArray.take(n: Int): List<Double> {
     require(n >= 0, "Requested element count $n is less than zero.")
     var count = 0
-    val realN = if (n > size()) size() else n
+    val realN = Math.min(n, size())
     val list = ArrayList<Double>(realN)
     for (item in this) {
         if (count++ == realN)
@@ -1061,7 +1047,7 @@ public fun DoubleArray.take(n: Int): List<Double> {
 public fun FloatArray.take(n: Int): List<Float> {
     require(n >= 0, "Requested element count $n is less than zero.")
     var count = 0
-    val realN = if (n > size()) size() else n
+    val realN = Math.min(n, size())
     val list = ArrayList<Float>(realN)
     for (item in this) {
         if (count++ == realN)
@@ -1077,7 +1063,7 @@ public fun FloatArray.take(n: Int): List<Float> {
 public fun IntArray.take(n: Int): List<Int> {
     require(n >= 0, "Requested element count $n is less than zero.")
     var count = 0
-    val realN = if (n > size()) size() else n
+    val realN = Math.min(n, size())
     val list = ArrayList<Int>(realN)
     for (item in this) {
         if (count++ == realN)
@@ -1093,7 +1079,7 @@ public fun IntArray.take(n: Int): List<Int> {
 public fun LongArray.take(n: Int): List<Long> {
     require(n >= 0, "Requested element count $n is less than zero.")
     var count = 0
-    val realN = if (n > size()) size() else n
+    val realN = Math.min(n, size())
     val list = ArrayList<Long>(realN)
     for (item in this) {
         if (count++ == realN)
@@ -1109,7 +1095,7 @@ public fun LongArray.take(n: Int): List<Long> {
 public fun ShortArray.take(n: Int): List<Short> {
     require(n >= 0, "Requested element count $n is less than zero.")
     var count = 0
-    val realN = if (n > size()) size() else n
+    val realN = Math.min(n, size())
     val list = ArrayList<Short>(realN)
     for (item in this) {
         if (count++ == realN)
@@ -1119,29 +1105,13 @@ public fun ShortArray.take(n: Int): List<Short> {
     return list
 }
 
-/**
- * Returns a list containing first [n] elements
- */
-public fun <T> Collection<T>.take(n: Int): List<T> {
-    require(n >= 0, "Requested element count $n is less than zero.")
-    var count = 0
-    val realN = if (n > size()) size() else n
-    val list = ArrayList<T>(realN)
-    for (item in this) {
-        if (count++ == realN)
-            break;
-        list.add(item)
-    }
-    return list
-}
-
 /**
  * Returns a list containing first [n] elements
  */
 public fun <T> Iterable<T>.take(n: Int): List<T> {
     require(n >= 0, { "Requested element count $n is less than zero." })
     var count = 0
-    val list = ArrayList<T>(n)
+    val list = ArrayList<T>(Math.min(n, collectionSizeOrDefault(n)))
     for (item in this) {
         if (count++ == n)
             break
@@ -1172,6 +1142,7 @@ public fun <T> Stream<T>.take(n: Int): Stream<T> {
  * Returns a string containing the first [n] characters from this string, or the entire string if this string is shorter
  */
 public fun String.take(n: Int): String {
+    require(n >= 0, { "Requested element count $n is less than zero." })
     return substring(0, Math.min(n, length()))
 }
 
@@ -1181,7 +1152,7 @@ public fun String.take(n: Int): String {
 public fun <T> Array<out T>.takeLast(n: Int): List<T> {
     require(n >= 0, { "Requested element count $n is less than zero." })
     val size = size()
-    val realN = if (n > size) size else n
+    val realN = Math.min(n, size)
     val list = ArrayList<T>(realN)
     for (index in size - realN .. size - 1)
         list.add(this[index])
@@ -1194,7 +1165,7 @@ public fun <T> Array<out T>.takeLast(n: Int): List<T> {
 public fun BooleanArray.takeLast(n: Int): List<Boolean> {
     require(n >= 0, { "Requested element count $n is less than zero." })
     val size = size()
-    val realN = if (n > size) size else n
+    val realN = Math.min(n, size)
     val list = ArrayList<Boolean>(realN)
     for (index in size - realN .. size - 1)
         list.add(this[index])
@@ -1207,7 +1178,7 @@ public fun BooleanArray.takeLast(n: Int): List<Boolean> {
 public fun ByteArray.takeLast(n: Int): List<Byte> {
     require(n >= 0, { "Requested element count $n is less than zero." })
     val size = size()
-    val realN = if (n > size) size else n
+    val realN = Math.min(n, size)
     val list = ArrayList<Byte>(realN)
     for (index in size - realN .. size - 1)
         list.add(this[index])
@@ -1220,7 +1191,7 @@ public fun ByteArray.takeLast(n: Int): List<Byte> {
 public fun CharArray.takeLast(n: Int): List<Char> {
     require(n >= 0, { "Requested element count $n is less than zero." })
     val size = size()
-    val realN = if (n > size) size else n
+    val realN = Math.min(n, size)
     val list = ArrayList<Char>(realN)
     for (index in size - realN .. size - 1)
         list.add(this[index])
@@ -1233,7 +1204,7 @@ public fun CharArray.takeLast(n: Int): List<Char> {
 public fun DoubleArray.takeLast(n: Int): List<Double> {
     require(n >= 0, { "Requested element count $n is less than zero." })
     val size = size()
-    val realN = if (n > size) size else n
+    val realN = Math.min(n, size)
     val list = ArrayList<Double>(realN)
     for (index in size - realN .. size - 1)
         list.add(this[index])
@@ -1246,7 +1217,7 @@ public fun DoubleArray.takeLast(n: Int): List<Double> {
 public fun FloatArray.takeLast(n: Int): List<Float> {
     require(n >= 0, { "Requested element count $n is less than zero." })
     val size = size()
-    val realN = if (n > size) size else n
+    val realN = Math.min(n, size)
     val list = ArrayList<Float>(realN)
     for (index in size - realN .. size - 1)
         list.add(this[index])
@@ -1259,7 +1230,7 @@ public fun FloatArray.takeLast(n: Int): List<Float> {
 public fun IntArray.takeLast(n: Int): List<Int> {
     require(n >= 0, { "Requested element count $n is less than zero." })
     val size = size()
-    val realN = if (n > size) size else n
+    val realN = Math.min(n, size)
     val list = ArrayList<Int>(realN)
     for (index in size - realN .. size - 1)
         list.add(this[index])
@@ -1272,7 +1243,7 @@ public fun IntArray.takeLast(n: Int): List<Int> {
 public fun LongArray.takeLast(n: Int): List<Long> {
     require(n >= 0, { "Requested element count $n is less than zero." })
     val size = size()
-    val realN = if (n > size) size else n
+    val realN = Math.min(n, size)
     val list = ArrayList<Long>(realN)
     for (index in size - realN .. size - 1)
         list.add(this[index])
@@ -1285,7 +1256,7 @@ public fun LongArray.takeLast(n: Int): List<Long> {
 public fun ShortArray.takeLast(n: Int): List<Short> {
     require(n >= 0, { "Requested element count $n is less than zero." })
     val size = size()
-    val realN = if (n > size) size else n
+    val realN = Math.min(n, size)
     val list = ArrayList<Short>(realN)
     for (index in size - realN .. size - 1)
         list.add(this[index])
@@ -1298,7 +1269,7 @@ public fun ShortArray.takeLast(n: Int): List<Short> {
 public fun <T> List<T>.takeLast(n: Int): List<T> {
     require(n >= 0, { "Requested element count $n is less than zero." })
     val size = size()
-    val realN = if (n > size) size else n
+    val realN = Math.min(n, size)
     val list = ArrayList<T>(realN)
     for (index in size - realN .. size - 1)
         list.add(this[index])
@@ -1309,7 +1280,9 @@ public fun <T> List<T>.takeLast(n: Int): List<T> {
  * Returns a string containing the last [n] characters from this string, or the entire string if this string is shorter
  */
 public fun String.takeLast(n: Int): String {
-    return substring(length() - Math.min(n, length()), length())
+    require(n >= 0, { "Requested element count $n is less than zero." })
+    val length = length()
+    return substring(length - Math.min(n, length), length)
 }
 
 /**
diff --git a/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt b/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt
index 5abba8d820d..6845d4428f3 100644
--- a/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt
+++ b/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt
@@ -12,8 +12,24 @@ fun filtering(): List<GenericFunction> {
         body {
             """
             require(n >= 0, { "Requested element count $n is less than zero." })
+            val list: ArrayList<T>
+            if (this is Collection<*>) {
+                val resultSize = size() - n
+                if (resultSize <= 0)
+                    return emptyList()
+
+                list = ArrayList<T>(resultSize)
+                if (this is List<T>) {
+                    for (index in n..size() - 1) {
+                        list.add(this[index])
+                    }
+                    return list
+                }
+            }
+            else {
+                list = ArrayList<T>()
+            }
             var count = 0
-            val list = ArrayList<T>()
             for (item in this) {
                 if (count++ >= n) list.add(item)
             }
@@ -34,22 +50,7 @@ fun filtering(): List<GenericFunction> {
         body(Strings) { "return substring(Math.min(n, length()))" }
         returns(Strings) { "String" }
 
-        body(Collections) {
-            """
-            require(n >= 0, { "Requested element count $n is less than zero." })
-            if (n >= size())
-                return emptyList()
-
-            var count = 0
-            val list = ArrayList<T>(size() - n)
-            for (item in this) {
-                if (count++ >= n) list.add(item)
-            }
-            return list
-            """
-        }
-
-        body(Lists, ArraysOfObjects, ArraysOfPrimitives) {
+        body(ArraysOfObjects, ArraysOfPrimitives) {
             """
             require(n >= 0, { "Requested element count $n is less than zero." })
             if (n >= size())
@@ -72,7 +73,7 @@ fun filtering(): List<GenericFunction> {
             """
             require(n >= 0, { "Requested element count $n is less than zero." })
             var count = 0
-            val list = ArrayList<T>(n)
+            val list = ArrayList<T>(Math.min(n, collectionSizeOrDefault(n)))
             for (item in this) {
                 if (count++ == n)
                     break
@@ -83,7 +84,12 @@ fun filtering(): List<GenericFunction> {
         }
 
         doc(Strings) { "Returns a string containing the first [n] characters from this string, or the entire string if this string is shorter"}
-        body(Strings) { "return substring(0, Math.min(n, length()))" }
+        body(Strings) {
+            """
+            require(n >= 0, { "Requested element count $n is less than zero." })
+            return substring(0, Math.min(n, length()))
+            """
+        }
         returns(Strings) { "String" }
 
         doc(Sequences) { "Returns a sequence containing first *n* elements" }
@@ -95,12 +101,11 @@ fun filtering(): List<GenericFunction> {
             """
         }
 
-        include(Collections)
-        body(Collections, ArraysOfObjects, ArraysOfPrimitives) {
+        body(ArraysOfObjects, ArraysOfPrimitives) {
             """
             require(n >= 0, "Requested element count $n is less than zero.")
             var count = 0
-            val realN = if (n > size()) size() else n
+            val realN = Math.min(n, size())
             val list = ArrayList<T>(realN)
             for (item in this) {
                 if (count++ == realN)
@@ -119,14 +124,20 @@ fun filtering(): List<GenericFunction> {
         returns("List<T>")
 
         doc(Strings) { "Returns a string containing the last [n] characters from this string, or the entire string if this string is shorter"}
-        body(Strings) { "return substring(length() - Math.min(n, length()), length())" }
+        body(Strings) {
+            """
+            require(n >= 0, { "Requested element count $n is less than zero." })
+            val length = length()
+            return substring(length - Math.min(n, length), length)
+            """
+        }
         returns(Strings) { "String" }
 
         body(Lists, ArraysOfObjects, ArraysOfPrimitives) {
             """
             require(n >= 0, { "Requested element count $n is less than zero." })
             val size = size()
-            val realN = if (n > size) size else n
+            val realN = Math.min(n, size)
             val list = ArrayList<T>(realN)
             for (index in size - realN .. size - 1)
                 list.add(this[index])
@@ -366,7 +377,7 @@ fun filtering(): List<GenericFunction> {
         returns("List<T>")
         body {
             """
-            val list = ArrayList<T>()
+            val list = ArrayList<T>(indices.collectionSizeOrDefault(10))
             for (index in indices) {
                 list.add(get(index))
             }
@@ -378,7 +389,7 @@ fun filtering(): List<GenericFunction> {
         returns(Strings) { "String" }
         body(Strings) {
             """
-            val result = StringBuilder()
+            val result = StringBuilder(indices.collectionSizeOrDefault(10))
             for (i in indices) {
                 result.append(get(i))
             }