use index to track memoized sequence iterator status

2 files changed, 36 insertions, 9 deletions

diff --git a/common-util/src/main/kotlin/com/google/devtools/ksp/MemoizedSequence.kt b/common-util/src/main/kotlin/com/google/devtools/ksp/MemoizedSequence.kt
index 49bfcecc..9fa09323 100644
--- a/common-util/src/main/kotlin/com/google/devtools/ksp/MemoizedSequence.kt
+++ b/common-util/src/main/kotlin/com/google/devtools/ksp/MemoizedSequence.kt
@@ -3,29 +3,29 @@ package com.google.devtools.ksp
 // TODO: garbage collect underlying sequence after exhaust.
 class MemoizedSequence<T>(sequence: Sequence<T>) : Sequence<T> {
 
-    private val cache = mutableListOf<T>()
+    private val cache = arrayListOf<T>()
 
     private val iter: Iterator<T> by lazy {
         sequence.iterator()
     }
 
-    private inner class CachedIterator(val iterCache: Iterator<T>, var buffer: Int) : Iterator<T> {
+    private inner class CachedIterator() : Iterator<T> {
+        var idx = 0
         override fun hasNext(): Boolean {
-            return buffer > 0 || iter.hasNext()
+            return idx < cache.size || iter.hasNext()
         }
 
         override fun next(): T {
-            if (buffer > 0) {
-                buffer --
-                return iterCache.next()
+            if (idx == cache.size) {
+                cache.add(iter.next())
             }
-            val value = iter.next()
-            cache.add(value)
+            val value = cache[idx]
+            idx += 1
             return value
         }
     }
 
     override fun iterator(): Iterator<T> {
-        return CachedIterator(cache.iterator(), cache.size)
+        return CachedIterator()
     }
 }
diff --git a/common-util/src/test/kotlin/com/google/devtools/ksp/MemoizedSequenceTest.kt b/common-util/src/test/kotlin/com/google/devtools/ksp/MemoizedSequenceTest.kt
new file mode 100644
index 00000000..52ac4b2a
--- /dev/null
+++ b/common-util/src/test/kotlin/com/google/devtools/ksp/MemoizedSequenceTest.kt
@@ -0,0 +1,27 @@
+package com.google.devtools.ksp
+
+import org.junit.Assert
+import org.junit.Test
+
+class MemoizedSequenceTest {
+    @Test
+    fun testConcurrentRead() {
+        val memoized = MemoizedSequence(
+            sequenceOf(1, 2, 3, 4, 5, 6)
+        )
+        val s1 = memoized.iterator()
+        val s2 = memoized.iterator()
+        val s1read = mutableListOf<Int>()
+        val s2read = mutableListOf<Int>()
+        while (s1.hasNext() || s2.hasNext()) {
+            if (s1.hasNext()) {
+                s1read.add(s1.next())
+            }
+            if (s2.hasNext()) {
+                s2read.add(s2.next())
+            }
+        }
+        Assert.assertEquals(listOf(1, 2, 3, 4, 5, 6), s1read)
+        Assert.assertEquals(listOf(1, 2, 3, 4, 5, 6), s2read)
+    }
+}