Merge pull request #7732 from brunnerant/tuple-split

Extend Tuple API

Author: nicolasstucki

bench-run/inputs/drop.in

@@ -0,0 +1 @@
+size:0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50

bench-run/inputs/split.in

@@ -0,0 +1 @@
+size:0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50

bench-run/inputs/take.in

@@ -0,0 +1 @@
+size:0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50

bench-run/src/main/scala/dotty/tools/benchmarks/tuples/TupleOps.scala

@@ -1,11 +1,13 @@
 package dotty.tools.benchmarks.tuples
 
 import org.openjdk.jmh.annotations._
+import scala.util.Random
 
 @State(Scope.Thread)
 class TupleOps {
   var tuple1: Tuple = _
   var tuple2: Tuple = _
+  var tuple3: Tuple = _
 
   @Setup
   def setup(): Unit = {
@@ -16,6 +18,9 @@ class TupleOps {
     tuple2 = ()
     for (i <- 1 until 10)
       tuple2 = s"elem$i" *: tuple2
+
+    val rand = new Random(12345)
+    tuple3 = Tuple.fromArray(rand.shuffle(1 to 15).toArray)
   }
 
   def tupleFlatMap(tuple: Tuple, f: [A] => A => Tuple): Tuple = {
@@ -34,6 +39,22 @@ class TupleOps {
     tailRecReverse(tuple, ())
   }
 
+  def tupleMerge(tuple1: Tuple, tuple2: Tuple): Tuple = (tuple1, tuple2) match {
+    case (_, ()) => tuple1
+    case ((), _) => tuple2
+    case (x *: xs, y *: ys) =>
+      if (x.asInstanceOf[Int] <= y.asInstanceOf[Int]) x *: tupleMerge(xs, tuple2)
+      else y *: tupleMerge(tuple1, ys)
+  }
+
+  def tupleMergeSort(tuple: Tuple): Tuple =
+    if (tuple.size <= 1) tuple
+    else {
+      val (tuple1, tuple2) = tuple.splitAt(tuple.size / 2)
+      val (sorted1, sorted2) = (tupleMergeSort(tuple1), tupleMergeSort(tuple2))
+      tupleMerge(sorted1, sorted2)
+    }
+
   @Benchmark
   def reverse(): Tuple = {
     tupleReverse(tuple1)
@@ -43,4 +64,9 @@ class TupleOps {
   def flatMap(): Tuple = {
     tupleFlatMap(tuple2, [A] => (x: A) => (x, x))
   }
+
+  @Benchmark
+  def mergeSort(): Tuple = {
+    tupleMergeSort(tuple3)
+  }
 }

bench-run/src/main/scala/tuples/Drop.scala

@@ -0,0 +1,34 @@
+package dotty.tools.benchmarks.tuples
+
+import org.openjdk.jmh.annotations._
+import scala.runtime.DynamicTuple
+
+@State(Scope.Thread)
+class Drop {
+  @Param(Array("0"))
+  var size: Int = _
+  var tuple: Tuple = _
+  var array: Array[Object] = _
+  var half: Int = _
+
+  @Setup
+  def setup(): Unit = {
+    tuple = ()
+    half = size / 2
+
+    for (i <- 1 to size)
+      tuple = "elem" *: tuple
+
+    array = new Array[Object](size)
+  }
+
+  @Benchmark
+  def tupleDrop(): Tuple = {
+    DynamicTuple.dynamicDrop(tuple, half)
+  }
+
+  @Benchmark
+  def arrayDrop(): Array[Object] = {
+    array.drop(half)
+  }
+}

bench-run/src/main/scala/tuples/Split.scala

@@ -0,0 +1,34 @@
+package dotty.tools.benchmarks.tuples
+
+import org.openjdk.jmh.annotations._
+import scala.runtime.DynamicTuple
+
+@State(Scope.Thread)
+class Split {
+  @Param(Array("0"))
+  var size: Int = _
+  var tuple: Tuple = _
+  var array: Array[Object] = _
+  var half: Int = _
+
+  @Setup
+  def setup(): Unit = {
+    tuple = ()
+    half = size / 2
+
+    for (i <- 1 to size)
+      tuple = "elem" *: tuple
+
+    array = new Array[Object](size)
+  }
+
+  @Benchmark
+  def tupleSplit(): (Tuple, Tuple) = {
+    DynamicTuple.dynamicSplitAt(tuple, half)
+  }
+
+  @Benchmark
+  def arraySplit(): (Array[Object], Array[Object]) = {
+    array.splitAt(half)
+  }
+}

bench-run/src/main/scala/tuples/Take.scala

@@ -0,0 +1,34 @@
+package dotty.tools.benchmarks.tuples
+
+import org.openjdk.jmh.annotations._
+import scala.runtime.DynamicTuple
+
+@State(Scope.Thread)
+class Take {
+  @Param(Array("0"))
+  var size: Int = _
+  var tuple: Tuple = _
+  var array: Array[Object] = _
+  var half: Int = _
+
+  @Setup
+  def setup(): Unit = {
+    tuple = ()
+    half = size / 2
+
+    for (i <- 1 to size)
+      tuple = "elem" *: tuple
+
+    array = new Array[Object](size)
+  }
+
+  @Benchmark
+  def tupleTake(): Tuple = {
+    DynamicTuple.dynamicTake(tuple, half)
+  }
+
+  @Benchmark
+  def arrayTake(): Array[Object] = {
+    array.take(half)
+  }
+}

compiler/test/dotc/run-test-pickling.blacklist

@@ -6,6 +6,10 @@ t7374
 tuples1.scala
 tuples1a.scala
 tuples1b.scala
+tuple-ops.scala
+tuple-take.scala
+tuple-drop.scala
+tuple-zip.scala
 typeclass-derivation1.scala
 typeclass-derivation2.scala
 typeclass-derivation2a.scala
@@ -19,4 +23,3 @@ t10889
 macros-in-same-project1
 i5257.scala
 enum-java
-tuple-ops.scala

library/src/scala/Tuple.scala

@@ -43,13 +43,33 @@ sealed trait Tuple extends Any {
   inline def zip[This >: this.type <: Tuple, T2 <: Tuple](t2: T2): Zip[This, T2] =
     DynamicTuple.dynamicZip(this, t2)
 
- /** Called on a tuple `(a1, ..., an)`, returns a new tuple `(f(a1), ..., f(an))`.
-  *  The result is typed as `(F[A1], ..., F[An])` if the tuple type is fully known.
-  *  If the tuple is of the form `a1 *: ... *: Tuple` (that is, the tail is not known
-  *  to be the cons type.
-  */
+  /** Called on a tuple `(a1, ..., an)`, returns a new tuple `(f(a1), ..., f(an))`.
+   *  The result is typed as `(F[A1], ..., F[An])` if the tuple type is fully known.
+   *  If the tuple is of the form `a1 *: ... *: Tuple` (that is, the tail is not known
+   *  to be the cons type.
+   */
   inline def map[F[_]](f: [t] => t => F[t]): Map[this.type, F] =
     DynamicTuple.dynamicMap(this, f)
+
+  /** Given a tuple `(a1, ..., am)`, returns the tuple `(a1, ..., an)` consisting
+   *  of its first n elements.
+   */
+  inline def take[This >: this.type <: Tuple](n: Int): Take[This, n.type] =
+    DynamicTuple.dynamicTake[This, n.type](this, n)
+
+
+  /** Given a tuple `(a1, ..., am)`, returns the tuple `(an+1, ..., am)` consisting
+   *  all its elements except the first n ones.
+   */
+  inline def drop[This >: this.type <: Tuple](n: Int): Drop[This, n.type] =
+    DynamicTuple.dynamicDrop[This, n.type](this, n)
+
+  /** Given a tuple `(a1, ..., am)`, returns a pair of the tuple `(a1, ..., an)`
+   *  consisting of the first n elements, and the tuple `(an+1, ..., am)` consisting
+   *  of the remaining elements.
+   */
+  inline def splitAt[This >: this.type <: Tuple](n: Int): Split[This, n.type] =
+    DynamicTuple.dynamicSplitAt[This, n.type](this, n)
 }
 
 object Tuple {
@@ -116,6 +136,29 @@ object Tuple {
    */
   type IsMappedBy[F[_]] = [X <: Tuple] =>> X =:= Map[InverseMap[X, F], F]
 
+  /** Transforms a tuple `(T1, ..., Tn)` into `(T1, ..., Ti)`. */
+  type Take[T <: Tuple, N <: Int] <: Tuple = N match {
+    case 0 => Unit
+    case S[n1] => T match {
+      case Unit => Unit
+      case x *: xs => x *: Take[xs, n1]
+    }
+  }
+
+  /** Transforms a tuple `(T1, ..., Tn)` into `(Ti+1, ..., Tn)`. */
+  type Drop[T <: Tuple, N <: Int] <: Tuple = N match {
+    case 0 => T
+    case S[n1] => T match {
+      case Unit => Unit
+      case x *: xs => Drop[xs, n1]
+    }
+  }
+
+  /** Splits a tuple (T1, ..., Tn) into a pair of two tuples `(T1, ..., Ti)` and
+   * `(Ti+1, ..., Tn)`.
+   */
+  type Split[T <: Tuple, N <: Int] = (Take[T, N], Drop[T, N])
+
   /** Convert an array into a tuple of unknown arity and types */
   def fromArray[T](xs: Array[T]): Tuple = {
     val xs2 = xs match {

library/src/scala/runtime/DynamicTuple.scala

@@ -1,6 +1,6 @@
 package scala.runtime
 
-import scala.Tuple.{ Concat, Size, Head, Tail, Elem, Zip, Map }
+import scala.Tuple.{ Concat, Size, Head, Tail, Elem, Zip, Map, Take, Drop, Split }
 
 object DynamicTuple {
   inline val MaxSpecialized = 22
@@ -380,7 +380,7 @@ object DynamicTuple {
     case _ => specialCaseTail(self)
   }
 
-  def dynamicApply[This <: NonEmptyTuple, N <: Int] (self: This, n: Int): Elem[This, N] = {
+  def dynamicApply[This <: NonEmptyTuple, N <: Int] (self: This, n: N): Elem[This, N] = {
     type Result = Elem[This, N]
     val res = (self: Any) match {
       case self: Unit => throw new IndexOutOfBoundsException(n.toString)
@@ -473,6 +473,78 @@ object DynamicTuple {
         .asInstanceOf[Map[This, F]]
   }
 
+  def dynamicTake[This <: Tuple, N <: Int](self: This, n: N): Take[This, N] = {
+    if (n < 0) throw new IndexOutOfBoundsException(n.toString)
+    val actualN = Math.min(n, self.size)
+
+    type Result = Take[This, N]
+
+    if (actualN == 0) ().asInstanceOf[Result]
+    else {
+      val arr = (self: Any) match {
+        case xxl: TupleXXL =>
+          xxl.elems.asInstanceOf[Array[Object]].take(actualN)
+        case _ =>
+          val arr = new Array[Object](actualN)
+          self.asInstanceOf[Product].productIterator.asInstanceOf[Iterator[Object]]
+            .copyToArray(arr, 0, actualN)
+          arr
+      }
+
+      dynamicFromIArray(arr.asInstanceOf[IArray[Object]]).asInstanceOf[Result]
+    }
+  }
+
+  def dynamicDrop[This <: Tuple, N <: Int](self: This, n: N): Drop[This, N] = {
+    if (n < 0) throw new IndexOutOfBoundsException(n.toString)
+    val size = self.size
+    val actualN = Math.min(n, size)
+    val rem = size - actualN
+
+    type Result = Drop[This, N]
+
+    if (rem == 0) ().asInstanceOf[Result]
+    else {
+      val arr = (self: Any) match {
+        case xxl: TupleXXL =>
+          xxl.elems.asInstanceOf[Array[Object]].drop(actualN)
+        case _ =>
+          val arr = new Array[Object](rem)
+          self.asInstanceOf[Product].productIterator.asInstanceOf[Iterator[Object]]
+            .drop(actualN).copyToArray(arr, 0, rem)
+          arr
+      }
+
+      dynamicFromIArray(arr.asInstanceOf[IArray[Object]]).asInstanceOf[Result]
+    }
+  }
+
+  def dynamicSplitAt[This <: Tuple, N <: Int](self: This, n: N): Split[This, N] = {
+    if (n < 0) throw new IndexOutOfBoundsException(n.toString)
+    val size = self.size
+    val actualN = Math.min(n, size)
+
+    type Result = Split[This, N]
+
+    val (arr1, arr2) = (self: Any) match {
+      case () => (Array.empty[Object], Array.empty[Object])
+      case xxl: TupleXXL =>
+        xxl.elems.asInstanceOf[Array[Object]].splitAt(actualN)
+      case _ =>
+        val arr1 = new Array[Object](actualN)
+        val arr2 = new Array[Object](size - actualN)
+        val it = self.asInstanceOf[Product].productIterator.asInstanceOf[Iterator[Object]]
+        it.copyToArray(arr1, 0, actualN)
+        it.copyToArray(arr2, 0, size - actualN)
+        (arr1, arr2)
+    }
+
+    (
+      dynamicFromIArray(arr1.asInstanceOf[IArray[Object]]),
+      dynamicFromIArray(arr2.asInstanceOf[IArray[Object]])
+    ).asInstanceOf[Result]
+  }
+
   def consIterator(head: Any, tail: Tuple): Iterator[Any] =
     Iterator.single(head) ++ tail.asInstanceOf[Product].productIterator
 

tests/run/tuple-drop.check

@@ -0,0 +1,14 @@
+()
+()
+()
+(1,2,3,4,5)
+(2,3,4,5)
+(4,5)
+()
+()
+(11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35)
+(12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35)
+(21,22,23,24,25,26,27,28,29,30,31,32,33,34,35)
+(31,32,33,34,35)
+()
+()

tests/run/tuple-drop.scala

@@ -0,0 +1,23 @@
+
+object Test extends App {
+  val emptyTuple: Tuple = ()
+  val tuple: Tuple = ("1", "2", "3", "4", "5")
+  val tupleXXL: Tuple = ("11", "12", "13", "14", "15", "16", "17", "18", "19", "20", "21", "22", "23", "24", "25", "26", "27", "28", "29", "30", "31", "32", "33", "34", "35")
+
+  println(emptyTuple.drop(0))
+  println(emptyTuple.drop(1))
+  println(emptyTuple.drop(10))
+
+  println(tuple.drop(0))
+  println(tuple.drop(1))
+  println(tuple.drop(3))
+  println(tuple.drop(5))
+  println(tuple.drop(10))
+
+  println(tupleXXL.drop(0))
+  println(tupleXXL.drop(1))
+  println(tupleXXL.drop(10))
+  println(tupleXXL.drop(20))
+  println(tupleXXL.drop(25))
+  println(tupleXXL.drop(30))
+}