{-|
Module      : Data.LruCache.IO
Copyright   : (c) Moritz Kiefer, 2016
              (c) Jasper Van der Jeugt, 2015
License     : BSD3
Maintainer  : moritz.kiefer@purelyfunctional.org
Convenience module for the common case of caching results of IO actions.
See 'Data.LruCache.IO.Finalizer' if you want to run finalizers
automatically when cache entries are evicted
-}
module Data.LruCache.IO
  ( LruHandle(..)
  , cached
  , newLruHandle
  , StripedLruHandle(..)
  , stripedCached
  , newStripedLruHandle
  ) where

import           Control.Applicative ((<$>))
import           Data.Hashable (Hashable, hash)
import           Data.IORef (IORef, atomicModifyIORef', newIORef)
import           Data.Vector (Vector)
import qualified Data.Vector as Vector
import           Prelude hiding (lookup)

import           Data.LruCache

-- | Store a LRU cache in an 'IORef to be able to conveniently update it.
newtype LruHandle k v = LruHandle (IORef (LruCache k v))

-- | Create a new LRU cache of the given size.
newLruHandle :: Int -> IO (LruHandle k v)
newLruHandle :: forall k v. Int -> IO (LruHandle k v)
newLruHandle Int
capacity = IORef (LruCache k v) -> LruHandle k v
forall k v. IORef (LruCache k v) -> LruHandle k v
LruHandle (IORef (LruCache k v) -> LruHandle k v)
-> IO (IORef (LruCache k v)) -> IO (LruHandle k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> LruCache k v -> IO (IORef (LruCache k v))
forall a. a -> IO (IORef a)
newIORef (Int -> LruCache k v
forall k v. Int -> LruCache k v
empty Int
capacity)

-- | Return the cached result of the action or, in the case of a cache
-- miss, execute the action and insert it in the cache.
cached :: (Hashable k, Ord k) => LruHandle k v -> k -> IO v -> IO v
cached :: forall k v.
(Hashable k, Ord k) =>
LruHandle k v -> k -> IO v -> IO v
cached (LruHandle IORef (LruCache k v)
ref) k
k IO v
io =
  do Maybe v
lookupRes <- IORef (LruCache k v)
-> (LruCache k v -> (LruCache k v, Maybe v)) -> IO (Maybe v)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (LruCache k v)
ref ((LruCache k v -> (LruCache k v, Maybe v)) -> IO (Maybe v))
-> (LruCache k v -> (LruCache k v, Maybe v)) -> IO (Maybe v)
forall a b. (a -> b) -> a -> b
$ \LruCache k v
c ->
       case k -> LruCache k v -> Maybe (v, LruCache k v)
forall k v.
(Hashable k, Ord k) =>
k -> LruCache k v -> Maybe (v, LruCache k v)
lookup k
k LruCache k v
c of
         Maybe (v, LruCache k v)
Nothing      -> (LruCache k v
c,  Maybe v
forall a. Maybe a
Nothing)
         Just (v
v, LruCache k v
c') -> (LruCache k v
c', v -> Maybe v
forall a. a -> Maybe a
Just v
v)
     case Maybe v
lookupRes of
       Just v
v  -> v -> IO v
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return v
v
       Maybe v
Nothing ->
         do v
v <- IO v
io
            IORef (LruCache k v)
-> (LruCache k v -> (LruCache k v, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (LruCache k v)
ref ((LruCache k v -> (LruCache k v, ())) -> IO ())
-> (LruCache k v -> (LruCache k v, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \LruCache k v
c -> (k -> v -> LruCache k v -> LruCache k v
forall k v.
(Hashable k, Ord k) =>
k -> v -> LruCache k v -> LruCache k v
insert k
k v
v LruCache k v
c, ())
            v -> IO v
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return v
v

-- | Using a stripe of multiple handles can improve the performance in
-- the case of concurrent accesses since several handles can be
-- accessed in parallel.
newtype StripedLruHandle k v = StripedLruHandle (Vector (LruHandle k v))

-- | Create a new 'StripedHandle' with the given number of stripes and
-- the given capacity for each stripe.
newStripedLruHandle :: Int -> Int -> IO (StripedLruHandle k v)
newStripedLruHandle :: forall k v. Int -> Int -> IO (StripedLruHandle k v)
newStripedLruHandle Int
numStripes Int
capacityPerStripe =
  Vector (LruHandle k v) -> StripedLruHandle k v
forall k v. Vector (LruHandle k v) -> StripedLruHandle k v
StripedLruHandle (Vector (LruHandle k v) -> StripedLruHandle k v)
-> IO (Vector (LruHandle k v)) -> IO (StripedLruHandle k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO (LruHandle k v) -> IO (Vector (LruHandle k v))
forall (m :: * -> *) a. Monad m => Int -> m a -> m (Vector a)
Vector.replicateM Int
numStripes (Int -> IO (LruHandle k v)
forall k v. Int -> IO (LruHandle k v)
newLruHandle Int
capacityPerStripe)

-- | Striped version of 'cached'.
stripedCached ::
  (Hashable k, Ord k) =>
  StripedLruHandle k v ->
  k ->
  IO v ->
  IO v
stripedCached :: forall k v.
(Hashable k, Ord k) =>
StripedLruHandle k v -> k -> IO v -> IO v
stripedCached (StripedLruHandle Vector (LruHandle k v)
v) k
k =
    LruHandle k v -> k -> IO v -> IO v
forall k v.
(Hashable k, Ord k) =>
LruHandle k v -> k -> IO v -> IO v
cached (Vector (LruHandle k v)
v Vector (LruHandle k v) -> Int -> LruHandle k v
forall a. Vector a -> Int -> a
Vector.! Int
idx) k
k
  where
    idx :: Int
idx = k -> Int
forall a. Hashable a => a -> Int
hash k
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Vector (LruHandle k v) -> Int
forall a. Vector a -> Int
Vector.length Vector (LruHandle k v)
v