{-# LANGUAGE
BangPatterns
, NamedFieldPuns
, ScopedTypeVariables
, RankNTypes
#-}
module Data.TimeMap
(
TimeMap
,
newTimeMap
, insert
, insertWithTime
, update
, updateWithTime
, adjust
, adjustWithTime
, delete
, touch
,
lookup
, timeOf
, ageOf
, keys
, elems
, toList
, size
, null
,
filter
, filterWithKey
, filterSince
, filterFromNow
,
takeSince
, takeFromNow
) where
import Prelude hiding (lookup, null, filter)
import Data.Time (UTCTime, NominalDiffTime, addUTCTime, diffUTCTime, getCurrentTime)
import Data.Hashable (Hashable (..))
import Data.Maybe (fromMaybe, fromJust, catMaybes)
import qualified Data.Map.Strict as Map
import qualified Data.HashSet as HS
import qualified Data.TimeMap.Internal as MM
import qualified STMContainers.Map as HT
import qualified Focus as F
import qualified ListT as L
import Control.Monad (forM, void)
import Control.Concurrent.STM (STM, atomically, TVar, writeTVar, readTVar, modifyTVar', modifyTVar, newTVar)
data TimeIndexed a = TimeIndexed
{ indexedTime :: {-# UNPACK #-} !UTCTime
, indexedValue :: a
}
data TimeMap k a = TimeMap
{ timeMap :: !(TVar (MM.MultiMap UTCTime k))
, keysMap :: !(HT.Map k (TimeIndexed a))
}
newTimeMap :: STM (TimeMap k a)
newTimeMap = TimeMap <$> newTVar MM.empty
<*> HT.new
insert :: Hashable k
=> Eq k
=> k -> a -> TimeMap k a -> IO ()
insert k x xs = do
now <- getCurrentTime
atomically (insertWithTime now k x xs)
{-# INLINEABLE insert #-}
insertWithTime :: forall k a
. Hashable k
=> Eq k
=> UTCTime -> k -> a -> TimeMap k a -> STM ()
insertWithTime now k x TimeMap{timeMap,keysMap} =
HT.focus go k keysMap
where
go :: Maybe (TimeIndexed a) -> STM ((), F.Decision (TimeIndexed a))
go mx = do
modifyTVar timeMap $
let changeOld = case mx of
Nothing -> id
Just (TimeIndexed oldTime _) ->
MM.remove oldTime k
in MM.insert now k . changeOld
pure ((), F.Replace (TimeIndexed now x))
{-# INLINEABLE insertWithTime #-}
lookup :: Hashable k
=> Eq k
=> k -> TimeMap k a -> STM (Maybe a)
lookup k TimeMap{keysMap} =
(\mx' -> indexedValue <$> mx') <$> HT.lookup k keysMap
{-# INLINEABLE lookup #-}
keys :: Hashable k
=> Eq k
=> TimeMap k a -> STM (HS.HashSet k)
keys TimeMap{timeMap} = MM.elems <$> readTVar timeMap
{-# INLINEABLE keys #-}
elems :: TimeMap k a -> STM [a]
elems TimeMap{keysMap} = L.toList $ (indexedValue . snd) <$> HT.stream keysMap
toList :: Hashable k
=> Eq k
=> TimeMap k a -> STM [(k, a)]
toList TimeMap{keysMap,timeMap} = do
keys' <- (HS.toList . MM.elems) <$> readTVar timeMap
forM keys' $ \k -> do
mVal <- HT.lookup k keysMap
pure (k, indexedValue (fromJust mVal))
size :: TimeMap k a -> STM Int
size xs = length <$> elems xs
null :: TimeMap k a -> STM Bool
null xs = HT.null (keysMap xs)
timeOf :: Hashable k
=> Eq k
=> k -> TimeMap k a -> STM (Maybe UTCTime)
timeOf k xs = do
mx <- HT.lookup k (keysMap xs)
pure (indexedTime <$> mx)
{-# INLINEABLE timeOf #-}
ageOf :: Hashable k
=> Eq k
=> k -> TimeMap k a -> IO (Maybe NominalDiffTime)
ageOf k xs = do
now <- getCurrentTime
mt <- atomically (timeOf k xs)
pure (diffUTCTime now <$> mt)
{-# INLINEABLE ageOf #-}
update :: Hashable k
=> Eq k
=> (a -> Maybe a) -> k -> TimeMap k a -> IO ()
update p k xs = do
now <- getCurrentTime
atomically (updateWithTime now p k xs)
{-# INLINEABLE update #-}
updateWithTime :: forall k a
. Hashable k
=> Eq k
=> UTCTime -> (a -> Maybe a) -> k -> TimeMap k a -> STM ()
updateWithTime now p k TimeMap{keysMap,timeMap} =
HT.focus go k keysMap
where
go :: Maybe (TimeIndexed a) -> STM ((), F.Decision (TimeIndexed a))
go Nothing = pure ((), F.Keep)
go (Just (TimeIndexed oldTime y)) =
let (action,minsert) =
case p y of
Nothing -> (F.Remove , MM.remove oldTime k)
Just y' -> (F.Replace (TimeIndexed now y'), id)
in do modifyTVar timeMap (MM.insert now k . minsert)
pure ((), action)
{-# INLINEABLE updateWithTime #-}
adjust :: Hashable k
=> Eq k
=> (a -> a) -> k -> TimeMap k a -> IO ()
adjust f k xs = do
now <- getCurrentTime
atomically (adjustWithTime now f k xs)
{-# INLINEABLE adjust #-}
adjustWithTime :: forall k a
. Hashable k
=> Eq k
=> UTCTime -> (a -> a) -> k -> TimeMap k a -> STM ()
adjustWithTime now f k TimeMap{keysMap,timeMap} =
HT.focus go k keysMap
where
go :: Maybe (TimeIndexed a) -> STM ((), F.Decision (TimeIndexed a))
go Nothing = pure ((), F.Keep)
go (Just (TimeIndexed oldTime y)) = do
modifyTVar timeMap (MM.insert now k . MM.remove oldTime k)
pure ((), F.Replace $ TimeIndexed now $ f y)
{-# INLINEABLE adjustWithTime #-}
delete :: Hashable k
=> Eq k
=> k -> TimeMap k a -> STM ()
delete k TimeMap{timeMap,keysMap} = HT.focus go k keysMap
where
go mx = do
case mx of
Nothing -> pure ()
Just (TimeIndexed oldTime _) ->
modifyTVar' timeMap (MM.remove oldTime k)
pure ((), F.Remove)
{-# INLINEABLE delete #-}
touch :: Hashable k
=> Eq k
=> k -> TimeMap k a -> IO ()
touch = adjust id
{-# INLINEABLE touch #-}
filter :: Hashable k
=> Eq k
=> (a -> Bool) -> TimeMap k a -> STM ()
filter p = filterWithKey (const p)
{-# INLINEABLE filter #-}
filterWithKey :: forall k a
. Hashable k
=> Eq k
=> (k -> a -> Bool) -> TimeMap k a -> STM ()
filterWithKey p xs = do
ks <- (HS.toList . MM.elems) <$> readTVar (timeMap xs)
mapM_ go ks
where
go :: k -> STM ()
go k = HT.focus go' k (keysMap xs)
where
go' :: Maybe (TimeIndexed a) -> STM ((), F.Decision (TimeIndexed a))
go' (Just (TimeIndexed _ x))
| p k x = pure ((), F.Keep)
| otherwise = pure ((), F.Remove)
go' Nothing = pure ((), F.Keep)
{-# INLINEABLE filterWithKey #-}
takeSince :: Hashable k
=> Eq k
=> UTCTime
-> TimeMap k a
-> STM [(k, a)]
takeSince t TimeMap{timeMap,keysMap} = do
ts <- readTVar timeMap
let (toCut, mx, result) = Map.splitLookup t ts
toRemove = MM.elems toCut `HS.union` fromMaybe HS.empty mx
writeTVar timeMap result
taken <- fmap catMaybes $ forM (HS.toList toRemove) $ \k -> do
mX <- HT.lookup k keysMap
case mX of
Nothing -> pure Nothing
Just (TimeIndexed _ x) -> do
HT.delete k keysMap
pure (Just (k, x))
pure taken
{-# INLINEABLE takeSince #-}
takeFromNow :: Hashable k
=> Eq k
=> NominalDiffTime
-> TimeMap k a
-> IO [(k, a)]
takeFromNow t xs = do
now <- getCurrentTime
atomically (takeSince (addUTCTime (negate t) now) xs)
{-# INLINEABLE takeFromNow #-}
filterSince :: Hashable k
=> Eq k
=> UTCTime
-> TimeMap k a
-> STM ()
filterSince t = void . takeSince t
{-# INLINEABLE filterSince #-}
filterFromNow :: Hashable k
=> Eq k
=> NominalDiffTime
-> TimeMap k a
-> IO ()
filterFromNow t = void . takeFromNow t
{-# INLINEABLE filterFromNow #-}