-- |
-- Module      : Test.Speculate.Utils.Memoize
-- Copyright   : (c) 2016-2024 Rudy Matela
-- License     : 3-Clause BSD  (see the file LICENSE)
-- Maintainer  : Rudy Matela <rudy@matela.com.br>
--
-- This module is part of Speculate.
--
-- Memoization module.
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 :: Int
defaultMemory = Int
2520 -- 2^3 * 3^2 * 5 * 7

{- those don't work, GHC wont cache them
memoize :: (Listable a, Ord a) => (a -> b) -> a -> b
memoize = memoizeFor defaultMemory

memoizeFor :: (Listable a, Ord a) => Int -> (a -> b) -> a -> b
memoizeFor n f x = fromMaybe (f x) (M.lookup x m)
  where
  m = memoryFor n f
-}

withMemory :: Ord a => (a -> b) -> Map a b -> a -> b
withMemory :: forall a b. Ord a => (a -> b) -> Map a b -> a -> b
withMemory a -> b
f Map a b
m a
x = b -> Maybe b -> b
forall a. a -> Maybe a -> a
fromMaybe (a -> b
f a
x) (a -> Map a b -> Maybe b
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
x Map a b
m)

withMemory2 :: (Ord a, Ord b) => (a -> b -> c) -> Map (a,b) c -> a -> b -> c
withMemory2 :: forall a b c.
(Ord a, Ord b) =>
(a -> b -> c) -> Map (a, b) c -> a -> b -> c
withMemory2 a -> b -> c
f Map (a, b) c
m = ((a, b) -> c) -> a -> b -> c
forall a b c. ((a, b) -> c) -> a -> b -> c
curry ((a -> b -> c) -> (a, b) -> c
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> b -> c
f ((a, b) -> c) -> Map (a, b) c -> (a, b) -> c
forall a b. Ord a => (a -> b) -> Map a b -> a -> b
`withMemory` Map (a, b) c
m)

memory :: (Listable a, Ord a) => (a -> b) -> Map a b
memory :: forall a b. (Listable a, Ord a) => (a -> b) -> Map a b
memory = Int -> (a -> b) -> Map a b
forall a b. (Listable a, Ord a) => Int -> (a -> b) -> Map a b
memoryFor Int
defaultMemory

memory2 :: (Listable a, Listable b, Ord a, Ord b) => (a -> b -> c) -> Map (a,b) c
memory2 :: forall a b c.
(Listable a, Listable b, Ord a, Ord b) =>
(a -> b -> c) -> Map (a, b) c
memory2 = Int -> (a -> b -> c) -> Map (a, b) c
forall a b c.
(Listable a, Listable b, Ord a, Ord b) =>
Int -> (a -> b -> c) -> Map (a, b) c
memory2For Int
defaultMemory

memoryFor :: (Listable a, Ord a) => Int -> (a -> b) -> Map a b
memoryFor :: forall a b. (Listable a, Ord a) => Int -> (a -> b) -> Map a b
memoryFor Int
n a -> b
f = ((a, b) -> Map a b -> Map a b) -> Map a b -> [(a, b)] -> Map a b
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((a -> b -> Map a b -> Map a b) -> (a, b) -> Map a b -> Map a b
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> b -> Map a b -> Map a b
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert) Map a b
forall k a. Map k a
M.empty ([(a, b)] -> Map a b)
-> ([(a, b)] -> [(a, b)]) -> [(a, b)] -> Map a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [(a, b)] -> [(a, b)]
forall a. Int -> [a] -> [a]
take Int
n ([(a, b)] -> Map a b) -> [(a, b)] -> Map a b
forall a b. (a -> b) -> a -> b
$ (a -> (a, b)) -> [a] -> [(a, b)]
forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> (a
x, a -> b
f a
x)) [a]
forall a. Listable a => [a]
list

memory2For :: (Listable a, Listable b, Ord a, Ord b)
           => Int -> (a -> b -> c) -> Map (a,b) c
memory2For :: forall a b c.
(Listable a, Listable b, Ord a, Ord b) =>
Int -> (a -> b -> c) -> Map (a, b) c
memory2For Int
n a -> b -> c
f = Int -> ((a, b) -> c) -> Map (a, b) c
forall a b. (Listable a, Ord a) => Int -> (a -> b) -> Map a b
memoryFor Int
n ((a -> b -> c) -> (a, b) -> c
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> b -> c
f)