Optimize ArrayOps and ArraySeq sorting

- Return the original collection (for immutable.ArraySeq) or a straight
  clone for length <= 1.
- Use Java’s native array sorting methods when sorting primitive arrays
  with the default Ordering.
- Copy primitive arrays into a boxed representation for sorting large
  arrays with a non-default Ordering using Java’s native array sorting.
  ArrayOps copies the result back into an unboxed array,
  immutable.ArraySeq wraps the boxed array directly. Scala’s custom
  in-place sorting implementation is still faster for small arrays in
  ArrayOps (where we’d have to copy back when using the Java version).

Author: szeiger

src/library/scala/collection/ArrayOps.scala

@@ -1,12 +1,14 @@
 package scala
 package collection
 
-import java.lang.Math.{min, max}
+import java.lang.Math.{max, min}
+import java.util.Arrays
 
 import mutable.ArrayBuilder
 import immutable.Range
 import scala.reflect.ClassTag
 import scala.math.Ordering
+import scala.util.Sorting
 import scala.Predef.{ // unimport all array-related implicit conversions to avoid triggering them accidentally
   genericArrayOps => _,
   booleanArrayOps => _,
@@ -145,6 +147,10 @@ object ArrayOps {
     }
   }
 
+  /** The cut-off point for the array size after which we switch from `Sorting.stableSort` to
+    * an implementation that copies the data to a boxed representation for use with `Arrays.sort`.
+    */
+  private final val MaxStableSortLength = 300
 }
 
 /** This class serves as a wrapper for `Array`s with many of the operations found in
@@ -458,28 +464,40 @@ final class ArrayOps[A](val xs: Array[A]) extends AnyVal {
     */
   def sorted[B >: A](implicit ord: Ordering[B]): Array[A] = {
     val len = xs.length
-    if(xs.getClass.getComponentType.isPrimitive && len > 1) {
-      // need to copy into a boxed representation to use Java's Arrays.sort
-      val a = new Array[AnyRef](len)
-      var i = 0
-      while(i < len) {
-        a(i) = xs(i).asInstanceOf[AnyRef]
-        i += 1
-      }
-      java.util.Arrays.sort(a, ord.asInstanceOf[Ordering[AnyRef]])
-      val res = new Array[A](len)
-      i = 0
-      while(i < len) {
-        res(i) = a(i).asInstanceOf[A]
-        i += 1
-      }
-      res
+    def boxed = if(len < ArrayOps.MaxStableSortLength) {
+      val a = xs.clone()
+      Sorting.stableSort(xs)(ord.asInstanceOf[Ordering[A]])
+      a
     } else {
-      val copy = slice(0, len)
-      if(len > 1)
-        java.util.Arrays.sort(copy.asInstanceOf[Array[AnyRef]], ord.asInstanceOf[Ordering[AnyRef]])
-      copy
+      val a = Array.copyAs[AnyRef](xs, len)(ClassTag.AnyRef)
+      Arrays.sort(a, ord.asInstanceOf[Ordering[AnyRef]])
+      Array.copyAs[A](a, len)
+      a
     }
+    if(len <= 1) xs.clone()
+    else ((xs: Array[_]) match {
+      case xs: Array[AnyRef] =>
+        val a = Arrays.copyOf(xs, len); Arrays.sort(a, ord.asInstanceOf[Ordering[AnyRef]]); a
+      case xs: Array[Int] =>
+        if(ord eq Ordering.Int) { val a = Arrays.copyOf(xs, len); Arrays.sort(a); a }
+        else boxed
+      case xs: Array[Long] =>
+        if(ord eq Ordering.Long) { val a = Arrays.copyOf(xs, len); Arrays.sort(a); a }
+        else boxed
+      case xs: Array[Char] =>
+        if(ord eq Ordering.Char) { val a = Arrays.copyOf(xs, len); Arrays.sort(a); a }
+        else boxed
+      case xs: Array[Byte] =>
+        if(ord eq Ordering.Byte) { val a = Arrays.copyOf(xs, len); Arrays.sort(a); a }
+        else boxed
+      case xs: Array[Short] =>
+        if(ord eq Ordering.Short) { val a = Arrays.copyOf(xs, len); Arrays.sort(a); a }
+        else boxed
+      case xs: Array[Boolean] =>
+        if(ord eq Ordering.Boolean) { val a = Arrays.copyOf(xs, len); Sorting.stableSort(a); a }
+        else boxed
+      case xs => boxed
+    }).asInstanceOf[Array[A]]
   }
 
   /** Sorts this array according to a comparison function.

src/library/scala/collection/immutable/ArraySeq.scala

@@ -7,6 +7,7 @@ import scala.collection.mutable.{ArrayBuffer, ArrayBuilder, Builder, ArraySeq =>
 import scala.collection.{ArrayOps, ClassTagSeqFactory, SeqFactory, StrictOptimizedClassTagSeqFactory}
 import scala.collection.IterableOnce
 import scala.annotation.unchecked.uncheckedVariance
+import scala.util.Sorting
 import scala.util.hashing.MurmurHash3
 import scala.reflect.ClassTag
 import scala.runtime.ScalaRunTime
@@ -147,6 +148,14 @@ sealed abstract class ArraySeq[+A]
   override protected[this] def writeReplace(): AnyRef = this
 
   override protected final def applyPreferredMaxLength: Int = Int.MaxValue
+
+  override def sorted[B >: A](implicit ord: Ordering[B]): ArraySeq[A] =
+    if(unsafeArray.length <= 1) this
+    else {
+      val a = Array.copyAs[AnyRef](unsafeArray, length)(ClassTag.AnyRef)
+      Arrays.sort(a, ord.asInstanceOf[Ordering[AnyRef]])
+      new ArraySeq.ofRef[AnyRef](a).asInstanceOf[ArraySeq[A]]
+    }
 }
 
 /**
@@ -231,6 +240,14 @@ object ArraySeq extends StrictOptimizedClassTagSeqFactory[ArraySeq] { self =>
           that.unsafeArray.asInstanceOf[Array[AnyRef]])
       case _ => super.equals(that)
     }
+    override def sorted[B >: T](implicit ord: Ordering[B]): ArraySeq.ofRef[T] = {
+      if(unsafeArray.length <= 1) this
+      else {
+        val a = unsafeArray.clone()
+        Arrays.sort(a, ord.asInstanceOf[Ordering[T]])
+        new ArraySeq.ofRef(a)
+      }
+    }
   }
 
   @SerialVersionUID(3L)
@@ -244,6 +261,13 @@ object ArraySeq extends StrictOptimizedClassTagSeqFactory[ArraySeq] { self =>
       case that: ofByte => Arrays.equals(unsafeArray, that.unsafeArray)
       case _ => super.equals(that)
     }
+    override def sorted[B >: Byte](implicit ord: Ordering[B]): ArraySeq[Byte] =
+      if(length <= 1) this
+      else if(ord eq Ordering.Byte) {
+        val a = unsafeArray.clone()
+        Arrays.sort(a)
+        new ArraySeq.ofByte(a)
+      } else super.sorted[B]
   }
 
   @SerialVersionUID(3L)
@@ -257,6 +281,13 @@ object ArraySeq extends StrictOptimizedClassTagSeqFactory[ArraySeq] { self =>
       case that: ofShort => Arrays.equals(unsafeArray, that.unsafeArray)
       case _ => super.equals(that)
     }
+    override def sorted[B >: Short](implicit ord: Ordering[B]): ArraySeq[Short] =
+      if(length <= 1) this
+      else if(ord eq Ordering.Short) {
+        val a = unsafeArray.clone()
+        Arrays.sort(a)
+        new ArraySeq.ofShort(a)
+      } else super.sorted[B]
   }
 
   @SerialVersionUID(3L)
@@ -270,6 +301,13 @@ object ArraySeq extends StrictOptimizedClassTagSeqFactory[ArraySeq] { self =>
       case that: ofChar => Arrays.equals(unsafeArray, that.unsafeArray)
       case _ => super.equals(that)
     }
+    override def sorted[B >: Char](implicit ord: Ordering[B]): ArraySeq[Char] =
+      if(length <= 1) this
+      else if(ord eq Ordering.Char) {
+        val a = unsafeArray.clone()
+        Arrays.sort(a)
+        new ArraySeq.ofChar(a)
+      } else super.sorted[B]
 
     override def addString(sb: StringBuilder, start: String, sep: String, end: String): StringBuilder =
       (new MutableArraySeq.ofChar(unsafeArray)).addString(sb, start, sep, end)
@@ -286,6 +324,13 @@ object ArraySeq extends StrictOptimizedClassTagSeqFactory[ArraySeq] { self =>
       case that: ofInt => Arrays.equals(unsafeArray, that.unsafeArray)
       case _ => super.equals(that)
     }
+    override def sorted[B >: Int](implicit ord: Ordering[B]): ArraySeq[Int] =
+      if(length <= 1) this
+      else if(ord eq Ordering.Int) {
+        val a = unsafeArray.clone()
+        Arrays.sort(a)
+        new ArraySeq.ofInt(a)
+      } else super.sorted[B]
   }
 
   @SerialVersionUID(3L)
@@ -299,6 +344,13 @@ object ArraySeq extends StrictOptimizedClassTagSeqFactory[ArraySeq] { self =>
       case that: ofLong => Arrays.equals(unsafeArray, that.unsafeArray)
       case _ => super.equals(that)
     }
+    override def sorted[B >: Long](implicit ord: Ordering[B]): ArraySeq[Long] =
+      if(length <= 1) this
+      else if(ord eq Ordering.Long) {
+        val a = unsafeArray.clone()
+        Arrays.sort(a)
+        new ArraySeq.ofLong(a)
+      } else super.sorted[B]
   }
 
   @SerialVersionUID(3L)
@@ -338,6 +390,13 @@ object ArraySeq extends StrictOptimizedClassTagSeqFactory[ArraySeq] { self =>
       case that: ofBoolean => Arrays.equals(unsafeArray, that.unsafeArray)
       case _ => super.equals(that)
     }
+    override def sorted[B >: Boolean](implicit ord: Ordering[B]): ArraySeq[Boolean] =
+      if(length <= 1) this
+      else if(ord eq Ordering.Boolean) {
+        val a = unsafeArray.clone()
+        Sorting.stableSort(a)
+        new ArraySeq.ofBoolean(a)
+      } else super.sorted[B]
   }
 
   @SerialVersionUID(3L)

src/library/scala/collection/mutable/ArraySeq.scala

@@ -78,6 +78,9 @@ sealed abstract class ArraySeq[T]
       super.equals(other)
   }
 
+  override def sorted[B >: T](implicit ord: Ordering[B]): ArraySeq[T] =
+    ArraySeq.make(array.asInstanceOf[Array[Any]].sorted(ord.asInstanceOf[Ordering[Any]])).asInstanceOf[ArraySeq[T]]
+
   override def sortInPlace[B >: T]()(implicit ord: Ordering[B]): this.type = {
     if (length > 1) scala.util.Sorting.stableSort(array.asInstanceOf[Array[B]])
     this

test/benchmarks/src/main/scala/scala/collection/immutable/ArraySeqBenchmark.scala

@@ -0,0 +1,71 @@
+package scala.collection.immutable
+
+import java.util.concurrent.TimeUnit
+import java.util.Arrays
+
+import org.openjdk.jmh.annotations._
+import org.openjdk.jmh.infra.Blackhole
+
+import scala.reflect.ClassTag
+
+@BenchmarkMode(Array(Mode.AverageTime))
+@Fork(2)
+@Threads(1)
+@Warmup(iterations = 10)
+@Measurement(iterations = 10)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@State(Scope.Benchmark)
+class ArraySeqBenchmark {
+
+  @Param(Array("0", "1", "10", "1000", "10000"))
+  var size: Int = _
+  var integersS: ArraySeq[Int] = _
+  var stringsS: ArraySeq[String] = _
+
+  @Setup(Level.Trial) def initNumbers: Unit = {
+    val integers = (1 to size).toList
+    val strings = integers.map(_.toString)
+    integersS = ArraySeq.unsafeWrapArray(integers.toArray)
+    stringsS = ArraySeq.unsafeWrapArray(strings.toArray)
+  }
+
+  @Benchmark def sortedStringOld(bh: Blackhole): Unit =
+    bh.consume(oldSorted(stringsS))
+
+  @Benchmark def sortedIntOld(bh: Blackhole): Unit =
+    bh.consume(oldSorted(integersS))
+
+  @Benchmark def sortedIntCustomOld(bh: Blackhole): Unit =
+    bh.consume(oldSorted(integersS)(Ordering.Int.reverse, implicitly))
+
+  @Benchmark def sortedStringNew(bh: Blackhole): Unit =
+    bh.consume(stringsS.sorted)
+
+  @Benchmark def sortedIntNew(bh: Blackhole): Unit =
+    bh.consume(integersS.sorted)
+
+  @Benchmark def sortedIntCustomNew(bh: Blackhole): Unit =
+    bh.consume(integersS.sorted(Ordering.Int.reverse))
+
+  private[this] def oldSorted[A](seq: ArraySeq[A])(implicit ord: Ordering[A], tag: ClassTag[A]): ArraySeq[A] = {
+    val len = seq.length
+    val b = ArraySeq.newBuilder[A](tag)
+    if (len == 1) b ++= seq.toIterable
+    else if (len > 1) {
+      b.sizeHint(len)
+      val arr = new Array[AnyRef](len)
+      var i = 0
+      for (x <- seq) {
+        arr(i) = x.asInstanceOf[AnyRef]
+        i += 1
+      }
+      java.util.Arrays.sort(arr, ord.asInstanceOf[Ordering[Object]])
+      i = 0
+      while (i < arr.length) {
+        b += arr(i).asInstanceOf[A]
+        i += 1
+      }
+    }
+    b.result()
+  }
+}

test/benchmarks/src/main/scala/scala/collection/mutable/ArrayOpsBenchmark.scala

@@ -1,10 +1,13 @@
 package scala.collection.mutable
 
 import java.util.concurrent.TimeUnit
+import java.util.Arrays
 
 import org.openjdk.jmh.annotations._
 import org.openjdk.jmh.infra.Blackhole
 
+import scala.reflect.ClassTag
+
 @BenchmarkMode(Array(Mode.AverageTime))
 @Fork(2)
 @Threads(1)
@@ -14,17 +17,19 @@ import org.openjdk.jmh.infra.Blackhole
 @State(Scope.Benchmark)
 class ArrayOpsBenchmark {
 
-  @Param(Array("10", "1000", "10000"))
+  @Param(Array("0", "1", "10", "1000", "10000"))
   var size: Int = _
   var integers: List[Int] = _
   var strings: List[String] = _
   var integersA: Array[Int] = _
+  var stringsA: Array[String] = _
 
 
   @Setup(Level.Trial) def initNumbers: Unit = {
     integers = (1 to size).toList
     strings = integers.map(_.toString)
     integersA = integers.toArray
+    stringsA = strings.toArray
   }
 
   @Benchmark def appendInteger(bh: Blackhole): Unit = {
@@ -66,4 +71,49 @@ class ArrayOpsBenchmark {
   @Benchmark def foldSum(bh: Blackhole): Unit = {
     bh.consume(integersA.fold(0){ (a,b) => a + b })
   }
+
+  @Benchmark def sortedStringOld(bh: Blackhole): Unit =
+    bh.consume(oldSorted(stringsA))
+
+  @Benchmark def sortedIntOld(bh: Blackhole): Unit =
+    bh.consume(oldSorted(integersA))
+
+  @Benchmark def sortedIntCustomOld(bh: Blackhole): Unit =
+    bh.consume(oldSorted(integersA)(Ordering.Int.reverse))
+
+  @Benchmark def sortedStringNew(bh: Blackhole): Unit =
+    bh.consume(stringsA.sorted)
+
+  @Benchmark def sortedIntNew(bh: Blackhole): Unit =
+    bh.consume(integersA.sorted)
+
+  @Benchmark def sortedIntCustomNew(bh: Blackhole): Unit =
+    bh.consume(integersA.sorted(Ordering.Int.reverse))
+
+  def oldSorted[A, B >: A](xs: Array[A])(implicit ord: Ordering[B]): Array[A] = {
+    implicit def ct = ClassTag[A](xs.getClass.getComponentType)
+    val len = xs.length
+    if(xs.getClass.getComponentType.isPrimitive && len > 1) {
+      // need to copy into a boxed representation to use Java's Arrays.sort
+      val a = new Array[AnyRef](len)
+      var i = 0
+      while(i < len) {
+        a(i) = xs(i).asInstanceOf[AnyRef]
+        i += 1
+      }
+      Arrays.sort(a, ord.asInstanceOf[Ordering[AnyRef]])
+      val res = new Array[A](len)
+      i = 0
+      while(i < len) {
+        res(i) = a(i).asInstanceOf[A]
+        i += 1
+      }
+      res
+    } else {
+      val copy = xs.slice(0, len)
+      if(len > 1)
+        Arrays.sort(copy.asInstanceOf[Array[AnyRef]], ord.asInstanceOf[Ordering[AnyRef]])
+      copy
+    }
+  }
 }