Provide stable sorting when JS engine sorting doesn't look stable

ilya-g · ilya-g · commit 51eb21d44b2e · 2019-01-11T20:53:29.000+03:00
#KT-12473
diff --git a/libraries/stdlib/js/irRuntime/generated/_ArraysJs.kt b/libraries/stdlib/js/irRuntime/generated/_ArraysJs.kt
@@ -1128,8 +1128,7 @@ public actual fun IntArray.sort(): Unit {
  * Sorts the array in-place.
  */
 public actual fun LongArray.sort(): Unit {
-    if (size > 1)
-        sort { a: Long, b: Long -> a.compareTo(b) }
+    if (size > 1) sort { a: Long, b: Long -> a.compareTo(b) }
 }
 
 /**
@@ -1171,16 +1170,14 @@ public actual fun CharArray.sort(): Unit {
  * Sorts the array in-place according to the natural order of its elements.
  */
 public actual fun <T : Comparable<T>> Array<out T>.sort(): Unit {
-    if (size > 1)
-        sort { a: T, b: T -> a.compareTo(b) }
+    if (size > 1) sortArray(this)
 }
 
 /**
  * Sorts the array in-place according to the order specified by the given [comparison] function.
  */
-@kotlin.internal.InlineOnly
-public inline fun <T> Array<out T>.sort(noinline comparison: (a: T, b: T) -> Int): Unit {
-    asDynamic().sort(comparison)
+public fun <T> Array<out T>.sort(comparison: (a: T, b: T) -> Int): Unit {
+    if (size > 1) sortArrayWith(this, comparison)
 }
 
 /**
@@ -1243,8 +1240,7 @@ public inline fun CharArray.sort(noinline comparison: (a: Char, b: Char) -> Int)
  * Sorts the array in-place according to the order specified by the given [comparator].
  */
 public actual fun <T> Array<out T>.sortWith(comparator: Comparator<in T>): Unit {
-    if (size > 1)
-        sort { a, b -> comparator.compare(a, b) }
+    if (size > 1) sortArrayWith(this, comparator)
 }
 
 /**
diff --git a/libraries/stdlib/js/src/generated/_ArraysJs.kt b/libraries/stdlib/js/src/generated/_ArraysJs.kt
@@ -1159,8 +1159,7 @@ public actual fun IntArray.sort(): Unit {
  * Sorts the array in-place.
  */
 public actual fun LongArray.sort(): Unit {
-    if (size > 1)
-        sort { a: Long, b: Long -> a.compareTo(b) }
+    if (size > 1) sort { a: Long, b: Long -> a.compareTo(b) }
 }
 
 /**
@@ -1207,16 +1206,14 @@ public actual fun CharArray.sort(): Unit {
  * Sorts the array in-place according to the natural order of its elements.
  */
 public actual fun <T : Comparable<T>> Array<out T>.sort(): Unit {
-    if (size > 1)
-        sort { a: T, b: T -> a.compareTo(b) }
+    if (size > 1) sortArray(this)
 }
 
 /**
  * Sorts the array in-place according to the order specified by the given [comparison] function.
  */
-@kotlin.internal.InlineOnly
-public inline fun <T> Array<out T>.sort(noinline comparison: (a: T, b: T) -> Int): Unit {
-    asDynamic().sort(comparison)
+public fun <T> Array<out T>.sort(comparison: (a: T, b: T) -> Int): Unit {
+    if (size > 1) sortArrayWith(this, comparison)
 }
 
 /**
@@ -1279,8 +1276,7 @@ public inline fun CharArray.sort(noinline comparison: (a: Char, b: Char) -> Int)
  * Sorts the array in-place according to the order specified by the given [comparator].
  */
 public actual fun <T> Array<out T>.sortWith(comparator: Comparator<in T>): Unit {
-    if (size > 1)
-        sort { a, b -> comparator.compare(a, b) }
+    if (size > 1) sortArrayWith(this, comparator)
 }
 
 /**
diff --git a/libraries/stdlib/js/src/kotlin/collections/ArraySorting.kt b/libraries/stdlib/js/src/kotlin/collections/ArraySorting.kt
@@ -0,0 +1,104 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package kotlin.collections
+
+internal fun <T> sortArrayWith(array: Array<out T>, comparison: (T, T) -> Int) {
+    if (getStableSortingIsSupported()) {
+        array.asDynamic().sort(comparison)
+    } else {
+        mergeSort(array.unsafeCast<Array<T>>(), 0, array.lastIndex, Comparator(comparison))
+    }
+}
+
+internal fun <T> sortArrayWith(array: Array<out T>, comparator: Comparator<in T>) {
+    if (getStableSortingIsSupported()) {
+        val comparison = { a: T, b: T -> comparator.compare(a, b) }
+        array.asDynamic().sort(comparison)
+    } else {
+        mergeSort(array.unsafeCast<Array<T>>(), 0, array.lastIndex, comparator)
+    }
+}
+
+internal fun <T : Comparable<T>> sortArray(array: Array<out T>) {
+    if (getStableSortingIsSupported()) {
+        val comparison = { a: T, b: T -> a.compareTo(b) }
+        array.asDynamic().sort(comparison)
+    } else {
+        mergeSort(array.unsafeCast<Array<T>>(), 0, array.lastIndex, naturalOrder())
+    }
+}
+
+private var _stableSortingIsSupported: Boolean? = null
+private fun getStableSortingIsSupported(): Boolean {
+    _stableSortingIsSupported?.let { return it }
+    _stableSortingIsSupported = false
+
+    val array = js("[]").unsafeCast<Array<Int>>()
+    // known implementations may use stable sort for arrays of up to 512 elements
+    // so we create slightly more elements to test stability
+    for (index in 0 until 600) array.asDynamic().push(index)
+    val comparison = { a: Int, b: Int -> (a and 3) - (b and 3) }
+    array.asDynamic().sort(comparison)
+    for (index in 1 until array.size) {
+        val a = array[index - 1]
+        val b = array[index]
+        if ((a and 3) == (b and 3) && a >= b) return false
+    }
+    _stableSortingIsSupported = true
+    return true
+}
+
+
+private fun <T> mergeSort(array: Array<T>, start: Int, endInclusive: Int, comparator: Comparator<in T>) {
+    val buffer = arrayOfNulls<Any?>(array.size).unsafeCast<Array<T>>()
+    val result = mergeSort(array, buffer, start, endInclusive, comparator)
+    if (result !== array) {
+        result.forEachIndexed { i, v -> array[i] = v }
+    }
+}
+
+// Both start and end are inclusive indices.
+private fun <T> mergeSort(array: Array<T>, buffer: Array<T>, start: Int, end: Int, comparator: Comparator<in T>): Array<T> {
+    if (start == end) {
+        return array
+    }
+
+    val median = (start + end) / 2
+    val left = mergeSort(array, buffer, start, median, comparator)
+    val right = mergeSort(array, buffer, median + 1, end, comparator)
+
+    val target = if (left === buffer) array else buffer
+
+    // Merge.
+    var leftIndex = start
+    var rightIndex = median + 1
+    for (i in start..end) {
+        when {
+            leftIndex <= median && rightIndex <= end -> {
+                val leftValue = left[leftIndex]
+                val rightValue = right[rightIndex]
+
+                if (comparator.compare(leftValue, rightValue) <= 0) {
+                    target[i] = leftValue
+                    leftIndex++
+                } else {
+                    target[i] = rightValue
+                    rightIndex++
+                }
+            }
+            leftIndex <= median -> {
+                target[i] = left[leftIndex]
+                leftIndex++
+            }
+            else /* rightIndex <= end */ -> {
+                target[i] = right[rightIndex]
+                rightIndex++
+            }
+        }
+    }
+
+    return target
+}
diff --git a/libraries/stdlib/test/collections/ArraysTest.kt b/libraries/stdlib/test/collections/ArraysTest.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
  * that can be found in the license/LICENSE.txt file.
  */
 
@@ -1412,6 +1412,21 @@ class ArraysTest {
         }
     }
 
+    @Test fun sortStable() {
+        val keyRange = 'A'..'D'
+        for (size in listOf(10, 100, 2000)) {
+            val array = Array(size) { index -> Sortable(keyRange.random(), index) }
+
+            array.sortedArray().assertStableSorted()
+            array.sortedArrayDescending().assertStableSorted(descending = true)
+
+            array.sort()
+            array.assertStableSorted()
+            array.sortDescending()
+            array.assertStableSorted(descending = true)
+        }
+    }
+
     @Test fun sortByInPlace() {
         val data = arrayOf("aa" to 20, "ab" to 3, "aa" to 3)
         data.sortBy { it.second }
@@ -1431,6 +1446,22 @@ class ArraysTest {
         assertEquals(listOf(1, 2, 0), indices.sortedBy { values[it] })
     }
 
+    @Test fun sortByStable() {
+        val keyRange = 'A'..'D'
+        for (size in listOf(10, 100, 2000)) {
+            val array = Array(size) { index -> Sortable(keyRange.random(), index) }
+
+            array.sortedBy { it.key }.iterator().assertStableSorted()
+            array.sortedByDescending { it.key }.iterator().assertStableSorted(descending = true)
+
+            array.sortBy { it.key }
+            array.assertStableSorted()
+
+            array.sortByDescending { it.key }
+            array.assertStableSorted(descending = true)
+        }
+    }
+
     @Test fun sortedNullableBy() {
         fun String.nullIfEmpty() = if (this.isEmpty()) null else this
         arrayOf(null, "").let {
@@ -1457,6 +1488,20 @@ class ArraysTest {
     }
 }
 
+private data class Sortable<K : Comparable<K>>(val key: K, val index: Int) : Comparable<Sortable<K>> {
+    override fun compareTo(other: Sortable<K>): Int = this.key.compareTo(other.key)
+}
+
+private fun <K : Comparable<K>> Array<out Sortable<K>>.assertStableSorted(descending: Boolean = false) =
+    iterator().assertStableSorted(descending = descending)
+
+private fun <K : Comparable<K>> Iterator<Sortable<K>>.assertStableSorted(descending: Boolean = false) {
+    assertSorted { a, b ->
+        val relation = a.key.compareTo(b.key)
+        (if (descending) relation > 0 else relation < 0) || relation == 0 && a.index < b.index
+    }
+}
+
 private class ArraySortedChecker<A, T>(val array: A, val comparator: Comparator<in T>) {
     public fun <R> checkSorted(sorted: A.() -> R, sortedDescending: A.() -> R, iterator: R.() -> Iterator<T>) {
         array.sorted().iterator().assertSorted { a, b -> comparator.compare(a, b) <= 0 }
diff --git a/libraries/tools/kotlin-stdlib-gen/src/templates/Arrays.kt b/libraries/tools/kotlin-stdlib-gen/src/templates/Arrays.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
  * that can be found in the license/LICENSE.txt file.
  */
 
@@ -884,10 +884,7 @@ object ArrayOps : TemplateGroupBase() {
         returns("Unit")
         on(Platform.JS) {
             body {
-                """
-                if (size > 1)
-                    sort { a: T, b: T -> a.compareTo(b) }
-                """
+                """if (size > 1) sortArray(this)"""
             }
             specialFor(ArraysOfPrimitives) {
                 if (primitive != PrimitiveType.Long) {
@@ -898,6 +895,10 @@ object ArrayOps : TemplateGroupBase() {
                     on(Backend.IR) {
                         body { "this.asDynamic().sort()" }
                     }
+                } else {
+                    body {
+                        """if (size > 1) sort { a: T, b: T -> a.compareTo(b) }"""
+                    }
                 }
             }
         }
@@ -939,26 +940,29 @@ object ArrayOps : TemplateGroupBase() {
         }
         on(Platform.JS) {
             body {
-                """
-                if (size > 1)
-                    sort { a, b -> comparator.compare(a, b) }
-                """
+                """if (size > 1) sortArrayWith(this, comparator)"""
             }
         }
         on(Platform.Native) {
             body { """if (size > 1) kotlin.util.sortArrayWith(this, 0, size, comparator)""" }
         }
     }
 
-    val f_sort_comparison = fn("sort(noinline comparison: (a: T, b: T) -> Int)") {
+    val f_sort_comparison = fn("sort(comparison: (a: T, b: T) -> Int)") {
         platforms(Platform.JS)
         include(ArraysOfObjects, ArraysOfPrimitives)
         exclude(PrimitiveType.Boolean)
     } builder {
-        inlineOnly()
         returns("Unit")
         doc { "Sorts the array in-place according to the order specified by the given [comparison] function." }
-        body { "asDynamic().sort(comparison)" }
+        specialFor(ArraysOfPrimitives) {
+            inlineOnly()
+            signature("sort(noinline comparison: (a: T, b: T) -> Int)")
+            body { "asDynamic().sort(comparison)" }
+        }
+        specialFor(ArraysOfObjects) {
+            body { """if (size > 1) sortArrayWith(this, comparison)""" }
+        }
     }
 
     val f_sort_objects = fn("sort()") {
