Scala example source code file (TrieIterator.scala)

This example Scala source code file (TrieIterator.scala) is included in my "Source Code Warehouse" project. The intent of this project is to help you more easily find Scala source code examples by using tags.

All credit for the original source code belongs to scala-lang.org; I'm just trying to make examples easier to find. (For my Scala work, see my Scala examples and tutorials.)

Scala tags/keywords

annotation, anyref, array, hashmapcollision1, hashtriemap, hashtrieset, int, iterator, splititerators, t, trieiterator

The TrieIterator.scala Scala example source code

/*                     __                                               *\
**     ________ ___   / /  ___     Scala API                            **
**    / __/ __// _ | / /  / _ |    (c) 2003-2013, LAMP/EPFL             **
**  __\ \/ /__/ __ |/ /__/ __ |    http://scala-lang.org/               **
** /____/\___/_/ |_/____/_/ | |                                         **
**                          |/                                          **
\*                                                                      */

package scala
package collection
package immutable

import HashMap.{ HashTrieMap, HashMapCollision1, HashMap1 }
import HashSet.{ HashTrieSet, HashSetCollision1, HashSet1 }
import scala.annotation.unchecked.{ uncheckedVariance => uV }
import scala.annotation.tailrec

/** Abandons any pretense of type safety for speed.  You can't say I
 *  didn't try: see r23934.
 */
private[collection] abstract class TrieIterator[+T](elems: Array[Iterable[T]]) extends AbstractIterator[T] {
  outer =>

  private[immutable] def getElem(x: AnyRef): T

  def initDepth                                     = 0
  def initArrayStack: Array[Array[Iterable[T @uV]]] = new Array[Array[Iterable[T]]](6)
  def initPosStack                                  = new Array[Int](6)
  def initArrayD: Array[Iterable[T @uV]]            = elems
  def initPosD                                      = 0
  def initSubIter: Iterator[T]                      = null // to traverse collision nodes

  private[this] var depth                                     = initDepth
  private[this] var arrayStack: Array[Array[Iterable[T @uV]]] = initArrayStack
  private[this] var posStack                                  = initPosStack
  private[this] var arrayD: Array[Iterable[T @uV]]            = initArrayD
  private[this] var posD                                      = initPosD
  private[this] var subIter                                   = initSubIter

  private[this] def getElems(x: Iterable[T]): Array[Iterable[T]] = (x match {
    case x: HashTrieMap[_, _] => x.elems
    case x: HashTrieSet[_]    => x.elems
  }).asInstanceOf[Array[Iterable[T]]]

  private[this] def collisionToArray(x: Iterable[T]): Array[Iterable[T]] = (x match {
    case x: HashMapCollision1[_, _] => x.kvs.map(x => HashMap(x)).toArray
    case x: HashSetCollision1[_]    => x.ks.map(x => HashSet(x)).toArray
  }).asInstanceOf[Array[Iterable[T]]]

  private[this] type SplitIterators = ((Iterator[T], Int), Iterator[T])

  private def isTrie(x: AnyRef) = x match {
    case _: HashTrieMap[_,_] | _: HashTrieSet[_] => true
    case _                                       => false
  }
  private def isContainer(x: AnyRef) = x match {
    case _: HashMap1[_, _] | _: HashSet1[_] => true
    case _                                  => false
  }

  final class DupIterator(xs: Array[Iterable[T]]) extends {
    override val initDepth                                     = outer.depth
    override val initArrayStack: Array[Array[Iterable[T @uV]]] = outer.arrayStack
    override val initPosStack                                  = outer.posStack
    override val initArrayD: Array[Iterable[T @uV]]            = outer.arrayD
    override val initPosD                                      = outer.posD
    override val initSubIter                                   = outer.subIter
  } with TrieIterator[T](xs) {
    final override def getElem(x: AnyRef): T = outer.getElem(x)
  }

  def dupIterator: TrieIterator[T] = new DupIterator(elems)

  private[this] def newIterator(xs: Array[Iterable[T]]) = new TrieIterator(xs) {
    final override def getElem(x: AnyRef): T = outer.getElem(x)
  }

  private[this] def iteratorWithSize(arr: Array[Iterable[T]]): (Iterator[T], Int) =
    (newIterator(arr), arr.map(_.size).sum)

  private[this] def arrayToIterators(arr: Array[Iterable[T]]): SplitIterators = {
    val (fst, snd) = arr.splitAt(arr.length / 2)

    (iteratorWithSize(snd), newIterator(fst))
  }
  private[this] def splitArray(ad: Array[Iterable[T]]): SplitIterators =
    if (ad.length > 1) arrayToIterators(ad)
    else ad(0) match {
      case _: HashMapCollision1[_, _] | _: HashSetCollision1[_] =>
        arrayToIterators(collisionToArray(ad(0)))
      case _ =>
        splitArray(getElems(ad(0)))
    }

  def hasNext = (subIter ne null) || depth >= 0
  def next(): T = {
    if (subIter ne null) {
      val el = subIter.next()
      if (!subIter.hasNext)
        subIter = null
      el
    } else
      next0(arrayD, posD)
  }

  @tailrec private[this] def next0(elems: Array[Iterable[T]], i: Int): T = {
    if (i == elems.length-1) { // reached end of level, pop stack
      depth -= 1
      if (depth >= 0) {
        arrayD = arrayStack(depth)
        posD = posStack(depth)
        arrayStack(depth) = null
      } else {
        arrayD = null
        posD = 0
      }
    } else
      posD += 1

    val m = elems(i)

    // Note: this block is over twice as fast written this way as it is
    // as a pattern match.  Haven't started looking into why that is, but
    // it's pretty sad the pattern matcher is that much slower.
    if (isContainer(m))
      getElem(m) // push current pos onto stack and descend
    else if (isTrie(m)) {
      if (depth >= 0) {
        arrayStack(depth) = arrayD
        posStack(depth) = posD
      }
      depth += 1
      arrayD = getElems(m)
      posD = 0
      next0(getElems(m), 0)
    }
    else {
      subIter = m.iterator
      next()
    }
    // The much slower version:
    //
    // m match {
    //   case _: HashMap1[_, _] | _: HashSet1[_] =>
    //     getElem(m) // push current pos onto stack and descend
    //   case _: HashTrieMap[_,_] | _: HashTrieSet[_] =>
    //     if (depth >= 0) {
    //       arrayStack(depth) = arrayD
    //       posStack(depth) = posD
    //     }
    //     depth += 1
    //     arrayD = getElems(m)
    //     posD = 0
    //     next0(getElems(m), 0)
    //   case _ =>
    //     subIter = m.iterator
    //     next
    // }
  }

  // assumption: contains 2 or more elements
  // splits this iterator into 2 iterators
  // returns the 1st iterator, its number of elements, and the second iterator
  def split: SplitIterators = {
    // 0) simple case: no elements have been iterated - simply divide arrayD
    if (arrayD != null && depth == 0 && posD == 0)
      return splitArray(arrayD)

    // otherwise, some elements have been iterated over
    // 1) collision case: if we have a subIter, we return subIter and elements after it
    if (subIter ne null) {
      val buff = subIter.toBuffer
      subIter = null
      ((buff.iterator, buff.length), this)
    }
    else {
      // otherwise find the topmost array stack element
      if (depth > 0) {
        // 2) topmost comes before (is not) arrayD
        //    steal a portion of top to create a new iterator
        if (posStack(0) == arrayStack(0).length - 1) {
          // 2a) only a single entry left on top
          // this means we have to modify this iterator - pop topmost
          val snd = Array[Iterable[T]](arrayStack(0).last)
          val szsnd = snd(0).size
          // modify this - pop
          depth -= 1
          1 until arrayStack.length foreach (i => arrayStack(i - 1) = arrayStack(i))
          arrayStack(arrayStack.length - 1) = Array[Iterable[T]](null)
          posStack = posStack.tail ++ Array[Int](0)
          // we know that `this` is not empty, since it had something on the arrayStack and arrayStack elements are always non-empty
          ((newIterator(snd), szsnd), this)
        } else {
          // 2b) more than a single entry left on top
          val (fst, snd) = arrayStack(0).splitAt(arrayStack(0).length - (arrayStack(0).length - posStack(0) + 1) / 2)
          arrayStack(0) = fst
          (iteratorWithSize(snd), this)
        }
      } else {
        // 3) no topmost element (arrayD is at the top)
        //    steal a portion of it and update this iterator
        if (posD == arrayD.length - 1) {
          // 3a) positioned at the last element of arrayD
          val m = arrayD(posD)
          arrayToIterators(
            if (isTrie(m)) getElems(m)
            else collisionToArray(m)
          )
        }
        else {
          // 3b) arrayD has more free elements
          val (fst, snd) = arrayD.splitAt(arrayD.length - (arrayD.length - posD + 1) / 2)
          arrayD = fst
          (iteratorWithSize(snd), this)
        }
      }
    }
  }
}