{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-|
Module      : Data.LruCache
Copyright   : (c) Moritz Kiefer, 2016
              (c) Jasper Van der Jeugt, 2015
License     : BSD3
Maintainer  : moritz.kiefer@purelyfunctional.org
Pure API to an LRU cache.
-}
module Data.LruCache
  ( LruCache
  , empty
  , insert
  , insertView
  , lookup
  ) where

import qualified Data.HashPSQ as HashPSQ
import           Data.Hashable (Hashable)
import           Data.List.Compat (sortOn)
import           Data.Maybe (isNothing)
import           Prelude hiding (lookup)

import           Data.LruCache.Internal

-- | Create an empty 'LruCache' of the given size.
empty :: Int -> LruCache k v
empty :: forall k v. Int -> LruCache k v
empty Int
capacity
  | Int
capacity Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = [Char] -> LruCache k v
forall a. HasCallStack => [Char] -> a
error [Char]
"LruCache.empty: capacity < 1"
  | Bool
otherwise    =
      LruCache 
        { lruCapacity :: Int
lruCapacity = Int
capacity
        , lruSize :: Int
lruSize     = Int
0
        , lruTick :: Priority
lruTick     = Priority
0
        , lruQueue :: HashPSQ k Priority v
lruQueue    = HashPSQ k Priority v
forall k p v. HashPSQ k p v
HashPSQ.empty
        }

-- | Restore 'LruCache' invariants returning the evicted element if any.
trim' :: (Hashable k, Ord k) => LruCache k v -> (Maybe (k, v), LruCache k v)
trim' :: forall k v.
(Hashable k, Ord k) =>
LruCache k v -> (Maybe (k, v), LruCache k v)
trim' LruCache k v
c
  | LruCache k v -> Priority
forall k v. LruCache k v -> Priority
lruTick LruCache k v
c Priority -> Priority -> Bool
forall a. Eq a => a -> a -> Bool
== Priority
forall a. Bounded a => a
maxBound     =
      -- It is not physically possible to have that many elements but
      -- the clock could potentially get here. We then simply decrease
      -- all priorities in O(nlogn) and start over.
      let queue' :: HashPSQ k Priority v
queue' = [(k, Priority, v)] -> HashPSQ k Priority v
forall k p v.
(Hashable k, Ord k, Ord p) =>
[(k, p, v)] -> HashPSQ k p v
HashPSQ.fromList ([(k, Priority, v)] -> HashPSQ k Priority v)
-> (HashPSQ k Priority v -> [(k, Priority, v)])
-> HashPSQ k Priority v
-> HashPSQ k Priority v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(k, Priority, v)] -> [(k, Priority, v)]
forall k v. [(k, Priority, v)] -> [(k, Priority, v)]
compress ([(k, Priority, v)] -> [(k, Priority, v)])
-> (HashPSQ k Priority v -> [(k, Priority, v)])
-> HashPSQ k Priority v
-> [(k, Priority, v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashPSQ k Priority v -> [(k, Priority, v)]
forall k p v.
(Hashable k, Ord k, Ord p) =>
HashPSQ k p v -> [(k, p, v)]
HashPSQ.toList (HashPSQ k Priority v -> HashPSQ k Priority v)
-> HashPSQ k Priority v -> HashPSQ k Priority v
forall a b. (a -> b) -> a -> b
$ LruCache k v -> HashPSQ k Priority v
forall k v. LruCache k v -> HashPSQ k Priority v
lruQueue LruCache k v
c
      in LruCache k v -> (Maybe (k, v), LruCache k v)
forall k v.
(Hashable k, Ord k) =>
LruCache k v -> (Maybe (k, v), LruCache k v)
trim' (LruCache k v -> (Maybe (k, v), LruCache k v))
-> LruCache k v -> (Maybe (k, v), LruCache k v)
forall a b. (a -> b) -> a -> b
$!
           LruCache k v
c { lruTick = fromIntegral (lruSize c)
             , lruQueue = queue'
             }
  | LruCache k v -> Int
forall k v. LruCache k v -> Int
lruSize LruCache k v
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> LruCache k v -> Int
forall k v. LruCache k v -> Int
lruCapacity LruCache k v
c = 
      let Just (k
k, Priority
_, v
v) = HashPSQ k Priority v -> Maybe (k, Priority, v)
forall k p v.
(Hashable k, Ord k, Ord p) =>
HashPSQ k p v -> Maybe (k, p, v)
HashPSQ.findMin (LruCache k v -> HashPSQ k Priority v
forall k v. LruCache k v -> HashPSQ k Priority v
lruQueue LruCache k v
c)
          c' :: LruCache k v
c' = LruCache k v
c  { lruSize  = lruSize c - 1
                  , lruQueue = HashPSQ.deleteMin (lruQueue c)
                  }
      in LruCache k v
-> (Maybe (k, v), LruCache k v) -> (Maybe (k, v), LruCache k v)
forall a b. a -> b -> b
seq LruCache k v
c' ((k, v) -> Maybe (k, v)
forall a. a -> Maybe a
Just (k
k, v
v), LruCache k v
c')
  | Bool
otherwise                 = (Maybe (k, v)
forall a. Maybe a
Nothing, LruCache k v
c)

compress :: [(k,Priority,v)] -> [(k,Priority,v)]
compress :: forall k v. [(k, Priority, v)] -> [(k, Priority, v)]
compress [(k, Priority, v)]
q =
  let sortedQ :: [(k, Priority, v)]
sortedQ = ((k, Priority, v) -> Priority)
-> [(k, Priority, v)] -> [(k, Priority, v)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (\(k
_,Priority
p,v
_) -> Priority
p) [(k, Priority, v)]
q
  in ((k, Priority, v) -> Priority -> (k, Priority, v))
-> [(k, Priority, v)] -> [Priority] -> [(k, Priority, v)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(k
k,Priority
_,v
v) Priority
p -> (k
k,Priority
p,v
v)) [(k, Priority, v)]
sortedQ [Priority
1..]

-- TODO benchmark to see if this is actually faster than snd . trim'
-- | Restore 'LruCache' invariants. For performance reasons this is
-- not @snd . trim'@.
trim :: (Hashable k, Ord k) => LruCache k v -> LruCache k v
trim :: forall k v. (Hashable k, Ord k) => LruCache k v -> LruCache k v
trim LruCache k v
c
  | LruCache k v -> Priority
forall k v. LruCache k v -> Priority
lruTick LruCache k v
c Priority -> Priority -> Bool
forall a. Eq a => a -> a -> Bool
== Priority
forall a. Bounded a => a
maxBound     = Int -> LruCache k v
forall k v. Int -> LruCache k v
empty (LruCache k v -> Int
forall k v. LruCache k v -> Int
lruCapacity LruCache k v
c)
  | LruCache k v -> Int
forall k v. LruCache k v -> Int
lruSize LruCache k v
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> LruCache k v -> Int
forall k v. LruCache k v -> Int
lruCapacity LruCache k v
c = 
      LruCache k v
c  { lruSize  = lruSize c - 1
         , lruQueue = HashPSQ.deleteMin (lruQueue c)
         }
  | Bool
otherwise                 = LruCache k v
c

-- | Insert an element into the 'LruCache'.
insert :: (Hashable k, Ord k) => k -> v -> LruCache k v -> LruCache k v
insert :: forall k v.
(Hashable k, Ord k) =>
k -> v -> LruCache k v -> LruCache k v
insert k
key v
val LruCache k v
c =
  LruCache k v -> LruCache k v
forall k v. (Hashable k, Ord k) => LruCache k v -> LruCache k v
trim (LruCache k v -> LruCache k v) -> LruCache k v -> LruCache k v
forall a b. (a -> b) -> a -> b
$!
  let (Maybe (Priority, v)
mbOldVal,HashPSQ k Priority v
queue) = k
-> Priority
-> v
-> HashPSQ k Priority v
-> (Maybe (Priority, v), HashPSQ k Priority v)
forall k p v.
(Hashable k, Ord k, Ord p) =>
k -> p -> v -> HashPSQ k p v -> (Maybe (p, v), HashPSQ k p v)
HashPSQ.insertView k
key (LruCache k v -> Priority
forall k v. LruCache k v -> Priority
lruTick LruCache k v
c) v
val (LruCache k v -> HashPSQ k Priority v
forall k v. LruCache k v -> HashPSQ k Priority v
lruQueue LruCache k v
c)
  in LruCache k v
c  { lruSize  = if isNothing mbOldVal
                     then lruSize c + 1
                     else lruSize c
        , lruTick  = lruTick c + 1
        , lruQueue = queue
        }

-- | Insert an element into the 'LruCache' returning the evicted
-- element if any.
--
-- When the logical clock reaches its maximum value and all values are
-- evicted 'Nothing' is returned.
insertView :: (Hashable k, Ord k) => k -> v -> LruCache k v -> (Maybe (k, v), LruCache k v)
insertView :: forall k v.
(Hashable k, Ord k) =>
k -> v -> LruCache k v -> (Maybe (k, v), LruCache k v)
insertView k
key v
val LruCache k v
cache =
  let (Maybe (Priority, v)
mbOldVal,HashPSQ k Priority v
queue) = 
        k
-> Priority
-> v
-> HashPSQ k Priority v
-> (Maybe (Priority, v), HashPSQ k Priority v)
forall k p v.
(Hashable k, Ord k, Ord p) =>
k -> p -> v -> HashPSQ k p v -> (Maybe (p, v), HashPSQ k p v)
HashPSQ.insertView k
key (LruCache k v -> Priority
forall k v. LruCache k v -> Priority
lruTick LruCache k v
cache) v
val (LruCache k v -> HashPSQ k Priority v
forall k v. LruCache k v -> HashPSQ k Priority v
lruQueue LruCache k v
cache)
  in LruCache k v -> (Maybe (k, v), LruCache k v)
forall k v.
(Hashable k, Ord k) =>
LruCache k v -> (Maybe (k, v), LruCache k v)
trim' (LruCache k v -> (Maybe (k, v), LruCache k v))
-> LruCache k v -> (Maybe (k, v), LruCache k v)
forall a b. (a -> b) -> a -> b
$! LruCache k v
cache 
       { lruSize  = if isNothing mbOldVal
                    then lruSize cache + 1
                    else lruSize cache
       , lruTick  = lruTick cache + 1
       , lruQueue = queue
       }

-- | Lookup an element in an 'LruCache' and mark it as the least
-- recently accessed.
lookup :: (Hashable k, Ord k) => k -> LruCache k v -> Maybe (v, LruCache k v)
lookup :: forall k v.
(Hashable k, Ord k) =>
k -> LruCache k v -> Maybe (v, LruCache k v)
lookup k
k LruCache k v
c = 
  case (Maybe (Priority, v) -> (Maybe v, Maybe (Priority, v)))
-> k -> HashPSQ k Priority v -> (Maybe v, HashPSQ k Priority v)
forall k p v b.
(Hashable k, Ord k, Ord p) =>
(Maybe (p, v) -> (b, Maybe (p, v)))
-> k -> HashPSQ k p v -> (b, HashPSQ k p v)
HashPSQ.alter Maybe (Priority, v) -> (Maybe v, Maybe (Priority, v))
forall {a} {b}. Maybe (a, b) -> (Maybe b, Maybe (Priority, b))
lookupAndBump k
k (LruCache k v -> HashPSQ k Priority v
forall k v. LruCache k v -> HashPSQ k Priority v
lruQueue LruCache k v
c) of
    (Maybe v
Nothing, HashPSQ k Priority v
_) -> Maybe (v, LruCache k v)
forall a. Maybe a
Nothing
    (Just v
x, HashPSQ k Priority v
q)  ->
      let !c' :: LruCache k v
c' = LruCache k v -> LruCache k v
forall k v. (Hashable k, Ord k) => LruCache k v -> LruCache k v
trim (LruCache k v -> LruCache k v) -> LruCache k v -> LruCache k v
forall a b. (a -> b) -> a -> b
$ LruCache k v
c {lruTick = lruTick c + 1, lruQueue = q}
      in (v, LruCache k v) -> Maybe (v, LruCache k v)
forall a. a -> Maybe a
Just (v
x, LruCache k v
c')
  where
    lookupAndBump :: Maybe (a, b) -> (Maybe b, Maybe (Priority, b))
lookupAndBump Maybe (a, b)
Nothing       = (Maybe b
forall a. Maybe a
Nothing, Maybe (Priority, b)
forall a. Maybe a
Nothing)
    lookupAndBump (Just (a
_, b
x)) = (b -> Maybe b
forall a. a -> Maybe a
Just b
x,  (Priority, b) -> Maybe (Priority, b)
forall a. a -> Maybe a
Just ((LruCache k v -> Priority
forall k v. LruCache k v -> Priority
lruTick LruCache k v
c), b
x))