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
) where
import Prelude hiding (lookup, null, filter)
import Data.Time (UTCTime, NominalDiffTime, addUTCTime, diffUTCTime, getCurrentTime)
import Data.Hashable (Hashable (..))
import Data.Maybe (fromMaybe, fromJust)
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)
import Control.Concurrent.STM (STM, atomically, TVar, writeTVar, readTVar, modifyTVar', modifyTVar, newTVar)
data TimeIndexed a = TimeIndexed
{ indexedTime :: !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
insertWithTime :: ( Hashable k
, Eq k
) => UTCTime -> k -> a -> TimeMap k a -> STM ()
insertWithTime now k x xs =
HT.focus go k (keysMap xs)
where
go mx = do
modifyTVar (timeMap xs) $
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))
lookup :: ( Hashable k
, Eq k
) => k -> TimeMap k a -> STM (Maybe a)
lookup k xs =
(\mx' -> indexedValue <$> mx') <$> HT.lookup k (keysMap xs)
keys :: ( Hashable k
, Eq k
) => TimeMap k a -> STM (HS.HashSet k)
keys xs = MM.elems <$> readTVar (timeMap xs)
elems :: TimeMap k a -> STM [a]
elems xs = L.toList $ (indexedValue . snd) <$> HT.stream (keysMap xs)
toList :: ( Hashable k
, Eq k
) => TimeMap k a -> STM [(k, a)]
toList xs = do
keys' <- (HS.toList . MM.elems) <$> readTVar (timeMap xs)
forM keys' $ \k -> do
mVal <- HT.lookup k (keysMap xs)
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
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
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
updateWithTime :: ( Hashable k
, Eq k
) => UTCTime -> (a -> Maybe a) -> k -> TimeMap k a -> STM ()
updateWithTime now p k xs =
HT.focus go k (keysMap xs)
where
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 xs) (MM.insert now k . minsert)
pure ((), action)
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
adjustWithTime :: ( Hashable k
, Eq k
) => UTCTime -> (a -> a) -> k -> TimeMap k a -> STM ()
adjustWithTime now f k xs =
HT.focus go k (keysMap xs)
where
go Nothing = pure ((), F.Keep)
go (Just (TimeIndexed oldTime y)) = do
modifyTVar (timeMap xs) (MM.insert now k . MM.remove oldTime k)
pure ((), F.Replace (TimeIndexed now $! f y))
delete :: ( Hashable k
, Eq k
) => k -> TimeMap k a -> STM ()
delete k xs = HT.focus go k (keysMap xs)
where
go mx = do
case mx of
Nothing -> pure ()
Just (TimeIndexed oldTime _) ->
modifyTVar' (timeMap xs) (MM.remove oldTime k)
pure ((), F.Remove)
touch :: ( Hashable k
, Eq k
) => k -> TimeMap k a -> IO ()
touch = adjust id
filter :: ( Hashable k
, Eq k
) => (a -> Bool) -> TimeMap k a -> STM ()
filter p = filterWithKey (const p)
filterWithKey :: ( 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 = HT.focus go' k (keysMap xs)
where
go' (Just (TimeIndexed _ x))
| p k x = pure ((), F.Keep)
| otherwise = pure ((), F.Remove)
go' Nothing = pure ((), F.Keep)
filterSince :: ( Hashable k
, Eq k
) => UTCTime
-> TimeMap k a
-> STM ()
filterSince t xs = do
ts <- readTVar (timeMap xs)
let (toCut, mx, result) = Map.splitLookup t ts
found = fromMaybe HS.empty mx
toRemove = MM.elems toCut `HS.union` found
writeTVar (timeMap xs) result
mapM_ (\k -> HT.delete k $ keysMap xs) $! HS.toList toRemove
filterFromNow :: ( Hashable k
, Eq k
) => NominalDiffTime
-> TimeMap k a
-> IO ()
filterFromNow t xs = do
now <- getCurrentTime
atomically $ (filterSince $! addUTCTime (negate t) now) xs