{-# 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 = (forall a. Num a => a -> a -> a
+ POSIXTime
ma) forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) = forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map POSIXTime k v)
m (\Map POSIXTime k v
m -> forall k g a.
(Ord k, Ord g) =>
g -> k -> a -> Map g k a -> Map g k a
Map.insert (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
_) = forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map POSIXTime k v)
m (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
maxAge :: POSIXTime
mapVar :: TVar (Map POSIXTime k v)
maxAge :: forall k a. Map k a -> POSIXTime
mapVar :: forall k a. Map k a -> TVar (Map POSIXTime k a)
..} = forall a. a -> Maybe a -> a
fromMaybe v
d forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k g a. (Ord k, Ord g) => k -> Map g k a -> Maybe a
Map.lookup k
k forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
maxAge :: POSIXTime
mapVar :: TVar (Map POSIXTime k a)
maxAge :: forall k a. Map k a -> POSIXTime
mapVar :: forall k a. Map k a -> TVar (Map POSIXTime k a)
..} = forall (t :: * -> *) a. Foldable t => t a -> [a]
toList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
maxAge :: POSIXTime
mapVar :: TVar (Map POSIXTime k v)
maxAge :: forall k a. Map k a -> POSIXTime
mapVar :: forall k a. Map k a -> TVar (Map POSIXTime k a)
..} = IO POSIXTime
getPOSIXTime forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \POSIXTime
t -> forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map POSIXTime k v)
mapVar (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) = forall s a. TVar s -> (s -> (a, s)) -> STM a
stateTVar TVar (Map POSIXTime k a)
mv (\Map POSIXTime k a
m -> 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 (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 b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask 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 <- forall a. a -> IO (TVar a)
newTVarIO (forall g k a. g -> Map g k a
Map.new POSIXTime
now)
  Weak (TVar (Map POSIXTime k a))
wVar <- forall a. TVar a -> IO () -> IO (Weak (TVar a))
mkWeakTVar TVar (Map POSIXTime k a)
var forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  let m :: Map k a
m = forall k a. TVar (Map POSIXTime k a) -> POSIXTime -> Map k a
Map TVar (Map POSIXTime k a)
var POSIXTime
maxAge
  forall a. Async a -> IO ()
link forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall a. IO a -> IO a
restore forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO (Async a)
async forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
Monad m =>
m (Maybe a) -> (a -> m b) -> m ()
whileJust_ (forall v. Weak v -> IO (Maybe v)
deRefWeak Weak (TVar (Map POSIXTime k a))
wVar) (\TVar (Map POSIXTime k a)
m' -> forall {k} {a}. Ord k => TVar (Map POSIXTime k a) -> IO ()
tick' TVar (Map POSIXTime k a)
m' forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> IO ()
threadDelay Int
1000000)
  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 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \POSIXTime
t -> forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map POSIXTime k a)
m (forall k g a. (Ord k, Ord g) => g -> Map g k a -> Map g k a
Map.newGen POSIXTime
t)