MemoizedDeepRecursiveFunction

arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/MemoizedDeepRecursiveFunction.kt

@@ -0,0 +1,28 @@
+package arrow.core
+
+import arrow.atomic.Atomic
+import arrow.atomic.loop
+
+public fun <T, R> MemoizedDeepRecursiveFunction(
+  block: suspend DeepRecursiveScope<T, R>.(T) -> R
+): DeepRecursiveFunction<T, R> {
+  val cache = Atomic(emptyMap<T, R>())
+  return DeepRecursiveFunction { x ->
+    when (x) {
+      in cache.get() -> cache.get().getValue(x)
+      else -> {
+        val result = block(x)
+        cache.loop { old ->
+          when (x) {
+            in old ->
+              return@DeepRecursiveFunction old.getValue(x)
+            else -> {
+              if (cache.compareAndSet(old, old + Pair(x, result)))
+                return@DeepRecursiveFunction result
+            }
+          }
+        }
+      }
+    }
+  }
+}

arrow-libs/core/arrow-core/src/commonTest/kotlin/arrow/core/MemoizationTest.kt

@@ -215,6 +215,21 @@ class MemoizationTest : StringSpec({
     memoized(1, 2, 3, 4, 5) shouldBe null
     runs shouldBe 1
   }
+
+  "Recursive memoization" {
+    var runs = 0
+    val memoizedDeepRecursiveFibonacci: DeepRecursiveFunction<Int, Int> =
+      MemoizedDeepRecursiveFunction { n ->
+        when (n) {
+          0 -> 0.also { runs++ }
+          1 -> 1
+          else -> callRecursive(n - 1) + callRecursive(n - 2)
+        }
+      }
+    val result = memoizedDeepRecursiveFibonacci(5)
+    result shouldBe 5
+    runs shouldBe 1
+  }
 })
 
 private fun consecSumResult(n: Int): Int = (n * (n + 1)) / 2