{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}

module Control.Memo.Fix where

import Control.Monad.State
#if MIN_VERSION_mtl(2,3,0)
import Data.Function (fix)
#endif
import Data.Functor (($>))
import qualified Data.IntMap.Strict as IM
import qualified Data.Map.Strict as M
import qualified Data.Vector as V

{- |
>>> :set -XLambdaCase
>>> :{
let memoFib = \fib -> \case
      0 -> 0
      1 -> 1
      i -> fib (i - 1) + fib (i - 2)
 in memoFix 100 memoFib 50
:}
12586269025
-}
memoFix ::
  -- | memo size
  Int ->
  ((Int -> a) -> Int -> a) ->
  Int ->
  a
memoFix :: forall a. Int -> ((Int -> a) -> Int -> a) -> Int -> a
memoFix Int
n (Int -> a) -> Int -> a
f = ((Int -> a) -> Int -> a) -> Int -> a
forall a. (a -> a) -> a
fix (((Int -> a) -> Int -> a) -> Int -> a)
-> ((Int -> a) -> Int -> a) -> Int -> a
forall a b. (a -> b) -> a -> b
$ \Int -> a
memo -> (Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
n ((Int -> a) -> Int -> a
f Int -> a
memo) V.!)

{- |
>>> :set -XLambdaCase
>>> :{
let memoFibM :: (Monad m) => (Int -> m Integer) -> Int -> m Integer
    memoFibM = \fib -> \case
      0 -> pure 0
      1 -> pure 1
      i -> (+) <$> fib (i - 1) <*> fib (i - 2)
 in memoFixMap memoFibM 50
:}
12586269025
-}
memoFixMap ::
  (Ord k) =>
  ((k -> State (M.Map k a) a) -> k -> State (M.Map k a) a) ->
  k ->
  a
memoFixMap :: forall k a.
Ord k =>
((k -> State (Map k a) a) -> k -> State (Map k a) a) -> k -> a
memoFixMap (k -> State (Map k a) a) -> k -> State (Map k a) a
f k
k = (State (Map k a) a -> Map k a -> a)
-> Map k a -> State (Map k a) a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip State (Map k a) a -> Map k a -> a
forall s a. State s a -> s -> a
evalState Map k a
forall k a. Map k a
M.empty (State (Map k a) a -> a) -> State (Map k a) a -> a
forall a b. (a -> b) -> a -> b
$ do
  (((k -> State (Map k a) a) -> k -> State (Map k a) a)
 -> k -> State (Map k a) a)
-> k
-> ((k -> State (Map k a) a) -> k -> State (Map k a) a)
-> State (Map k a) a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((k -> State (Map k a) a) -> k -> State (Map k a) a)
-> k -> State (Map k a) a
forall a. (a -> a) -> a
fix k
k (((k -> State (Map k a) a) -> k -> State (Map k a) a)
 -> State (Map k a) a)
-> ((k -> State (Map k a) a) -> k -> State (Map k a) a)
-> State (Map k a) a
forall a b. (a -> b) -> a -> b
$ \k -> State (Map k a) a
memo k
x -> do
    (Map k a -> Maybe a) -> StateT (Map k a) Identity (Maybe a)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (k -> Map k a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
x) StateT (Map k a) Identity (Maybe a)
-> (Maybe a -> State (Map k a) a) -> State (Map k a) a
forall a b.
StateT (Map k a) Identity a
-> (a -> StateT (Map k a) Identity b)
-> StateT (Map k a) Identity b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Just a
fx -> a -> State (Map k a) a
forall a. a -> StateT (Map k a) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
fx
      Maybe a
Nothing ->
        (k -> State (Map k a) a) -> k -> State (Map k a) a
f k -> State (Map k a) a
memo k
x State (Map k a) a -> (a -> State (Map k a) a) -> State (Map k a) a
forall a b.
StateT (Map k a) Identity a
-> (a -> StateT (Map k a) Identity b)
-> StateT (Map k a) Identity b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
fx ->
          (Map k a -> Map k a) -> StateT (Map k a) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' (k -> a -> Map k a -> Map k a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
x a
fx) StateT (Map k a) Identity () -> a -> State (Map k a) a
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> a
fx

{- |
>>> :set -XLambdaCase
>>> :{
let memoFibM :: (Monad m) => (Int -> m Integer) -> Int -> m Integer
    memoFibM = \fib -> \case
      0 -> pure 0
      1 -> pure 1
      i -> (+) <$> fib (i - 1) <*> fib (i - 2)
 in memoFixIntMap memoFibM 50
:}
12586269025
-}
memoFixIntMap ::
  ((Int -> State (IM.IntMap a) a) -> Int -> State (IM.IntMap a) a) ->
  Int ->
  a
memoFixIntMap :: forall a.
((Int -> State (IntMap a) a) -> Int -> State (IntMap a) a)
-> Int -> a
memoFixIntMap (Int -> State (IntMap a) a) -> Int -> State (IntMap a) a
f Int
n = (State (IntMap a) a -> IntMap a -> a)
-> IntMap a -> State (IntMap a) a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip State (IntMap a) a -> IntMap a -> a
forall s a. State s a -> s -> a
evalState IntMap a
forall a. IntMap a
IM.empty (State (IntMap a) a -> a) -> State (IntMap a) a -> a
forall a b. (a -> b) -> a -> b
$ do
  (((Int -> State (IntMap a) a) -> Int -> State (IntMap a) a)
 -> Int -> State (IntMap a) a)
-> Int
-> ((Int -> State (IntMap a) a) -> Int -> State (IntMap a) a)
-> State (IntMap a) a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> State (IntMap a) a) -> Int -> State (IntMap a) a)
-> Int -> State (IntMap a) a
forall a. (a -> a) -> a
fix Int
n (((Int -> State (IntMap a) a) -> Int -> State (IntMap a) a)
 -> State (IntMap a) a)
-> ((Int -> State (IntMap a) a) -> Int -> State (IntMap a) a)
-> State (IntMap a) a
forall a b. (a -> b) -> a -> b
$ \Int -> State (IntMap a) a
memo Int
x -> do
    (IntMap a -> Maybe a) -> StateT (IntMap a) Identity (Maybe a)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
x) StateT (IntMap a) Identity (Maybe a)
-> (Maybe a -> State (IntMap a) a) -> State (IntMap a) a
forall a b.
StateT (IntMap a) Identity a
-> (a -> StateT (IntMap a) Identity b)
-> StateT (IntMap a) Identity b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Just a
fx -> a -> State (IntMap a) a
forall a. a -> StateT (IntMap a) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
fx
      Maybe a
Nothing ->
        (Int -> State (IntMap a) a) -> Int -> State (IntMap a) a
f Int -> State (IntMap a) a
memo Int
x State (IntMap a) a
-> (a -> State (IntMap a) a) -> State (IntMap a) a
forall a b.
StateT (IntMap a) Identity a
-> (a -> StateT (IntMap a) Identity b)
-> StateT (IntMap a) Identity b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
fx ->
          (IntMap a -> IntMap a) -> StateT (IntMap a) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' (Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
x a
fx) StateT (IntMap a) Identity () -> a -> State (IntMap a) a
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> a
fx