PrefixMap.scala

package com.rockymadden.collection.mutable

import scala.annotation.tailrec
import scala.collection.generic.CanBuildFrom
import scala.collection.mutable.{Builder, Map, MapBuilder, MapLike}

final class PrefixMap[A] extends Map[String, A] with MapLike[String, A, PrefixMap[A]] {
	private val suffixes: Map[Char, PrefixMap[A]] = Map.empty
	private var value: Option[A] = None

	def +=(pair: (String, A)): this.type = {
		update(pair._1, pair._2)
		this
	}

	def -=(key: String): this.type = {
		remove(key)
		this
	}

	override def empty = new PrefixMap[A]

	def get(key: String): Option[A] =
		if (key.isEmpty) value else suffixes.get(key(0)).flatMap(_.get(key.substring(1)))

	def iterator: Iterator[(String, A)] =
		value.iterator.map(("", _)) ++
			suffixes.iterator.flatMap { case (c, m) => m.iterator.map { case (k, v) => (c +: k, v) } }

	override def remove(key: String): Option[A] =
		if (key.isEmpty) {
			val previous = value
			value = None
			previous
		} else suffixes.get(key(0)).flatMap(_.remove(key.substring(1)))

	override def update(key: String, value: A): Unit = withPrefix(key).value = Some(value)

	@tailrec
	def withPrefix(key: String): PrefixMap[A] =
		if (key.isEmpty) this
		else {
			val leading = key(0)
			if (!suffixes.contains(leading)) suffixes += (leading -> empty)
			suffixes(leading).withPrefix(key.substring(1))
		}
}

object PrefixMap {
	def apply[A](pairs: (String, A)*): PrefixMap[A] = {
		val map: PrefixMap[A] = empty
		pairs.foreach(map += _)
		map
	}

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

	def empty[A] = new PrefixMap[A]

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