{-# LANGUAGE RecordWildCards #-}

module Network.Control.LRUCache (
    -- * LRU cache
    LRUCache,
    empty,
    insert,
    delete,
    lookup,
    lookup',

    -- * IO
    LRUCacheRef,
    newLRUCacheRef,
    cached,
    cached',

    -- * Internal
    empty',
) where

import Data.IORef (IORef, atomicModifyIORef', newIORef)
import Data.Int (Int64)
import Data.OrdPSQ (OrdPSQ)
import qualified Data.OrdPSQ as PSQ
import Prelude hiding (lookup)

----------------------------------------------------------------

type Priority = Int64

-- | Sized cache based on least recently used.
data LRUCache k v = LRUCache
    { forall k v. LRUCache k v -> Int
lcLimit :: Int
    -- ^ The maximum number of elements in the queue
    , forall k v. LRUCache k v -> Priority
lcTick :: Priority
    -- ^ The next logical time
    , forall k v. LRUCache k v -> OrdPSQ k Priority v
lcQueue :: OrdPSQ k Priority v
    }
    deriving (LRUCache k v -> LRUCache k v -> Bool
(LRUCache k v -> LRUCache k v -> Bool)
-> (LRUCache k v -> LRUCache k v -> Bool) -> Eq (LRUCache k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Ord k, Eq v) => LRUCache k v -> LRUCache k v -> Bool
$c== :: forall k v. (Ord k, Eq v) => LRUCache k v -> LRUCache k v -> Bool
== :: LRUCache k v -> LRUCache k v -> Bool
$c/= :: forall k v. (Ord k, Eq v) => LRUCache k v -> LRUCache k v -> Bool
/= :: LRUCache k v -> LRUCache k v -> Bool
Eq, Int -> LRUCache k v -> ShowS
[LRUCache k v] -> ShowS
LRUCache k v -> String
(Int -> LRUCache k v -> ShowS)
-> (LRUCache k v -> String)
-> ([LRUCache k v] -> ShowS)
-> Show (LRUCache k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show v, Show k) => Int -> LRUCache k v -> ShowS
forall k v. (Show v, Show k) => [LRUCache k v] -> ShowS
forall k v. (Show v, Show k) => LRUCache k v -> String
$cshowsPrec :: forall k v. (Show v, Show k) => Int -> LRUCache k v -> ShowS
showsPrec :: Int -> LRUCache k v -> ShowS
$cshow :: forall k v. (Show v, Show k) => LRUCache k v -> String
show :: LRUCache k v -> String
$cshowList :: forall k v. (Show v, Show k) => [LRUCache k v] -> ShowS
showList :: [LRUCache k v] -> ShowS
Show)

----------------------------------------------------------------

-- | Empty 'LRUCache'. /O(1)/
empty
    :: Int
    -- ^ The size of 'LRUCache'.
    -> LRUCache k v
empty :: forall k v. Int -> LRUCache k v
empty Int
capacity =
    LRUCache
        { lcLimit :: Int
lcLimit = Int
capacity
        , lcTick :: Priority
lcTick = Priority
0
        , lcQueue :: OrdPSQ k Priority v
lcQueue = OrdPSQ k Priority v
forall k p v. OrdPSQ k p v
PSQ.empty
        }

-- | Empty 'LRUCache'. /O(1)/
empty'
    :: Int
    -- ^ The size of 'LRUCache'.
    -> Int64
    -- ^ Counter
    -> LRUCache k v
empty' :: forall k v. Int -> Priority -> LRUCache k v
empty' Int
capacity Priority
tick =
    LRUCache
        { lcLimit :: Int
lcLimit = Int
capacity
        , lcTick :: Priority
lcTick = Priority
tick
        , lcQueue :: OrdPSQ k Priority v
lcQueue = OrdPSQ k Priority v
forall k p v. OrdPSQ k p v
PSQ.empty
        }

----------------------------------------------------------------

trim :: Ord k => LRUCache k v -> LRUCache k v
trim :: forall k v. Ord k => LRUCache k v -> LRUCache k v
trim c :: LRUCache k v
c@LRUCache{Int
Priority
OrdPSQ k Priority v
lcLimit :: forall k v. LRUCache k v -> Int
lcTick :: forall k v. LRUCache k v -> Priority
lcQueue :: forall k v. LRUCache k v -> OrdPSQ k Priority v
lcLimit :: Int
lcTick :: Priority
lcQueue :: OrdPSQ k Priority v
..}
    | Priority
lcTick Priority -> Priority -> Bool
forall a. Eq a => a -> a -> Bool
== Priority
forall a. Bounded a => a
maxBound =
        let siz :: Priority
siz = Int -> Priority
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Priority) -> Int -> Priority
forall a b. (a -> b) -> a -> b
$ OrdPSQ k Priority v -> Int
forall k p v. OrdPSQ k p v -> Int
PSQ.size OrdPSQ k Priority v
lcQueue
            diff :: Priority
diff = (Priority
forall a. Bounded a => a
maxBound :: Priority) Priority -> Priority -> Priority
forall a. Num a => a -> a -> a
- Priority
siz
            psq :: OrdPSQ k Priority v
psq = (k -> Priority -> v -> (Priority, v))
-> OrdPSQ k Priority v -> OrdPSQ k Priority v
forall k p q v w.
(k -> p -> v -> (q, w)) -> OrdPSQ k p v -> OrdPSQ k q w
PSQ.unsafeMapMonotonic (\k
_ Priority
p v
v -> (Priority
p Priority -> Priority -> Priority
forall a. Num a => a -> a -> a
- Priority
diff, v
v)) OrdPSQ k Priority v
lcQueue
         in LRUCache
                { lcLimit :: Int
lcLimit = Int
lcLimit
                , lcTick :: Priority
lcTick = Priority
siz
                , lcQueue :: OrdPSQ k Priority v
lcQueue = OrdPSQ k Priority v
psq
                }
    | OrdPSQ k Priority v -> Int
forall k p v. OrdPSQ k p v -> Int
PSQ.size OrdPSQ k Priority v
lcQueue Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
lcLimit = LRUCache k v
c{lcQueue = PSQ.deleteMin lcQueue}
    | Bool
otherwise = LRUCache k v
c

----------------------------------------------------------------

-- | Inserting. /O(log n)/
insert :: Ord k => k -> v -> LRUCache k v -> LRUCache k v
insert :: forall k v. Ord k => k -> v -> LRUCache k v -> LRUCache k v
insert k
key v
val c :: LRUCache k v
c@LRUCache{Int
Priority
OrdPSQ k Priority v
lcLimit :: forall k v. LRUCache k v -> Int
lcTick :: forall k v. LRUCache k v -> Priority
lcQueue :: forall k v. LRUCache k v -> OrdPSQ k Priority v
lcLimit :: Int
lcTick :: Priority
lcQueue :: OrdPSQ k Priority v
..} = LRUCache k v -> LRUCache k v
forall k v. Ord k => LRUCache k v -> LRUCache k v
trim LRUCache k v
c'
  where
    queue :: OrdPSQ k Priority v
queue = k -> Priority -> v -> OrdPSQ k Priority v -> OrdPSQ k Priority v
forall k p v.
(Ord k, Ord p) =>
k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
PSQ.insert k
key Priority
lcTick v
val OrdPSQ k Priority v
lcQueue
    c' :: LRUCache k v
c' = LRUCache k v
c{lcTick = lcTick + 1, lcQueue = queue}

----------------------------------------------------------------

-- | Deleting. /O(log n)/
delete :: Ord k => k -> LRUCache k v -> LRUCache k v
delete :: forall k v. Ord k => k -> LRUCache k v -> LRUCache k v
delete k
k c :: LRUCache k v
c@LRUCache{Int
Priority
OrdPSQ k Priority v
lcLimit :: forall k v. LRUCache k v -> Int
lcTick :: forall k v. LRUCache k v -> Priority
lcQueue :: forall k v. LRUCache k v -> OrdPSQ k Priority v
lcLimit :: Int
lcTick :: Priority
lcQueue :: OrdPSQ k Priority v
..} = LRUCache k v
c{lcQueue = q}
  where
    q :: OrdPSQ k Priority v
q = k -> OrdPSQ k Priority v -> OrdPSQ k Priority v
forall k p v. (Ord k, Ord p) => k -> OrdPSQ k p v -> OrdPSQ k p v
PSQ.delete k
k OrdPSQ k Priority v
lcQueue

----------------------------------------------------------------

-- | Looking up. /O(log n)/
lookup :: Ord k => k -> LRUCache k v -> Maybe v
lookup :: forall k v. Ord k => k -> LRUCache k v -> Maybe v
lookup k
k LRUCache{Int
Priority
OrdPSQ k Priority v
lcLimit :: forall k v. LRUCache k v -> Int
lcTick :: forall k v. LRUCache k v -> Priority
lcQueue :: forall k v. LRUCache k v -> OrdPSQ k Priority v
lcLimit :: Int
lcTick :: Priority
lcQueue :: OrdPSQ k Priority v
..} = (Priority, v) -> v
forall a b. (a, b) -> b
snd ((Priority, v) -> v) -> Maybe (Priority, v) -> Maybe v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> k -> OrdPSQ k Priority v -> Maybe (Priority, v)
forall k p v. Ord k => k -> OrdPSQ k p v -> Maybe (p, v)
PSQ.lookup k
k OrdPSQ k Priority v
lcQueue

-- | Looking up and changing priority. /O(log n)/
lookup' :: Ord k => k -> LRUCache k v -> Maybe (v, LRUCache k v)
lookup' :: forall k v. Ord k => k -> LRUCache k v -> Maybe (v, LRUCache k v)
lookup' k
k c :: LRUCache k v
c@LRUCache{Int
Priority
OrdPSQ k Priority v
lcLimit :: forall k v. LRUCache k v -> Int
lcTick :: forall k v. LRUCache k v -> Priority
lcQueue :: forall k v. LRUCache k v -> OrdPSQ k Priority v
lcLimit :: Int
lcTick :: Priority
lcQueue :: OrdPSQ k Priority v
..} = case (Maybe (Priority, v) -> (Maybe v, Maybe (Priority, v)))
-> k -> OrdPSQ k Priority v -> (Maybe v, OrdPSQ k Priority v)
forall k p v b.
(Ord k, Ord p) =>
(Maybe (p, v) -> (b, Maybe (p, v)))
-> k -> OrdPSQ k p v -> (b, OrdPSQ k p v)
PSQ.alter Maybe (Priority, v) -> (Maybe v, Maybe (Priority, v))
forall {a} {b}. Maybe (a, b) -> (Maybe b, Maybe (Priority, b))
lookupAndBump k
k OrdPSQ k Priority v
lcQueue of
    (Maybe v
Nothing, OrdPSQ k Priority v
_) -> Maybe (v, LRUCache k v)
forall a. Maybe a
Nothing
    (Just v
v, OrdPSQ k Priority v
q) ->
        let c' :: LRUCache k v
c' = LRUCache k v -> LRUCache k v
forall k v. 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{lcTick = lcTick + 1, lcQueue = q}
         in (v, LRUCache k v) -> Maybe (v, LRUCache k v)
forall a. a -> Maybe a
Just (v
v, 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)
    -- setting its priority to lcTick
    lookupAndBump (Just (a
_p, b
v)) = (b -> Maybe b
forall a. a -> Maybe a
Just b
v, (Priority, b) -> Maybe (Priority, b)
forall a. a -> Maybe a
Just (Priority
lcTick, b
v))

----------------------------------------------------------------

newtype LRUCacheRef k v = LRUCacheRef (IORef (LRUCache k v))

newLRUCacheRef :: Int -> IO (LRUCacheRef k v)
newLRUCacheRef :: forall k v. Int -> IO (LRUCacheRef k v)
newLRUCacheRef Int
capacity = IORef (LRUCache k v) -> LRUCacheRef k v
forall k v. IORef (LRUCache k v) -> LRUCacheRef k v
LRUCacheRef (IORef (LRUCache k v) -> LRUCacheRef k v)
-> IO (IORef (LRUCache k v)) -> IO (LRUCacheRef 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)

cached' :: Ord k => LRUCacheRef k v -> k -> IO (Maybe v)
cached' :: forall k v. Ord k => LRUCacheRef k v -> k -> IO (Maybe v)
cached' (LRUCacheRef IORef (LRUCache k v)
ref) k
k = do
    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. 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)

cached :: Ord k => LRUCacheRef k v -> k -> IO v -> IO (v, Bool)
cached :: forall k v. Ord k => LRUCacheRef k v -> k -> IO v -> IO (v, Bool)
cached (LRUCacheRef 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. 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, Bool) -> IO (v, Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (v
v, Bool
True)
        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. Ord k => k -> v -> LRUCache k v -> LRUCache k v
insert k
k v
v LRUCache k v
c, ())
            (v, Bool) -> IO (v, Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (v
v, Bool
False)