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 = unsafePerformIO (newIORef Map.empty)

{-# NOINLINE doubleCache #-}
doubleCache :: IORef (Map Double Tree)
doubleCache = unsafePerformIO (newIORef Map.empty)

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

memo :: (Typeable a, Show a) => (a -> Tree) -> (a -> Tree)
memo fn =
  case listToMaybe (catMaybes caches) of
    Just ret -> ret
    Nothing  -> memoAny fn
  where
    caches = [try intCache, try doubleCache]
    try cache = cast . memoUsing cache =<< cast fn

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

memoAny :: (Typeable a, Show a) => (a -> Tree) -> (a -> Tree)
memoAny fn a = unsafePerformIO $
  atomicModifyIORef anyCache $ \m ->
    let newVal = fn a
        notFound =
          (Map.insert (typeOf a, show a) newVal m, newVal) in
    case Map.lookup (typeOf a, show a) m of
      Nothing -> notFound
      Just t  -> (m, 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