Add missing IArray operations

Author: nicolasstucki

library/src/scala/IArray.scala

@@ -1,7 +1,8 @@
 package scala
 import reflect.ClassTag
 
-import scala.collection.immutable
+import scala.collection._
+import scala.collection.mutable.Buffer
 
 opaque type IArray[+T] = Array[_ <: T]
 
@@ -39,15 +40,10 @@ object IArray:
   extension (arr: IArray[Object]) def length: Int = arr.asInstanceOf[Array[Object]].length
   extension [T](arr: IArray[T]) def length: Int = arr.asInstanceOf[Array[T]].length
 
-  /** Returns this array concatenated with the given array. */
-  extension [T](arr: IArray[T]) def ++ [U >: T: ClassTag](that: IArray[U]): IArray[U] =
-    genericArrayOps(arr) ++ that
 
   /** Tests whether this array contains a given value as an element. */
   extension [T](arr: IArray[T]) def contains(elem: T): Boolean =
-    // `genericArrayOps(arr).contains(elem)` does not work because `elem` does not have type `arr.T`
-    // but we can use `exists` instead, which is how `ArrayOps#contains` itself is implemented:
-    genericArrayOps(arr).exists(_ == elem)
+    genericArrayOps(arr).contains(elem.asInstanceOf)
 
   /** Copy elements of this array to another array. */
   extension [T](arr: IArray[T]) def copyToArray[U >: T](xs: Array[U]): Int =
@@ -100,7 +96,7 @@ object IArray:
 
   /** Flattens a two-dimensional array by concatenating all its rows
     * into a single array. */
-  extension [T](arr: IArray[T]) def flatten[U: ClassTag](using T => Iterable[U]): IArray[U] =
+  extension [T](arr: IArray[T]) def flatten[U](using asIterable: T => Iterable[U], ct: ClassTag[U]): IArray[U] =
     genericArrayOps(arr).flatten
 
   /** Folds the elements of this array using the specified associative binary operator. */
@@ -135,9 +131,6 @@ object IArray:
 
   /** Finds index of first occurrence of some value in this array after or at some start index. */
   extension [T](arr: IArray[T]) def indexOf(elem: T, from: Int = 0): Int =
-    // `asInstanceOf` needed because `elem` does not have type `arr.T`
-    // We could use `arr.iterator.indexOf(elem, from)` or `arr.indexWhere(_ == elem, from)`
-    // but these would incur some overhead.
     genericArrayOps(arr).indexOf(elem.asInstanceOf, from)
 
   /** Finds index of the first element satisfying some predicate after or at some start index. */
@@ -170,7 +163,6 @@ object IArray:
 
   /** Finds index of last occurrence of some value in this array before or at a given end index. */
   extension [T](arr: IArray[T]) def lastIndexOf(elem: T, end: Int = arr.length - 1): Int =
-    // see: same issue in `indexOf`
     genericArrayOps(arr).lastIndexOf(elem.asInstanceOf, end)
 
   /** Finds index of last element satisfying some predicate before or at given end index. */
@@ -236,10 +228,6 @@ object IArray:
   extension [T](arr: IArray[T]) def splitAt(n: Int): (IArray[T], IArray[T]) =
     genericArrayOps(arr).splitAt(n)
 
-  /** Tests whether this array starts with the given array. */
-  extension [T](arr: IArray[T]) def startsWith[U >: T](that: IArray[U], offset: Int = 0): Boolean =
-    genericArrayOps(arr).startsWith(that)
-
   /** The rest of the array without its first element. */
   extension [T](arr: IArray[T]) def tail: IArray[T] =
     genericArrayOps(arr).tail
@@ -260,15 +248,133 @@ object IArray:
   extension [T](arr: IArray[T]) def toArray: Array[T] =
     arr.clone.asInstanceOf[Array[T]]
 
-  /** Converts an array of pairs into an array of first elements and an array of second elements. */
-  extension [U: ClassTag, V: ClassTag](arr: IArray[(U, V)]) def unzip: (IArray[U], IArray[V]) =
-    genericArrayOps(arr).unzip
-
-  /** Returns an array formed from this array and another iterable collection
-    * by combining corresponding elements in pairs.
-    * If one of the two collections is longer than the other, its remaining elements are ignored. */
-  extension [T](arr: IArray[T]) def zip[U](that: IArray[U]): IArray[(T, U)] =
-    genericArrayOps(arr).zip(that)
+  extension [T](arr: IArray[T])
+    def ++[U >: T: ClassTag](suffix: IArray[U]): IArray[U] = genericArrayOps(arr) ++ suffix.toSeq
+    def ++[U >: T: ClassTag](suffix: IterableOnce[U]): IArray[U] = genericArrayOps(arr) ++ suffix
+    def :+ [U >: T: ClassTag](x: U): IArray[U] = genericArrayOps(arr) :+ x
+    def :++ [U >: T: ClassTag](suffix: IArray[U]): IArray[U] = genericArrayOps(arr) :++ suffix
+    def :++ [U >: T: ClassTag](suffix: IterableOnce[U]): IArray[U] = genericArrayOps(arr) :++ suffix
+    def addString(b: mutable.StringBuilder): mutable.StringBuilder = arr.toSeq.addString(b)
+    def addString(b: mutable.StringBuilder, sep: String): mutable.StringBuilder = arr.toSeq.addString(b, sep)
+    def addString(b: mutable.StringBuilder, start: String, sep: String, end: String): mutable.StringBuilder = arr.toSeq.addString(b, start, sep, end)
+    def appended[U >: T: ClassTag](x: U): IArray[U] = genericArrayOps(arr).appended(x)
+    def appendedAll[U >: T: ClassTag](suffix: IArray[U]): IArray[U] = genericArrayOps(arr).appendedAll(suffix)
+    def appendedAll[U >: T: ClassTag](suffix: IterableOnce[U]): IArray[U] = genericArrayOps(arr).appendedAll(suffix)
+    def collect[U: ClassTag](pf: PartialFunction[T, U]): IArray[U] = genericArrayOps(arr).collect(pf)
+    def collectFirst[U](f: PartialFunction[T, U]): Option[U] = genericArrayOps(arr).collectFirst(f)
+    def combinations(n: Int): Iterator[IArray[T]] = genericArrayOps(arr).combinations(n)
+    def concat[U >: T: ClassTag](suffix: IArray[U]): IArray[U] = genericArrayOps(arr).concat(suffix)
+    def concat[U >: T: ClassTag](suffix: IterableOnce[U]): IArray[U] = genericArrayOps(arr).concat(suffix)
+    def containsSlice[U](that: IArray[U]): Boolean = arr.toSeq.containsSlice(that.toSeq)
+    def containsSlice[U](that: Seq[U]): Boolean = arr.toSeq.containsSlice(that)
+    def corresponds[U](that: IArray[U])(p: (T, U) => Boolean): Boolean = arr.toSeq.corresponds(that.toSeq)(p)
+    def corresponds[U](that: IterableOnce[U])(p: (T, U) => Boolean): Boolean = arr.toSeq.corresponds(that)(p)
+    def diff[U >: T](that: IArray[U]): IArray[T] = genericArrayOps(arr).diff(that.toSeq)
+    def diff[U >: T](that: Seq[U]): IArray[T] = genericArrayOps(arr).diff(that)
+    def distinct: IArray[T] = genericArrayOps(arr).distinct
+    def distinctBy[U](f: T => U): IArray[T] = genericArrayOps(arr).distinctBy(f)
+    def empty: immutable.ArraySeq[T] = arr.toSeq.empty
+    def startsWith[U >: T](that: IArray[U]): Boolean = genericArrayOps(arr).startsWith(that, 0)
+    def startsWith[U >: T](that: IArray[U], offset: Int): Boolean = genericArrayOps(arr).startsWith(that, offset)
+    def startsWith[U >: T](that: IterableOnce[U]): Boolean = genericArrayOps(arr).startsWith(that, 0)
+    def startsWith[U >: T](that: IterableOnce[U], offset: Int): Boolean = genericArrayOps(arr).startsWith(that, offset)
+    def endsWith[U >: T](that: IArray[U]): Boolean = genericArrayOps(arr).endsWith(that)
+    def endsWith[U >: T](that: Iterable[U]): Boolean = genericArrayOps(arr).endsWith(that)
+    def findLast(p: T => Boolean): Option[T] = arr.toSeq.findLast(p)
+    def groupBy[K](f: T => K): immutable.Map[K, IArray[T]] = genericArrayOps(arr).groupBy(f)
+    def groupMap[K, U: ClassTag](key: T => K)(f: T => U): immutable.Map[K, IArray[U]] = genericArrayOps(arr).groupMap(key)(f)
+    def groupMapReduce[K, U](key: (T) => K)(f: (T) => U)(reduce: (U, U) => U): immutable.Map[K, U] = arr.toSeq.groupMapReduce(key)(f)(reduce)
+    def grouped(size: Int): Iterator[IArray[T]] = genericArrayOps(arr).grouped(size)
+    def indexOfSlice[U >: T](that: IArray[U]): Int = arr.toSeq.indexOfSlice(that)
+    def indexOfSlice[U >: T](that: Seq[U]): Int = arr.toSeq.indexOfSlice(that)
+    def indexOfSlice[U >: T](that: IArray[U], from: Int): Int = arr.toSeq.indexOfSlice(that, from)
+    def indexOfSlice[U >: T](that: Seq[U], from: Int): Int = arr.toSeq.indexOfSlice(that, from)
+    def inits: Iterator[IArray[T]] = genericArrayOps(arr).inits
+    def intersect[U >: T](that: IArray[U]): IArray[T] = genericArrayOps(arr).intersect(that)
+    def intersect[U >: T](that: Seq[U]): IArray[T] = genericArrayOps(arr).intersect(that)
+    def isTraversableAgain: Boolean = arr.toSeq.isTraversableAgain
+    def knownSize: Int = arr.length
+    def lastIndexOfSlice[U >: T](that: IArray[U]): Int = arr.toSeq.lastIndexOfSlice(that)
+    def lastIndexOfSlice[U >: T](that: Seq[U]): Int = arr.toSeq.lastIndexOfSlice(that)
+    def lastIndexOfSlice[U >: T](that: IArray[U], end: Int): Int = arr.toSeq.lastIndexOfSlice(that, end)
+    def lastIndexOfSlice[U >: T](that: Seq[U], end: Int): Int = arr.toSeq.lastIndexOfSlice(that, end)
+    def lazyZip[U](that: IArray[U]): LazyZip2[T, U, IArray[T]] = genericArrayOps(arr).lazyZip[U](that).asInstanceOf[LazyZip2[T, U, IArray[T]]]
+    def lazyZip[U](that: Iterable[U]): LazyZip2[T, U, IArray[T]] = genericArrayOps(arr).lazyZip[U](that).asInstanceOf[LazyZip2[T, U, IArray[T]]]
+    def lengthCompare(len: Int): Int = genericArrayOps(arr).lengthCompare(len)
+    def lengthIs: IterableOps.SizeCompareOps = arr.toSeq.lengthIs
+    def max[U >: T](using math.Ordering[U]): T = arr.toSeq.max[U]
+    def maxBy[U](f: T => U)(using math.Ordering[U]): T = arr.toSeq.maxBy(f)
+    def maxByOption[U](f: T => U)(using math.Ordering[U]): Option[T] = arr.toSeq.maxByOption(f)
+    def maxOption[U >: T](using math.Ordering[U]): Option[U] = arr.toSeq.maxOption[U]
+    def min[U >: T](using math.Ordering[U]): T = arr.toSeq.min[U]
+    def minBy[U](f: T => U)(using math.Ordering[U]): T = arr.toSeq.minBy(f)
+    def minByOption[U](f: T => U)(using math.Ordering[U]): Option[T] = arr.toSeq.minByOption(f)
+    def minOption[U >: T](using math.Ordering[U]): Option[U] = arr.toSeq.minOption[U]
+    def mkString: String = arr.toSeq.mkString
+    def mkString(sep: String): String = arr.toSeq.mkString(sep)
+    def mkString(start: String, sep: String, end: String): String = arr.toSeq.mkString(start, sep, end)
+    def padTo[U >: T: ClassTag](len: Int, elem: U): IArray[U] = genericArrayOps(arr).padTo(len, elem)
+    def partitionMap[T1: ClassTag, T2: ClassTag](f: T => Either[T1, T2]): (IArray[T1], IArray[T2]) = genericArrayOps(arr).partitionMap(f)
+    def patch[U >: T: ClassTag](from: Int, other: IterableOnce[U], replaced: Int): IArray[U] = genericArrayOps(arr).patch(from, other, replaced)
+    def permutations: Iterator[IArray[T]] = genericArrayOps(arr).permutations
+    def prepended[U >: T: ClassTag](x: U): IArray[U] = genericArrayOps(arr).prepended(x)
+    def prependedAll[U >: T: ClassTag](prefix: IterableOnce[U]): IArray[U] = genericArrayOps(arr).prependedAll(prefix)
+    def product[U >: T](using math.Numeric[U]): U = arr.toSeq.product[U]
+    def reduce[U >: T](op: (U, U) => U): U = arr.toSeq.reduce(op)
+    def reduceLeft[U >: T](op: (U, T) => U): U = arr.toSeq.reduceLeft(op)
+    def reduceRight[U >: T](op: (T, U) => U): U = arr.toSeq.reduceRight(op)
+    def reverseIterator: Iterator[T] = genericArrayOps(arr).reverseIterator
+    def sameElements[U >: T](that: IArray[U]): Boolean = arr.toSeq.sameElements(that)
+    def sameElements[U >: T](that: IterableOnce[U]): Boolean = arr.toSeq.sameElements(that)
+    def search[U >: T](elem: U)(using Ordering[U]): Searching.SearchResult = arr.toSeq.search(elem)
+    def search[U >: T](elem: U, from: Int, to: Int)(using Ordering[U]): Searching.SearchResult = arr.toSeq.search(elem, from, to)
+    def segmentLength(p: (T) => Boolean, from: Int): Int = arr.toSeq.segmentLength(p, from)
+    def segmentLength(p: (T) => Boolean): Int = arr.toSeq.segmentLength(p)
+    def sizeCompare(that: IArray[Any]): Int = arr.toSeq.sizeCompare(that)
+    def sizeCompare(that: Iterable[_]): Int = arr.toSeq.sizeCompare(that)
+    def sizeCompare(otherSize: Int): Int = genericArrayOps(arr).sizeCompare(otherSize)
+    def sizeIs: IterableOps.SizeCompareOps = arr.toSeq.sizeIs
+    def sliding(size: Int, step: Int = 1): Iterator[IArray[T]] = genericArrayOps(arr).sliding(size, step)
+    def stepper[S <: Stepper[_]](using StepperShape[T, S]): S = genericArrayOps(arr).stepper[S]
+    def sum[U >: T](using math.Numeric[U]): U = arr.toSeq.sum[U]
+    def tails: Iterator[IArray[T]] = genericArrayOps(arr).tails
+    def tapEach[U](f: (T) => U): IArray[T] =
+      arr.toSeq.foreach(f)
+      arr
+    def to[C1](factory: Factory[T, C1]): C1 = arr.toSeq.to(factory)
+    def toBuffer[U >: T]: Buffer[U] = arr.toSeq.toBuffer[U]
+    def toIndexedSeq: immutable.IndexedSeq[T] = arr.toSeq.toIndexedSeq
+    def toIterable: Iterable[T] = arr.toSeq.toIterable
+    def toList: List[T] = arr.toSeq.toList
+    def toSet: Set[T] = arr.toSeq.toSet
+    def toVector: Vector[T] = arr.toSeq.toVector
+    def unzip[T1, T2](using asPair: T => (T1, T2), ct1: ClassTag[T1], ct2: ClassTag[T2]): (IArray[T1], IArray[T2]) = genericArrayOps(arr).unzip
+    def unzip3[T1, T2, T3](using asTriple: T => (T1, T2, T3), ct1: ClassTag[T1], ct2: ClassTag[T2], ct3: ClassTag[T3]): (IArray[T1], IArray[T2], IArray[T3]) = genericArrayOps(arr).unzip3
+    def updated[U >: T: ClassTag](index: Int, elem: U): IArray[U] = genericArrayOps(arr).updated(index, elem)
+    def view: SeqView[T] = genericArrayOps(arr).view
+    def withFilter(p: T => Boolean): WithFilter[T] = new WithFilter(p, arr)
+    def zip[U](that: IArray[U]): IArray[(T, U)] = genericArrayOps(arr).zip(that)
+    def zip[U](that: IterableOnce[U]): IArray[(T, U)] = genericArrayOps(arr).zip(that)
+    def zipAll[T1 >: T, U](that: IArray[U], thisElem: T1, thatElem: U): IArray[(T1, U)] = genericArrayOps(arr).zipAll(that, thisElem, thatElem)
+    def zipAll[T1 >: T, U](that: Iterable[U], thisElem: T1, thatElem: U): IArray[(T1, U)] = genericArrayOps(arr).zipAll(that, thisElem, thatElem)
+    def zipWithIndex: IArray[(T, Int)] = genericArrayOps(arr).zipWithIndex
+  end extension
+
+  extension [T](arr: IArray[T])
+    def transpose[U](implicit asArray: T => IArray[U]): IArray[IArray[U]] =
+      genericArrayOps(arr).transpose(using asArray.asInstanceOf[T => Array[U]])
+
+  extension [T, U >: T: ClassTag](prefix: IterableOnce[T])
+    def ++:(arr: IArray[U]): IArray[U] = genericArrayOps(arr).prependedAll(prefix)
+
+  extension [T, U >: T: ClassTag](prefix: IArray[T])
+    def ++:(arr: IArray[U]): IArray[U] = genericArrayOps(arr).prependedAll(prefix)
+
+  extension [T, U >: T: ClassTag](x: T)
+    def +:(arr: IArray[U]): IArray[U] = genericArrayOps(arr).prepended(x)
+
+  extension [T1, T2](arr: IArray[(T1, T2)])
+    def toMap: Map[T1, T2] = arr.toSeq.toMap
 
   /** Conversion from IArray to immutable.ArraySeq */
   extension [T](arr: IArray[T]) def toSeq: immutable.ArraySeq[T] =
@@ -363,16 +469,32 @@ object IArray:
   /** An immutable array with given elements. */
   def apply(x: Unit, xs: Unit*): IArray[Unit] = Array(x, xs: _*)
 
+  /** Build an array from the iterable collection.
+   *
+   *  {{{
+   *  scala> val a = Array.from(Seq(1, 5))
+   *  val a: Array[Int] = Array(1, 5)
+   *
+   *  scala> val b = Array.from(Range(1, 5))
+   *  val b: Array[Int] = Array(1, 2, 3, 4)
+   *  }}}
+   *
+   *  @param  it the iterable collection
+   *  @return    an array consisting of elements of the iterable collection
+   */
+  def from[A : ClassTag](it: IterableOnce[A]): Array[A] =
+    Array.from(it)
+
+  def newBuilder[T](using t: ClassTag[T]): scala.collection.immutable.iarray.IArrayBuilder[T] =
+    scala.collection.immutable.iarray.IArrayBuilder.make[T]
+
   /** Concatenates all arrays into a single immutable array.
    *
    *  @param xss the given immutable arrays
    *  @return   the array created from concatenating `xss`
    */
   def concat[T: ClassTag](xss: IArray[T]*): IArray[T] =
-    // `Array.concat` should arguably take in a `Seq[Array[_ <: T]]`,
-    // but since it currently takes a `Seq[Array[T]]` we have to perform a cast,
-    // knowing tacitly that `concat` is not going to do the wrong thing.
-    Array.concat[T](xss.asInstanceOf[Seq[Array[T]]]: _*)
+    Array.concat[T](xss.asInstanceOf[immutable.Seq[Array[T]]]: _*)
 
   /** Returns an immutable array that contains the results of some element computation a number
    *  of times. Each element is determined by a separate computation.
@@ -391,7 +513,6 @@ object IArray:
    *  @param   elem the element computation
    */
   def fill[T: ClassTag](n1: Int, n2: Int)(elem: => T): IArray[IArray[T]] =
-    // We cannot avoid a cast here as Array.fill creates inner arrays out of our control:
     Array.fill(n1, n2)(elem)
 
   /** Returns a three-dimensional immutable array that contains the results of some element computation a number
@@ -512,6 +633,17 @@ object IArray:
    */
   def iterate[T: ClassTag](start: T, len: Int)(f: T => T): IArray[T] = Array.iterate(start, len)(f)
 
+  /** Compare two arrays per element.
+   *
+   *  A more efficient version of `xs.sameElements(ys)`.
+   *
+   *  @param xs an array of AnyRef
+   *  @param ys an array of AnyRef
+   *  @return true if corresponding elements are equal
+   */
+   def equals(xs: IArray[AnyRef], ys: IArray[AnyRef]): Boolean =
+    Array.equals(xs.asInstanceOf[Array[AnyRef]], ys.asInstanceOf[Array[AnyRef]])
+
   /** Returns a decomposition of the array into a sequence. This supports
    *  a pattern match like `{ case IArray(x,y,z) => println('3 elements')}`.
    *
@@ -521,4 +653,66 @@ object IArray:
   def unapplySeq[T](x: IArray[T]): Array.UnapplySeqWrapper[_ <: T] =
     Array.unapplySeq(x)
 
-end IArray
+  /** A lazy filtered array. No filtering is applied until one of `foreach`, `map` or `flatMap` is called. */
+  class WithFilter[T](p: T => Boolean, xs: IArray[T]):
+
+    /** Apply `f` to each element for its side effects.
+      * Note: [U] parameter needed to help scalac's type inference.
+      */
+    def foreach[U](f: T => U): Unit = {
+      val len = xs.length
+      var i = 0
+      while(i < len) {
+        val x = xs(i)
+        if(p(x)) f(x)
+        i += 1
+      }
+    }
+
+    /** Builds a new array by applying a function to all elements of this array.
+      *
+      *  @param f      the function to apply to each element.
+      *  @tparam U     the element type of the returned array.
+      *  @return       a new array resulting from applying the given function
+      *                `f` to each element of this array and collecting the results.
+      */
+    def map[U: ClassTag](f: T => U): IArray[U] = {
+      val b = scala.collection.immutable.iarray.IArrayBuilder.make[U]
+      var i = 0
+      while (i < xs.length) {
+        val x = xs(i)
+        if(p(x)) b += f(x)
+        i = i + 1
+      }
+      b.result()
+    }
+
+    /** Builds a new array by applying a function to all elements of this array
+      * and using the elements of the resulting collections.
+      *
+      *  @param f      the function to apply to each element.
+      *  @tparam U     the element type of the returned array.
+      *  @return       a new array resulting from applying the given collection-valued function
+      *                `f` to each element of this array and concatenating the results.
+      */
+    def flatMap[U: ClassTag](f: T => IterableOnce[U]): IArray[U] = {
+      val b = scala.collection.immutable.iarray.IArrayBuilder.make[U]
+      var i = 0
+      while(i < xs.length) {
+        val x = xs(i)
+        if(p(x)) b ++= f(xs(i))
+        i += 1
+      }
+      b.result()
+    }
+
+    def flatMap[BS, U](f: T => BS)(using asIterable: BS => Iterable[U], m: ClassTag[U]): IArray[U] =
+      flatMap[U](x => asIterable(f(x)))
+
+    /** Creates a new non-strict filter which combines this filter with the given predicate. */
+    def withFilter(q: T => Boolean): WithFilter[T] = new WithFilter[T](a => p(a) && q(a), xs)
+
+  end WithFilter
+
+
+end IArray