Pimping scalaz Memo

Question

I like the new scalaz Memo functionality but find it lacks 2 things: 1) it hides the underlying Map, which I need access to--at least a List of all the values, and 2) I want a version that's implemented using a val scala.collection.concurrent.TrieMap, which I read somewhere is preferable than a var Map.

I'm not yet an implicit wizard. Is there a way to pimp this Memo class to add versions that support this capability or am I going to have to cut/paste into a distinct, new class?

Answer 1

This can be accomplished with the built-in Memo.memo function. Memo.memo creates a Memo instance from a function F => K => V. This also lets you get access to the underlying trie easily. For example:

scala> def trieMemo[A, B](trie: collection.concurrent.TrieMap[A, B]) = 
         Memo.memo[A, B](f => k => trie.getOrElseUpdate(k, f(k)))
trieMemo: [A, B](trie: scala.collection.concurrent.TrieMap[A,B])scalaz.Memo[A,B]

scala> val trie = collection.concurrent.TrieMap[Int, Int]()
trie: scala.collection.concurrent.TrieMap[Int,Int] = TrieMap()

scala> val f = trieMemo(trie)(n => n * n)
f: Int => Int = <function1>

scala> f(5)
res0: Int = 25

scala> f(10)
res1: Int = 100

scala> trie
res2: scala.collection.concurrent.TrieMap[Int,Int] = TrieMap(5 -> 25, 10 -> 100)