A lazily sorted vector

gistfile1.scala

import scala.collection.GenSeqLike
import scala.collection.mutable.ArraySeq

/**
 * A lazily sorted vector.
 *
 * Some operations may force some sorting, so return a LazySortedVector as well as the expected result.
 */
sealed trait LazySortedVector[E] {
  /** Vector length. */
  def length: Int
  /** Fetch an element from the vector by index in the sorted result. */
  def apply(i: Int): (LazySortedVector[E], E)
  /** Fetch the elements in this vector `from` (inclusive) `to` (exclusive).
   *
   * The elements may be returned out of order, but will include exactly those values that would be returned in an
   * in-order iteration.
   *
   * If the indexes are out of bounds (`from<0` or `to>=length`), return the slice after clamping `from` and `to`
   * into range.
   */
  def slice(from: Int, to: Int): (LazySortedVector[E], Traversable[E])
  /** Perform one sorting step on the vector, if unsorted. */
  def pivot: LazySortedVector[E]
}

/**
 * A fully sorted block.
 *
 * The operations here are in effect the null ops.
 */
case class Sorted[E](elems: ArraySeq[E]) extends LazySortedVector[E] {
  def length = elems.length
  def apply(i: Int) = (this, elems.apply(i))
  def slice(from: Int, to: Int): (LazySortedVector[E], Iterable[E]) = {
    def clamp(i: Int) = if(i < 0) 0 else if (i > length) length else i
    (this, elems.slice(clamp(from), clamp(to)))
  }
  def pivot: LazySortedVector[E] = this
}

/**
 * An unsorted block.
 */
case class Unsorted[E](elems: ArraySeq[E], ord: Ordering[E]) extends LazySortedVector[E] {
  def length: Int = elems.length
  def apply(i: Int) = pivot.apply(i) // we have to pivot to make sure that `i` is looked up in sorted index space
  def slice(from: Int, to: Int): (LazySortedVector[E], Traversable[E]) = // try to avoid forcing a sort if the slice covers this range
    if(from <= 0 && to >= elems.length) (this, elems) // optimize for the covering case
    else pivot.slice(from, to) // otherwise, pivot and slice
 
  def pivot: LazySortedVector[E] = {
    val pivot = elems.head
    val rest = elems.tail
    val (lower, higher) = rest.partition(ord.lt(_, pivot))
    Pivoted(Unsorted(lower, ord), pivot, Unsorted(higher, ord))
  }
}

/**
 * A pair of blocks that contain values less than and greater than a pivot value.
 */
case class Pivoted[E](lower: LazySortedVector[E], pivotV: E, higher: LazySortedVector[E]) extends LazySortedVector[E] {
  def length = lower.length + 1 + higher.length

  def apply(i: Int): (LazySortedVector[E], E) = i match {
    case l if l < lower.length =>
      val (pl, ei) = lower.pivot.apply(l)
      mkFromPivoted(pl, higher) -> ei
    case p if p == lower.length =>
      this -> pivotV
    case h if h > lower.length =>
      val (ph, ei) = higher.pivot.apply(h - (lower.length + 1))
      mkFromPivoted(lower, ph) -> ei
  }

  def slice(from: Int, to: Int): (LazySortedVector[E], Traversable[E]) = {
    val partitionI = lower.length
    def inHigher(i: Int): Int = i - partitionI - 1 // transform `i` relative to indexes in `higher`

    if(to < partitionI) { // only touches the `lower` block
      val (pl, ls) = lower.slice(from, to)
      mkFromPivoted(pl, higher) -> ls
    } else
    if(from > partitionI) { // only touches the `higher` block
      val (ph, hs) = higher.slice(inHigher(from), inHigher(to))
      mkFromPivoted(lower, ph) -> hs
    } else
    if(from == partitionI) { // includes the partition element and `higher`
      val (ph, hs) = higher.slice(inHigher(from), inHigher(to))
      mkFromPivoted(lower, ph) -> (List(pivotV) ++: hs)
    } else
    if(to == partitionI) { // includes the partition element and `higher`
      val (pl, ls) = lower.slice(from, to)
      mkFromPivoted(pl, higher) -> (ls ++ List(pivotV))
    } else { // overlaps both `lower` and `higher`
      val (pl, ls) = lower.slice(from, to)
      val (ph, hs) = higher.slice(inHigher(from), inHigher(to))
      mkFromPivoted(pl, ph) -> (ls ++ List(pivotV) ++ hs)
    }
  }

  def pivot: LazySortedVector[E] = this

  /** Try to return an optimal representation of this object if the lower and higher blocks where to be replaced. */
  protected def mkFromPivoted(pl: LazySortedVector[E], ph: LazySortedVector[E]): LazySortedVector[E] = (pl, ph) match {
    case (l, h) if (l == lower && h == higher) => // no change - return this to avoid allocating objects
      this
    case (Sorted(sl), Sorted(sh)) => // both are sorted, so we should join this all up to make a larger sorted block
      Sorted((sl ++ List(pivotV) ++ sh))
    case (_, _) => // otherwise, 
      Pivoted(pl, pivotV, ph)
  }
}