module Data.TimeMap
(
TimeMap
,
newTimeMap
, insert
, adjust
, delete
,
lookup
, timeOf
, ageOf
, keys
, elems
, 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)
import qualified Data.Map 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.Concurrent.STM
data TimeMap k a = TimeMap
{ timeMap :: TVar (MM.MultiMap UTCTime k)
, keysMap :: HT.Map k (UTCTime, 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 $ HT.focus (go now) k (keysMap xs)
where
go now mx = do
modifyTVar (timeMap xs) $
let changeOld = case mx of
Nothing -> id
Just (oldTime,_) -> MM.remove oldTime k
in MM.insert now k . changeOld
pure ((), F.Replace (now, x))
lookup :: ( Hashable k
, Eq k
) => k -> TimeMap k a -> STM (Maybe a)
lookup k xs = do
mx <- HT.lookup k (keysMap xs)
pure (snd <$> mx)
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 $ (snd . snd) <$> HT.stream (keysMap 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 (fst <$> 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)
adjust :: ( Hashable k
, Eq k
) => (a -> a) -> k -> TimeMap k a -> IO ()
adjust f k xs = do
now <- getCurrentTime
atomically $ HT.focus (go now) k (keysMap xs)
where
go _ Nothing = pure ((), F.Keep)
go now (Just (oldTime, y)) = do
modifyTVar (timeMap xs) (MM.insert now k . MM.remove oldTime k)
pure ((), F.Replace (now, f y))
delete :: ( Hashable k
, Eq k
) => k -> TimeMap k a -> STM ()
delete k xs = do
HT.focus go k (keysMap xs)
where
go mx = do
case mx of
Nothing -> pure ()
Just (oldTime,_) -> modifyTVar' (timeMap xs) (MM.remove oldTime k)
pure ((), F.Remove)
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 (_,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