From fafefc2820062e13bba446c4a3d7b233ba5983d9 Mon Sep 17 00:00:00 2001
From: kageru <kageru@encode.moe>
Date: Wed, 9 Oct 2019 12:01:53 +0200
Subject: [PATCH] Add Sequence.peekable()

---
 .../moe/kageru/sekwences/PeekableSequence.kt  | 37 +++++++++++++++++++
 .../kotlin/moe/kageru/sekwences/Sekwences.kt  |  9 +++++
 .../{SekwencesTest.kt => CoalesceTest.kt}     |  2 +-
 .../moe/kageru/sekwences/PeekableTest.kt      | 24 ++++++++++++
 4 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 src/main/kotlin/moe/kageru/sekwences/PeekableSequence.kt
 rename src/test/kotlin/moe/kageru/sekwences/{SekwencesTest.kt => CoalesceTest.kt} (97%)
 create mode 100644 src/test/kotlin/moe/kageru/sekwences/PeekableTest.kt

diff --git a/src/main/kotlin/moe/kageru/sekwences/PeekableSequence.kt b/src/main/kotlin/moe/kageru/sekwences/PeekableSequence.kt
new file mode 100644
index 0000000..cb155f2
--- /dev/null
+++ b/src/main/kotlin/moe/kageru/sekwences/PeekableSequence.kt
@@ -0,0 +1,37 @@
+package moe.kageru.sekwences
+
+class PeekableSequence<T>(source: Sequence<T>) : Sequence<T> {
+    private val iterator: PeekableIterator<T> = PeekableIterator(source.iterator())
+
+    override fun iterator(): Iterator<T> = iterator
+
+    fun peek(): T? = iterator.peek()
+}
+
+internal class PeekableIterator<T>(private val source: Iterator<T>) : Iterator<T> {
+    private var peeked: T? = null
+    override fun hasNext(): Boolean {
+        return source.hasNext() || peeked != null
+    }
+
+    override fun next(): T {
+        return if (peeked == null) {
+            source.next()
+        } else {
+            val ret = peeked
+            peeked = null
+            ret!!
+        }
+    }
+
+    fun peek(): T? {
+        return when {
+            peeked != null -> peeked
+            hasNext() -> {
+                peeked = next()
+                peeked!!
+            }
+            else -> null
+        }
+    }
+}
diff --git a/src/main/kotlin/moe/kageru/sekwences/Sekwences.kt b/src/main/kotlin/moe/kageru/sekwences/Sekwences.kt
index 6284abc..3f08af5 100644
--- a/src/main/kotlin/moe/kageru/sekwences/Sekwences.kt
+++ b/src/main/kotlin/moe/kageru/sekwences/Sekwences.kt
@@ -18,3 +18,12 @@ fun <T> Sequence<T>.coalesce(merger: (T, T) -> T?): Sequence<T> {
         emptySequence()
     }
 }
+
+/**
+ * Modifies the current [Sequence] by adding a peek() method
+ * that returns the next element or null without advancing the sequence.
+ * Calling peek() multiple times without next() will return the same value.
+ */
+fun <T> Sequence<T>.peekable(): PeekableSequence<T> {
+    return PeekableSequence(this)
+}
diff --git a/src/test/kotlin/moe/kageru/sekwences/SekwencesTest.kt b/src/test/kotlin/moe/kageru/sekwences/CoalesceTest.kt
similarity index 97%
rename from src/test/kotlin/moe/kageru/sekwences/SekwencesTest.kt
rename to src/test/kotlin/moe/kageru/sekwences/CoalesceTest.kt
index 07b9404..4e6ec9d 100644
--- a/src/test/kotlin/moe/kageru/sekwences/SekwencesTest.kt
+++ b/src/test/kotlin/moe/kageru/sekwences/CoalesceTest.kt
@@ -3,7 +3,7 @@ package moe.kageru.sekwences
 import io.kotlintest.shouldBe
 import io.kotlintest.specs.ShouldSpec
 
-class SekwencesTest : ShouldSpec({
+class CoalesceTest : ShouldSpec({
     "should merge numbers" {
         val input = sequenceOf(1, 2, 3, 4, 5, 6, 4)
         val output = input.coalesce { first, second ->
diff --git a/src/test/kotlin/moe/kageru/sekwences/PeekableTest.kt b/src/test/kotlin/moe/kageru/sekwences/PeekableTest.kt
new file mode 100644
index 0000000..082ad9d
--- /dev/null
+++ b/src/test/kotlin/moe/kageru/sekwences/PeekableTest.kt
@@ -0,0 +1,24 @@
+package moe.kageru.sekwences
+
+import io.kotlintest.shouldBe
+import io.kotlintest.specs.ShouldSpec
+
+class PeekableTest : ShouldSpec({
+    "peek should return correct values" {
+        val input = sequenceOf(1, 2, 3, 4).peekable()
+        input.peek() shouldBe 1
+        input.peek() shouldBe input.peek()
+        input.iterator().next() shouldBe 1
+        input.peek() shouldBe 2
+        input.iterator().next() shouldBe 2
+        input.iterator().next() shouldBe 3
+        input.iterator().next() shouldBe 4
+        input.peek() shouldBe null
+    }
+
+    "sequence should operate normally" {
+        val input = sequenceOf(1, 2, 3, 4).peekable()
+        input.peek() shouldBe 1
+        input.map { it * it }.toList() shouldBe listOf(1, 4, 9, 16)
+    }
+})