Vector.scala

package scala.collection.immutable;

private[immutable] object VectorUtils{
  val delta = 5;
  val ratio = 2;
  def rotateL[T](x : Vector[T], t : T, y : VectorBin[T]) = 
    if (y.left.size < ratio * y.right.size) singleL(x, t, y)
    else doubleL(x, t, y);

  def rotateR[T](x : VectorBin[T], t : T,  y : Vector[T]) = 
    if (x.right.size < ratio * x.left.size) singleR(x, t, y)
    else doubleR(x, t, y);

  def singleL[T](x : Vector[T], t : T, y : VectorBin[T]) = VectorBin(VectorBin(x, t, y.left), y.t, y.right);
  def singleR[T](x : VectorBin[T], t : T,  y : Vector[T]) = VectorBin(x.left, x.t, VectorBin(x.right, t, y));

  def doubleL[T](x : Vector[T], t1 : T, y : Vector[T]) = y match {
    case VectorBin(VectorBin(z1, t2, z2), t3, z3) => VectorBin(VectorBin(x, t1, z1), t2, VectorBin(z2, t3, z3))
    case _ => error ("The impossible happened")
  }

  def doubleR[T](x : Vector[T], t3 : T, z4 : Vector[T]) = x match {
    case VectorBin(z1, t1, VectorBin(z2, t2, z3)) => VectorBin(VectorBin(z1, t1, z2), t2, VectorBin(z3, t3, z4));
    case _ => error ("The impossible happened")
  }
}

import VectorUtils._;

object Vector{

  def singleton[T](t : T) = VectorBin[T](VectorTip, t, VectorTip);
  def empty[T] : Vector[T] = VectorTip;

  def fromIterable[T](ts : Iterable[T]) : Vector[T] = ts match {
    case (x : Vector[T]) => x
    case x => x.foldLeft(empty[T])(_ :+ _)
  }

  def apply[T](ts : T*) : Vector[T] = fromIterable(ts);

}

sealed abstract class Vector[+T] extends Seq[T]{
  import Vector._;
  override def isEmpty = size == 0;

  override def stringPrefix = "Vector"

  def length = size;

  def lookup(f : T => Int) : T = this match {
    case VectorTip => error ("Key not found");
    case VectorBin(left, t, right) => {
      val i = f(t);

      if (i == 0) t;
      else if (i < 0) left.lookup(f);
      else right.lookup(f);
    }
  }
 

  /**
   * Perform a shallow balancing of this tree - assumes its subtrees are balanced.
   */
  def balance : Vector[T];

  /**
   * Deep balancing of the tree
   */
  def join : Vector[T] = this match {
    case VectorTip => VectorTip;
    case VectorBin(VectorTip, t, right) => t +: right;
    case VectorBin(left, t, VectorTip) => left :+ t;
    case VectorBin(left@VectorBin(_, _, _), t, right@VectorBin(_, _, _)) => 
      if (delta * left.size <= right.size) VectorBin(VectorBin(left, t, right.left).join, right.t, right.right).balance
      else if (delta * right.size <= left.size) VectorBin(left.left, left.t, VectorBin(left.right, t, right).join).balance
      else this;
  }

  def +:[S >: T](t : S) : Vector[S] = this match {
    case VectorTip => singleton(t);
    case VectorBin(l, t2, r) => VectorBin(t +: l, t2, r).balance;
  }

  def :+[S >: T](t : S) : Vector[S] = this match {
    case VectorTip => singleton(t);
    case VectorBin(l, t2, r) => VectorBin(l, t2, r :+ t).balance;
  }

  override def filter(f : T => Boolean) : Vector[T] = this match {
    case VectorTip => VectorTip;
    case VectorBin(l, t, r) => {
      if (!f(t)) l.filter(f) ++ r.filter(f);
      else {
        val lf = l.filter(f);
        val rf = r.filter(f);
        if (lf.size + rf.size + 1 < size) VectorBin(lf, t, rf).balance
        else this;
      }
    }
  }

  override def foreach(f : T => Unit) = this match {
    case VectorTip => ();
    case VectorBin(l, t, r) => {l.foreach(f); f(t); r.foreach(f);}
  }

  def delete(i : Int) : Vector[T] = {
    inBounds(i);
    (this : @unchecked) match {
      case VectorTip => VectorTip;
      case VectorBin(l, t, r) => {
        if (i < l.size) {
          VectorBin(l.delete(i), t, r).balance;        
        } else if (i == l.size) l glue r 
        else {
          VectorBin(l, t, r.delete(i - l.size - 1)).balance;
        }
      }
    }
  }

  def removeRightmost : (Vector[T], T) = this match {
    case VectorBin(l, t, VectorTip) => (l, t);
    case VectorBin(l, t, r) => { val (r2, rightmost) = r.removeRightmost; (VectorBin(l, t, r2).balance, rightmost) }
    case VectorTip => error ("Tried to remove rightmost from tip");
  }

  def removeLeftmost : (T, Vector[T]) = this match {
    case VectorBin(VectorTip, t, r) => (t, r);
    case VectorBin(l, t, r) => { val (leftmost, l2) = l.removeLeftmost; (leftmost, VectorBin(l2, t, r).balance) }
    case VectorTip => error ("Tried to remove leftmost from tip")
  }

  /**
   * Glues that tree to the right of this, performing no balancing between them,
   */
  private def glue[S >: T](that : Vector[S]) = (this, that) match {
    case (_, VectorTip) => this;
    case (VectorTip, _) => that;
    case (l, r) => if (l.size > r.size) { val (l2, t) = l.removeRightmost; VectorBin(l2, t, r).balance; }
                   else { val (t, r2) = r.removeLeftmost; VectorBin(l, t, r2).balance; }
  }


  override def ++[S >: T](that : Iterable[S]) : Vector[S] = this ++ Vector.fromIterable(that);
  def ++[S >: T](that : Vector[S]) : Vector[S] = glue(that).join 

  def inBounds(i : Int) = if (i < 0 || i >= size) error("Index " + i + " out of bounds");

  final def apply(i : Int) : T = {
    inBounds(i);
    (this : @unchecked) match {
      case VectorBin(left, t, right) => {
        if (i < left.size) left(i)
        else if (i == left.size) t;
        else right(i - left.size - 1);
      }
    } 
  }

  def update[S >: T](i : Int, x : S) : Vector[S] = {
    inBounds(i);
    (this : @unchecked) match {
      case VectorBin(l, y, r) => {
        if (i < l.size) VectorBin(l.update(i, x), y, r)
        else if (i == l.size) VectorBin(l, x, r);
        else VectorBin(l, y, r.update(i - l.size - 1, x))
      }
    }
  }

  def split(i : Int) : (Vector[T], T, Vector[T]) = {
    inBounds(i);
    (this : @unchecked) match {
      case VectorBin(l, t, r) => {
        if (i < l.size) {
          val (ll, lt, lr) = l.split(i);          
          (ll, lt, (lr :+ t) ++ r);
        }
        else if (i == l.size) {
          (l, t, r)
        } else {
          val (rl, rt, rr) = r.split(i - l.size - 1);
          ((l :+ t) ++ rl, rt, rr)
        }
      }
    }
  }
}

private[immutable] case object VectorTip extends Vector[Nothing]{
  override def size = 0;
  def elements = Iterator.empty;
  def balance = this;
}

private[immutable] case class VectorBin[T](left : Vector[T], t : T,  right : Vector[T]) extends Vector[T]{
  import Vector._;
  override val size = left.size + right.size + 1;

  def elements = left.elements ++ Iterator.single(t) ++ right.elements

  def balance = if (left.size + right.size <= 1) VectorBin(left, t, right);
    // The following casts are safe, because VectorTips can not have non-zero size.
    else if (right.size >= delta * left.size) rotateL(left, t, right.asInstanceOf[VectorBin[T]]);
    else if (left.size >= delta * right.size) rotateR(left.asInstanceOf[VectorBin[T]], t, right); 
    else this; 

}