Immutable PrefixMap

gistfile1.scala

import scala.collection.generic.CanBuildFrom
import scala.collection.immutable.MapLike
import scala.collection.mutable

/**
 * Immutable version of the PrefixMap demonstrated in the Scala collections
 * guide here: http://www.scala-lang.org/docu/files/collections-api/collections-impl_6.html
 * 
 * This version also has a smarter remove method (doesn't leave behind dead nodes with empty values)
 */
object PrefixMap {
  def empty[A]: PrefixMap[A] = new PrefixMap[A](None, Map.empty)

  def apply[A](kvs: (String, A)*): PrefixMap[A] =
    kvs.foldLeft(empty[A]) { _ + _ }

  def newBuilder[A]: mutable.Builder[(String, A), PrefixMap[A]] =
    new mutable.MapBuilder[String, A, PrefixMap[A]](empty)

  implicit def cbf[A]: CanBuildFrom[PrefixMap[_], (String, A), PrefixMap[A]] =
    new CanBuildFrom[PrefixMap[_], (String, A), PrefixMap[A]] {
      def apply(from: PrefixMap[_]) = newBuilder[A]
      def apply() = newBuilder[A]
    }
}

class PrefixMap[+A](value: Option[A], children: Map[Char, PrefixMap[A]])
  extends Map[String, A] with MapLike[String, A, PrefixMap[A]] {

  override def get(s: String): Option[A] = {
    if (s.isEmpty) value
    else children.get(s.head).flatMap { _.get(s.tail) }
  }

  override def +[B >: A](kv: (String, B)): PrefixMap[B] = {
    val (s, v) = kv
    if (s.isEmpty) new PrefixMap(Some(v), children)
    else {
      val first = s.head
      children.get(first) match {
        case Some(m) =>
          val updatedMap = m + (s.tail -> v)
          val newChildren = children.updated(first, updatedMap)
          new PrefixMap(value, newChildren)
        case None =>
          val updatedMap = PrefixMap.empty + (s.tail -> v)
          new PrefixMap(value, children.updated(first, updatedMap))
      }
    }
  }

  override def -(s: String): PrefixMap[A] = {
    if (s.isEmpty) new PrefixMap(None, children)
    else {
      val first = s.head
      children.get(first) match {
        case Some(m) =>
          val updatedMap = m - s.tail
          // after updating, prune this part of the tree
          // if it's now empty
          val newChildren =
            if (updatedMap.isEmpty) children - first
            else children.updated(first, updatedMap)
          new PrefixMap(value, newChildren)
        case None => this
      }
    }
  }

  override def isEmpty: Boolean =
    value.isEmpty && children.values.forall { _.isEmpty }

  override def empty: PrefixMap[A] = new PrefixMap(None, Map.empty)

  override def iterator: Iterator[(String, A)] = {
    (for (v <- value.iterator) yield ("", v)) ++
      (for {
        (chr, m) <- children.iterator
        (s, v) <- m.iterator
      } yield (chr +: s, v))
  }

  def withPrefix(s: String): PrefixMap[A] = {
    if (s.isEmpty) this
    else {
      children.get(s.head) match {
        case Some(m) => m.withPrefix(s.tail)
        case None => empty
      }
    }
  }
}