/core/src/main/scala/scalaz/Memo.scala

 1package scalaz
 2
 3import reflect.ClassTag
 4
 5/** A function memoization strategy.  See companion for various
 6  * instances employing various strategies.
 7  */
 8sealed abstract class Memo[@specialized(Int) K, @specialized(Int, Long, Double) V] {
 9  def apply(z: K => V): K => V
10}
11
12
13sealed abstract class MemoInstances {
14}
15
16/** @define immuMapNote As this memo uses a single var, it's
17  * thread-safe. */
18object Memo extends MemoInstances {
19  def memo[@specialized(Int) K, @specialized(Int, Long, Double) V](f: (K => V) => K => V): Memo[K, V] = new Memo[K, V] {
20    def apply(z: K => V) = f(z)
21  }
22
23  def nilMemo[@specialized(Int) K, @specialized(Int, Long, Double) V]: Memo[K, V] = memo[K, V](z => z)
24
25  private class ArrayMemo[V >: Null : ClassTag](n: Int) extends Memo[Int, V] {
26    override def apply(f: (Int) => V) = {
27      val a = Need(new Array[V](n))
28      k => if (k < 0 || k >= n) f(k) else {
29        val t = a.value(k)
30        if (t == null) {
31          val v = f(k)
32          a.value(k) = v
33          v
34        } else t
35      }
36    }
37  }
38
39  private class DoubleArrayMemo(n: Int, sentinel: Double) extends Memo[Int, Double] {
40    override def apply(f: (Int) => Double) = {
41      val a = Need {
42        if (sentinel == 0d) {
43          new Array[Double](n)
44        } else {
45          Array.fill(n)(sentinel)
46        }
47      }
48      k => if (k < 0 || k >= n) f(k) else {
49        val t = a.value(k)
50        if (t == sentinel || (sentinel.isNaN && t.isNaN)) {
51          val v = f(k)
52          a.value(k) = v
53          v
54        } else t
55      }
56    }
57  }
58
59  /** Cache results in an `n`-long array. */
60  def arrayMemo[V >: Null : ClassTag](n: Int): Memo[Int, V] = new ArrayMemo(n)
61
62  /** As with `arrayMemo`, but memoizing double results !=
63    * `sentinel`.
64    */
65  def doubleArrayMemo(n: Int, sentinel: Double = 0d): Memo[Int, Double] = new DoubleArrayMemo(n, sentinel)
66
67  import scala.collection.mutable
68
69  private def mutableMapMemo[K, V](a: mutable.Map[K, V]): Memo[K, V] = memo[K, V](f => k => a.getOrElseUpdate(k, f(k)))
70
71  /** Cache results in a [[scala.collection.mutable.HashMap]].
72    * Nonsensical if `K` lacks a meaningful `hashCode` and
73    * `java.lang.Object.equals`.
74    */
75  def mutableHashMapMemo[K, V]: Memo[K, V] = mutableMapMemo(new mutable.HashMap[K, V])
76
77  import scala.collection.immutable
78
79  def immutableMapMemo[K, V](m: immutable.Map[K, V]): Memo[K, V] = {
80    var a = m
81
82    memo[K, V](f =>
83      k => a.getOrElse(k, {
84        val v = f(k)
85        a = a updated(k, v)
86        v
87      }))
88  }
89
90  /** Cache results in a hash map.  Nonsensical unless `K` has
91    * a meaningful `hashCode` and `java.lang.Object.equals`.
92    * $immuMapNote
93    */
94  def immutableHashMapMemo[K, V]: Memo[K, V] = immutableMapMemo(immutable.HashMap.empty[K, V])
95
96  /** Cache results in a tree map. $immuMapNote */
97  def immutableTreeMapMemo[K: scala.Ordering, V]: Memo[K, V] = immutableMapMemo(new immutable.TreeMap[K, V])
98}