Added mapWhileNotNull & mapWhileCheck (#1768)

korlibs · Jul 5, 2023 · b70656a · b70656a
1 parent eba15d4
commit b70656a
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 5 deletions.
diff --git a/kds/src/commonMain/kotlin/korlibs/datastructure/_Extensions.kt b/kds/src/commonMain/kotlin/korlibs/datastructure/_Extensions.kt
@@ -14,6 +14,20 @@ inline fun mapWhileInt(cond: (index: Int) -> Boolean, gen: (Int) -> Int): IntArr
 inline fun mapWhileFloat(cond: (index: Int) -> Boolean, gen: (Int) -> Float): FloatArray = FloatArrayList().apply { while (cond(this.size)) this += gen(this.size) }.toFloatArray()
 inline fun mapWhileDouble(cond: (index: Int) -> Boolean, gen: (Int) -> Double): DoubleArray = DoubleArrayList().apply { while (cond(this.size)) this += gen(this.size) }.toDoubleArray()
 
+inline fun <reified T> mapWhileNotNull(gen: (Int) -> T?): List<T> = arrayListOf<T>().apply {
+    while (true) {
+        this += gen(this.size) ?: break
+    }
+}
+
+inline fun <reified T> mapWhileCheck(check: (T) -> Boolean, gen: (Int) -> T): List<T> = arrayListOf<T>().apply {
+    while (true) {
+        val res = gen(this.size)
+        if (!check(res)) break
+        this += res
+    }
+}
+
 fun <T> List<T>.getCyclic(index: Int): T = this[index umod this.size]
 fun <T> List<T>.getCyclicOrNull(index: Int): T? = if (this.isEmpty()) null else this.getOrNull(index umod this.size)
 fun <T> Array<T>.getCyclic(index: Int) = this[index umod this.size]

diff --git a/kds/src/commonTest/kotlin/korlibs/datastructure/extensions/MapWhileTest.kt b/kds/src/commonTest/kotlin/korlibs/datastructure/extensions/MapWhileTest.kt
@@ -1,10 +1,6 @@
 package korlibs.datastructure.extensions
 
-import korlibs.datastructure.mapWhile
-import korlibs.datastructure.mapWhileArray
-import korlibs.datastructure.mapWhileDouble
-import korlibs.datastructure.mapWhileFloat
-import korlibs.datastructure.mapWhileInt
+import korlibs.datastructure.*
 import kotlin.test.Test
 import kotlin.test.assertEquals
 
@@ -23,4 +19,14 @@ class MapWhileTest {
         val iterator = listOf(1, 2, 3).iterator()
         assertEquals(listOf(1, 2, 3), mapWhile({ iterator.hasNext() }) { iterator.next()})
     }
+
+    @Test
+    fun testNotNull() {
+        assertEquals(listOf(0, 1, 2), mapWhileNotNull { if (it >= 3) null else it })
+    }
+
+    @Test
+    fun testNotCheck() {
+        assertEquals(listOf(0, 1000, 2000), mapWhileCheck({ it < 3000 }) { it * 1000 })
+    }
 }