HAMT: valuesIterator

durban · Nov 1, 2024 · ba419b8 · ba419b8
1 parent 59e2701
commit ba419b8
Show file tree

Hide file tree

Showing 3 changed files with 203 additions and 1 deletion.
diff --git a/mcas/shared/src/main/scala/dev/tauri/choam/internal/mcas/AbstractHamt.scala b/mcas/shared/src/main/scala/dev/tauri/choam/internal/mcas/AbstractHamt.scala
@@ -20,8 +20,9 @@ package internal
 package mcas
 
 import scala.util.hashing.MurmurHash3
+import scala.collection.AbstractIterator
 
-private[mcas] abstract class AbstractHamt[K <: Hamt.HasHash, V <: Hamt.HasKey[K], E, T1, T2, H <: AbstractHamt[K, V, E, T1, T2, H]] protected[mcas] () { this: H =>
+private[mcas] abstract class AbstractHamt[K <: Hamt.HasHash, V <: Hamt.HasKey[K], E, T1, T2, H <: AbstractHamt[K, V, E, T1, T2, H]] protected[mcas] () { self: H =>
 
   protected def newArray(size: Int): Array[E]
 
@@ -46,6 +47,77 @@ private[mcas] abstract class AbstractHamt[K <: Hamt.HasHash, V <: Hamt.HasKey[K]
     this.forAllInternal(tok)
   }
 
+  /** Only call on the root! */
+  final def valuesIterator: Iterator[V] = {
+    new AbstractIterator[V] {
+
+      private[this] val arrays =
+        new Array[Array[AnyRef]](11)
+
+      private[this] val indices =
+        new Array[Int](11)
+
+      private[this] var depth: Int =
+        0
+
+      private[this] var loadedNext: V =
+        nullOf[V]
+
+      locally {
+        this.arrays(0) = self.contentsArr
+        this.indices(0) = -1
+        this.loadNext()
+      }
+
+      final override def hasNext: Boolean = {
+        !isNull(this.loadedNext)
+      }
+
+      final override def next(): V = {
+        if (this.hasNext) {
+          val res = this.loadedNext
+          this.loadedNext = nullOf[V]
+          this.loadNext()
+          res
+        } else {
+          throw new NoSuchElementException
+        }
+      }
+
+      @tailrec
+      private[this] final def loadNext(): Unit = {
+        val d = this.depth
+        val idx = this.indices(d) + 1
+        this.indices(d) = idx
+        val arr = this.arrays(d)
+        if (idx < arr.length) {
+          arr(idx) match {
+            case null =>
+              // skip empty slot:
+              this.loadNext()
+            case node: AbstractHamt[_, _, _, _, _, _] =>
+              // descend:
+              val newDepth = d + 1
+              this.depth = newDepth
+              this.indices(newDepth) = -1
+              this.arrays(newDepth) = node.contentsArr
+              this.loadNext()
+            case a =>
+              // found it:
+              this.loadedNext = a.asInstanceOf[V]
+          }
+        } else {
+          // ascend:
+          val newDepth = d - 1
+          this.depth = newDepth
+          if (newDepth >= 0) {
+            this.loadNext()
+          } // else: at end
+        }
+      }
+    }
+  }
+
   final def toString(pre: String, post: String): String = {
     val sb = new java.lang.StringBuilder(pre)
     val _ = this.toStringInternal(sb, first = true)

diff --git a/mcas/shared/src/test/scala/dev/tauri/choam/internal/mcas/HamtSpec.scala b/mcas/shared/src/test/scala/dev/tauri/choam/internal/mcas/HamtSpec.scala
@@ -569,6 +569,71 @@ final class HamtSpec extends ScalaCheckSuite with MUnitUtils with PropertyHelper
     // but copying to an array detects it:
     assertEquals(hamt.toArray((), flag = false, nullIfBlue = true), null)
   }
+
+  test("valuesIterator examples") {
+    val c0 = 0x0L
+    val c1 = 0x8000000000000000L
+    val h0 = LongHamt.empty
+    val it00 = h0.valuesIterator
+    assert(!it00.hasNext)
+    assert(!it00.hasNext)
+    try it00.next() catch { case _: NoSuchElementException => () }
+    val it01 = h0.valuesIterator
+    val it02 = h0.valuesIterator
+    assert(!it01.hasNext)
+    assert(!it02.hasNext)
+    assert(!it00.hasNext)
+    val h1 = h0.inserted(Val(c0))
+    val it10 = h1.valuesIterator
+    val it11 = h1.valuesIterator
+    assert(it10.hasNext)
+    assert(it10.hasNext)
+    assertEquals(it10.next(), Val(c0))
+    assert(!it10.hasNext)
+    assert(!it10.hasNext)
+    assert(it11.hasNext)
+    assertEquals(it11.next(), Val(c0))
+    assert(!it10.hasNext)
+    assert(!it11.hasNext)
+    val h2 = h1.inserted(Val(c1))
+    val it20 = h2.valuesIterator
+    val it21 = h2.valuesIterator
+    assert(it20.hasNext)
+    assert(it21.hasNext)
+    assertEquals(it20.next(), Val(c0))
+    assert(it20.hasNext)
+    assert(it21.hasNext)
+    assertEquals(it20.next(), Val(c1))
+    assert(!it20.hasNext)
+    assert(!it20.hasNext)
+    assert(it21.hasNext)
+    assertEquals(it21.next(), Val(c0))
+    assertEquals(it21.next(), Val(c1))
+    assert(!it20.hasNext)
+    assert(!it20.hasNext)
+    assert(!it21.hasNext)
+    assert(!it21.hasNext)
+  }
+
+  property("valuesIterator (default generator)") {
+    forAll { (seed: Long, nums: Set[Long]) =>
+      testValuesIterator(seed, nums)
+    }
+  }
+
+  property("valuesIterator (RIG generator)") {
+    myForAll { (seed: Long, nums: Set[Long]) =>
+      testValuesIterator(seed, nums)
+    }
+  }
+
+  private def testValuesIterator(seed: Long, nums: Set[Long]): Unit = {
+    val rng = new Random(seed)
+    val h = hamtFromList(rng.shuffle(nums.toList))
+    val expected = h.toArray.toList
+    val actual = h.valuesIterator.toList
+    assertEquals(actual, expected)
+  }
 }
 
 object HamtSpec {

diff --git a/mcas/shared/src/test/scala/dev/tauri/choam/internal/mcas/MutHamtSpec.scala b/mcas/shared/src/test/scala/dev/tauri/choam/internal/mcas/MutHamtSpec.scala
@@ -698,6 +698,71 @@ final class MutHamtSpec extends ScalaCheckSuite with MUnitUtils with PropertyHel
     assertEquals(immutableAct.size, immutableExp.size)
     assertEquals(immutableAct.toArray.toList, immutableExp.toArray.toList)
   }
+
+  test("valuesIterator examples") {
+    val c0 = 0x0L
+    val c1 = 0x8000000000000000L
+    val h = LongMutHamt.newEmpty()
+    val it00 = h.valuesIterator
+    assert(!it00.hasNext)
+    assert(!it00.hasNext)
+    try it00.next() catch { case _: NoSuchElementException => () }
+    val it01 = h.valuesIterator
+    val it02 = h.valuesIterator
+    assert(!it01.hasNext)
+    assert(!it02.hasNext)
+    assert(!it00.hasNext)
+    h.insert(Val(c0))
+    val it10 = h.valuesIterator
+    val it11 = h.valuesIterator
+    assert(it10.hasNext)
+    assert(it10.hasNext)
+    assertEquals(it10.next(), Val(c0))
+    assert(!it10.hasNext)
+    assert(!it10.hasNext)
+    assert(it11.hasNext)
+    assertEquals(it11.next(), Val(c0))
+    assert(!it10.hasNext)
+    assert(!it11.hasNext)
+    h.insert(Val(c1))
+    val it20 = h.valuesIterator
+    val it21 = h.valuesIterator
+    assert(it20.hasNext)
+    assert(it21.hasNext)
+    assertEquals(it20.next(), Val(c0))
+    assert(it20.hasNext)
+    assert(it21.hasNext)
+    assertEquals(it20.next(), Val(c1))
+    assert(!it20.hasNext)
+    assert(!it20.hasNext)
+    assert(it21.hasNext)
+    assertEquals(it21.next(), Val(c0))
+    assertEquals(it21.next(), Val(c1))
+    assert(!it20.hasNext)
+    assert(!it20.hasNext)
+    assert(!it21.hasNext)
+    assert(!it21.hasNext)
+  }
+
+  property("valuesIterator (default generator)") {
+    forAll { (seed: Long, nums: Set[Long]) =>
+      testValuesIterator(seed, nums)
+    }
+  }
+
+  property("valuesIterator (RIG generator)") {
+    myForAll { (seed: Long, nums: Set[Long]) =>
+      testValuesIterator(seed, nums)
+    }
+  }
+
+  private def testValuesIterator(seed: Long, nums: Set[Long]): Unit = {
+    val rng = new Random(seed)
+    val h = mutHamtFromList(rng.shuffle(nums.toList))
+    val expected = h.toArray.toList
+    val actual = h.valuesIterator.toList
+    assertEquals(actual, expected)
+  }
 }
 
 object MutHamtSpec {