Make plus operation to be strict binary or right external binary for Sets and Arrays.

libraries/tools/kotlin-stdlib-gen/src/templates/Generators.kt

@@ -7,27 +7,51 @@ fun generators(): List<GenericFunction> {
 
     templates add f("plus(element: T)") {
         exclude(Strings)
-        doc { "Returns a list containing all elements of original collection and then the given [element]." }
+        doc { "Returns a list containing all elements of the original collection and then the given [element]." }
         returns("List<T>")
+        returns("SELF", ArraysOfObjects, ArraysOfPrimitives, Sets, Sequences)
         body {
             """
-            val answer = ArrayList<T>(collectionSizeOrNull()?.let { it + 1 } ?: 10)
-            for (thisElement in this) answer.add(thisElement)
+            if (this is Collection) return this.plus(element)
+            val answer = ArrayList<T>()
+            answer.addAll(this)
             answer.add(element)
             return answer
             """
         }
-        body(ArraysOfObjects, ArraysOfPrimitives) {
+        body(Collections) {
             """
-            val answer = ArrayList<T>(size()+1)
-            for (thisElement in this) answer.add(thisElement)
+            val answer = ArrayList<T>(size() + 1)
+            answer.addAll(this)
             answer.add(element)
             return answer
             """
         }
 
-        doc(Sequences) { "Returns a sequence containing all elements of original sequence and then the given [element]." }
-        returns(Sequences) { "Sequence<T>" }
+
+        doc(ArraysOfObjects, ArraysOfPrimitives) { "Returns an array containing all elements of the original array and then the given [element]." }
+        body(ArraysOfObjects, ArraysOfPrimitives) {
+            """
+            val answer = this.copyOf(size() + 1)
+            answer[size()] = element
+            return answer as SELF
+            """
+        }
+
+        // TODO: use build scope function when available
+        // TODO: use immutable sets when available
+        // TODO: precalculate size
+        doc(Sets) { "Returns a set containing all elements of the original set and then the given [element]." }
+        body(Sets) {
+            """
+            val copyOfSet = LinkedHashSet<T>(mapCapacity(size() + 1))
+            copyOfSet.addAll(this)
+            copyOfSet.add(element)
+            return copyOfSet
+            """
+        }
+
+        doc(Sequences) { "Returns a sequence containing all elements of the original sequence and then the given [element]." }
         body(Sequences) {
             """
             return sequenceOf(this, sequenceOf(element)).flatten()
@@ -36,57 +60,120 @@ fun generators(): List<GenericFunction> {
     }
 
     templates add f("plus(collection: Iterable<T>)") {
-        exclude(Strings, Sequences)
-        doc { "Returns a list containing all elements of original collection and then all elements of the given [collection]." }
+        only(Iterables, Collections, Sets, Sequences)
+        doc { "Returns a list containing all elements of the original collection and then all elements of the given [collection]." }
         returns("List<T>")
+        returns("SELF", ArraysOfObjects, ArraysOfPrimitives, Sets, Sequences)
         body {
             """
-            val answer = ArrayList<T>(collectionSizeOrDefault(10) + collection.collectionSizeOrDefault(10))
-            for (thisElement in this) answer.add(thisElement)
+            if (this is Collection) return this.plus(collection)
+            val answer = ArrayList<T>(0)
+            answer.addAll(this)
+            answer.addAll(collection)
+            return answer
+            """
+        }
+        body(Collections) {
+            """
+            if (collection is Collection) return this.plus(collection)
+            val answer = ArrayList<T>(this)
             answer.addAll(collection)
             return answer
             """
         }
 
-        body(ArraysOfObjects, ArraysOfPrimitives) {
+        // TODO: try to precalculate size
+        // TODO: use immutable set builder when available
+        body(Sets) {
             """
-            val answer = ArrayList<T>(size() + collection.collectionSizeOrDefault(10))
-            for (thisElement in this) answer.add(thisElement)
+            val result = LinkedHashSet<T>(mapCapacity(collection.collectionSizeOrNull()?.let { this.size() + it } ?: this.size() * 2))
+            result.addAll(this)
+            result.addAll(collection)
+            return result
+            """
+        }
+
+        doc(Sequences) { "Returns a sequence containing all elements of original sequence and then all elements of the given [collection]" }
+        body(Sequences) {
+            """
+            return sequenceOf(this, collection.asSequence()).flatten()
+            """
+        }
+    }
+
+    templates add f("plus(collection: Collection<T>)") {
+        only(Collections, ArraysOfPrimitives, ArraysOfObjects)
+        returns("List<T>")
+        returns("SELF", ArraysOfObjects, ArraysOfPrimitives)
+        body(Collections) {
+            """
+            val answer = ArrayList<T>(this.size() + collection.size())
+            answer.addAll(this)
             answer.addAll(collection)
             return answer
             """
         }
+        // ?
+        body(ArraysOfObjects, ArraysOfPrimitives) {
+            """
+            val thisSize = size()
+            val answer = this.copyOf(thisSize + collection.size())
+            collection.forEachIndexed { i, element ->
+                answer[thisSize + i] = element
+            }
+            return answer as SELF
+            """
+        }
     }
 
     templates add f("plus(array: Array<out T>)") {
-        exclude(Strings, Sequences)
+        only(Iterables, Collections, Sets, Sequences)
         doc { "Returns a list containing all elements of original collection and then all elements of the given [array]." }
         returns("List<T>")
+        returns("SELF", Sets, Sequences)
         body {
             """
-            val answer = ArrayList<T>(collectionSizeOrDefault(10) + array.size())
-            for (thisElement in this) answer.add(thisElement)
+            if (this is Collection) return this.plus(array)
+            val answer = ArrayList<T>()
+            answer.addAll(this)
             answer.addAll(array)
             return answer
             """
         }
-        body(ArraysOfObjects, ArraysOfPrimitives) {
+        body(Collections) {
             """
-            val answer = ArrayList<T>(size() + array.size())
-            for (thisElement in this) answer.add(thisElement)
+            val answer = ArrayList<T>(this.size() + array.size())
+            answer.addAll(this)
             answer.addAll(array)
             return answer
             """
         }
+        body(Sets) {
+            """
+            val result = LinkedHashSet<T>(mapCapacity(this.size() + array.size()))
+            result.addAll(this)
+            result.addAll(array)
+            return result
+            """
+        }
+        body(Sequences) {
+            """
+            return this.plus(array.asList())
+            """
+        }
     }
 
-    templates add f("plus(collection: Iterable<T>)") {
-        only(Sequences)
-        doc { "Returns a sequence containing all elements of original sequence and then all elements of the given [collection]." }
-        returns("Sequence<T>")
+    templates add f("plus(collection: SELF)") {
+        only(ArraysOfObjects, ArraysOfPrimitives)
+        returns("SELF")
         body {
             """
-            return sequenceOf(this, collection.asSequence()).flatten()
+            val thisSize = size()
+            val answer = this.copyOf(thisSize + collection.size())
+            collection.forEachIndexed { i, element ->
+                answer[thisSize + i] = element
+            }
+            return answer as SELF
             """
         }
     }