commonMain.arrow.core.memoization.kt Maven / Gradle / Ivy

Go to download
@file:JvmName("Memoization")

package arrow.core

import arrow.atomic.Atomic
import arrow.atomic.loop
import arrow.atomic.value
import kotlin.jvm.JvmName

/**
 * Memoizes the given **pure** function so that invocations with the same arguments will only execute the function once.
 *
 * ```kotlin
 * import arrow.core.memoize
 * fun someWorkIntensiveFunction(someParam: Int): String = "$someParam"
 *
 * fun main() {
 *   //sampleStart
 *   val memoizedF = ::someWorkIntensiveFunction.memoize()
 *
 *   // The first invocation will store the argument and result in a cache inside the `memoizedF` reference.
 *   val value1 = memoizedF(42)
 *   // This second invocation won't really call the `someWorkIntensiveFunction` function
 *   //but retrieve the result from the previous invocation instead.
 *   val value2 = memoizedF(42)
 *
 *   //sampleEnd
 *   println("$value1 $value2")
 * }
 * ```
 * 
 *
 * Note that calling this function with the same parameters in parallel might cause the function to be executed twice.
 */
public fun  (() -> R).memoize(): () -> R {
  val m = MemoizedHandler<() -> R, MemoizeKey0, R>(this@memoize)
  return { m(MemoizeKey0(0)) }
}

/**
 * @see memoize
 */
public fun  ((P1) -> R).memoize(): (P1) -> R {
  val m = MemoizedHandler<((P1) -> R), MemoizeKey1, R>(this@memoize)
  return { p1 -> m(MemoizeKey1(p1)) }
}

/**
 * @see memoize
 */
public fun  ((P1, P2) -> R).memoize(): (P1, P2) -> R {
  val m = MemoizedHandler<((P1, P2) -> R), MemoizeKey2, R>(this@memoize)
  return { p1: P1, p2: P2 -> m(MemoizeKey2(p1, p2)) }
}

/**
 * @see memoize
 */
public fun  ((P1, P2, P3) -> R).memoize(): (P1, P2, P3) -> R {
  val m = MemoizedHandler<((P1, P2, P3) -> R), MemoizeKey3, R>(this@memoize)
  return { p1: P1, p2: P2, p3: P3 -> m(MemoizeKey3(p1, p2, p3)) }
}

/**
 * @see memoize
 */
public fun  ((P1, P2, P3, P4) -> R).memoize(): (P1, P2, P3, P4) -> R {
  val m = MemoizedHandler<((P1, P2, P3, P4) -> R), MemoizeKey4, R>(this@memoize)
  return { p1: P1, p2: P2, p3: P3, p4: P4 -> m(MemoizeKey4(p1, p2, p3, p4)) }
}

/**
 * @see memoize
 */
public fun  ((P1, P2, P3, P4, P5) -> R).memoize(): (P1, P2, P3, P4, P5) -> R {
  val m = MemoizedHandler<((P1, P2, P3, P4, P5) -> R), MemoizeKey5, R>(this@memoize)
  return { p1: P1, p2: P2, p3: P3, p4: P4, p5: P5 -> m(MemoizeKey5(p1, p2, p3, p4, p5)) }
}

private interface MemoizedCall {
  operator fun invoke(f: F): R
}

private data class MemoizeKey0(val p1: Byte) : MemoizedCall<() -> R, R> {
  override fun invoke(f: () -> R): R = f()
}

private data class MemoizeKey1(val p1: P1) : MemoizedCall<(P1) -> R, R> {
  override fun invoke(f: (P1) -> R) = f(p1)
}

private data class MemoizeKey2(val p1: P1, val p2: P2) : MemoizedCall<(P1, P2) -> R, R> {
  override fun invoke(f: (P1, P2) -> R) = f(p1, p2)
}

private data class MemoizeKey3(val p1: P1, val p2: P2, val p3: P3) :
  MemoizedCall<(P1, P2, P3) -> R, R> {
  override fun invoke(f: (P1, P2, P3) -> R) = f(p1, p2, p3)
}

private data class MemoizeKey4(val p1: P1, val p2: P2, val p3: P3, val p4: P4) :
  MemoizedCall<(P1, P2, P3, P4) -> R, R> {
  override fun invoke(f: (P1, P2, P3, P4) -> R) = f(p1, p2, p3, p4)
}

private data class MemoizeKey5(
  val p1: P1,
  val p2: P2,
  val p3: P3,
  val p4: P4,
  val p5: P5
) : MemoizedCall<(P1, P2, P3, P4, P5) -> R, R> {
  override fun invoke(f: (P1, P2, P3, P4, P5) -> R) = f(p1, p2, p3, p4, p5)
}

private class MemoizedHandler, out R>(val f: F) {
  private val cache = Atomic(emptyMap())

  operator fun invoke(k: K): R = when (k) {
    in cache.value -> cache.value.getValue(k)
    else -> {
      val b = k(f)
      cache.loop { old ->
        when (k) {
          in old ->
            return@invoke old.getValue(k)
          else -> {
            if (cache.compareAndSet(old, old + Pair(k, b)))
              return@invoke b
          }
        }
      }
    }
  }
}