TimedMap.hs

-- |
-- Module:     Control.Wire.TimedMap
-- Copyright:  (c) 2011 Ertugrul Soeylemez
-- License:    BSD3
-- Maintainer: Ertugrul Soeylemez <es@ertes.de>
--
-- This module implements a map, where each key has a timestamp.  It
-- maintains a timestamp index allowing you delete oldest entries
-- quickly.

module Control.Wire.TimedMap
    ( -- * Timed map
      TimedMap(..),

      -- * Operations
      -- ** Construct
      tmEmpty,
      -- ** Read
      tmFindWithDefault,
      tmLookup,
      -- ** Modify
      tmInsert,
      tmLimitAge,
      tmLimitSize
    )
    where

import qualified Data.Map as M
import qualified Data.Set as S
import Data.Map (Map)
import Data.Set (Set)


-- | A timed map is a regular map with timestamps and a timestamp index.

data TimedMap t k a =
    TimedMap {
      tmMap   :: Map k (a, t),  -- ^ Underlying map with timestamps.
      tmTimes :: Map t (Set k)  -- ^ Timestamp index.
    }
    deriving Show


-- | Find a value with default.

tmFindWithDefault ::
    Ord k
    => a               -- ^ Default, if key is not found.
    -> k               -- ^ Key to look up.
    -> TimedMap t k a  -- ^ Map to query.
    -> a               -- ^ Retrieved or default value.
tmFindWithDefault x0 k = M.findWithDefault x0 k . fmap fst . tmMap


-- | The empty timed map.

tmEmpty :: TimedMap t k a
tmEmpty = TimedMap M.empty M.empty


-- | Insert a value into the map.

tmInsert ::
    (Ord k, Ord t)
    => t               -- ^ Timestamp.
    -> k               -- ^ Key.
    -> a               -- ^ Value.
    -> TimedMap t k a  -- ^ Original map.
    -> TimedMap t k a  -- ^ Map with the value added.
tmInsert t k x (TimedMap xs' ts'') =
    TimedMap xs ts

    where
    xs = M.insert k (x, t) xs'
    ts = M.insertWith S.union t (S.singleton k) ts'
    ts' =
        case M.lookup k xs' of
          Nothing      -> ts''
          Just (_, t') ->
              M.update (\s' -> let s = S.delete k s' in
                               if S.null s then Nothing else Just s)
                       t' ts''


-- | Delete all items older than the specified timestamp.

tmLimitAge :: (Ord t, Ord k) => t -> TimedMap t k a -> TimedMap t k a
tmLimitAge minT (TimedMap xs' ts') = TimedMap xs ts
    where
    xs = xs' M.\\ delMap
    ts = maybe id (M.insert minT) tsCur tsYounger

    (tsOlder, tsCur, tsYounger) = M.splitLookup minT ts'
    delMap =
        M.fromDistinctAscList . map (, ()) .
        S.toAscList . S.unions . M.elems $ tsOlder


-- | Delete at least as many oldest items as necessary to limit the
-- map's size to the given value.  If you have multiple keys with the
-- same timestamp, this function can delete more keys than necessary.

tmLimitSize :: Ord k => Int -> TimedMap t k a -> TimedMap t k a
tmLimitSize n tm@(TimedMap xs ts') =
    if n >= 0 && M.size xs > n
      then tmLimitSize n $ TimedMap (xs M.\\ delMap) ts
      else tm

    where
    delMap = M.fromDistinctAscList . map (, ()) . S.toAscList $ delKeys
    ((_, delKeys), ts) = M.deleteFindMin ts'


-- | Look up the value for the given key.

tmLookup :: Ord k => k -> TimedMap t k a -> Maybe a
tmLookup k (TimedMap xs _) = fmap fst (M.lookup k xs)