Lazy Monoid Map

gistfile1.scala

import com.twitter.algebird.Monoid
import scala.collection.immutable.MapLike

/**
 * LazyMonoidMap is an immutable lazy map that wraps a collection of maps
 * that all have the same types of keys/values. Values must have a corresponding
 * monoid so that the values can be appropriately summed for a particular key that
 * exists in more than one underlying map. This data structure is useful for
 * those situations where you already have a collection of maps loaded in
 * memory and you don't want to create extra garbage by eagerly merging them
 * together.
 */
class LazyMonoidMap[A, B: Monoid](maps: Seq[Map[A, B]])
  extends Map[A, B] with MapLike[A, B, LazyMonoidMap[A,B]] {

  /** union of keys across all underlying maps */
  private lazy val allKeys: Set[A] = maps.iterator.flatMap { _.keySet }.toSet

  override lazy val size: Int = allKeys.size

  override def empty = LazyMonoidMap.empty[A, B]

  override def iterator: Iterator[(A, B)] = allKeys.iterator.map { k => (k, get(k).get) }

  override def foreach[U](f: ((A, B)) => U): Unit = iterator.foreach(f)

  override def get(key: A): Option[B] = {
    val matches = maps.iterator.flatMap { _.get(key) }
    if (matches.isEmpty) None else Some(Monoid.sum(matches))
  }

  override def updated[B1 >: B](key: A, value: B1): LazyMonoidMap[A, B1] =
    throw new UnsupportedOperationException("read only")

  override def +[B1 >: B](kv: (A, B1)): LazyMonoidMap[A, B1] =
    throw new UnsupportedOperationException("read only")

  override def +[B1 >: B](elem1: (A, B1), elem2: (A, B1), elems: (A, B1)*): LazyMonoidMap[A, B1] =
    throw new UnsupportedOperationException("read only")

  override def -(key: A): LazyMonoidMap[A, B] = new LazyMonoidMap(maps.map { _ - key })
}

object LazyMonoidMap {
  def empty[A, B: Monoid]: LazyMonoidMap[A, B] =
    new LazyMonoidMap[A, B](Seq.empty)

  def apply[A, B: Monoid](maps: Map[A, B]*): LazyMonoidMap[A, B] =
    new LazyMonoidMap[A, B](maps)
}