{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
module Haxl.Core.Memo
(
cachedComputation
, preCacheComputation
, memo
, memoFingerprint
, MemoFingerprintKey(..)
, memoize, memoize1, memoize2
, memoUnique
, MemoVar
, newMemo
, newMemoWith
, prepareMemo
, runMemo
) where
import Control.Exception as Exception hiding (throw)
import Data.IORef
import qualified Data.HashMap.Strict as HashMap
import Data.Text (Text)
import Data.Typeable
import Data.Hashable
import Data.Int
import Data.Word
import GHC.Prim (Addr#)
import Haxl.Core.Exception
import Haxl.Core.DataCache as DataCache
import Haxl.Core.Flags
import Haxl.Core.Monad
import Haxl.Core.Stats
import Haxl.Core.Profile
import Haxl.Core.Util (trace_)
cachedComputation
:: forall req u w a.
( Eq (req a)
, Hashable (req a)
, Typeable (req a))
=> req a -> GenHaxl u w a -> GenHaxl u w a
cachedComputation req haxl = GenHaxl $ \env@Env{..} -> do
mbRes <- DataCache.lookup req memoCache
case mbRes of
Just (DataCacheItem ivar k) -> do
ifProfiling flags $ do
incrementMemoHitCounterFor env k True
unHaxl (getIVarWithWrites ivar) env
Nothing -> do
ivar <- newIVar
k <- nextCallId env
DataCache.insertNotShowable req (DataCacheItem ivar k) memoCache
execMemoNowProfiled env haxl ivar k
preCacheComputation
:: forall req u w a.
( Eq (req a)
, Hashable (req a)
, Typeable (req a))
=> req a -> GenHaxl u w a -> GenHaxl u w a
preCacheComputation req haxl = GenHaxl $ \env@Env{..} -> do
mbRes <- DataCache.lookup req memoCache
case mbRes of
Just _ -> return $ Throw $ toException $ InvalidParameter
"preCacheComputation: key is already cached"
Nothing -> do
ivar <- newIVar
k <- nextCallId env
DataCache.insertNotShowable req (DataCacheItem ivar k) memoCache
execMemoNowProfiled env haxl ivar k
newtype MemoVar u w a = MemoVar (IORef (MemoStatus u w a))
data MemoStatus u w a
= MemoEmpty
| MemoReady (GenHaxl u w a) CallId
| MemoRun {-# UNPACK #-} !(IVar u w a) {-# UNPACK #-} !CallId
newMemo :: GenHaxl u w (MemoVar u w a)
newMemo = unsafeLiftIO $ MemoVar <$> newIORef MemoEmpty
prepareMemo :: MemoVar u w a -> GenHaxl u w a -> GenHaxl u w ()
prepareMemo (MemoVar memoRef) memoCmp
= GenHaxl $ \env -> do
k <- nextCallId env
writeIORef memoRef (MemoReady memoCmp k)
return (Done ())
newMemoWith :: GenHaxl u w a -> GenHaxl u w (MemoVar u w a)
newMemoWith memoCmp = do
memoVar <- newMemo
prepareMemo memoVar memoCmp
return memoVar
runMemo :: MemoVar u w a -> GenHaxl u w a
runMemo (MemoVar memoRef) = GenHaxl $ \env -> do
stored <- readIORef memoRef
case stored of
MemoEmpty -> trace_ "MemoEmpty " $
raise $ CriticalError "Attempting to run empty memo."
MemoReady cont k -> trace_ "MemoReady" $ do
ivar <- newIVar
writeIORef memoRef (MemoRun ivar k)
execMemoNowProfiled env cont ivar k
MemoRun ivar k -> trace_ "MemoRun" $ do
ifProfiling (flags env) $ do
incrementMemoHitCounterFor env k True
unHaxl (getIVarWithWrites ivar) env
execMemoNowProfiled
:: Env u w
-> GenHaxl u w a
-> IVar u w a
-> CallId
-> IO (Result u w a)
execMemoNowProfiled envOuter cont ivar cid = if report (flags envOuter) < 4
then execMemoNow envOuter cont ivar
else do
incrementMemoHitCounterFor envOuter cid False
unHaxl
(collectMemoData 0 $ GenHaxl $ \e -> execMemoNow e cont ivar)
envOuter
where
addStats :: Env u w -> Int64 -> IO ()
addStats env acc = modifyIORef' (statsRef env) $ \(Stats s) ->
Stats (MemoCall cid acc : s)
collectMemoData :: Int64 -> GenHaxl u w a -> GenHaxl u w a
collectMemoData acc f = GenHaxl $ \env -> do
a0 <- getAllocationCounter
r <- unHaxl f env{memoKey=cid}
a1 <- getAllocationCounter
let newTotal = acc + (a0 - a1)
ret <- case r of
Done a -> do addStats env newTotal; return (Done a)
Throw e -> do addStats env newTotal; return (Throw e)
Blocked ivar k ->
return (Blocked ivar (Cont (collectMemoData newTotal (toHaxl k))))
setAllocationCounter a1
return ret
execMemoNow :: Env u w -> GenHaxl u w a -> IVar u w a -> IO (Result u w a)
execMemoNow env cont ivar = do
wlogs <- newIORef NilWrites
let
!menv = env { writeLogsRef = wlogs }
r <- Exception.try $ unHaxl cont menv
case r of
Left e -> trace_ ("execMemoNow: Left " ++ show e) $ do
rethrowAsyncExceptions e
putIVar ivar (ThrowIO e) env
throwIO e
Right (Done a) -> trace_ "execMemoNow: Done" $ do
wt <- readIORef wlogs
putIVar ivar (Ok a wt) env
mbModifyWLRef wt (writeLogsRef env)
return (Done a)
Right (Throw ex) -> trace_ ("execMemoNow: Throw" ++ show ex) $ do
wt <- readIORef wlogs
putIVar ivar (ThrowHaxl ex wt) env
mbModifyWLRef wt (writeLogsRef env)
return (Throw ex)
Right (Blocked ivar' cont) -> trace_ "execMemoNow: Blocked" $ do
addJob menv (toHaxl cont) ivar ivar'
return (Blocked ivar (Cont (getIVarWithWrites ivar)))
newtype MemoVar1 u w a b = MemoVar1 (IORef (MemoStatus1 u w a b))
newtype MemoVar2 u w a b c = MemoVar2 (IORef (MemoStatus2 u w a b c))
data MemoStatus1 u w a b
= MemoEmpty1
| MemoTbl1 (a -> GenHaxl u w b) (HashMap.HashMap a (MemoVar u w b))
data MemoStatus2 u w a b c
= MemoEmpty2
| MemoTbl2
(a -> b -> GenHaxl u w c)
(HashMap.HashMap a (HashMap.HashMap b (MemoVar u w c)))
newMemo1 :: GenHaxl u w (MemoVar1 u w a b)
newMemo1 = unsafeLiftIO $ MemoVar1 <$> newIORef MemoEmpty1
newMemoWith1 :: (a -> GenHaxl u w b) -> GenHaxl u w (MemoVar1 u w a b)
newMemoWith1 f = newMemo1 >>= \r -> prepareMemo1 r f >> return r
prepareMemo1 :: MemoVar1 u w a b -> (a -> GenHaxl u w b) -> GenHaxl u w ()
prepareMemo1 (MemoVar1 r) f
= unsafeLiftIO $ writeIORef r (MemoTbl1 f HashMap.empty)
runMemo1 :: (Eq a, Hashable a) => MemoVar1 u w a b -> a -> GenHaxl u w b
runMemo1 (MemoVar1 r) k = unsafeLiftIO (readIORef r) >>= \case
MemoEmpty1 -> throw $ CriticalError "Attempting to run empty memo."
MemoTbl1 f h -> case HashMap.lookup k h of
Nothing -> do
x <- newMemoWith (f k)
unsafeLiftIO $ writeIORef r (MemoTbl1 f (HashMap.insert k x h))
runMemo x
Just v -> runMemo v
newMemo2 :: GenHaxl u w (MemoVar2 u w a b c)
newMemo2 = unsafeLiftIO $ MemoVar2 <$> newIORef MemoEmpty2
newMemoWith2 :: (a -> b -> GenHaxl u w c) -> GenHaxl u w (MemoVar2 u w a b c)
newMemoWith2 f = newMemo2 >>= \r -> prepareMemo2 r f >> return r
prepareMemo2 :: MemoVar2 u w a b c -> (a -> b -> GenHaxl u w c) -> GenHaxl u w ()
prepareMemo2 (MemoVar2 r) f
= unsafeLiftIO $ writeIORef r (MemoTbl2 f HashMap.empty)
runMemo2 :: (Eq a, Hashable a, Eq b, Hashable b)
=> MemoVar2 u w a b c
-> a -> b -> GenHaxl u w c
runMemo2 (MemoVar2 r) k1 k2 = unsafeLiftIO (readIORef r) >>= \case
MemoEmpty2 -> throw $ CriticalError "Attempting to run empty memo."
MemoTbl2 f h1 -> case HashMap.lookup k1 h1 of
Nothing -> do
v <- newMemoWith (f k1 k2)
unsafeLiftIO $ writeIORef r
(MemoTbl2 f (HashMap.insert k1 (HashMap.singleton k2 v) h1))
runMemo v
Just h2 -> case HashMap.lookup k2 h2 of
Nothing -> do
v <- newMemoWith (f k1 k2)
unsafeLiftIO $ writeIORef r
(MemoTbl2 f (HashMap.insert k1 (HashMap.insert k2 v h2) h1))
runMemo v
Just v -> runMemo v
memo
:: (Typeable a, Typeable k, Hashable k, Eq k)
=> k -> GenHaxl u w a -> GenHaxl u w a
memo key = cachedComputation (MemoKey key)
{-# RULES
"memo/Text" memo = memoText :: (Typeable a) =>
Text -> GenHaxl u w a -> GenHaxl u w a
#-}
{-# NOINLINE memo #-}
memoUnique
:: (Typeable a, Typeable k, Hashable k, Eq k)
=> MemoFingerprintKey a -> Text -> k -> GenHaxl u w a -> GenHaxl u w a
memoUnique fp label key = withLabel label . memo (fp, key)
{-# NOINLINE memoUnique #-}
data MemoKey k a where
MemoKey :: (Typeable k, Hashable k, Eq k) => k -> MemoKey k a
deriving Typeable
deriving instance Eq (MemoKey k a)
instance Hashable (MemoKey k a) where
hashWithSalt s (MemoKey t) = hashWithSalt s t
data MemoTextKey a where
MemoText :: Text -> MemoTextKey a
deriving Typeable
deriving instance Eq (MemoTextKey a)
instance Hashable (MemoTextKey a) where
hashWithSalt s (MemoText t) = hashWithSalt s t
memoText :: (Typeable a) => Text -> GenHaxl u w a -> GenHaxl u w a
memoText key = withLabel key . cachedComputation (MemoText key)
data MemoFingerprintKey a where
MemoFingerprintKey
:: {-# UNPACK #-} !Word64
-> {-# UNPACK #-} !Word64
-> Addr# -> Addr#
-> MemoFingerprintKey a
deriving Typeable
deriving instance Eq (MemoFingerprintKey a)
instance Hashable (MemoFingerprintKey a) where
hashWithSalt s (MemoFingerprintKey x _ _ _) =
hashWithSalt s (fromIntegral x :: Int)
{-# NOINLINE memoFingerprint #-}
memoFingerprint
:: Typeable a => MemoFingerprintKey a -> GenHaxl u w a -> GenHaxl u w a
memoFingerprint key@(MemoFingerprintKey _ _ mnPtr nPtr) =
withFingerprintLabel mnPtr nPtr . cachedComputation key
memoize :: GenHaxl u w a -> GenHaxl u w (GenHaxl u w a)
memoize a = runMemo <$> newMemoWith a
memoize1 :: (Eq a, Hashable a)
=> (a -> GenHaxl u w b)
-> GenHaxl u w (a -> GenHaxl u w b)
memoize1 f = runMemo1 <$> newMemoWith1 f
memoize2 :: (Eq a, Hashable a, Eq b, Hashable b)
=> (a -> b -> GenHaxl u w c)
-> GenHaxl u w (a -> b -> GenHaxl u w c)
memoize2 f = runMemo2 <$> newMemoWith2 f