{-# LANGUAGE RecordWildCards #-}
module Data.Map.Strict.Decaying
  ( Map,
    new,
    insert,
    delete,
    elems,
    findWithDefault,
    updateLookupWithKey,

    -- * Testing Hooks
    tick,
  )
where

import           Control.Concurrent.Async    (async, cancel, link)
import           Control.Concurrent.STM      (STM, atomically)
import           Control.Concurrent.STM.TVar (TVar, mkWeakTVar, modifyTVar', newTVarIO, readTVar, stateTVar, writeTVar)
import           Control.Exception           (mask)
import           Control.Monad               (forever, join, void, (<=<))
import           Control.Monad.Loops         (whileJust_)
import           Data.Foldable               (toList)
import qualified Data.Map.Strict.Expiring    as Map
import           Data.Maybe                  (fromMaybe, mapMaybe)
import           Data.Time                   (NominalDiffTime)
import           Data.Time.Clock.POSIX       (POSIXTime, getPOSIXTime)
import           GHC.Conc                    (threadDelay)
import           System.Mem.Weak             (deRefWeak)

data Map k a = Map {
  forall k a. Map k a -> TVar (Map POSIXTime k a)
mapVar :: TVar (Map.Map POSIXTime k a),
  forall k a. Map k a -> POSIXTime
maxAge :: NominalDiffTime
}

expiry :: NominalDiffTime -> Map.Map POSIXTime k a -> POSIXTime
expiry :: forall k a. POSIXTime -> Map POSIXTime k a -> POSIXTime
expiry POSIXTime
ma = (POSIXTime -> POSIXTime -> POSIXTime
forall a. Num a => a -> a -> a
+ POSIXTime
ma) (POSIXTime -> POSIXTime)
-> (Map POSIXTime k a -> POSIXTime)
-> Map POSIXTime k a
-> POSIXTime
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map POSIXTime k a -> POSIXTime
forall g k a. Map g k a -> g
Map.generation

insert :: Ord k => k -> v -> Map k v -> STM ()
insert :: forall k v. Ord k => k -> v -> Map k v -> STM ()
insert k
k v
v (Map TVar (Map POSIXTime k v)
m POSIXTime
ma) = TVar (Map POSIXTime k v)
-> (Map POSIXTime k v -> Map POSIXTime k v) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map POSIXTime k v)
m (\Map POSIXTime k v
m -> POSIXTime -> k -> v -> Map POSIXTime k v -> Map POSIXTime k v
forall k g a.
(Ord k, Ord g) =>
g -> k -> a -> Map g k a -> Map g k a
Map.insert (POSIXTime -> Map POSIXTime k v -> POSIXTime
forall k a. POSIXTime -> Map POSIXTime k a -> POSIXTime
expiry POSIXTime
ma Map POSIXTime k v
m) k
k v
v Map POSIXTime k v
m)

delete :: Ord k => k -> Map k v -> STM ()
delete :: forall k v. Ord k => k -> Map k v -> STM ()
delete k
k (Map TVar (Map POSIXTime k v)
m POSIXTime
_) = TVar (Map POSIXTime k v)
-> (Map POSIXTime k v -> Map POSIXTime k v) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map POSIXTime k v)
m (k -> Map POSIXTime k v -> Map POSIXTime k v
forall k g a. (Ord k, Ord g) => k -> Map g k a -> Map g k a
Map.delete k
k)

findWithDefault :: Ord k => v -> k -> Map k v -> STM v
findWithDefault :: forall k v. Ord k => v -> k -> Map k v -> STM v
findWithDefault v
d k
k Map{TVar (Map POSIXTime k v)
POSIXTime
mapVar :: forall k a. Map k a -> TVar (Map POSIXTime k a)
maxAge :: forall k a. Map k a -> POSIXTime
mapVar :: TVar (Map POSIXTime k v)
maxAge :: POSIXTime
..} = v -> Maybe v -> v
forall a. a -> Maybe a -> a
fromMaybe v
d (Maybe v -> v)
-> (Map POSIXTime k v -> Maybe v) -> Map POSIXTime k v -> v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> Map POSIXTime k v -> Maybe v
forall k g a. (Ord k, Ord g) => k -> Map g k a -> Maybe a
Map.lookup k
k (Map POSIXTime k v -> v) -> STM (Map POSIXTime k v) -> STM v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Map POSIXTime k v) -> STM (Map POSIXTime k v)
forall a. TVar a -> STM a
readTVar TVar (Map POSIXTime k v)
mapVar

-- | All visible records.
elems :: Map k a -> STM [a]
elems :: forall k a. Map k a -> STM [a]
elems Map{TVar (Map POSIXTime k a)
POSIXTime
mapVar :: forall k a. Map k a -> TVar (Map POSIXTime k a)
maxAge :: forall k a. Map k a -> POSIXTime
mapVar :: TVar (Map POSIXTime k a)
maxAge :: POSIXTime
..} = Map POSIXTime k a -> [a]
forall a. Map POSIXTime k a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Map POSIXTime k a -> [a]) -> STM (Map POSIXTime k a) -> STM [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Map POSIXTime k a) -> STM (Map POSIXTime k a)
forall a. TVar a -> STM a
readTVar TVar (Map POSIXTime k a)
mapVar

tick :: Ord k => Map k v -> IO ()
tick :: forall k v. Ord k => Map k v -> IO ()
tick Map{TVar (Map POSIXTime k v)
POSIXTime
mapVar :: forall k a. Map k a -> TVar (Map POSIXTime k a)
maxAge :: forall k a. Map k a -> POSIXTime
mapVar :: TVar (Map POSIXTime k v)
maxAge :: POSIXTime
..} = IO POSIXTime
getPOSIXTime IO POSIXTime -> (POSIXTime -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \POSIXTime
t -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Map POSIXTime k v)
-> (Map POSIXTime k v -> Map POSIXTime k v) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map POSIXTime k v)
mapVar (POSIXTime -> Map POSIXTime k v -> Map POSIXTime k v
forall k g a. (Ord k, Ord g) => g -> Map g k a -> Map g k a
Map.newGen POSIXTime
t)

updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> STM (Maybe a)
updateLookupWithKey :: forall k a.
Ord k =>
(k -> a -> Maybe a) -> k -> Map k a -> STM (Maybe a)
updateLookupWithKey k -> a -> Maybe a
f k
k (Map TVar (Map POSIXTime k a)
mv POSIXTime
ma) = TVar (Map POSIXTime k a)
-> (Map POSIXTime k a -> (Maybe a, Map POSIXTime k a))
-> STM (Maybe a)
forall s a. TVar s -> (s -> (a, s)) -> STM a
stateTVar TVar (Map POSIXTime k a)
mv (\Map POSIXTime k a
m -> POSIXTime
-> (k -> a -> Maybe a)
-> k
-> Map POSIXTime k a
-> (Maybe a, Map POSIXTime k a)
forall g k a.
(Ord g, Ord k) =>
g -> (k -> a -> Maybe a) -> k -> Map g k a -> (Maybe a, Map g k a)
Map.updateLookupWithKey (POSIXTime -> Map POSIXTime k a -> POSIXTime
forall k a. POSIXTime -> Map POSIXTime k a -> POSIXTime
expiry POSIXTime
ma Map POSIXTime k a
m) k -> a -> Maybe a
f k
k Map POSIXTime k a
m)

new :: Ord k => NominalDiffTime -> IO (Map k a)
new :: forall k a. Ord k => POSIXTime -> IO (Map k a)
new POSIXTime
maxAge = ((forall a. IO a -> IO a) -> IO (Map k a)) -> IO (Map k a)
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO (Map k a)) -> IO (Map k a))
-> ((forall a. IO a -> IO a) -> IO (Map k a)) -> IO (Map k a)
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
  POSIXTime
now <- IO POSIXTime
getPOSIXTime
  TVar (Map POSIXTime k a)
var <- Map POSIXTime k a -> IO (TVar (Map POSIXTime k a))
forall a. a -> IO (TVar a)
newTVarIO (POSIXTime -> Map POSIXTime k a
forall g k a. g -> Map g k a
Map.new POSIXTime
now)
  Weak (TVar (Map POSIXTime k a))
wVar <- TVar (Map POSIXTime k a)
-> IO () -> IO (Weak (TVar (Map POSIXTime k a)))
forall a. TVar a -> IO () -> IO (Weak (TVar a))
mkWeakTVar TVar (Map POSIXTime k a)
var (IO () -> IO (Weak (TVar (Map POSIXTime k a))))
-> IO () -> IO (Weak (TVar (Map POSIXTime k a)))
forall a b. (a -> b) -> a -> b
$ () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  let m :: Map k a
m = TVar (Map POSIXTime k a) -> POSIXTime -> Map k a
forall k a. TVar (Map POSIXTime k a) -> POSIXTime -> Map k a
Map TVar (Map POSIXTime k a)
var POSIXTime
maxAge
  Async () -> IO ()
forall a. Async a -> IO ()
link (Async () -> IO ())
-> (IO (Async ()) -> IO (Async ())) -> IO (Async ()) -> IO ()
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< IO (Async ()) -> IO (Async ())
forall a. IO a -> IO a
restore (IO (Async ()) -> IO ()) -> IO (Async ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ IO (Maybe (TVar (Map POSIXTime k a)))
-> (TVar (Map POSIXTime k a) -> IO ()) -> IO ()
forall (m :: * -> *) a b.
Monad m =>
m (Maybe a) -> (a -> m b) -> m ()
whileJust_ (Weak (TVar (Map POSIXTime k a))
-> IO (Maybe (TVar (Map POSIXTime k a)))
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak (TVar (Map POSIXTime k a))
wVar) (\TVar (Map POSIXTime k a)
m' -> TVar (Map POSIXTime k a) -> IO ()
forall {k} {a}. Ord k => TVar (Map POSIXTime k a) -> IO ()
tick' TVar (Map POSIXTime k a)
m' IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> IO ()
threadDelay Int
1000000)
  Map k a -> IO (Map k a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map k a
m
  where
    tick' :: TVar (Map POSIXTime k a) -> IO ()
tick' TVar (Map POSIXTime k a)
m = IO POSIXTime
getPOSIXTime IO POSIXTime -> (POSIXTime -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \POSIXTime
t -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Map POSIXTime k a)
-> (Map POSIXTime k a -> Map POSIXTime k a) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map POSIXTime k a)
m (POSIXTime -> Map POSIXTime k a -> Map POSIXTime k a
forall k g a. (Ord k, Ord g) => g -> Map g k a -> Map g k a
Map.newGen POSIXTime
t)