{-# LANGUAGE BangPatterns #-}
{-|
Module      : Data.LruCache.IO.Finalizer
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 
when finalizers have to be run when cache entries are evicted.
-}
module Data.LruCache.IO.Finalizer
  ( LruHandle(..)
  , newLruHandle
  , cached
  , StripedLruHandle(..)
  , newStripedLruHandle
  , stripedCached
  ) where

import           Control.Applicative ((<$>))
import           Data.Foldable (traverse_)
import           Data.Hashable (Hashable, hash)
import           Data.IORef (IORef, atomicModifyIORef', newIORef)
import           Data.Tuple (swap)
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, v -> IO ())))

-- | 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, v -> IO ())) -> LruHandle k v
forall k v. IORef (LruCache k (v, v -> IO ())) -> LruHandle k v
LruHandle (IORef (LruCache k (v, v -> IO ())) -> LruHandle k v)
-> IO (IORef (LruCache k (v, v -> IO ()))) -> IO (LruHandle k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> LruCache k (v, v -> IO ())
-> IO (IORef (LruCache k (v, v -> IO ())))
forall a. a -> IO (IORef a)
newIORef (Int -> LruCache k (v, v -> IO ())
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 ->
  (v -> IO ()) {- ^ finalizer -} ->
  IO v
cached :: forall k v.
(Hashable k, Ord k) =>
LruHandle k v -> k -> IO v -> (v -> IO ()) -> IO v
cached (LruHandle IORef (LruCache k (v, v -> IO ()))
ref) k
k IO v
io v -> IO ()
finalizer =
  do Maybe (v, v -> IO ())
lookupRes <- IORef (LruCache k (v, v -> IO ()))
-> (LruCache k (v, v -> IO ())
    -> (LruCache k (v, v -> IO ()), Maybe (v, v -> IO ())))
-> IO (Maybe (v, v -> IO ()))
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (LruCache k (v, v -> IO ()))
ref ((LruCache k (v, v -> IO ())
  -> (LruCache k (v, v -> IO ()), Maybe (v, v -> IO ())))
 -> IO (Maybe (v, v -> IO ())))
-> (LruCache k (v, v -> IO ())
    -> (LruCache k (v, v -> IO ()), Maybe (v, v -> IO ())))
-> IO (Maybe (v, v -> IO ()))
forall a b. (a -> b) -> a -> b
$ \LruCache k (v, v -> IO ())
c ->
       case k
-> LruCache k (v, v -> IO ())
-> Maybe ((v, v -> IO ()), LruCache k (v, v -> IO ()))
forall k v.
(Hashable k, Ord k) =>
k -> LruCache k v -> Maybe (v, LruCache k v)
lookup k
k LruCache k (v, v -> IO ())
c of
         Maybe ((v, v -> IO ()), LruCache k (v, v -> IO ()))
Nothing      -> (LruCache k (v, v -> IO ())
c,  Maybe (v, v -> IO ())
forall a. Maybe a
Nothing)
         Just ((v, v -> IO ())
v, LruCache k (v, v -> IO ())
c') -> (LruCache k (v, v -> IO ())
c', (v, v -> IO ()) -> Maybe (v, v -> IO ())
forall a. a -> Maybe a
Just (v, v -> IO ())
v)
     case Maybe (v, v -> IO ())
lookupRes of
       Just (!v
v,v -> IO ()
_)  -> v -> IO v
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return v
v
       Maybe (v, v -> IO ())
Nothing      ->
         do v
v <- IO v
io
            Maybe (k, (v, v -> IO ()))
evicted <- IORef (LruCache k (v, v -> IO ()))
-> (LruCache k (v, v -> IO ())
    -> (LruCache k (v, v -> IO ()), Maybe (k, (v, v -> IO ()))))
-> IO (Maybe (k, (v, v -> IO ())))
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (LruCache k (v, v -> IO ()))
ref ((LruCache k (v, v -> IO ())
  -> (LruCache k (v, v -> IO ()), Maybe (k, (v, v -> IO ()))))
 -> IO (Maybe (k, (v, v -> IO ()))))
-> (LruCache k (v, v -> IO ())
    -> (LruCache k (v, v -> IO ()), Maybe (k, (v, v -> IO ()))))
-> IO (Maybe (k, (v, v -> IO ())))
forall a b. (a -> b) -> a -> b
$ \LruCache k (v, v -> IO ())
c -> 
              (Maybe (k, (v, v -> IO ())), LruCache k (v, v -> IO ()))
-> (LruCache k (v, v -> IO ()), Maybe (k, (v, v -> IO ())))
forall a b. (a, b) -> (b, a)
swap (k
-> (v, v -> IO ())
-> LruCache k (v, v -> IO ())
-> (Maybe (k, (v, v -> IO ())), LruCache k (v, v -> IO ()))
forall k v.
(Hashable k, Ord k) =>
k -> v -> LruCache k v -> (Maybe (k, v), LruCache k v)
insertView k
k (v
v,v -> IO ()
finalizer) LruCache k (v, v -> IO ())
c)
            ((k, (v, v -> IO ())) -> IO ())
-> Maybe (k, (v, v -> IO ())) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (\(k
_,(v
v',v -> IO ()
finalize')) -> v -> IO ()
finalize' v
v') Maybe (k, (v, v -> IO ()))
evicted
            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 ->
  (v -> IO ()) {- ^ finalizer -} ->
  IO v
stripedCached :: forall k v.
(Hashable k, Ord k) =>
StripedLruHandle k v -> k -> IO v -> (v -> IO ()) -> IO v
stripedCached (StripedLruHandle Vector (LruHandle k v)
v) k
k =
    LruHandle k v -> k -> IO v -> (v -> IO ()) -> IO v
forall k v.
(Hashable k, Ord k) =>
LruHandle k v -> k -> IO v -> (v -> IO ()) -> 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