ReducingIterator.scala

ReducingIterator.scala

class ReducingIterator[A] private (iterator: Iterator[A], shouldGroup: (A, A) => Boolean, reduce: (A, A) => A) extends Iterator[A] {
  // Basic implementation robbed from https://github.com/fulcrumgenomics/fgbio/blob/27b0bc20f8baf26f555bb056073bb8de528c4ef3/src/main/scala/com/fulcrumgenomics/util/BetterBufferedIterator.scala#L36

  private var buffer: Option[A] = maybeNext

  /** Returns a Some(A) if there is a next in the iterator else None. */
  private def maybeNext = if (this.iterator.hasNext) Some(this.iterator.next()) else None

  /** True if head() and hasNext() will return an item, false otherwise. */
  override def hasNext: Boolean = this.buffer.nonEmpty

  private def advance(): Unit = { this.buffer = maybeNext }

  override def next: A = {
    accumulate(buffer.get)
  }

  @annotation.tailrec
  private def accumulate(acc: A): A = {
    advance()
    buffer match {
      case None => acc
      case Some(b) if !shouldGroup(acc, b) => acc
      case Some(b) => accumulate(reduce(acc, b))
    }
  }
}

object ReducingIterator {
  def apply[A](iterator: Iterator[A])(shouldGroup: (A, A) => Boolean, reduce: (A, A) => A): Iterator[A] = {
    new ReducingIterator(iterator, shouldGroup, reduce)
  }

  def keyValue[A, K, V](iterator: Iterator[A])(key: A => K, value: A => V, shouldGroup: (K, K) => Boolean, reduce: (V, V) => V, setV: (A, V) => A): Iterator[A] = {
    val shouldGroupA: (A, A) => Boolean = (a, b) => shouldGroup(key(a), key(b))
    val reduceA: (A, A) => A = (a, b) => setV(a, reduce(value(a), value(b)))
    new ReducingIterator[A](iterator, shouldGroupA, reduceA)
  }

  def tuple[K, V](iterator: Iterator[(K, V)])(shouldGroup: (K, K) => Boolean, reduce: (V, V) => V): Iterator[(K, V)] = {
    keyValue[(K, V), K, V](iterator)(_._1, _._2, shouldGroup, reduce, (t, v) => t._1 -> v)
  }

  def main(a: Array[String]): Unit = {
    val a = ReducingIterator(List(1, 1, 2, 3, 5, 5, 10, 99).iterator)(_ == _, _ + _).toList
    println("a = " + a)
    assert(a == List(4, 3, 20, 99))

    val b = ReducingIterator(List(1, 1, 2).iterator)(_ == _, _ + _).toList
    println("b = " + b)
    assert(b == List(4))

    val c = ReducingIterator.tuple(List(1 -> 1, 1 -> 3, 1 -> 5, 2 -> 2, 3 -> 4).iterator)(_ == _, _ + _).toList
    println("c = " + c)
    assert(c == List((1, 9), (2, 2), (3, 4)))

    val d = ReducingIterator.keyValue[(Int, Int), Int, Int](List(1 -> 1, 1 -> 3, 1 -> 5, 2 -> 2, 3 -> 4).iterator)(_._1, _._2, _ == _, _ + _, _._1 -> _).toList
    println("d = " + d)
    assert(d == List((1, 9), (2, 2), (3, 4)))
  }
}