{-# LANGUAGE BangPatterns              #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE RankNTypes                #-}
{-|
Copyright   : Written by David Himmelstrup
License     : Unlicense
Maintainer  : lemmih@gmail.com
Stability   : experimental
Portability : POSIX

With animations defined as SVG images over time, it is unfortunately
quite easy to write inefficient code where expensive expressions are
re-evaluated for every frame even if nothing has changed. This get
around this issue, this module defines a global key-value table that
can store expensive expressions such that they are evaluated only once.

There is currently no way to clear values from the table and it is your
own responsibility to not consume all available memory.

-}
module Reanimate.Memo
  ( Key(..)
  , memo
  ) where

import           Data.Dynamic          (Dynamic, fromDynamic, toDyn)
import           Data.IORef            (IORef, atomicModifyIORef', newIORef)
import qualified Data.Map              as Map
import           Data.Maybe            (fromMaybe)
import           Data.Typeable         (Typeable, cast, typeOf)
import           System.IO.Unsafe      (unsafePerformIO)
import           System.Mem.StableName (StableName, eqStableName, hashStableName, makeStableName)

data DynamicName = forall a. DynamicName !(StableName a) | forall a. (Eq a, Ord a, Typeable a) => DynamicKey a
instance Eq DynamicName where
  DynamicName StableName a
a == :: DynamicName -> DynamicName -> Bool
== DynamicName StableName a
b = StableName a -> StableName a -> Bool
forall a b. StableName a -> StableName b -> Bool
eqStableName StableName a
a StableName a
b
  DynamicKey a
a == DynamicKey a
b =
    case a -> Maybe a
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast a
a of
      Maybe a
Nothing -> Bool
False
      Just a
a' -> a
a'a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==a
b
  DynamicName
_ == DynamicName
_ = Bool
False

instance Ord DynamicName where
  DynamicName StableName a
a compare :: DynamicName -> DynamicName -> Ordering
`compare` DynamicName StableName a
b =
    StableName a -> Int
forall a. StableName a -> Int
hashStableName StableName a
a Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` StableName a -> Int
forall a. StableName a -> Int
hashStableName StableName a
b
  DynamicName{} `compare` DynamicName
_ = Ordering
LT
  DynamicKey a
a `compare` DynamicKey a
b =
    case a -> Maybe a
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast a
a of
      Maybe a
Nothing -> a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf a
a TypeRep -> TypeRep -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf a
b
      Just a
a' -> a
a' a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` a
b
  DynamicName
_ `compare` DynamicName
_ = Ordering
GT

data CacheMap = CacheMap !(Map.Map DynamicName CacheMap) !(Map.Map DynamicName Dynamic)

emptyCacheMap :: CacheMap
emptyCacheMap :: CacheMap
emptyCacheMap = Map DynamicName CacheMap -> Map DynamicName Dynamic -> CacheMap
CacheMap Map DynamicName CacheMap
forall k a. Map k a
Map.empty Map DynamicName Dynamic
forall k a. Map k a
Map.empty

-- sizeCacheMap :: CacheMap -> Int
-- sizeCacheMap (CacheMap sub vals) =
--   sum (map sizeCacheMap (Map.elems sub)) + Map.size vals

cacheMapLookup :: [DynamicName] -> CacheMap -> Maybe Dynamic
cacheMapLookup :: [DynamicName] -> CacheMap -> Maybe Dynamic
cacheMapLookup [] CacheMap
_ = Maybe Dynamic
forall a. Maybe a
Nothing
cacheMapLookup [DynamicName
k] (CacheMap Map DynamicName CacheMap
_ Map DynamicName Dynamic
vals) = DynamicName -> Map DynamicName Dynamic -> Maybe Dynamic
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup DynamicName
k Map DynamicName Dynamic
vals
cacheMapLookup (DynamicName
k:[DynamicName]
ks) (CacheMap Map DynamicName CacheMap
sub Map DynamicName Dynamic
_) =
  [DynamicName] -> CacheMap -> Maybe Dynamic
cacheMapLookup [DynamicName]
ks (CacheMap -> Maybe Dynamic) -> Maybe CacheMap -> Maybe Dynamic
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< DynamicName -> Map DynamicName CacheMap -> Maybe CacheMap
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup DynamicName
k Map DynamicName CacheMap
sub

cacheMapInsert :: [DynamicName] -> Dynamic -> CacheMap -> CacheMap
cacheMapInsert :: [DynamicName] -> Dynamic -> CacheMap -> CacheMap
cacheMapInsert [] Dynamic
_ CacheMap
m = CacheMap
m
cacheMapInsert [DynamicName
k] Dynamic
v (CacheMap Map DynamicName CacheMap
sub Map DynamicName Dynamic
vals) = Map DynamicName CacheMap -> Map DynamicName Dynamic -> CacheMap
CacheMap Map DynamicName CacheMap
sub (DynamicName
-> Dynamic -> Map DynamicName Dynamic -> Map DynamicName Dynamic
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert DynamicName
k Dynamic
v Map DynamicName Dynamic
vals)
cacheMapInsert (DynamicName
k:[DynamicName]
ks) Dynamic
v (CacheMap Map DynamicName CacheMap
sub Map DynamicName Dynamic
vals) =
  Map DynamicName CacheMap -> Map DynamicName Dynamic -> CacheMap
CacheMap ((Maybe CacheMap -> Maybe CacheMap)
-> DynamicName
-> Map DynamicName CacheMap
-> Map DynamicName CacheMap
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter Maybe CacheMap -> Maybe CacheMap
fn DynamicName
k Map DynamicName CacheMap
sub) Map DynamicName Dynamic
vals
  where
    fn :: Maybe CacheMap -> Maybe CacheMap
fn = CacheMap -> Maybe CacheMap
forall a. a -> Maybe a
Just (CacheMap -> Maybe CacheMap)
-> (Maybe CacheMap -> CacheMap) -> Maybe CacheMap -> Maybe CacheMap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [DynamicName] -> Dynamic -> CacheMap -> CacheMap
cacheMapInsert [DynamicName]
ks Dynamic
v (CacheMap -> CacheMap)
-> (Maybe CacheMap -> CacheMap) -> Maybe CacheMap -> CacheMap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CacheMap -> Maybe CacheMap -> CacheMap
forall a. a -> Maybe a -> a
fromMaybe CacheMap
emptyCacheMap

{-# NOINLINE cacheMap #-}
-- FIXME: There should be a way to clear the cache.
cacheMap :: IORef CacheMap
cacheMap :: IORef CacheMap
cacheMap = IO (IORef CacheMap) -> IORef CacheMap
forall a. IO a -> a
unsafePerformIO (CacheMap -> IO (IORef CacheMap)
forall a. a -> IO (IORef a)
newIORef CacheMap
emptyCacheMap)

-- | Keys can either by any boxed object with identity (to be compared with
--   StableNames) or a primitive type with an Eq instance.
data Key = forall a. Key !a | forall a. (Typeable a, Eq a, Ord a) => KeyPrim !a

fromKey :: Key -> IO DynamicName
fromKey :: Key -> IO DynamicName
fromKey (Key a
val)     = StableName a -> DynamicName
forall a. StableName a -> DynamicName
DynamicName (StableName a -> DynamicName)
-> IO (StableName a) -> IO DynamicName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> IO (StableName a)
forall a. a -> IO (StableName a)
makeStableName a
val
fromKey (KeyPrim a
val) = DynamicName -> IO DynamicName
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> DynamicName
forall a. (Eq a, Ord a, Typeable a) => a -> DynamicName
DynamicKey a
val)

-- | Cache expensive value in global store. You must guarantee that the
--   key uniquely determines the value.
memo :: Typeable a => [Key] -> a -> a
memo :: [Key] -> a -> a
memo ![Key]
k a
v = IO a -> a
forall a. IO a -> a
unsafePerformIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ do
  [DynamicName]
keys <- (Key -> IO DynamicName) -> [Key] -> IO [DynamicName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Key -> IO DynamicName
fromKey [Key]
k
  IORef CacheMap -> (CacheMap -> (CacheMap, a)) -> IO a
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef CacheMap
cacheMap ((CacheMap -> (CacheMap, a)) -> IO a)
-> (CacheMap -> (CacheMap, a)) -> IO a
forall a b. (a -> b) -> a -> b
$ \CacheMap
m ->
    case Dynamic -> Maybe a
forall a. Typeable a => Dynamic -> Maybe a
fromDynamic (Dynamic -> Maybe a) -> Maybe Dynamic -> Maybe a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [DynamicName] -> CacheMap -> Maybe Dynamic
cacheMapLookup [DynamicName]
keys CacheMap
m of
      Just a
v' -> (CacheMap
m, a
v')
      Maybe a
Nothing -> ([DynamicName] -> Dynamic -> CacheMap -> CacheMap
cacheMapInsert [DynamicName]
keys (a -> Dynamic
forall a. Typeable a => a -> Dynamic
toDyn a
v) CacheMap
m, a
v)