Peek as much as you'd like, memory allowing

FullyPeekableIterator.scala

/*
 * The MIT License
 *
 * Copyright © 2022 Clint Valentine
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
package io.cvbio.collection

import com.fulcrumgenomics.FgBioDef.yieldAndThen

import scala.collection.{AbstractIterator, BufferedIterator, mutable}

/** An iterator that is capable of buffering forward any number of elements. */
trait FullyPeekableIterator[+A] extends BufferedIterator[A] {

  /** Lift an element from this iterator without advancing the iterator. Stores peeked elements in memory. */
  def lift(index: Int): Option[A]

  /** Lift many elements from this iterator without advancing the iterator. Stores peeked elements in memory. */
  def liftMany(start: Int, end: Int): Seq[Option[A]]

  /** Peek elements while the predicate remains true. */
  def peekWhile(p: A => Boolean): FullyPeekableIterator[A]

  /** Return this iterator as it is already buffered. */
  override def buffered: this.type = this
}

/** Companion object for [[FullyPeekableIterator]]. */
object FullyPeekableIterator {

  /** Additional methods on the base Scala iterator class. */
  implicit class FullyPeekableIteratorImpl[A](private val iter: Iterator[A]) {

    /** Create a fully peekable iterator from this iterator. */
    def fullyPeekable: FullyPeekableIterator[A] = {
      new AbstractIterator[A] with FullyPeekableIterator[A] {
        private val queue: mutable.Queue[A] = new mutable.Queue[A]()

        /** If this iterator has another item or not. */
        def hasNext: Boolean = queue.nonEmpty || iter.hasNext

        /** The next item in this iterator without advancing the iterator. */
        def head: A = headOption.getOrElse(Iterator.empty.next())

        /** Lift an element from this iterator without advancing the iterator. */
        override def headOption: Option[A] = lift(0)

        /** The known size of this iterator. If there is no known size, return <= -1. */
        override def knownSize: Int = {
          val innerSize = iter.knownSize
          if (innerSize >= 0) innerSize + queue.knownSize else innerSize
        }

        /** Lift an element from this iterator without advancing the iterator. */
        def lift(index: Int): Option[A] = {
          while (queue.length <= index + 1 && iter.hasNext) queue.enqueue(iter.next())
          queue.lift(index)
        }

        /** Lift many elements from this iterator without advancing the iterator. */
        def liftMany(start: Int, end: Int): Seq[Option[A]] = Range(start, end).inclusive.map(lift)

        /** Peek elements while the predicate remains true. */
        def peekWhile(p: A => Boolean): FullyPeekableIterator[A] = {
          val self: FullyPeekableIterator[A] = this
          new AbstractIterator[A] {
            private var index: Int = 0
            override def hasNext: Boolean = self.lift(index).exists(p)
            override def next(): A = yieldAndThen(self.lift(index))(index += 1).getOrElse(Iterator.empty.next())
          }.fullyPeekable
        }

        /** Returns an iterator over contiguous elements that match the predicate without discarding excess. */
        override def takeWhile(p: A => Boolean): FullyPeekableIterator[A] = {
          val self: FullyPeekableIterator[A] = this
          new AbstractIterator[A] {
            def hasNext: Boolean = self.headOption.exists(p)
            def next(): A = self.next()
          }.fullyPeekable
        }

        /** Drops items while they match the predicate without discarding excess. */
        override def dropWhile(p: A => Boolean): this.type = {
          while (headOption.exists(p)) next()
          this
        }

        /** The next element in the iterator which will advance the iterator. */
        def next(): A = if (queue.nonEmpty) queue.dequeue() else iter.next()
      }
    }
  }
}

FullyPeekableIteratorTest.scala

/*
 * The MIT License
 *
 * Copyright © 2022 Clint Valentine
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
package io.cvbio.collection

import com.fulcrumgenomics.commons.CommonsDef._
import io.cvbio.collection.FullyPeekableIterator._
import io.cvbio.testing.UnitSpec
import org.scalatest.OptionValues._

/** Unit tests for [[FullyPeekableIterator]]. */
class FullyPeekableIteratorTest extends UnitSpec {

  "FullyPeekableIterator" should "iterate over elements like a normal iterator" in {
    val expected = Seq(1, 2, 3)
    val peekable = expected.iterator.fullyPeekable
    peekable.toSeq should contain theSameElementsInOrderAs expected
  }

  it should "also support operations that a normal buffered iterator has (without advancing)" in {
    val expected = Seq(1, 2, 3)
    val peekable = expected.iterator.fullyPeekable
    peekable.headOption.value shouldBe 1
    peekable.head shouldBe 1
    peekable.toSeq should contain theSameElementsInOrderAs Seq(1, 2, 3)
  }

  it should "lift a single element from the head of the iterator without advancing" in {
    val expected = Seq(1, 2, 3)
    val peekable = expected.iterator.fullyPeekable
    peekable.lift(0).value shouldBe 1
    peekable.toSeq should contain theSameElementsInOrderAs expected
  }

  it should "know it has a next element even if it has peeked that single element" in {
    val expected = Seq(1)
    val peekable = expected.iterator.fullyPeekable
    peekable.lift(0).value shouldBe 1
    peekable.hasNext shouldBe true
    peekable.next() shouldBe 1
    peekable.hasNext shouldBe false
  }

  it should "not have a known size when it wraps an iterator without a known size" in {
    val expected = Seq(1, 2, 3)
    val iterator = expected.iterator
    val peekable = iterator.fullyPeekable
    peekable.knownSize shouldBe iterator.knownSize
    peekable.headOption.value shouldBe 1 // Fills the underlying buffer with 1 element.
    peekable.knownSize shouldBe iterator.knownSize
  }

  it should "have a known size when it wraps an iterator with a known size" in {
    class SizedIterator[T](iter: Iterator[T]) extends Iterator[T] {
      private val forced            = iter.toSeq
      private var seen              = 0
      private val underlying        = forced.iterator
      override def knownSize: Int   = forced.length - seen
      override def hasNext: Boolean = underlying.hasNext
      override def next(): T        = yieldAndThen(underlying.next())(seen += 1)
    }

    val sized    = new SizedIterator(Seq(1, 2, 3).iterator)
    val peekable = sized.fullyPeekable
    peekable.knownSize shouldBe 3
    peekable.headOption.value shouldBe 1
    peekable.knownSize shouldBe 3
    peekable.next() shouldBe 1
    peekable.knownSize shouldBe 2
    peekable.next() shouldBe 2
    peekable.next() shouldBe 3
    peekable.knownSize shouldBe 0
  }

  it should "be able to lift many elements without advancing the iterator" in {
    val expected = Seq(1, 2, 3)
    val peekable = expected.iterator.fullyPeekable
    peekable.liftMany(0, 3) should contain theSameElementsInOrderAs Seq(Some(1), Some(2), Some(3), None)
    peekable.toSeq should contain theSameElementsInOrderAs expected
  }

  it should "be able to peek forward while a predicate is true without advancing the iterator" in {
    val expected = Seq(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
    val peekable = expected.iterator.fullyPeekable
    peekable.peekWhile(_ == 0).toSeq shouldBe empty
    peekable.peekWhile(_ <= 5).toSeq should contain theSameElementsInOrderAs Seq(1, 2, 3, 4, 5)
    peekable.next() shouldBe 1 // Advance past the first element.
    peekable.peekWhile(_ <= 5).toSeq should contain theSameElementsInOrderAs Seq(2, 3, 4, 5)
    peekable.toSeq should contain theSameElementsInOrderAs expected.drop(1)
  }

  it should "be able to take while a predicate is true without excessively advancing the iterator" in {
    val expected = Seq(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
    val peekable = expected.iterator.fullyPeekable
    peekable.takeWhile(_ == 0).toSeq shouldBe empty
    peekable.takeWhile(_ <= 5).toSeq should contain theSameElementsInOrderAs Seq(1, 2, 3, 4, 5)
    peekable.next() shouldBe 6 // Advance past the sixth element.
    peekable.takeWhile(_ => true).toSeq should contain theSameElementsInOrderAs Seq(7, 8, 9, 10)
  }

  it should "be able to drop while a predicate is true without excessively advancing the iterator" in {
    val expected = Seq(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
    val peekable = expected.iterator.fullyPeekable
    peekable.dropWhile(_ <= 5).toSeq should contain theSameElementsInOrderAs Seq(6, 7, 8, 9, 10)
  }
}