MemoizedDeepRecursiveFunction (#3091)

Co-authored-by: serras <[email protected]>
Co-authored-by: Simon Vergauwen <[email protected]>

Author: serras

arrow-libs/core/arrow-core/api/arrow-core.api

@@ -776,6 +776,10 @@ public final class arrow/core/Memoization {
 	public static final fun memoize (Lkotlin/jvm/functions/Function5;)Lkotlin/jvm/functions/Function5;
 }
 
+public final class arrow/core/MemoizedDeepRecursiveFunctionKt {
+	public static final fun MemoizedDeepRecursiveFunction (Lkotlin/jvm/functions/Function3;)Lkotlin/DeepRecursiveFunction;
+}
+
 public abstract interface class arrow/core/NonEmptyCollection : java/util/Collection, kotlin/jvm/internal/markers/KMappedMarker {
 	public abstract fun distinct ()Larrow/core/NonEmptyList;
 	public abstract fun distinctBy (Lkotlin/jvm/functions/Function1;)Larrow/core/NonEmptyList;

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

@@ -0,0 +1,38 @@
+package arrow.core
+
+import arrow.atomic.Atomic
+import arrow.atomic.loop
+
+/**
+ * Defines a recursive **pure** function that:
+ * - keeps its stack on the heap, which allows very deep recursive computations that do not use the actual call stack;
+ * - memoizes every call, which means that the function is execute only once per argument.
+ *
+ * [MemoizedDeepRecursiveFunction] takes one parameter of type [T] and returns a result of type [R].
+ * The [block] of code defines the body of a recursive function. In this block
+ * [callRecursive][DeepRecursiveScope.callRecursive] function can be used to make a recursive call
+ * to the declared function.
+ */
+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 (val v = cache.get()[x]) {
+      null -> {
+        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
+            }
+          }
+        }
+      }
+      else -> v
+    }
+  }
+}

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

@@ -215,6 +215,37 @@ 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
+  }
+
+  "Recursive memoization, run twice should be memoized" {
+    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 result1 = memoizedDeepRecursiveFibonacci(5)
+    val result2 = memoizedDeepRecursiveFibonacci(5)
+    result1 shouldBe result2
+    runs shouldBe 1
+  }
 })
 
 private fun consecSumResult(n: Int): Int = (n * (n + 1)) / 2