module Graphics.SvgTree.Memo
  ( memo
  ) where

import           Data.IORef             (IORef, atomicModifyIORef, newIORef)
import           Data.Map               (Map)
import qualified Data.Map               as Map
import           Data.Maybe             (catMaybes, listToMaybe)
import           Data.Typeable          (TypeRep, Typeable, cast, typeOf)
import           Graphics.SvgTree.Types (Tree)
import           System.IO.Unsafe       (unsafePerformIO)

{-# NOINLINE intCache #-}
intCache :: IORef (Map Int Tree)
intCache :: IORef (Map Int Tree)
intCache = IO (IORef (Map Int Tree)) -> IORef (Map Int Tree)
forall a. IO a -> a
unsafePerformIO (Map Int Tree -> IO (IORef (Map Int Tree))
forall a. a -> IO (IORef a)
newIORef Map Int Tree
forall k a. Map k a
Map.empty)

{-# NOINLINE doubleCache #-}
doubleCache :: IORef (Map Double Tree)
doubleCache :: IORef (Map Double Tree)
doubleCache = IO (IORef (Map Double Tree)) -> IORef (Map Double Tree)
forall a. IO a -> a
unsafePerformIO (Map Double Tree -> IO (IORef (Map Double Tree))
forall a. a -> IO (IORef a)
newIORef Map Double Tree
forall k a. Map k a
Map.empty)

{-# NOINLINE anyCache #-}
anyCache :: IORef (Map (TypeRep,String) Tree)
anyCache :: IORef (Map (TypeRep, String) Tree)
anyCache = IO (IORef (Map (TypeRep, String) Tree))
-> IORef (Map (TypeRep, String) Tree)
forall a. IO a -> a
unsafePerformIO (Map (TypeRep, String) Tree
-> IO (IORef (Map (TypeRep, String) Tree))
forall a. a -> IO (IORef a)
newIORef Map (TypeRep, String) Tree
forall k a. Map k a
Map.empty)

memo :: (Typeable a, Show a) => (a -> Tree) -> (a -> Tree)
memo :: (a -> Tree) -> a -> Tree
memo a -> Tree
fn =
  case [a -> Tree] -> Maybe (a -> Tree)
forall a. [a] -> Maybe a
listToMaybe ([Maybe (a -> Tree)] -> [a -> Tree]
forall a. [Maybe a] -> [a]
catMaybes [Maybe (a -> Tree)]
caches) of
    Just a -> Tree
ret -> a -> Tree
ret
    Maybe (a -> Tree)
Nothing  -> (a -> Tree) -> a -> Tree
forall a. (Typeable a, Show a) => (a -> Tree) -> a -> Tree
memoAny a -> Tree
fn
  where
    caches :: [Maybe (a -> Tree)]
caches = [IORef (Map Int Tree) -> Maybe (a -> Tree)
forall a b.
(Ord a, Typeable a, Typeable b) =>
IORef (Map a Tree) -> Maybe b
try IORef (Map Int Tree)
intCache, IORef (Map Double Tree) -> Maybe (a -> Tree)
forall a b.
(Ord a, Typeable a, Typeable b) =>
IORef (Map a Tree) -> Maybe b
try IORef (Map Double Tree)
doubleCache]
    try :: IORef (Map a Tree) -> Maybe b
try IORef (Map a Tree)
cache = (a -> Tree) -> Maybe b
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast ((a -> Tree) -> Maybe b)
-> ((a -> Tree) -> a -> Tree) -> (a -> Tree) -> Maybe b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IORef (Map a Tree) -> (a -> Tree) -> a -> Tree
forall a. Ord a => IORef (Map a Tree) -> (a -> Tree) -> a -> Tree
memoUsing IORef (Map a Tree)
cache ((a -> Tree) -> Maybe b) -> Maybe (a -> Tree) -> Maybe b
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (a -> Tree) -> Maybe (a -> Tree)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast a -> Tree
fn

memoUsing :: Ord a => IORef (Map a Tree) -> (a -> Tree) -> (a -> Tree)
memoUsing :: IORef (Map a Tree) -> (a -> Tree) -> a -> Tree
memoUsing IORef (Map a Tree)
cache a -> Tree
fn a
a = IO Tree -> Tree
forall a. IO a -> a
unsafePerformIO (IO Tree -> Tree) -> IO Tree -> Tree
forall a b. (a -> b) -> a -> b
$
  IORef (Map a Tree) -> (Map a Tree -> (Map a Tree, Tree)) -> IO Tree
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (Map a Tree)
cache ((Map a Tree -> (Map a Tree, Tree)) -> IO Tree)
-> (Map a Tree -> (Map a Tree, Tree)) -> IO Tree
forall a b. (a -> b) -> a -> b
$ \Map a Tree
m ->
    let newVal :: Tree
newVal = a -> Tree
fn a
a
        notFound :: (Map a Tree, Tree)
notFound =
          (a -> Tree -> Map a Tree -> Map a Tree
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
a Tree
newVal Map a Tree
m, Tree
newVal) in
    case a -> Map a Tree -> Maybe Tree
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
a Map a Tree
m of
      Maybe Tree
Nothing -> (Map a Tree, Tree)
notFound
      Just Tree
t  -> (Map a Tree
m, Tree
t)

memoAny :: (Typeable a, Show a) => (a -> Tree) -> (a -> Tree)
memoAny :: (a -> Tree) -> a -> Tree
memoAny a -> Tree
fn a
a = IO Tree -> Tree
forall a. IO a -> a
unsafePerformIO (IO Tree -> Tree) -> IO Tree -> Tree
forall a b. (a -> b) -> a -> b
$
  IORef (Map (TypeRep, String) Tree)
-> (Map (TypeRep, String) Tree
    -> (Map (TypeRep, String) Tree, Tree))
-> IO Tree
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (Map (TypeRep, String) Tree)
anyCache ((Map (TypeRep, String) Tree -> (Map (TypeRep, String) Tree, Tree))
 -> IO Tree)
-> (Map (TypeRep, String) Tree
    -> (Map (TypeRep, String) Tree, Tree))
-> IO Tree
forall a b. (a -> b) -> a -> b
$ \Map (TypeRep, String) Tree
m ->
    let newVal :: Tree
newVal = a -> Tree
fn a
a
        notFound :: (Map (TypeRep, String) Tree, Tree)
notFound =
          ((TypeRep, String)
-> Tree -> Map (TypeRep, String) Tree -> Map (TypeRep, String) Tree
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf a
a, a -> String
forall a. Show a => a -> String
show a
a) Tree
newVal Map (TypeRep, String) Tree
m, Tree
newVal) in
    case (TypeRep, String) -> Map (TypeRep, String) Tree -> Maybe Tree
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf a
a, a -> String
forall a. Show a => a -> String
show a
a) Map (TypeRep, String) Tree
m of
      Maybe Tree
Nothing -> (Map (TypeRep, String) Tree, Tree)
notFound
      Just Tree
t  -> (Map (TypeRep, String) Tree
m, Tree
t)

-- {-# INLINE memo #-}
-- memo :: (a -> b) -> (a -> b)
-- memo fn = unsafePerformIO $ do
--   ref <- newIORef Map.empty
--   return $ \a -> unsafePerformIO $ do
--     stableA <- makeStableName a
--     let key = hashStableName stableA
--     atomicModifyIORef ref $ \m ->
--       case Map.lookup key m of
--         -- Just (s,b) | s == stableA ->
--         --   (m, b)
--         _Nothing -> let !b = fn a in
--           (Map.insert key (stableA, b) m, b)
-- memo fn = unsafePerformIO $ do
--   ht <- HT.new :: IO (HT.BasicHashTable (StableName Any) Any)
--   return $ \a -> unsafePerformIO $ do
--     stableA <- makeStableName $ unsafeCoerce a
--     mbB <- HT.lookup ht stableA
--     case mbB of
--       Just b -> return (unsafeCoerce b)
--       Nothing -> do
--         let !b = fn a
--         HT.insert ht stableA (unsafeCoerce b)
--         return b