-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}

-- |
-- Memoization support. This module is provided for access to Haxl
-- internals only; most users should import "Haxl.Core" instead.
--
module Haxl.Core.Memo
  (
    -- * Basic memoization
    cachedComputation
  , preCacheComputation

    -- * High-level memoization
  , memo
  , memoFingerprint
  , MemoFingerprintKey(..)
  , memoize, memoize1, memoize2
  , memoUnique

    -- * Local memoization
  , 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_)

-- -----------------------------------------------------------------------------
-- Memoization

-- | 'cachedComputation' memoizes a Haxl computation.  The key is a
-- request.
--
-- /Note:/ These cached computations will /not/ be included in the output
-- of 'dumpCacheAsHaskell'.
--
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 a -> GenHaxl u w a -> GenHaxl u w a
cachedComputation req a
req GenHaxl u w a
haxl = (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall u w a. (Env u w -> IO (Result u w a)) -> GenHaxl u w a
GenHaxl ((Env u w -> IO (Result u w a)) -> GenHaxl u w a)
-> (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall a b. (a -> b) -> a -> b
$ \env :: Env u w
env@Env{u
CallId
Maybe (DataCacheLookup w)
TVar [CompleteReq u w]
IORef CallId
IORef Profile
IORef Stats
IORef ReqCountMap
IORef (RequestStore u)
IORef (JobList u w)
IORef (WriteTree w)
HaxlDataCache u w
Flags
StateStore
ProfileCurrent
dataCacheFetchFallback :: forall u w. Env u w -> Maybe (DataCacheLookup w)
writeLogsRefNoMemo :: forall u w. Env u w -> IORef (WriteTree w)
writeLogsRef :: forall u w. Env u w -> IORef (WriteTree w)
completions :: forall u w. Env u w -> TVar [CompleteReq u w]
submittedReqsRef :: forall u w. Env u w -> IORef ReqCountMap
runQueueRef :: forall u w. Env u w -> IORef (JobList u w)
reqStoreRef :: forall u w. Env u w -> IORef (RequestStore u)
states :: forall u w. Env u w -> StateStore
profRef :: forall u w. Env u w -> IORef Profile
profCurrent :: forall u w. Env u w -> ProfileCurrent
callIdRef :: forall u w. Env u w -> IORef CallId
statsBatchIdRef :: forall u w. Env u w -> IORef CallId
statsRef :: forall u w. Env u w -> IORef Stats
userEnv :: forall u w. Env u w -> u
flags :: forall u w. Env u w -> Flags
memoKey :: forall u w. Env u w -> CallId
memoCache :: forall u w. Env u w -> HaxlDataCache u w
dataCache :: forall u w. Env u w -> HaxlDataCache u w
dataCacheFetchFallback :: Maybe (DataCacheLookup w)
writeLogsRefNoMemo :: IORef (WriteTree w)
writeLogsRef :: IORef (WriteTree w)
completions :: TVar [CompleteReq u w]
submittedReqsRef :: IORef ReqCountMap
runQueueRef :: IORef (JobList u w)
reqStoreRef :: IORef (RequestStore u)
states :: StateStore
profRef :: IORef Profile
profCurrent :: ProfileCurrent
callIdRef :: IORef CallId
statsBatchIdRef :: IORef CallId
statsRef :: IORef Stats
userEnv :: u
flags :: Flags
memoKey :: CallId
memoCache :: HaxlDataCache u w
dataCache :: HaxlDataCache u w
..} -> do
  Maybe (DataCacheItem u w a)
mbRes <- req a -> HaxlDataCache u w -> IO (Maybe (DataCacheItem u w a))
forall (req :: * -> *) a (res :: * -> *).
Typeable (req a) =>
req a -> DataCache res -> IO (Maybe (res a))
DataCache.lookup req a
req HaxlDataCache u w
memoCache
  case Maybe (DataCacheItem u w a)
mbRes of
    Just (DataCacheItem IVar u w a
ivar CallId
k) -> do
      Flags -> IO () -> IO ()
forall (m :: * -> *) a. Monad m => Flags -> m a -> m ()
ifProfiling Flags
flags (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        Env u w -> CallId -> Bool -> IO ()
forall u w. Env u w -> CallId -> Bool -> IO ()
incrementMemoHitCounterFor Env u w
env CallId
k Bool
True
      GenHaxl u w a -> Env u w -> IO (Result u w a)
forall u w a. GenHaxl u w a -> Env u w -> IO (Result u w a)
unHaxl (IVar u w a -> GenHaxl u w a
forall u w a. IVar u w a -> GenHaxl u w a
getIVarWithWrites IVar u w a
ivar) Env u w
env
    Maybe (DataCacheItem u w a)
Nothing -> do
      IVar u w a
ivar <- IO (IVar u w a)
forall u w a. IO (IVar u w a)
newIVar
      CallId
k <- Env u w -> IO CallId
forall u w. Env u w -> IO CallId
nextCallId Env u w
env
      -- no need to incremenetMemoHitCounter as execMemo will do it
      req a -> DataCacheItem u w a -> HaxlDataCache u w -> IO ()
forall (req :: * -> *) a (res :: * -> *).
(Hashable (req a), Typeable (req a), Eq (req a)) =>
req a -> res a -> DataCache res -> IO ()
DataCache.insertNotShowable req a
req (IVar u w a -> CallId -> DataCacheItem u w a
forall u w a. IVar u w a -> CallId -> DataCacheItem u w a
DataCacheItem IVar u w a
ivar CallId
k) HaxlDataCache u w
memoCache
      Env u w
-> GenHaxl u w a -> IVar u w a -> CallId -> IO (Result u w a)
forall u w a.
Env u w
-> GenHaxl u w a -> IVar u w a -> CallId -> IO (Result u w a)
execMemoNowProfiled Env u w
env GenHaxl u w a
haxl IVar u w a
ivar CallId
k


-- | Like 'cachedComputation', but fails if the cache is already
-- populated.
--
-- Memoization can be (ab)used to "mock" a cached computation, by
-- pre-populating the cache with an alternative implementation. In
-- that case we don't want the operation to populate the cache to
-- silently succeed if the cache is already populated.
--
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 a -> GenHaxl u w a -> GenHaxl u w a
preCacheComputation req a
req GenHaxl u w a
haxl = (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall u w a. (Env u w -> IO (Result u w a)) -> GenHaxl u w a
GenHaxl ((Env u w -> IO (Result u w a)) -> GenHaxl u w a)
-> (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall a b. (a -> b) -> a -> b
$ \env :: Env u w
env@Env{u
CallId
Maybe (DataCacheLookup w)
TVar [CompleteReq u w]
IORef CallId
IORef Profile
IORef Stats
IORef ReqCountMap
IORef (RequestStore u)
IORef (JobList u w)
IORef (WriteTree w)
HaxlDataCache u w
Flags
StateStore
ProfileCurrent
dataCacheFetchFallback :: Maybe (DataCacheLookup w)
writeLogsRefNoMemo :: IORef (WriteTree w)
writeLogsRef :: IORef (WriteTree w)
completions :: TVar [CompleteReq u w]
submittedReqsRef :: IORef ReqCountMap
runQueueRef :: IORef (JobList u w)
reqStoreRef :: IORef (RequestStore u)
states :: StateStore
profRef :: IORef Profile
profCurrent :: ProfileCurrent
callIdRef :: IORef CallId
statsBatchIdRef :: IORef CallId
statsRef :: IORef Stats
userEnv :: u
flags :: Flags
memoKey :: CallId
memoCache :: HaxlDataCache u w
dataCache :: HaxlDataCache u w
dataCacheFetchFallback :: forall u w. Env u w -> Maybe (DataCacheLookup w)
writeLogsRefNoMemo :: forall u w. Env u w -> IORef (WriteTree w)
writeLogsRef :: forall u w. Env u w -> IORef (WriteTree w)
completions :: forall u w. Env u w -> TVar [CompleteReq u w]
submittedReqsRef :: forall u w. Env u w -> IORef ReqCountMap
runQueueRef :: forall u w. Env u w -> IORef (JobList u w)
reqStoreRef :: forall u w. Env u w -> IORef (RequestStore u)
states :: forall u w. Env u w -> StateStore
profRef :: forall u w. Env u w -> IORef Profile
profCurrent :: forall u w. Env u w -> ProfileCurrent
callIdRef :: forall u w. Env u w -> IORef CallId
statsBatchIdRef :: forall u w. Env u w -> IORef CallId
statsRef :: forall u w. Env u w -> IORef Stats
userEnv :: forall u w. Env u w -> u
flags :: forall u w. Env u w -> Flags
memoKey :: forall u w. Env u w -> CallId
memoCache :: forall u w. Env u w -> HaxlDataCache u w
dataCache :: forall u w. Env u w -> HaxlDataCache u w
..} -> do
  Maybe (DataCacheItem u w a)
mbRes <- req a -> HaxlDataCache u w -> IO (Maybe (DataCacheItem u w a))
forall (req :: * -> *) a (res :: * -> *).
Typeable (req a) =>
req a -> DataCache res -> IO (Maybe (res a))
DataCache.lookup req a
req HaxlDataCache u w
memoCache
  case Maybe (DataCacheItem u w a)
mbRes of
    Just DataCacheItem u w a
_ -> Result u w a -> IO (Result u w a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Result u w a -> IO (Result u w a))
-> Result u w a -> IO (Result u w a)
forall a b. (a -> b) -> a -> b
$ SomeException -> Result u w a
forall u w a. SomeException -> Result u w a
Throw (SomeException -> Result u w a) -> SomeException -> Result u w a
forall a b. (a -> b) -> a -> b
$ InvalidParameter -> SomeException
forall e. Exception e => e -> SomeException
toException (InvalidParameter -> SomeException)
-> InvalidParameter -> SomeException
forall a b. (a -> b) -> a -> b
$ Text -> InvalidParameter
InvalidParameter
      Text
"preCacheComputation: key is already cached"
    Maybe (DataCacheItem u w a)
Nothing -> do
      IVar u w a
ivar <- IO (IVar u w a)
forall u w a. IO (IVar u w a)
newIVar
      CallId
k <- Env u w -> IO CallId
forall u w. Env u w -> IO CallId
nextCallId Env u w
env
      req a -> DataCacheItem u w a -> HaxlDataCache u w -> IO ()
forall (req :: * -> *) a (res :: * -> *).
(Hashable (req a), Typeable (req a), Eq (req a)) =>
req a -> res a -> DataCache res -> IO ()
DataCache.insertNotShowable req a
req (IVar u w a -> CallId -> DataCacheItem u w a
forall u w a. IVar u w a -> CallId -> DataCacheItem u w a
DataCacheItem IVar u w a
ivar CallId
k) HaxlDataCache u w
memoCache
      Env u w
-> GenHaxl u w a -> IVar u w a -> CallId -> IO (Result u w a)
forall u w a.
Env u w
-> GenHaxl u w a -> IVar u w a -> CallId -> IO (Result u w a)
execMemoNowProfiled Env u w
env GenHaxl u w a
haxl IVar u w a
ivar CallId
k

-- -----------------------------------------------------------------------------
-- Memoization

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

-- | Create a new @MemoVar@ for storing a memoized computation. The created
-- @MemoVar@ is initially empty, not tied to any specific computation. Running
-- this memo (with @runMemo@) without preparing it first (with @prepareMemo@)
-- will result in an exception.
newMemo :: GenHaxl u w (MemoVar u w a)
newMemo :: GenHaxl u w (MemoVar u w a)
newMemo = IO (MemoVar u w a) -> GenHaxl u w (MemoVar u w a)
forall a u w. IO a -> GenHaxl u w a
unsafeLiftIO (IO (MemoVar u w a) -> GenHaxl u w (MemoVar u w a))
-> IO (MemoVar u w a) -> GenHaxl u w (MemoVar u w a)
forall a b. (a -> b) -> a -> b
$ IORef (MemoStatus u w a) -> MemoVar u w a
forall u w a. IORef (MemoStatus u w a) -> MemoVar u w a
MemoVar (IORef (MemoStatus u w a) -> MemoVar u w a)
-> IO (IORef (MemoStatus u w a)) -> IO (MemoVar u w a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MemoStatus u w a -> IO (IORef (MemoStatus u w a))
forall a. a -> IO (IORef a)
newIORef MemoStatus u w a
forall u w a. MemoStatus u w a
MemoEmpty

-- | Store a computation within a supplied @MemoVar@. Any memo stored within the
-- @MemoVar@ already (regardless of completion) will be discarded, in favor of
-- the supplied computation. A @MemoVar@ must be prepared before it is run.
prepareMemo :: MemoVar u w a -> GenHaxl u w a -> GenHaxl u w ()
prepareMemo :: MemoVar u w a -> GenHaxl u w a -> GenHaxl u w ()
prepareMemo (MemoVar IORef (MemoStatus u w a)
memoRef) GenHaxl u w a
memoCmp
  = (Env u w -> IO (Result u w ())) -> GenHaxl u w ()
forall u w a. (Env u w -> IO (Result u w a)) -> GenHaxl u w a
GenHaxl ((Env u w -> IO (Result u w ())) -> GenHaxl u w ())
-> (Env u w -> IO (Result u w ())) -> GenHaxl u w ()
forall a b. (a -> b) -> a -> b
$ \Env u w
env -> do
      CallId
k <- Env u w -> IO CallId
forall u w. Env u w -> IO CallId
nextCallId Env u w
env
      IORef (MemoStatus u w a) -> MemoStatus u w a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (MemoStatus u w a)
memoRef (GenHaxl u w a -> CallId -> MemoStatus u w a
forall u w a. GenHaxl u w a -> CallId -> MemoStatus u w a
MemoReady GenHaxl u w a
memoCmp CallId
k)
      Result u w () -> IO (Result u w ())
forall (m :: * -> *) a. Monad m => a -> m a
return (() -> Result u w ()
forall u w a. a -> Result u w a
Done ())

-- | Convenience function, combines @newMemo@ and @prepareMemo@.
newMemoWith :: GenHaxl u w a -> GenHaxl u w (MemoVar u w a)
newMemoWith :: GenHaxl u w a -> GenHaxl u w (MemoVar u w a)
newMemoWith GenHaxl u w a
memoCmp = do
  MemoVar u w a
memoVar <- GenHaxl u w (MemoVar u w a)
forall u w a. GenHaxl u w (MemoVar u w a)
newMemo
  MemoVar u w a -> GenHaxl u w a -> GenHaxl u w ()
forall u w a. MemoVar u w a -> GenHaxl u w a -> GenHaxl u w ()
prepareMemo MemoVar u w a
memoVar GenHaxl u w a
memoCmp
  MemoVar u w a -> GenHaxl u w (MemoVar u w a)
forall (m :: * -> *) a. Monad m => a -> m a
return MemoVar u w a
memoVar

-- | Continue the memoized computation within a given @MemoVar@.
-- Notes:
--
--   1. If the memo contains a complete result, return that result.
--   2. If the memo contains an in-progress computation, continue it as far as
--      possible for this round.
--   3. If the memo is empty (it was not prepared), throw an error.
--
-- For example, to memoize the computation @one@ given by:
--
-- > one :: Haxl Int
-- > one = return 1
--
-- use:
--
-- > do
-- >   oneMemo <- newMemoWith one
-- >   let memoizedOne = runMemo aMemo one
-- >   oneResult <- memoizedOne
--
-- To memoize mutually dependent computations such as in:
--
-- > h :: Haxl Int
-- > h = do
-- >   a <- f
-- >   b <- g
-- >   return (a + b)
-- >  where
-- >   f = return 42
-- >   g = succ <$> f
--
-- without needing to reorder them, use:
--
-- > h :: Haxl Int
-- > h = do
-- >   fMemoRef <- newMemo
-- >   gMemoRef <- newMemo
-- >
-- >   let f = runMemo fMemoRef
-- >       g = runMemo gMemoRef
-- >
-- >   prepareMemo fMemoRef $ return 42
-- >   prepareMemo gMemoRef $ succ <$> f
-- >
-- >   a <- f
-- >   b <- g
-- >   return (a + b)
--
runMemo :: MemoVar u w a -> GenHaxl u w a
runMemo :: MemoVar u w a -> GenHaxl u w a
runMemo (MemoVar IORef (MemoStatus u w a)
memoRef) = (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall u w a. (Env u w -> IO (Result u w a)) -> GenHaxl u w a
GenHaxl ((Env u w -> IO (Result u w a)) -> GenHaxl u w a)
-> (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall a b. (a -> b) -> a -> b
$ \Env u w
env -> do
  MemoStatus u w a
stored <- IORef (MemoStatus u w a) -> IO (MemoStatus u w a)
forall a. IORef a -> IO a
readIORef IORef (MemoStatus u w a)
memoRef
  case MemoStatus u w a
stored of
    -- Memo was not prepared first; throw an exception.
    MemoStatus u w a
MemoEmpty -> String -> IO (Result u w a) -> IO (Result u w a)
forall a. String -> a -> a
trace_ String
"MemoEmpty " (IO (Result u w a) -> IO (Result u w a))
-> IO (Result u w a) -> IO (Result u w a)
forall a b. (a -> b) -> a -> b
$
      Env u w -> CriticalError -> IO (Result u w a)
forall e u w a. Exception e => Env u w -> e -> IO (Result u w a)
raise Env u w
env (CriticalError -> IO (Result u w a))
-> CriticalError -> IO (Result u w a)
forall a b. (a -> b) -> a -> b
$ Text -> CriticalError
CriticalError Text
"Attempting to run empty memo."
    -- Memo has been prepared but not run yet
    MemoReady GenHaxl u w a
cont CallId
k -> String -> IO (Result u w a) -> IO (Result u w a)
forall a. String -> a -> a
trace_ String
"MemoReady" (IO (Result u w a) -> IO (Result u w a))
-> IO (Result u w a) -> IO (Result u w a)
forall a b. (a -> b) -> a -> b
$ do
      IVar u w a
ivar <- IO (IVar u w a)
forall u w a. IO (IVar u w a)
newIVar
      IORef (MemoStatus u w a) -> MemoStatus u w a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (MemoStatus u w a)
memoRef (IVar u w a -> CallId -> MemoStatus u w a
forall u w a. IVar u w a -> CallId -> MemoStatus u w a
MemoRun IVar u w a
ivar CallId
k)
      Env u w
-> GenHaxl u w a -> IVar u w a -> CallId -> IO (Result u w a)
forall u w a.
Env u w
-> GenHaxl u w a -> IVar u w a -> CallId -> IO (Result u w a)
execMemoNowProfiled Env u w
env GenHaxl u w a
cont IVar u w a
ivar CallId
k
    -- The memo has already been run, get (or wait for) for the result
    MemoRun IVar u w a
ivar CallId
k -> String -> IO (Result u w a) -> IO (Result u w a)
forall a. String -> a -> a
trace_ String
"MemoRun" (IO (Result u w a) -> IO (Result u w a))
-> IO (Result u w a) -> IO (Result u w a)
forall a b. (a -> b) -> a -> b
$ do
      Flags -> IO () -> IO ()
forall (m :: * -> *) a. Monad m => Flags -> m a -> m ()
ifProfiling (Env u w -> Flags
forall u w. Env u w -> Flags
flags Env u w
env) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        Env u w -> CallId -> Bool -> IO ()
forall u w. Env u w -> CallId -> Bool -> IO ()
incrementMemoHitCounterFor Env u w
env CallId
k Bool
True
      GenHaxl u w a -> Env u w -> IO (Result u w a)
forall u w a. GenHaxl u w a -> Env u w -> IO (Result u w a)
unHaxl (IVar u w a -> GenHaxl u w a
forall u w a. IVar u w a -> GenHaxl u w a
getIVarWithWrites IVar u w a
ivar) Env u w
env

execMemoNowProfiled
  :: Env u w
  -> GenHaxl u w a
  -> IVar u w a
  -> CallId
  -> IO (Result u w a)
execMemoNowProfiled :: Env u w
-> GenHaxl u w a -> IVar u w a -> CallId -> IO (Result u w a)
execMemoNowProfiled Env u w
envOuter GenHaxl u w a
cont IVar u w a
ivar CallId
cid =
  if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ ReportFlag -> ReportFlags -> Bool
testReportFlag ReportFlag
ReportProfiling (ReportFlags -> Bool) -> ReportFlags -> Bool
forall a b. (a -> b) -> a -> b
$ Flags -> ReportFlags
report (Flags -> ReportFlags) -> Flags -> ReportFlags
forall a b. (a -> b) -> a -> b
$ Env u w -> Flags
forall u w. Env u w -> Flags
flags Env u w
envOuter
  then Env u w -> GenHaxl u w a -> IVar u w a -> IO (Result u w a)
forall u w a.
Env u w -> GenHaxl u w a -> IVar u w a -> IO (Result u w a)
execMemoNow Env u w
envOuter GenHaxl u w a
cont IVar u w a
ivar
  else do
    Env u w -> CallId -> Bool -> IO ()
forall u w. Env u w -> CallId -> Bool -> IO ()
incrementMemoHitCounterFor Env u w
envOuter CallId
cid Bool
False
    GenHaxl u w a -> Env u w -> IO (Result u w a)
forall u w a. GenHaxl u w a -> Env u w -> IO (Result u w a)
unHaxl
      (Int64 -> GenHaxl u w a -> GenHaxl u w a
forall u w a. Int64 -> GenHaxl u w a -> GenHaxl u w a
collectMemoData Int64
0 (GenHaxl u w a -> GenHaxl u w a) -> GenHaxl u w a -> GenHaxl u w a
forall a b. (a -> b) -> a -> b
$ (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall u w a. (Env u w -> IO (Result u w a)) -> GenHaxl u w a
GenHaxl ((Env u w -> IO (Result u w a)) -> GenHaxl u w a)
-> (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall a b. (a -> b) -> a -> b
$ \Env u w
e -> Env u w -> GenHaxl u w a -> IVar u w a -> IO (Result u w a)
forall u w a.
Env u w -> GenHaxl u w a -> IVar u w a -> IO (Result u w a)
execMemoNow Env u w
e GenHaxl u w a
cont IVar u w a
ivar)
      Env u w
envOuter
  where
    addStats :: Env u w -> Int64 -> IO ()
    addStats :: Env u w -> Int64 -> IO ()
addStats Env u w
env Int64
acc = IORef Stats -> (Stats -> Stats) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Env u w -> IORef Stats
forall u w. Env u w -> IORef Stats
statsRef Env u w
env) ((Stats -> Stats) -> IO ()) -> (Stats -> Stats) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Stats [FetchStats]
s) ->
      [FetchStats] -> Stats
Stats (CallId -> Int64 -> FetchStats
MemoCall CallId
cid Int64
acc FetchStats -> [FetchStats] -> [FetchStats]
forall a. a -> [a] -> [a]
: [FetchStats]
s)
    collectMemoData :: Int64 -> GenHaxl u w a -> GenHaxl u w a
    collectMemoData :: Int64 -> GenHaxl u w a -> GenHaxl u w a
collectMemoData Int64
acc GenHaxl u w a
f = (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall u w a. (Env u w -> IO (Result u w a)) -> GenHaxl u w a
GenHaxl ((Env u w -> IO (Result u w a)) -> GenHaxl u w a)
-> (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall a b. (a -> b) -> a -> b
$ \Env u w
env -> do
      Int64
a0 <- IO Int64
getAllocationCounter
      Result u w a
r <- GenHaxl u w a -> Env u w -> IO (Result u w a)
forall u w a. GenHaxl u w a -> Env u w -> IO (Result u w a)
unHaxl GenHaxl u w a
f Env u w
env{memoKey :: CallId
memoKey=CallId
cid}
      Int64
a1 <- IO Int64
getAllocationCounter
      let newTotal :: Int64
newTotal = Int64
acc Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ (Int64
a0 Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
a1)
      Result u w a
ret <- case Result u w a
r of
        Done a
a -> do Env u w -> Int64 -> IO ()
forall u w. Env u w -> Int64 -> IO ()
addStats Env u w
env Int64
newTotal; Result u w a -> IO (Result u w a)
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Result u w a
forall u w a. a -> Result u w a
Done a
a)
        Throw SomeException
e -> do Env u w -> Int64 -> IO ()
forall u w. Env u w -> Int64 -> IO ()
addStats Env u w
env Int64
newTotal; Result u w a -> IO (Result u w a)
forall (m :: * -> *) a. Monad m => a -> m a
return (SomeException -> Result u w a
forall u w a. SomeException -> Result u w a
Throw SomeException
e)
        Blocked IVar u w b
ivar Cont u w a
k ->
          Result u w a -> IO (Result u w a)
forall (m :: * -> *) a. Monad m => a -> m a
return (IVar u w b -> Cont u w a -> Result u w a
forall u w a b. IVar u w b -> Cont u w a -> Result u w a
Blocked IVar u w b
ivar (GenHaxl u w a -> Cont u w a
forall u w a. GenHaxl u w a -> Cont u w a
Cont (Int64 -> GenHaxl u w a -> GenHaxl u w a
forall u w a. Int64 -> GenHaxl u w a -> GenHaxl u w a
collectMemoData Int64
newTotal (Cont u w a -> GenHaxl u w a
forall u w a. Cont u w a -> GenHaxl u w a
toHaxl Cont u w a
k))))
      Int64 -> IO ()
setAllocationCounter Int64
a1
      Result u w a -> IO (Result u w a)
forall (m :: * -> *) a. Monad m => a -> m a
return Result u w a
ret

execMemoNow :: Env u w -> GenHaxl u w a -> IVar u w a -> IO (Result u w a)
execMemoNow :: Env u w -> GenHaxl u w a -> IVar u w a -> IO (Result u w a)
execMemoNow Env u w
env GenHaxl u w a
cont IVar u w a
ivar = do
  IORef (WriteTree w)
wlogs <- WriteTree w -> IO (IORef (WriteTree w))
forall a. a -> IO (IORef a)
newIORef WriteTree w
forall w. WriteTree w
NilWrites
  let
    !menv :: Env u w
menv = Env u w
env { writeLogsRef :: IORef (WriteTree w)
writeLogsRef = IORef (WriteTree w)
wlogs }
    -- use an env with empty writes, so we can memoize the extra
    -- writes done as part of 'cont'
  Either SomeException (Result u w a)
r <- IO (Result u w a) -> IO (Either SomeException (Result u w a))
forall e a. Exception e => IO a -> IO (Either e a)
Exception.try (IO (Result u w a) -> IO (Either SomeException (Result u w a)))
-> IO (Result u w a) -> IO (Either SomeException (Result u w a))
forall a b. (a -> b) -> a -> b
$ GenHaxl u w a -> Env u w -> IO (Result u w a)
forall u w a. GenHaxl u w a -> Env u w -> IO (Result u w a)
unHaxl GenHaxl u w a
cont Env u w
menv

  case Either SomeException (Result u w a)
r of
    Left SomeException
e -> String -> IO (Result u w a) -> IO (Result u w a)
forall a. String -> a -> a
trace_ (String
"execMemoNow: Left " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SomeException -> String
forall a. Show a => a -> String
show SomeException
e) (IO (Result u w a) -> IO (Result u w a))
-> IO (Result u w a) -> IO (Result u w a)
forall a b. (a -> b) -> a -> b
$ do
      SomeException -> IO ()
rethrowAsyncExceptions SomeException
e
      IVar u w a -> ResultVal a w -> Env u w -> IO ()
forall u w a. IVar u w a -> ResultVal a w -> Env u w -> IO ()
putIVar IVar u w a
ivar (SomeException -> ResultVal a w
forall a w. SomeException -> ResultVal a w
ThrowIO SomeException
e) Env u w
env
      SomeException -> IO (Result u w a)
forall e a. Exception e => e -> IO a
throwIO SomeException
e
    Right (Done a
a) -> String -> IO (Result u w a) -> IO (Result u w a)
forall a. String -> a -> a
trace_ String
"execMemoNow: Done" (IO (Result u w a) -> IO (Result u w a))
-> IO (Result u w a) -> IO (Result u w a)
forall a b. (a -> b) -> a -> b
$ do
      WriteTree w
wt <- IORef (WriteTree w) -> IO (WriteTree w)
forall a. IORef a -> IO a
readIORef IORef (WriteTree w)
wlogs
      IVar u w a -> ResultVal a w -> Env u w -> IO ()
forall u w a. IVar u w a -> ResultVal a w -> Env u w -> IO ()
putIVar IVar u w a
ivar (a -> WriteTree w -> ResultVal a w
forall a w. a -> WriteTree w -> ResultVal a w
Ok a
a WriteTree w
wt) Env u w
env
      WriteTree w -> IORef (WriteTree w) -> IO ()
forall w. WriteTree w -> IORef (WriteTree w) -> IO ()
mbModifyWLRef WriteTree w
wt (Env u w -> IORef (WriteTree w)
forall u w. Env u w -> IORef (WriteTree w)
writeLogsRef Env u w
env)
      Result u w a -> IO (Result u w a)
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Result u w a
forall u w a. a -> Result u w a
Done a
a)
    Right (Throw SomeException
ex) -> String -> IO (Result u w a) -> IO (Result u w a)
forall a. String -> a -> a
trace_ (String
"execMemoNow: Throw" String -> String -> String
forall a. [a] -> [a] -> [a]
++ SomeException -> String
forall a. Show a => a -> String
show SomeException
ex) (IO (Result u w a) -> IO (Result u w a))
-> IO (Result u w a) -> IO (Result u w a)
forall a b. (a -> b) -> a -> b
$ do
      WriteTree w
wt <- IORef (WriteTree w) -> IO (WriteTree w)
forall a. IORef a -> IO a
readIORef IORef (WriteTree w)
wlogs
      IVar u w a -> ResultVal a w -> Env u w -> IO ()
forall u w a. IVar u w a -> ResultVal a w -> Env u w -> IO ()
putIVar IVar u w a
ivar (SomeException -> WriteTree w -> ResultVal a w
forall a w. SomeException -> WriteTree w -> ResultVal a w
ThrowHaxl SomeException
ex WriteTree w
wt) Env u w
env
      WriteTree w -> IORef (WriteTree w) -> IO ()
forall w. WriteTree w -> IORef (WriteTree w) -> IO ()
mbModifyWLRef WriteTree w
wt (Env u w -> IORef (WriteTree w)
forall u w. Env u w -> IORef (WriteTree w)
writeLogsRef Env u w
env)
      Result u w a -> IO (Result u w a)
forall (m :: * -> *) a. Monad m => a -> m a
return (SomeException -> Result u w a
forall u w a. SomeException -> Result u w a
Throw SomeException
ex)
    Right (Blocked IVar u w b
ivar' Cont u w a
cont) -> String -> IO (Result u w a) -> IO (Result u w a)
forall a. String -> a -> a
trace_ String
"execMemoNow: Blocked" (IO (Result u w a) -> IO (Result u w a))
-> IO (Result u w a) -> IO (Result u w a)
forall a b. (a -> b) -> a -> b
$ do
      -- We "block" this memoized computation in the new environment 'menv', so
      -- that when it finishes, we can store all the write logs from the env
      -- in the IVar.
      Env u w -> GenHaxl u w a -> IVar u w a -> IVar u w b -> IO ()
forall u w b a.
Env u w -> GenHaxl u w b -> IVar u w b -> IVar u w a -> IO ()
addJob Env u w
menv (Cont u w a -> GenHaxl u w a
forall u w a. Cont u w a -> GenHaxl u w a
toHaxl Cont u w a
cont) IVar u w a
ivar IVar u w b
ivar'
      -- Now we call @getIVarWithWrites@ to populate the writes in the original
      -- environment 'env'.
      Result u w a -> IO (Result u w a)
forall (m :: * -> *) a. Monad m => a -> m a
return (IVar u w a -> Cont u w a -> Result u w a
forall u w a b. IVar u w b -> Cont u w a -> Result u w a
Blocked IVar u w a
ivar (GenHaxl u w a -> Cont u w a
forall u w a. GenHaxl u w a -> Cont u w a
Cont (IVar u w a -> GenHaxl u w a
forall u w a. IVar u w a -> GenHaxl u w a
getIVarWithWrites IVar u w a
ivar)))

-- -----------------------------------------------------------------------------
-- 1-ary and 2-ary memo functions

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 :: GenHaxl u w (MemoVar1 u w a b)
newMemo1 = IO (MemoVar1 u w a b) -> GenHaxl u w (MemoVar1 u w a b)
forall a u w. IO a -> GenHaxl u w a
unsafeLiftIO (IO (MemoVar1 u w a b) -> GenHaxl u w (MemoVar1 u w a b))
-> IO (MemoVar1 u w a b) -> GenHaxl u w (MemoVar1 u w a b)
forall a b. (a -> b) -> a -> b
$ IORef (MemoStatus1 u w a b) -> MemoVar1 u w a b
forall u w a b. IORef (MemoStatus1 u w a b) -> MemoVar1 u w a b
MemoVar1 (IORef (MemoStatus1 u w a b) -> MemoVar1 u w a b)
-> IO (IORef (MemoStatus1 u w a b)) -> IO (MemoVar1 u w a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MemoStatus1 u w a b -> IO (IORef (MemoStatus1 u w a b))
forall a. a -> IO (IORef a)
newIORef MemoStatus1 u w a b
forall u w a b. MemoStatus1 u w a b
MemoEmpty1

newMemoWith1 :: (a -> GenHaxl u w b) -> GenHaxl u w (MemoVar1 u w a b)
newMemoWith1 :: (a -> GenHaxl u w b) -> GenHaxl u w (MemoVar1 u w a b)
newMemoWith1 a -> GenHaxl u w b
f = GenHaxl u w (MemoVar1 u w a b)
forall u w a b. GenHaxl u w (MemoVar1 u w a b)
newMemo1 GenHaxl u w (MemoVar1 u w a b)
-> (MemoVar1 u w a b -> GenHaxl u w (MemoVar1 u w a b))
-> GenHaxl u w (MemoVar1 u w a b)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \MemoVar1 u w a b
r -> MemoVar1 u w a b -> (a -> GenHaxl u w b) -> GenHaxl u w ()
forall u w a b.
MemoVar1 u w a b -> (a -> GenHaxl u w b) -> GenHaxl u w ()
prepareMemo1 MemoVar1 u w a b
r a -> GenHaxl u w b
f GenHaxl u w ()
-> GenHaxl u w (MemoVar1 u w a b) -> GenHaxl u w (MemoVar1 u w a b)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MemoVar1 u w a b -> GenHaxl u w (MemoVar1 u w a b)
forall (m :: * -> *) a. Monad m => a -> m a
return MemoVar1 u w a b
r

prepareMemo1 :: MemoVar1 u w a b -> (a -> GenHaxl u w b) -> GenHaxl u w ()
prepareMemo1 :: MemoVar1 u w a b -> (a -> GenHaxl u w b) -> GenHaxl u w ()
prepareMemo1 (MemoVar1 IORef (MemoStatus1 u w a b)
r) a -> GenHaxl u w b
f
  = IO () -> GenHaxl u w ()
forall a u w. IO a -> GenHaxl u w a
unsafeLiftIO (IO () -> GenHaxl u w ()) -> IO () -> GenHaxl u w ()
forall a b. (a -> b) -> a -> b
$ IORef (MemoStatus1 u w a b) -> MemoStatus1 u w a b -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (MemoStatus1 u w a b)
r ((a -> GenHaxl u w b)
-> HashMap a (MemoVar u w b) -> MemoStatus1 u w a b
forall u w a b.
(a -> GenHaxl u w b)
-> HashMap a (MemoVar u w b) -> MemoStatus1 u w a b
MemoTbl1 a -> GenHaxl u w b
f HashMap a (MemoVar u w b)
forall k v. HashMap k v
HashMap.empty)

runMemo1 :: (Eq a, Hashable a) => MemoVar1 u w a b -> a -> GenHaxl u w b
runMemo1 :: MemoVar1 u w a b -> a -> GenHaxl u w b
runMemo1 (MemoVar1 IORef (MemoStatus1 u w a b)
r) a
k = IO (MemoStatus1 u w a b) -> GenHaxl u w (MemoStatus1 u w a b)
forall a u w. IO a -> GenHaxl u w a
unsafeLiftIO (IORef (MemoStatus1 u w a b) -> IO (MemoStatus1 u w a b)
forall a. IORef a -> IO a
readIORef IORef (MemoStatus1 u w a b)
r) GenHaxl u w (MemoStatus1 u w a b)
-> (MemoStatus1 u w a b -> GenHaxl u w b) -> GenHaxl u w b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  MemoStatus1 u w a b
MemoEmpty1 -> CriticalError -> GenHaxl u w b
forall e u w a. Exception e => e -> GenHaxl u w a
throw (CriticalError -> GenHaxl u w b) -> CriticalError -> GenHaxl u w b
forall a b. (a -> b) -> a -> b
$ Text -> CriticalError
CriticalError Text
"Attempting to run empty memo."
  MemoTbl1 a -> GenHaxl u w b
f HashMap a (MemoVar u w b)
h -> case a -> HashMap a (MemoVar u w b) -> Maybe (MemoVar u w b)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup a
k HashMap a (MemoVar u w b)
h of
    Maybe (MemoVar u w b)
Nothing -> do
      MemoVar u w b
x <- GenHaxl u w b -> GenHaxl u w (MemoVar u w b)
forall u w a. GenHaxl u w a -> GenHaxl u w (MemoVar u w a)
newMemoWith (a -> GenHaxl u w b
f a
k)
      IO () -> GenHaxl u w ()
forall a u w. IO a -> GenHaxl u w a
unsafeLiftIO (IO () -> GenHaxl u w ()) -> IO () -> GenHaxl u w ()
forall a b. (a -> b) -> a -> b
$ IORef (MemoStatus1 u w a b) -> MemoStatus1 u w a b -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (MemoStatus1 u w a b)
r ((a -> GenHaxl u w b)
-> HashMap a (MemoVar u w b) -> MemoStatus1 u w a b
forall u w a b.
(a -> GenHaxl u w b)
-> HashMap a (MemoVar u w b) -> MemoStatus1 u w a b
MemoTbl1 a -> GenHaxl u w b
f (a
-> MemoVar u w b
-> HashMap a (MemoVar u w b)
-> HashMap a (MemoVar u w b)
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert a
k MemoVar u w b
x HashMap a (MemoVar u w b)
h))
      MemoVar u w b -> GenHaxl u w b
forall u w a. MemoVar u w a -> GenHaxl u w a
runMemo MemoVar u w b
x
    Just MemoVar u w b
v -> MemoVar u w b -> GenHaxl u w b
forall u w a. MemoVar u w a -> GenHaxl u w a
runMemo MemoVar u w b
v

newMemo2 :: GenHaxl u w (MemoVar2 u w a b c)
newMemo2 :: GenHaxl u w (MemoVar2 u w a b c)
newMemo2 = IO (MemoVar2 u w a b c) -> GenHaxl u w (MemoVar2 u w a b c)
forall a u w. IO a -> GenHaxl u w a
unsafeLiftIO (IO (MemoVar2 u w a b c) -> GenHaxl u w (MemoVar2 u w a b c))
-> IO (MemoVar2 u w a b c) -> GenHaxl u w (MemoVar2 u w a b c)
forall a b. (a -> b) -> a -> b
$ IORef (MemoStatus2 u w a b c) -> MemoVar2 u w a b c
forall u w a b c.
IORef (MemoStatus2 u w a b c) -> MemoVar2 u w a b c
MemoVar2 (IORef (MemoStatus2 u w a b c) -> MemoVar2 u w a b c)
-> IO (IORef (MemoStatus2 u w a b c)) -> IO (MemoVar2 u w a b c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MemoStatus2 u w a b c -> IO (IORef (MemoStatus2 u w a b c))
forall a. a -> IO (IORef a)
newIORef MemoStatus2 u w a b c
forall u w a b c. MemoStatus2 u w a b c
MemoEmpty2

newMemoWith2 :: (a -> b -> GenHaxl u w c) -> GenHaxl u w (MemoVar2 u w a b c)
newMemoWith2 :: (a -> b -> GenHaxl u w c) -> GenHaxl u w (MemoVar2 u w a b c)
newMemoWith2 a -> b -> GenHaxl u w c
f = GenHaxl u w (MemoVar2 u w a b c)
forall u w a b c. GenHaxl u w (MemoVar2 u w a b c)
newMemo2 GenHaxl u w (MemoVar2 u w a b c)
-> (MemoVar2 u w a b c -> GenHaxl u w (MemoVar2 u w a b c))
-> GenHaxl u w (MemoVar2 u w a b c)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \MemoVar2 u w a b c
r -> MemoVar2 u w a b c -> (a -> b -> GenHaxl u w c) -> GenHaxl u w ()
forall u w a b c.
MemoVar2 u w a b c -> (a -> b -> GenHaxl u w c) -> GenHaxl u w ()
prepareMemo2 MemoVar2 u w a b c
r a -> b -> GenHaxl u w c
f GenHaxl u w ()
-> GenHaxl u w (MemoVar2 u w a b c)
-> GenHaxl u w (MemoVar2 u w a b c)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MemoVar2 u w a b c -> GenHaxl u w (MemoVar2 u w a b c)
forall (m :: * -> *) a. Monad m => a -> m a
return MemoVar2 u w a b c
r

prepareMemo2 :: MemoVar2 u w a b c -> (a -> b -> GenHaxl u w c) -> GenHaxl u w ()
prepareMemo2 :: MemoVar2 u w a b c -> (a -> b -> GenHaxl u w c) -> GenHaxl u w ()
prepareMemo2 (MemoVar2 IORef (MemoStatus2 u w a b c)
r) a -> b -> GenHaxl u w c
f
  = IO () -> GenHaxl u w ()
forall a u w. IO a -> GenHaxl u w a
unsafeLiftIO (IO () -> GenHaxl u w ()) -> IO () -> GenHaxl u w ()
forall a b. (a -> b) -> a -> b
$ IORef (MemoStatus2 u w a b c) -> MemoStatus2 u w a b c -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (MemoStatus2 u w a b c)
r ((a -> b -> GenHaxl u w c)
-> HashMap a (HashMap b (MemoVar u w c)) -> MemoStatus2 u w a b c
forall u w a b c.
(a -> b -> GenHaxl u w c)
-> HashMap a (HashMap b (MemoVar u w c)) -> MemoStatus2 u w a b c
MemoTbl2 a -> b -> GenHaxl u w c
f HashMap a (HashMap b (MemoVar u w c))
forall k v. HashMap k v
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 u w a b c -> a -> b -> GenHaxl u w c
runMemo2 (MemoVar2 IORef (MemoStatus2 u w a b c)
r) a
k1 b
k2 = IO (MemoStatus2 u w a b c) -> GenHaxl u w (MemoStatus2 u w a b c)
forall a u w. IO a -> GenHaxl u w a
unsafeLiftIO (IORef (MemoStatus2 u w a b c) -> IO (MemoStatus2 u w a b c)
forall a. IORef a -> IO a
readIORef IORef (MemoStatus2 u w a b c)
r) GenHaxl u w (MemoStatus2 u w a b c)
-> (MemoStatus2 u w a b c -> GenHaxl u w c) -> GenHaxl u w c
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  MemoStatus2 u w a b c
MemoEmpty2 -> CriticalError -> GenHaxl u w c
forall e u w a. Exception e => e -> GenHaxl u w a
throw (CriticalError -> GenHaxl u w c) -> CriticalError -> GenHaxl u w c
forall a b. (a -> b) -> a -> b
$ Text -> CriticalError
CriticalError Text
"Attempting to run empty memo."
  MemoTbl2 a -> b -> GenHaxl u w c
f HashMap a (HashMap b (MemoVar u w c))
h1 -> case a
-> HashMap a (HashMap b (MemoVar u w c))
-> Maybe (HashMap b (MemoVar u w c))
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup a
k1 HashMap a (HashMap b (MemoVar u w c))
h1 of
    Maybe (HashMap b (MemoVar u w c))
Nothing -> do
      MemoVar u w c
v <- GenHaxl u w c -> GenHaxl u w (MemoVar u w c)
forall u w a. GenHaxl u w a -> GenHaxl u w (MemoVar u w a)
newMemoWith (a -> b -> GenHaxl u w c
f a
k1 b
k2)
      IO () -> GenHaxl u w ()
forall a u w. IO a -> GenHaxl u w a
unsafeLiftIO (IO () -> GenHaxl u w ()) -> IO () -> GenHaxl u w ()
forall a b. (a -> b) -> a -> b
$ IORef (MemoStatus2 u w a b c) -> MemoStatus2 u w a b c -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (MemoStatus2 u w a b c)
r
        ((a -> b -> GenHaxl u w c)
-> HashMap a (HashMap b (MemoVar u w c)) -> MemoStatus2 u w a b c
forall u w a b c.
(a -> b -> GenHaxl u w c)
-> HashMap a (HashMap b (MemoVar u w c)) -> MemoStatus2 u w a b c
MemoTbl2 a -> b -> GenHaxl u w c
f (a
-> HashMap b (MemoVar u w c)
-> HashMap a (HashMap b (MemoVar u w c))
-> HashMap a (HashMap b (MemoVar u w c))
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert a
k1 (b -> MemoVar u w c -> HashMap b (MemoVar u w c)
forall k v. Hashable k => k -> v -> HashMap k v
HashMap.singleton b
k2 MemoVar u w c
v) HashMap a (HashMap b (MemoVar u w c))
h1))
      MemoVar u w c -> GenHaxl u w c
forall u w a. MemoVar u w a -> GenHaxl u w a
runMemo MemoVar u w c
v
    Just HashMap b (MemoVar u w c)
h2 -> case b -> HashMap b (MemoVar u w c) -> Maybe (MemoVar u w c)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup b
k2 HashMap b (MemoVar u w c)
h2 of
      Maybe (MemoVar u w c)
Nothing -> do
        MemoVar u w c
v <- GenHaxl u w c -> GenHaxl u w (MemoVar u w c)
forall u w a. GenHaxl u w a -> GenHaxl u w (MemoVar u w a)
newMemoWith (a -> b -> GenHaxl u w c
f a
k1 b
k2)
        IO () -> GenHaxl u w ()
forall a u w. IO a -> GenHaxl u w a
unsafeLiftIO (IO () -> GenHaxl u w ()) -> IO () -> GenHaxl u w ()
forall a b. (a -> b) -> a -> b
$ IORef (MemoStatus2 u w a b c) -> MemoStatus2 u w a b c -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (MemoStatus2 u w a b c)
r
          ((a -> b -> GenHaxl u w c)
-> HashMap a (HashMap b (MemoVar u w c)) -> MemoStatus2 u w a b c
forall u w a b c.
(a -> b -> GenHaxl u w c)
-> HashMap a (HashMap b (MemoVar u w c)) -> MemoStatus2 u w a b c
MemoTbl2 a -> b -> GenHaxl u w c
f (a
-> HashMap b (MemoVar u w c)
-> HashMap a (HashMap b (MemoVar u w c))
-> HashMap a (HashMap b (MemoVar u w c))
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert a
k1 (b
-> MemoVar u w c
-> HashMap b (MemoVar u w c)
-> HashMap b (MemoVar u w c)
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert b
k2 MemoVar u w c
v HashMap b (MemoVar u w c)
h2) HashMap a (HashMap b (MemoVar u w c))
h1))
        MemoVar u w c -> GenHaxl u w c
forall u w a. MemoVar u w a -> GenHaxl u w a
runMemo MemoVar u w c
v
      Just MemoVar u w c
v -> MemoVar u w c -> GenHaxl u w c
forall u w a. MemoVar u w a -> GenHaxl u w a
runMemo MemoVar u w c
v

-- -----------------------------------------------------------------------------
-- A key type that can be used for memoizing computations by a Text key

-- | Memoize a computation using an arbitrary key.  The result will be
-- calculated once; the second and subsequent time it will be returned
-- immediately.  It is the caller's responsibility to ensure that for
-- every two calls @memo key haxl@, if they have the same @key@ then
-- they compute the same result.
memo
  :: (Typeable a, Typeable k, Hashable k, Eq k)
  => k -> GenHaxl u w a -> GenHaxl u w a
memo :: k -> GenHaxl u w a -> GenHaxl u w a
memo k
key = MemoKey k a -> GenHaxl u w a -> GenHaxl u w a
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 (k -> MemoKey k a
forall k a. (Typeable k, Hashable k, Eq k) => k -> MemoKey k a
MemoKey k
key)

{-# RULES
"memo/Text" memo = memoText :: (Typeable a) =>
            Text -> GenHaxl u w a -> GenHaxl u w a
 #-}

{-# NOINLINE memo #-}

-- | Memoize a computation using its location and a Fingerprint. This ensures
-- uniqueness across computations.
memoUnique
  :: (Typeable a, Typeable k, Hashable k, Eq k)
  => MemoFingerprintKey a -> Text -> k -> GenHaxl u w a -> GenHaxl u w a
memoUnique :: MemoFingerprintKey a -> Text -> k -> GenHaxl u w a -> GenHaxl u w a
memoUnique MemoFingerprintKey a
fp Text
label k
key = Text -> GenHaxl u w a -> GenHaxl u w a
forall u w a. Text -> GenHaxl u w a -> GenHaxl u w a
withLabel Text
label (GenHaxl u w a -> GenHaxl u w a)
-> (GenHaxl u w a -> GenHaxl u w a)
-> GenHaxl u w a
-> GenHaxl u w a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (MemoFingerprintKey a, k) -> GenHaxl u w a -> GenHaxl u w a
forall a k u w.
(Typeable a, Typeable k, Hashable k, Eq k) =>
k -> GenHaxl u w a -> GenHaxl u w a
memo (MemoFingerprintKey a
fp, k
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 :: CallId -> MemoKey k a -> CallId
hashWithSalt CallId
s (MemoKey k
t) = CallId -> k -> CallId
forall a. Hashable a => CallId -> a -> CallId
hashWithSalt CallId
s k
t

-- An optimised memo key for Text keys.  This is used automatically
-- when the key is Text, due to the RULES pragma above.

data MemoTextKey a where
  MemoText :: Text -> MemoTextKey a
  deriving Typeable

deriving instance Eq (MemoTextKey a)

instance Hashable (MemoTextKey a) where
  hashWithSalt :: CallId -> MemoTextKey a -> CallId
hashWithSalt CallId
s (MemoText Text
t) = CallId -> Text -> CallId
forall a. Hashable a => CallId -> a -> CallId
hashWithSalt CallId
s Text
t

memoText :: (Typeable a) => Text -> GenHaxl u w a -> GenHaxl u w a
memoText :: Text -> GenHaxl u w a -> GenHaxl u w a
memoText Text
key = Text -> GenHaxl u w a -> GenHaxl u w a
forall u w a. Text -> GenHaxl u w a -> GenHaxl u w a
withLabel Text
key (GenHaxl u w a -> GenHaxl u w a)
-> (GenHaxl u w a -> GenHaxl u w a)
-> GenHaxl u w a
-> GenHaxl u w a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MemoTextKey a -> GenHaxl u w a -> GenHaxl u w a
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 (Text -> MemoTextKey a
forall a. Text -> MemoTextKey a
MemoText Text
key)

-- | A memo key derived from a 128-bit MD5 hash.  Do not use this directly,
-- it is for use by automatically-generated memoization.
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 :: CallId -> MemoFingerprintKey a -> CallId
hashWithSalt CallId
s (MemoFingerprintKey Word64
x Word64
_ Addr#
_ Addr#
_) =
    CallId -> CallId -> CallId
forall a. Hashable a => CallId -> a -> CallId
hashWithSalt CallId
s (Word64 -> CallId
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
x :: Int)

-- This is optimised for cheap call sites: when we have a call
--
--   memoFingerprint (MemoFingerprintKey 1234 5678 "module"# "name"#) e
--
-- then the MemoFingerprintKey constructor will be statically
-- allocated (with two 64-bit fields and pointers to cstrings for the names),
-- and shared by all calls to memo. So the memo call will not allocate,
-- unlike memoText.
--
{-# NOINLINE memoFingerprint #-}
memoFingerprint
  :: Typeable a => MemoFingerprintKey a -> GenHaxl u w a -> GenHaxl u w a
memoFingerprint :: MemoFingerprintKey a -> GenHaxl u w a -> GenHaxl u w a
memoFingerprint key :: MemoFingerprintKey a
key@(MemoFingerprintKey Word64
_ Word64
_ Addr#
mnPtr Addr#
nPtr) =
  Addr# -> Addr# -> GenHaxl u w a -> GenHaxl u w a
forall u w a. Addr# -> Addr# -> GenHaxl u w a -> GenHaxl u w a
withFingerprintLabel Addr#
mnPtr Addr#
nPtr (GenHaxl u w a -> GenHaxl u w a)
-> (GenHaxl u w a -> GenHaxl u w a)
-> GenHaxl u w a
-> GenHaxl u w a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MemoFingerprintKey a -> GenHaxl u w a -> GenHaxl u w a
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 MemoFingerprintKey a
key

-- * Generic memoization machinery.

-- | Transform a Haxl computation into a memoized version of itself.
--
-- Given a Haxl computation, @memoize@ creates a version which stores its result
-- in a @MemoVar@ (which @memoize@ creates), and returns the stored result on
-- subsequent invocations. This permits the creation of local memos, whose
-- lifetimes are scoped to the current function, rather than the entire request.
memoize :: GenHaxl u w a -> GenHaxl u w (GenHaxl u w a)
memoize :: GenHaxl u w a -> GenHaxl u w (GenHaxl u w a)
memoize GenHaxl u w a
a = MemoVar u w a -> GenHaxl u w a
forall u w a. MemoVar u w a -> GenHaxl u w a
runMemo (MemoVar u w a -> GenHaxl u w a)
-> GenHaxl u w (MemoVar u w a) -> GenHaxl u w (GenHaxl u w a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GenHaxl u w a -> GenHaxl u w (MemoVar u w a)
forall u w a. GenHaxl u w a -> GenHaxl u w (MemoVar u w a)
newMemoWith GenHaxl u w a
a

-- | Transform a 1-argument function returning a Haxl computation into a
-- memoized version of itself.
--
-- Given a function @f@ of type @a -> GenHaxl u w b@, @memoize1@ creates a version
-- which memoizes the results of @f@ in a table keyed by its argument, and
-- returns stored results on subsequent invocations with the same argument.
--
-- e.g.:
--
-- > allFriends :: [Int] -> GenHaxl u w [Int]
-- > allFriends ids = do
-- >   memoizedFriendsOf <- memoize1 friendsOf
-- >   concat <$> mapM memoizeFriendsOf ids
--
-- The above implementation will not invoke the underlying @friendsOf@
-- repeatedly for duplicate values in @ids@.
memoize1 :: (Eq a, Hashable a)
         => (a -> GenHaxl u w b)
         -> GenHaxl u w (a -> GenHaxl u w b)
memoize1 :: (a -> GenHaxl u w b) -> GenHaxl u w (a -> GenHaxl u w b)
memoize1 a -> GenHaxl u w b
f = MemoVar1 u w a b -> a -> GenHaxl u w b
forall a u w b.
(Eq a, Hashable a) =>
MemoVar1 u w a b -> a -> GenHaxl u w b
runMemo1 (MemoVar1 u w a b -> a -> GenHaxl u w b)
-> GenHaxl u w (MemoVar1 u w a b)
-> GenHaxl u w (a -> GenHaxl u w b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> GenHaxl u w b) -> GenHaxl u w (MemoVar1 u w a b)
forall a u w b.
(a -> GenHaxl u w b) -> GenHaxl u w (MemoVar1 u w a b)
newMemoWith1 a -> GenHaxl u w b
f

-- | Transform a 2-argument function returning a Haxl computation, into a
-- memoized version of itself.
--
-- The 2-ary version of @memoize1@, see its documentation for details.
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 :: (a -> b -> GenHaxl u w c) -> GenHaxl u w (a -> b -> GenHaxl u w c)
memoize2 a -> b -> GenHaxl u w c
f = MemoVar2 u w a b c -> a -> b -> GenHaxl u w c
forall a b u w c.
(Eq a, Hashable a, Eq b, Hashable b) =>
MemoVar2 u w a b c -> a -> b -> GenHaxl u w c
runMemo2 (MemoVar2 u w a b c -> a -> b -> GenHaxl u w c)
-> GenHaxl u w (MemoVar2 u w a b c)
-> GenHaxl u w (a -> b -> GenHaxl u w c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> b -> GenHaxl u w c) -> GenHaxl u w (MemoVar2 u w a b c)
forall a b u w c.
(a -> b -> GenHaxl u w c) -> GenHaxl u w (MemoVar2 u w a b c)
newMemoWith2 a -> b -> GenHaxl u w c
f