{-# LANGUAGE DeriveFunctor   #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections   #-}
module Data.Map.Strict.Expiring (
    Map,
    new,
    generation,
    newGen,

    insert,
    delete,
    lookup,
    updateLookupWithKey,
    assocs,

    -- * for testing
    inspect
) where

import           Data.Foldable   (fold)
import qualified Data.Map.Strict as Map
import           Data.Maybe      (fromMaybe, mapMaybe)
import           Data.Set        (Set)
import qualified Data.Set        as Set
import           Prelude         hiding (lookup, map)

data Entry g a = Entry {
  forall g a. Entry g a -> a
value :: !a,
  forall g a. Entry g a -> g
gen   :: !g
} deriving ((forall a b. (a -> b) -> Entry g a -> Entry g b)
-> (forall a b. a -> Entry g b -> Entry g a) -> Functor (Entry g)
forall a b. a -> Entry g b -> Entry g a
forall a b. (a -> b) -> Entry g a -> Entry g b
forall g a b. a -> Entry g b -> Entry g a
forall g a b. (a -> b) -> Entry g a -> Entry g b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall g a b. (a -> b) -> Entry g a -> Entry g b
fmap :: forall a b. (a -> b) -> Entry g a -> Entry g b
$c<$ :: forall g a b. a -> Entry g b -> Entry g a
<$ :: forall a b. a -> Entry g b -> Entry g a
Functor, Int -> Entry g a -> ShowS
[Entry g a] -> ShowS
Entry g a -> String
(Int -> Entry g a -> ShowS)
-> (Entry g a -> String)
-> ([Entry g a] -> ShowS)
-> Show (Entry g a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall g a. (Show a, Show g) => Int -> Entry g a -> ShowS
forall g a. (Show a, Show g) => [Entry g a] -> ShowS
forall g a. (Show a, Show g) => Entry g a -> String
$cshowsPrec :: forall g a. (Show a, Show g) => Int -> Entry g a -> ShowS
showsPrec :: Int -> Entry g a -> ShowS
$cshow :: forall g a. (Show a, Show g) => Entry g a -> String
show :: Entry g a -> String
$cshowList :: forall g a. (Show a, Show g) => [Entry g a] -> ShowS
showList :: [Entry g a] -> ShowS
Show)

-- | A map of values that expire after a given generation.
data Map g k a = Map {
  -- | Primary store of values.
  forall g k a. Map g k a -> Map k (Entry g a)
map        :: !(Map.Map k (Entry g a)),
  -- | The current generation
  forall g k a. Map g k a -> g
generation :: !g,
  -- | A map of generations to keys that are expiring at that generation.
  forall g k a. Map g k a -> Map g (Set k)
aging      :: !(Map.Map g (Set k))
} deriving ((forall a b. (a -> b) -> Map g k a -> Map g k b)
-> (forall a b. a -> Map g k b -> Map g k a) -> Functor (Map g k)
forall a b. a -> Map g k b -> Map g k a
forall a b. (a -> b) -> Map g k a -> Map g k b
forall g k a b. a -> Map g k b -> Map g k a
forall g k a b. (a -> b) -> Map g k a -> Map g k b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall g k a b. (a -> b) -> Map g k a -> Map g k b
fmap :: forall a b. (a -> b) -> Map g k a -> Map g k b
$c<$ :: forall g k a b. a -> Map g k b -> Map g k a
<$ :: forall a b. a -> Map g k b -> Map g k a
Functor, Int -> Map g k a -> ShowS
[Map g k a] -> ShowS
Map g k a -> String
(Int -> Map g k a -> ShowS)
-> (Map g k a -> String)
-> ([Map g k a] -> ShowS)
-> Show (Map g k a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall g k a. (Show k, Show a, Show g) => Int -> Map g k a -> ShowS
forall g k a. (Show k, Show a, Show g) => [Map g k a] -> ShowS
forall g k a. (Show k, Show a, Show g) => Map g k a -> String
$cshowsPrec :: forall g k a. (Show k, Show a, Show g) => Int -> Map g k a -> ShowS
showsPrec :: Int -> Map g k a -> ShowS
$cshow :: forall g k a. (Show k, Show a, Show g) => Map g k a -> String
show :: Map g k a -> String
$cshowList :: forall g k a. (Show k, Show a, Show g) => [Map g k a] -> ShowS
showList :: [Map g k a] -> ShowS
Show)

instance Ord g => Foldable (Map g k) where
    foldMap :: forall m a. Monoid m => (a -> m) -> Map g k a -> m
foldMap a -> m
f = (Entry g a -> m) -> [Entry g a] -> m
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (a -> m
f (a -> m) -> (Entry g a -> a) -> Entry g a -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry g a -> a
forall g a. Entry g a -> a
value) ([Entry g a] -> m) -> (Map g k a -> [Entry g a]) -> Map g k a -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k (Entry g a) -> [Entry g a]
forall k a. Map k a -> [a]
Map.elems (Map k (Entry g a) -> [Entry g a])
-> (Map g k a -> Map k (Entry g a)) -> Map g k a -> [Entry g a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map g k a -> Map k (Entry g a)
forall g k a. Map g k a -> Map k (Entry g a)
map

-- | Make a new empty Map at the starting generation.
new :: g -> Map g k a
new :: forall g k a. g -> Map g k a
new g
g = Map k (Entry g a) -> g -> Map g (Set k) -> Map g k a
forall g k a. Map k (Entry g a) -> g -> Map g (Set k) -> Map g k a
Map Map k (Entry g a)
forall k a. Map k a
Map.empty g
g Map g (Set k)
forall k a. Map k a
Map.empty

-- | 𝑂(log𝑛). Assign the next generation and expire any data this new generation invalidates.
-- The generation may never decrease.  Attempts to decrease it are ignored.
newGen :: (Ord k, Ord g) => g -> Map g k a -> Map g k a
newGen :: forall k g a. (Ord k, Ord g) => g -> Map g k a -> Map g k a
newGen g
g Map g k a
m
    | g
g g -> g -> Bool
forall a. Ord a => a -> a -> Bool
> Map g k a -> g
forall g k a. Map g k a -> g
generation Map g k a
m = Map g k a -> Map g k a
forall g k a. (Ord g, Ord k) => Map g k a -> Map g k a
expire Map g k a
m { generation = g }
    | Bool
otherwise = Map g k a
m

-- | 𝑂(log𝑛). Insert a new value into the map to expire after the given generation.
-- alterF :: (Functor f, Ord k) => (Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
insert :: (Ord k, Ord g) => g -> k -> a -> Map g k a -> Map g k a
insert :: forall k g a.
(Ord k, Ord g) =>
g -> k -> a -> Map g k a -> Map g k a
insert g
g k
_ a
_ Map g k a
m | g
g g -> g -> Bool
forall a. Ord a => a -> a -> Bool
< Map g k a -> g
forall g k a. Map g k a -> g
generation Map g k a
m = Map g k a
m
insert g
g k
k a
v m :: Map g k a
m@Map{g
Map k (Entry g a)
Map g (Set k)
generation :: forall g k a. Map g k a -> g
map :: forall g k a. Map g k a -> Map k (Entry g a)
aging :: forall g k a. Map g k a -> Map g (Set k)
map :: Map k (Entry g a)
generation :: g
aging :: Map g (Set k)
..} = case (Maybe (Entry g a) -> (Maybe (Entry g a), Maybe (Entry g a)))
-> k -> Map k (Entry g a) -> (Maybe (Entry g a), Map k (Entry g a))
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
Map.alterF (, Entry g a -> Maybe (Entry g a)
forall a. a -> Maybe a
Just (a -> g -> Entry g a
forall g a. a -> g -> Entry g a
Entry a
v g
g)) k
k Map k (Entry g a)
map of
    (Just Entry g a
old, Map k (Entry g a)
m') -> Map g k a
m{map=m', aging = Map.insertWith (<>) g (Set.singleton k) (removeAging (gen old) k aging)}
    (Maybe (Entry g a)
Nothing, Map k (Entry g a)
m')  -> Map g k a
m{map=m', aging = Map.insertWith (<>) g (Set.singleton k) aging}

-- | 𝑂(log𝑛). Lookup and update.
-- The function returns changed value, if it is updated. Returns the original key value if the map entry is deleted.
updateLookupWithKey :: (Ord g, Ord k) => g -> (k -> a -> Maybe a) -> k -> Map g k a -> (Maybe a, Map g k a)
updateLookupWithKey :: forall g k a.
(Ord g, Ord k) =>
g -> (k -> a -> Maybe a) -> k -> Map g k a -> (Maybe a, Map g k a)
updateLookupWithKey g
g k -> a -> Maybe a
_ k
_ Map g k a
m | g
g g -> g -> Bool
forall a. Ord a => a -> a -> Bool
< Map g k a -> g
forall g k a. Map g k a -> g
generation Map g k a
m = (Maybe a
forall a. Maybe a
Nothing, Map g k a
m)
updateLookupWithKey g
g k -> a -> Maybe a
f k
k m :: Map g k a
m@Map{g
Map g (Set k)
Map k (Entry g a)
generation :: forall g k a. Map g k a -> g
map :: forall g k a. Map g k a -> Map k (Entry g a)
aging :: forall g k a. Map g k a -> Map g (Set k)
map :: Map k (Entry g a)
generation :: g
aging :: Map g (Set k)
..} = case (Maybe (Entry g a)
 -> ((Maybe (Entry g a), Maybe (Entry g a)), Maybe (Entry g a)))
-> k
-> Map k (Entry g a)
-> ((Maybe (Entry g a), Maybe (Entry g a)), Map k (Entry g a))
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
Map.alterF Maybe (Entry g a)
-> ((Maybe (Entry g a), Maybe (Entry g a)), Maybe (Entry g a))
forall {g}.
Maybe (Entry g a)
-> ((Maybe (Entry g a), Maybe (Entry g a)), Maybe (Entry g a))
f' k
k Map k (Entry g a)
map of
    ((Maybe (Entry g a)
Nothing, Maybe (Entry g a)
_), Map k (Entry g a)
m')   -> (Maybe a
forall a. Maybe a
Nothing, Map g k a
m)
    ((Just Entry g a
old, Maybe (Entry g a)
Nothing), Map k (Entry g a)
m')  -> (a -> Maybe a
forall a. a -> Maybe a
Just (Entry g a -> a
forall g a. Entry g a -> a
value Entry g a
old), Map g k a
m{map=m', aging = removeAging (gen old) k aging})
    ((Just Entry g a
old, Just Entry g a
new), Map k (Entry g a)
m') -> (a -> Maybe a
forall a. a -> Maybe a
Just (Entry g a -> a
forall g a. Entry g a -> a
value Entry g a
new), Map g k a
m{map=m', aging = Map.insertWith (<>) g (Set.singleton k) (removeAging (gen old) k aging)})
    where
        f' :: Maybe (Entry g a)
-> ((Maybe (Entry g a), Maybe (Entry g a)), Maybe (Entry g a))
f' Maybe (Entry g a)
Nothing = ((Maybe (Entry g a)
forall a. Maybe a
Nothing, Maybe (Entry g a)
forall a. Maybe a
Nothing), Maybe (Entry g a)
forall a. Maybe a
Nothing)
        f' (Just Entry g a
e) = case k -> a -> Maybe a
f k
k (Entry g a -> a
forall g a. Entry g a -> a
value Entry g a
e) of
            Maybe a
Nothing -> ((Entry g a -> Maybe (Entry g a)
forall a. a -> Maybe a
Just Entry g a
e, Maybe (Entry g a)
forall a. Maybe a
Nothing), Maybe (Entry g a)
forall a. Maybe a
Nothing)
            Just a
v  -> ((Entry g a -> Maybe (Entry g a)
forall a. a -> Maybe a
Just Entry g a
e, Entry g a -> Maybe (Entry g a)
forall a. a -> Maybe a
Just (a -> g -> Entry g a
forall g a. a -> g -> Entry g a
Entry a
v g
g)), Entry g a -> Maybe (Entry g a)
forall a. a -> Maybe a
Just (a -> g -> Entry g a
forall g a. a -> g -> Entry g a
Entry a
v g
g))

removeAging :: (Ord g, Ord k) => g -> k -> Map.Map g (Set k) -> Map.Map g (Set k)
removeAging :: forall g k.
(Ord g, Ord k) =>
g -> k -> Map g (Set k) -> Map g (Set k)
removeAging g
g k
k = (Set k -> Maybe (Set k)) -> g -> Map g (Set k) -> Map g (Set k)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update (Set k -> Maybe (Set k)
forall {a}. Set a -> Maybe (Set a)
nonNull (Set k -> Maybe (Set k))
-> (Set k -> Set k) -> Set k -> Maybe (Set k)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> Set k -> Set k
forall a. Ord a => a -> Set a -> Set a
Set.delete k
k) g
g
    where nonNull :: Set a -> Maybe (Set a)
nonNull Set a
s = if Set a -> Bool
forall a. Set a -> Bool
Set.null Set a
s then Maybe (Set a)
forall a. Maybe a
Nothing else Set a -> Maybe (Set a)
forall a. a -> Maybe a
Just Set a
s

-- | 𝑂(log𝑛). Lookup a value in the map.
-- This will not return any items that have expired.
lookup :: (Ord k, Ord g) => k -> Map g k a -> Maybe a
lookup :: forall k g a. (Ord k, Ord g) => k -> Map g k a -> Maybe a
lookup k
k = (Entry g a -> a) -> Maybe (Entry g a) -> Maybe a
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Entry g a -> a
forall g a. Entry g a -> a
value (Maybe (Entry g a) -> Maybe a)
-> (Map g k a -> Maybe (Entry g a)) -> Map g k a -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> Map k (Entry g a) -> Maybe (Entry g a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k (Map k (Entry g a) -> Maybe (Entry g a))
-> (Map g k a -> Map k (Entry g a))
-> Map g k a
-> Maybe (Entry g a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map g k a -> Map k (Entry g a)
forall g k a. Map g k a -> Map k (Entry g a)
map

-- | 𝑂(log𝑛). Delete an item.
delete :: (Ord k, Ord g) => k -> Map g k a -> Map g k a
delete :: forall k g a. (Ord k, Ord g) => k -> Map g k a -> Map g k a
delete k
k m :: Map g k a
m@Map{g
Map k (Entry g a)
Map g (Set k)
generation :: forall g k a. Map g k a -> g
map :: forall g k a. Map g k a -> Map k (Entry g a)
aging :: forall g k a. Map g k a -> Map g (Set k)
map :: Map k (Entry g a)
generation :: g
aging :: Map g (Set k)
..} = case k -> Map k (Entry g a) -> Maybe (Entry g a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k (Entry g a)
map of
  Maybe (Entry g a)
Nothing        -> Map g k a
m
  Just Entry{g
a
value :: forall g a. Entry g a -> a
gen :: forall g a. Entry g a -> g
value :: a
gen :: g
..} -> Map g k a
m { map = Map.delete k map, aging = removeAging gen k aging }

-- | 𝑂(𝑛). Return all current key/value associations.
assocs :: Ord g => Map g k a -> [(k,a)]
assocs :: forall g k a. Ord g => Map g k a -> [(k, a)]
assocs Map{g
Map g (Set k)
Map k (Entry g a)
generation :: forall g k a. Map g k a -> g
map :: forall g k a. Map g k a -> Map k (Entry g a)
aging :: forall g k a. Map g k a -> Map g (Set k)
map :: Map k (Entry g a)
generation :: g
aging :: Map g (Set k)
..} = (Entry g a -> a) -> (k, Entry g a) -> (k, a)
forall a b. (a -> b) -> (k, a) -> (k, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Entry g a -> a
forall g a. Entry g a -> a
value ((k, Entry g a) -> (k, a)) -> [(k, Entry g a)] -> [(k, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map k (Entry g a) -> [(k, Entry g a)]
forall k a. Map k a -> [(k, a)]
Map.assocs Map k (Entry g a)
map

-- | 𝑂(log𝑛).  Expire older generation items.
expire :: (Ord g, Ord k) => Map g k a -> Map g k a
expire :: forall g k a. (Ord g, Ord k) => Map g k a -> Map g k a
expire m :: Map g k a
m@Map{g
Map g (Set k)
Map k (Entry g a)
generation :: forall g k a. Map g k a -> g
map :: forall g k a. Map g k a -> Map k (Entry g a)
aging :: forall g k a. Map g k a -> Map g (Set k)
map :: Map k (Entry g a)
generation :: g
aging :: Map g (Set k)
..} = Map g k a
m{ map = map', aging = aging'}
    where
        (Map g (Set k)
todo, Maybe (Set k)
exact, Map g (Set k)
later) = g -> Map g (Set k) -> (Map g (Set k), Maybe (Set k), Map g (Set k))
forall k a. Ord k => k -> Map k a -> (Map k a, Maybe a, Map k a)
Map.splitLookup g
generation Map g (Set k)
aging
        aging' :: Map g (Set k)
aging' = Map g (Set k)
later Map g (Set k) -> Map g (Set k) -> Map g (Set k)
forall a. Semigroup a => a -> a -> a
<> Map g (Set k)
-> (Set k -> Map g (Set k)) -> Maybe (Set k) -> Map g (Set k)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Map g (Set k)
forall a. Monoid a => a
mempty (g -> Set k -> Map g (Set k)
forall k a. k -> a -> Map k a
Map.singleton g
generation) Maybe (Set k)
exact
        map' :: Map k (Entry g a)
map' = (k -> Map k (Entry g a) -> Map k (Entry g a))
-> Map k (Entry g a) -> Set k -> Map k (Entry g a)
forall a b. (a -> b -> b) -> b -> Set a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr k -> Map k (Entry g a) -> Map k (Entry g a)
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Map k (Entry g a)
map (Map g (Set k) -> Set k
forall m. Monoid m => Map g m -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold Map g (Set k)
todo)

-- | Inspect stored size for testing.
inspect :: Ord k => Map g k a -> (Int, g, Int)
inspect :: forall k g a. Ord k => Map g k a -> (Int, g, Int)
inspect Map{g
Map k (Entry g a)
Map g (Set k)
generation :: forall g k a. Map g k a -> g
map :: forall g k a. Map g k a -> Map k (Entry g a)
aging :: forall g k a. Map g k a -> Map g (Set k)
map :: Map k (Entry g a)
generation :: g
aging :: Map g (Set k)
..} = (Map k (Entry g a) -> Int
forall k a. Map k a -> Int
Map.size Map k (Entry g a)
map, g
generation, Set k -> Int
forall a. Set a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Set k -> Int) -> Set k -> Int
forall a b. (a -> b) -> a -> b
$ Map g (Set k) -> Set k
forall m. Monoid m => Map g m -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold Map g (Set k)
aging)