-- 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 ExistentialQuantification #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

-- |
-- A cache mapping data requests to their results.  This module is
-- provided for access to Haxl internals only; most users should not
-- need to import it.
--
module Haxl.Core.DataCache
  ( DataCache(..)
  , SubCache(..)
  , emptyDataCache
  , filter
  , insert
  , insertNotShowable
  , insertWithShow
  , lookup
  , showCache
  , readCache
  ) where

import Prelude hiding (lookup, filter)
import Control.Exception
import Unsafe.Coerce
import Data.Typeable
import Data.Hashable
import qualified Data.HashTable.IO as H

-- ---------------------------------------------------------------------------
-- DataCache

-- | A @'DataCache' res@ maps things of type @req a@ to @res a@, for
-- any @req@ and @a@ provided @req a@ is an instance of 'Typeable'. In
-- practice @req a@ will be a request type parameterised by its result.
--
newtype DataCache res = DataCache (HashTable TypeRep (SubCache res))

-- | The implementation is a two-level map: the outer level maps the
-- types of requests to 'SubCache', which maps actual requests to their
-- results.  So each 'SubCache' contains requests of the same type.
-- This works well because we only have to store the dictionaries for
-- 'Hashable' and 'Eq' once per request type.
--
data SubCache res =
  forall req a . (Hashable (req a), Eq (req a)) =>
       SubCache (req a -> String) (a -> String) !(HashTable (req a) (res a))
       -- NB. the inner HashMap is strict, to avoid building up
       -- a chain of thunks during repeated insertions.

type HashTable k v = H.BasicHashTable k v

-- | A new, empty 'DataCache'.
emptyDataCache :: IO (DataCache res)
emptyDataCache :: IO (DataCache res)
emptyDataCache = HashTable RealWorld SomeTypeRep (SubCache res) -> DataCache res
forall (res :: * -> *).
HashTable SomeTypeRep (SubCache res) -> DataCache res
DataCache (HashTable RealWorld SomeTypeRep (SubCache res) -> DataCache res)
-> IO (HashTable RealWorld SomeTypeRep (SubCache res))
-> IO (DataCache res)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (HashTable RealWorld SomeTypeRep (SubCache res))
forall (h :: * -> * -> * -> *) k v.
HashTable h =>
IO (IOHashTable h k v)
H.new

-- | Inserts a request-result pair into the 'DataCache'.
insert
  :: (Hashable (req a), Typeable (req a), Eq (req a), Show (req a), Show a)
  => req a
  -- ^ Request
  -> res a
  -- ^ Result
  -> DataCache res
  -> IO ()

insert :: req a -> res a -> DataCache res -> IO ()
insert = (req a -> String)
-> (a -> String) -> req a -> res a -> DataCache res -> IO ()
forall (req :: * -> *) a (res :: * -> *).
(Hashable (req a), Typeable (req a), Eq (req a)) =>
(req a -> String)
-> (a -> String) -> req a -> res a -> DataCache res -> IO ()
insertWithShow req a -> String
forall a. Show a => a -> String
show a -> String
forall a. Show a => a -> String
show

-- | Inserts a request-result pair into the 'DataCache', without
-- requiring Show instances of the request or the result.  The cache
-- cannot be subsequently used with `showCache`.
insertNotShowable
  :: (Hashable (req a), Typeable (req a), Eq (req a))
  => req a
  -- ^ Request
  -> res a
  -- ^ Result
  -> DataCache res
  -> IO ()

insertNotShowable :: req a -> res a -> DataCache res -> IO ()
insertNotShowable = (req a -> String)
-> (a -> String) -> req a -> res a -> DataCache res -> IO ()
forall (req :: * -> *) a (res :: * -> *).
(Hashable (req a), Typeable (req a), Eq (req a)) =>
(req a -> String)
-> (a -> String) -> req a -> res a -> DataCache res -> IO ()
insertWithShow req a -> String
forall a. a
notShowable a -> String
forall a. a
notShowable

-- | Inserts a request-result pair into the 'DataCache', with the given
-- functions used to show the request and result.
insertWithShow
  :: (Hashable (req a), Typeable (req a), Eq (req a))
  => (req a -> String)
  -- ^ Show function for request
  -> (a -> String)
  -- ^ Show function for result
  -> req a
  -- ^ Request
  -> res a
  -- ^ Result
  -> DataCache res
  -> IO ()

insertWithShow :: (req a -> String)
-> (a -> String) -> req a -> res a -> DataCache res -> IO ()
insertWithShow req a -> String
showRequest a -> String
showResult req a
request res a
result (DataCache HashTable SomeTypeRep (SubCache res)
m) =
  HashTable SomeTypeRep (SubCache res)
-> SomeTypeRep
-> (Maybe (SubCache res) -> IO (Maybe (SubCache res), ()))
-> IO ()
forall (h :: * -> * -> * -> *) k v a.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> (Maybe v -> IO (Maybe v, a)) -> IO a
H.mutateIO HashTable SomeTypeRep (SubCache res)
m (req a -> SomeTypeRep
forall a. Typeable a => a -> SomeTypeRep
typeOf req a
request) ((req a -> String)
-> (a -> String)
-> req a
-> res a
-> Maybe (SubCache res)
-> IO (Maybe (SubCache res), ())
forall (req :: * -> *) a (res :: * -> *).
(Hashable (req a), Typeable (req a), Eq (req a)) =>
(req a -> String)
-> (a -> String)
-> req a
-> res a
-> Maybe (SubCache res)
-> IO (Maybe (SubCache res), ())
mutate req a -> String
showRequest a -> String
showResult req a
request res a
result)

notShowable :: a
notShowable :: a
notShowable = String -> a
forall a. HasCallStack => String -> a
error String
"insertNotShowable"

-- | A mutation function for mutateIO. If the key doesn't exist in the top-level
-- cache, creates a new hashtable and inserts the request and result.
-- If the key exists, insert the request and result into the existing subcache,
-- replacing any existing mapping.
mutate :: (Hashable (req a), Typeable (req a), Eq (req a))
  => (req a -> String)
  -> (a -> String)
  -> req a
  -> res a
  -> Maybe (SubCache res)
  -> IO (Maybe (SubCache res), ())
mutate :: (req a -> String)
-> (a -> String)
-> req a
-> res a
-> Maybe (SubCache res)
-> IO (Maybe (SubCache res), ())
mutate req a -> String
showRequest a -> String
showResult req a
request res a
result Maybe (SubCache res)
Nothing = do
  HashTable RealWorld (req a) (res a)
newTable <- IO (HashTable RealWorld (req a) (res a))
forall (h :: * -> * -> * -> *) k v.
HashTable h =>
IO (IOHashTable h k v)
H.new
  IOHashTable HashTable (req a) (res a) -> req a -> res a -> IO ()
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> v -> IO ()
H.insert HashTable RealWorld (req a) (res a)
IOHashTable HashTable (req a) (res a)
newTable req a
request res a
result
  (Maybe (SubCache res), ()) -> IO (Maybe (SubCache res), ())
forall (m :: * -> *) a. Monad m => a -> m a
return (SubCache res -> Maybe (SubCache res)
forall a. a -> Maybe a
Just ((req a -> String)
-> (a -> String)
-> IOHashTable HashTable (req a) (res a)
-> SubCache res
forall (res :: * -> *) (req :: * -> *) a.
(Hashable (req a), Eq (req a)) =>
(req a -> String)
-> (a -> String) -> HashTable (req a) (res a) -> SubCache res
SubCache req a -> String
showRequest a -> String
showResult HashTable RealWorld (req a) (res a)
IOHashTable HashTable (req a) (res a)
newTable), ())
mutate req a -> String
_ a -> String
_ req a
request res a
result (Just sc :: SubCache res
sc@(SubCache req a -> String
_ a -> String
_ HashTable (req a) (res a)
oldTable)) = do
    HashTable (req a) (res a) -> req a -> res a -> IO ()
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> v -> IO ()
H.insert HashTable (req a) (res a)
oldTable (req a -> req a
forall a b. a -> b
unsafeCoerce req a
request) (res a -> res a
forall a b. a -> b
unsafeCoerce res a
result)
    (Maybe (SubCache res), ()) -> IO (Maybe (SubCache res), ())
forall (m :: * -> *) a. Monad m => a -> m a
return (SubCache res -> Maybe (SubCache res)
forall a. a -> Maybe a
Just SubCache res
sc, ())

-- | Looks up the cached result of a request.
lookup
  :: Typeable (req a)
  => req a
  -- ^ Request
  -> DataCache res
  -> IO (Maybe (res a))

lookup :: req a -> DataCache res -> IO (Maybe (res a))
lookup req a
req (DataCache HashTable SomeTypeRep (SubCache res)
m) = do
  Maybe (SubCache res)
mbRes <- HashTable SomeTypeRep (SubCache res)
-> SomeTypeRep -> IO (Maybe (SubCache res))
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> IO (Maybe v)
H.lookup HashTable SomeTypeRep (SubCache res)
m (req a -> SomeTypeRep
forall a. Typeable a => a -> SomeTypeRep
typeOf req a
req)
  case Maybe (SubCache res)
mbRes of
    Maybe (SubCache res)
Nothing -> Maybe (res a) -> IO (Maybe (res a))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (res a)
forall a. Maybe a
Nothing
    Just (SubCache req a -> String
_ a -> String
_ HashTable (req a) (res a)
sc) ->
      IO (Maybe (res a)) -> IO (Maybe (res a))
forall a b. a -> b
unsafeCoerce (HashTable (req a) (res a) -> req a -> IO (Maybe (res a))
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> IO (Maybe v)
H.lookup HashTable (req a) (res a)
sc (req a -> req a
forall a b. a -> b
unsafeCoerce req a
req))

filter
  :: forall res
  . (forall a. res a -> IO Bool)
  -> DataCache res
  -> IO (DataCache res)
filter :: (forall a. res a -> IO Bool) -> DataCache res -> IO (DataCache res)
filter forall a. res a -> IO Bool
pred (DataCache HashTable SomeTypeRep (SubCache res)
cache) = do
  [(SomeTypeRep, SubCache res)]
cacheList <- HashTable SomeTypeRep (SubCache res)
-> IO [(SomeTypeRep, SubCache res)]
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> IO [(k, v)]
H.toList HashTable SomeTypeRep (SubCache res)
cache
  [(SomeTypeRep, SubCache res)]
filteredCache <- (SomeTypeRep, SubCache res) -> IO (SomeTypeRep, SubCache res)
filterSubCache ((SomeTypeRep, SubCache res) -> IO (SomeTypeRep, SubCache res))
-> [(SomeTypeRep, SubCache res)]
-> IO [(SomeTypeRep, SubCache res)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
`mapM` [(SomeTypeRep, SubCache res)]
cacheList
  HashTable RealWorld SomeTypeRep (SubCache res) -> DataCache res
forall (res :: * -> *).
HashTable SomeTypeRep (SubCache res) -> DataCache res
DataCache (HashTable RealWorld SomeTypeRep (SubCache res) -> DataCache res)
-> IO (HashTable RealWorld SomeTypeRep (SubCache res))
-> IO (DataCache res)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(SomeTypeRep, SubCache res)]
-> IO (HashTable SomeTypeRep (SubCache res))
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
[(k, v)] -> IO (IOHashTable h k v)
H.fromList [(SomeTypeRep, SubCache res)]
filteredCache
  where
    filterSubCache
      :: (TypeRep, SubCache res)
      -> IO (TypeRep, SubCache res)
    filterSubCache :: (SomeTypeRep, SubCache res) -> IO (SomeTypeRep, SubCache res)
filterSubCache (SomeTypeRep
ty, SubCache req a -> String
showReq a -> String
showRes HashTable (req a) (res a)
hm) = do
      [(req a, res a)]
filteredList <- ([(req a, res a)] -> (req a, res a) -> IO [(req a, res a)])
-> [(req a, res a)]
-> HashTable (req a) (res a)
-> IO [(req a, res a)]
forall (h :: * -> * -> * -> *) a k v.
HashTable h =>
(a -> (k, v) -> IO a) -> a -> IOHashTable h k v -> IO a
H.foldM [(req a, res a)] -> (req a, res a) -> IO [(req a, res a)]
go [] HashTable (req a) (res a)
hm
      HashTable RealWorld (req a) (res a)
filteredSC <- [(req a, res a)] -> IO (HashTable (req a) (res a))
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
[(k, v)] -> IO (IOHashTable h k v)
H.fromList [(req a, res a)]
filteredList
      (SomeTypeRep, SubCache res) -> IO (SomeTypeRep, SubCache res)
forall (m :: * -> *) a. Monad m => a -> m a
return (SomeTypeRep
ty, (req a -> String)
-> (a -> String) -> HashTable (req a) (res a) -> SubCache res
forall (res :: * -> *) (req :: * -> *) a.
(Hashable (req a), Eq (req a)) =>
(req a -> String)
-> (a -> String) -> HashTable (req a) (res a) -> SubCache res
SubCache req a -> String
showReq a -> String
showRes HashTable RealWorld (req a) (res a)
HashTable (req a) (res a)
filteredSC)
      where
        go :: [(req a, res a)] -> (req a, res a) -> IO [(req a, res a)]
go [(req a, res a)]
res (req a
request, res a
rvar) = do
          Bool
predRes <- res a -> IO Bool
forall a. res a -> IO Bool
pred res a
rvar
          [(req a, res a)] -> IO [(req a, res a)]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(req a, res a)] -> IO [(req a, res a)])
-> [(req a, res a)] -> IO [(req a, res a)]
forall a b. (a -> b) -> a -> b
$ if Bool
predRes then (req a
request, res a
rvar)(req a, res a) -> [(req a, res a)] -> [(req a, res a)]
forall a. a -> [a] -> [a]
:[(req a, res a)]
res else [(req a, res a)]
res

-- | Dumps the contents of the cache, with requests and responses
-- converted to 'String's using the supplied show functions.  The
-- entries are grouped by 'TypeRep'.  Note that this will fail if
-- 'insertNotShowable' has been used to insert any entries.
showCache
  :: forall res
  .  DataCache res
  -> (forall a . res a -> IO (Maybe (Either SomeException a)))
  -> IO [(TypeRep, [(String, Either SomeException String)])]
showCache :: DataCache res
-> (forall a. res a -> IO (Maybe (Either SomeException a)))
-> IO [(SomeTypeRep, [(String, Either SomeException String)])]
showCache (DataCache HashTable SomeTypeRep (SubCache res)
cache) forall a. res a -> IO (Maybe (Either SomeException a))
readRes = ([(SomeTypeRep, [(String, Either SomeException String)])]
 -> (SomeTypeRep, SubCache res)
 -> IO [(SomeTypeRep, [(String, Either SomeException String)])])
-> [(SomeTypeRep, [(String, Either SomeException String)])]
-> HashTable SomeTypeRep (SubCache res)
-> IO [(SomeTypeRep, [(String, Either SomeException String)])]
forall (h :: * -> * -> * -> *) a k v.
HashTable h =>
(a -> (k, v) -> IO a) -> a -> IOHashTable h k v -> IO a
H.foldM [(SomeTypeRep, [(String, Either SomeException String)])]
-> (SomeTypeRep, SubCache res)
-> IO [(SomeTypeRep, [(String, Either SomeException String)])]
goSubCache [] HashTable SomeTypeRep (SubCache res)
cache
  where
    goSubCache
      :: [(TypeRep, [(String, Either SomeException String)])]
      -> (TypeRep, SubCache res)
      -> IO [(TypeRep, [(String, Either SomeException String)])]
    goSubCache :: [(SomeTypeRep, [(String, Either SomeException String)])]
-> (SomeTypeRep, SubCache res)
-> IO [(SomeTypeRep, [(String, Either SomeException String)])]
goSubCache [(SomeTypeRep, [(String, Either SomeException String)])]
res (SomeTypeRep
ty, SubCache req a -> String
showReq a -> String
showRes HashTable (req a) (res a)
hm) = do
      [(String, Either SomeException String)]
subCacheResult <- ([(String, Either SomeException String)]
 -> (req a, res a) -> IO [(String, Either SomeException String)])
-> [(String, Either SomeException String)]
-> HashTable (req a) (res a)
-> IO [(String, Either SomeException String)]
forall (h :: * -> * -> * -> *) a k v.
HashTable h =>
(a -> (k, v) -> IO a) -> a -> IOHashTable h k v -> IO a
H.foldM [(String, Either SomeException String)]
-> (req a, res a) -> IO [(String, Either SomeException String)]
go [] HashTable (req a) (res a)
hm
      [(SomeTypeRep, [(String, Either SomeException String)])]
-> IO [(SomeTypeRep, [(String, Either SomeException String)])]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(SomeTypeRep, [(String, Either SomeException String)])]
 -> IO [(SomeTypeRep, [(String, Either SomeException String)])])
-> [(SomeTypeRep, [(String, Either SomeException String)])]
-> IO [(SomeTypeRep, [(String, Either SomeException String)])]
forall a b. (a -> b) -> a -> b
$ (SomeTypeRep
ty, [(String, Either SomeException String)]
subCacheResult)(SomeTypeRep, [(String, Either SomeException String)])
-> [(SomeTypeRep, [(String, Either SomeException String)])]
-> [(SomeTypeRep, [(String, Either SomeException String)])]
forall a. a -> [a] -> [a]
:[(SomeTypeRep, [(String, Either SomeException String)])]
res
      where
        go :: [(String, Either SomeException String)]
-> (req a, res a) -> IO [(String, Either SomeException String)]
go [(String, Either SomeException String)]
res (req a
request, res a
rvar) = do
          Maybe (Either SomeException a)
maybe_r <- res a -> IO (Maybe (Either SomeException a))
forall a. res a -> IO (Maybe (Either SomeException a))
readRes res a
rvar
          [(String, Either SomeException String)]
-> IO [(String, Either SomeException String)]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(String, Either SomeException String)]
 -> IO [(String, Either SomeException String)])
-> [(String, Either SomeException String)]
-> IO [(String, Either SomeException String)]
forall a b. (a -> b) -> a -> b
$ case Maybe (Either SomeException a)
maybe_r of
            Maybe (Either SomeException a)
Nothing -> [(String, Either SomeException String)]
res
            Just (Left SomeException
e) -> (req a -> String
showReq req a
request, SomeException -> Either SomeException String
forall a b. a -> Either a b
Left SomeException
e) (String, Either SomeException String)
-> [(String, Either SomeException String)]
-> [(String, Either SomeException String)]
forall a. a -> [a] -> [a]
: [(String, Either SomeException String)]
res
            Just (Right a
result) ->
              (req a -> String
showReq req a
request, String -> Either SomeException String
forall a b. b -> Either a b
Right (a -> String
showRes a
result)) (String, Either SomeException String)
-> [(String, Either SomeException String)]
-> [(String, Either SomeException String)]
forall a. a -> [a] -> [a]
: [(String, Either SomeException String)]
res

-- | Dumps the contents of the cache responses to list
readCache
  :: forall res ret
  .  DataCache res
  -> (forall a . res a -> IO ret)
  -> IO [(TypeRep, [Either SomeException ret])]
readCache :: DataCache res
-> (forall a. res a -> IO ret)
-> IO [(SomeTypeRep, [Either SomeException ret])]
readCache (DataCache HashTable SomeTypeRep (SubCache res)
cache) forall a. res a -> IO ret
readRes = ([(SomeTypeRep, [Either SomeException ret])]
 -> (SomeTypeRep, SubCache res)
 -> IO [(SomeTypeRep, [Either SomeException ret])])
-> [(SomeTypeRep, [Either SomeException ret])]
-> HashTable SomeTypeRep (SubCache res)
-> IO [(SomeTypeRep, [Either SomeException ret])]
forall (h :: * -> * -> * -> *) a k v.
HashTable h =>
(a -> (k, v) -> IO a) -> a -> IOHashTable h k v -> IO a
H.foldM [(SomeTypeRep, [Either SomeException ret])]
-> (SomeTypeRep, SubCache res)
-> IO [(SomeTypeRep, [Either SomeException ret])]
goSubCache [] HashTable SomeTypeRep (SubCache res)
cache
  where
    goSubCache
      :: [(TypeRep, [Either SomeException ret])]
      -> (TypeRep, SubCache res)
      -> IO [(TypeRep, [Either SomeException ret])]
    goSubCache :: [(SomeTypeRep, [Either SomeException ret])]
-> (SomeTypeRep, SubCache res)
-> IO [(SomeTypeRep, [Either SomeException ret])]
goSubCache [(SomeTypeRep, [Either SomeException ret])]
res (SomeTypeRep
ty, SubCache req a -> String
_showReq a -> String
_showRes HashTable (req a) (res a)
hm) = do
      [Either SomeException ret]
subCacheResult <- ([Either SomeException ret]
 -> (req a, res a) -> IO [Either SomeException ret])
-> [Either SomeException ret]
-> HashTable (req a) (res a)
-> IO [Either SomeException ret]
forall (h :: * -> * -> * -> *) a k v.
HashTable h =>
(a -> (k, v) -> IO a) -> a -> IOHashTable h k v -> IO a
H.foldM [Either SomeException ret]
-> (req a, res a) -> IO [Either SomeException ret]
go [] HashTable (req a) (res a)
hm
      [(SomeTypeRep, [Either SomeException ret])]
-> IO [(SomeTypeRep, [Either SomeException ret])]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(SomeTypeRep, [Either SomeException ret])]
 -> IO [(SomeTypeRep, [Either SomeException ret])])
-> [(SomeTypeRep, [Either SomeException ret])]
-> IO [(SomeTypeRep, [Either SomeException ret])]
forall a b. (a -> b) -> a -> b
$ (SomeTypeRep
ty, [Either SomeException ret]
subCacheResult)(SomeTypeRep, [Either SomeException ret])
-> [(SomeTypeRep, [Either SomeException ret])]
-> [(SomeTypeRep, [Either SomeException ret])]
forall a. a -> [a] -> [a]
:[(SomeTypeRep, [Either SomeException ret])]
res
      where
        go :: [Either SomeException ret]
-> (req a, res a) -> IO [Either SomeException ret]
go [Either SomeException ret]
res (req a
_request, res a
rvar) = do
          Either SomeException ret
r <- IO ret -> IO (Either SomeException ret)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO ret -> IO (Either SomeException ret))
-> IO ret -> IO (Either SomeException ret)
forall a b. (a -> b) -> a -> b
$ res a -> IO ret
forall a. res a -> IO ret
readRes res a
rvar
          [Either SomeException ret] -> IO [Either SomeException ret]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Either SomeException ret] -> IO [Either SomeException ret])
-> [Either SomeException ret] -> IO [Either SomeException ret]
forall a b. (a -> b) -> a -> b
$ Either SomeException ret
r Either SomeException ret
-> [Either SomeException ret] -> [Either SomeException ret]
forall a. a -> [a] -> [a]
: [Either SomeException ret]
res