module Test.Speculate.Utils.Memoize
( memory, memory2
, memoryFor, memory2For
, withMemory, withMemory2
)
where
import qualified Data.Map as M
import Data.Map (Map)
import Test.LeanCheck (Listable(..))
import Data.Maybe (fromMaybe)
defaultMemory :: Int
defaultMemory = 2520
withMemory :: Ord a => (a -> b) -> Map a b -> a -> b
withMemory f m x = fromMaybe (f x) (M.lookup x m)
withMemory2 :: (Ord a, Ord b) => (a -> b -> c) -> Map (a,b) c -> a -> b -> c
withMemory2 f m = curry (uncurry f `withMemory` m)
memory :: (Listable a, Ord a) => (a -> b) -> Map a b
memory = memoryFor defaultMemory
memory2 :: (Listable a, Listable b, Ord a, Ord b) => (a -> b -> c) -> Map (a,b) c
memory2 = memory2For defaultMemory
memoryFor :: (Listable a, Ord a) => Int -> (a -> b) -> Map a b
memoryFor n f = foldr (uncurry M.insert) M.empty . take n $ map (\x -> (x, f x)) list
memory2For :: (Listable a, Listable b, Ord a, Ord b)
=> Int -> (a -> b -> c) -> Map (a,b) c
memory2For n f = memoryFor n (uncurry f)