CountingTrie Attempt

package org.gdget.util

import scala.collection.{immutable, mutable}

/** Description of Class
  *
  * @author hugofirth
  */
class CountingTrie[A] extends mutable.Map[IndexedSeq[A],Long] with mutable.MapLike[IndexedSeq[A],Long, CountingTrie[A]]{
  type TokenString = IndexedSeq[A]


  var suffixes: immutable.Map[A,CountingTrie[A]] = Map.empty
  var value: Option[Long] = None
  var total: Long = 0L

  def withPrefix(prefix: TokenString): CountingTrie[A] = {
    if (prefix.isEmpty) this
    else {
      val leading = prefix(0)
      suffixes get leading match {
        case None =>
          suffixes = suffixes + (leading -> empty)
        case _ =>
      }
      suffixes(leading) withPrefix (prefix takeRight prefix.size-1)
    }
  }

  override def get(key: TokenString): Option[Long] = {
    if(key.isEmpty) value else suffixes get key(0) flatMap( _.get(key takeRight key.size-1) )
  }

  override def update(key: TokenString, value: Long):Unit = withPrefix(key).value = Some(value)

  override def remove(key: TokenString): Option[Long] = {
    if(key.isEmpty) {
      val previous = value; value = None; previous
    } else {
      suffixes get key(0) flatMap( _.remove(key takeRight key.size-1) )
    }
  }

  override def +=(kv: (TokenString, Long)): CountingTrie[A] = { update(kv._1, kv._2); this }

  override def -=(key: TokenString): CountingTrie[A] = { remove(key); this }

  override def iterator: Iterator[(TokenString, Long)] = {
    (for(v <- value.iterator) yield (IndexedSeq.empty[A], v)) ++
      (for {
        (chr,m) <- suffixes.iterator
        (ts, v) <- m.iterator
      } yield (chr +: ts, v))
  }

  override def empty = new CountingTrie[A]

}

object CountingTrie {
  def apply[A](tokens: IndexedSeq[A]*): CountingTrie[A] = {
    val trie = new CountingTrie[A]; tokens foreach { tokenString: IndexedSeq[A] =>
      trie += ((tokenString, 0L))
    }
    trie
  }
}