{-# 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 -> 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
<$ :: forall a b. a -> Entry g b -> Entry g a
$c<$ :: forall g a b. a -> Entry g b -> Entry g a
fmap :: forall a b. (a -> b) -> Entry g a -> Entry g b
$cfmap :: forall g a b. (a -> b) -> Entry g a -> Entry g b
Functor, Int -> Entry g a -> ShowS
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
showList :: [Entry g a] -> ShowS
$cshowList :: forall g a. (Show a, Show g) => [Entry g a] -> ShowS
show :: Entry g a -> String
$cshow :: forall g a. (Show a, Show g) => Entry g a -> String
showsPrec :: Int -> Entry g a -> ShowS
$cshowsPrec :: forall g a. (Show a, Show g) => Int -> 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 -> 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
<$ :: forall a b. a -> Map g k b -> Map g k a
$c<$ :: forall g k a b. a -> Map g k b -> Map g k a
fmap :: forall a b. (a -> b) -> Map g k a -> Map g k b
$cfmap :: forall g k a b. (a -> b) -> Map g k a -> Map g k b
Functor, Int -> Map g k a -> ShowS
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
showList :: [Map g k a] -> ShowS
$cshowList :: forall g k a. (Show k, Show a, Show g) => [Map g k a] -> ShowS
show :: Map g k a -> String
$cshow :: forall g k a. (Show k, Show a, Show g) => Map g k a -> String
showsPrec :: Int -> Map g k a -> ShowS
$cshowsPrec :: forall g k a. (Show k, Show a, Show g) => Int -> 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 = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (a -> m
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall g a. Entry g a -> a
value) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [a]
Map.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall g k a. Map k (Entry g a) -> g -> Map g (Set k) -> Map g k a
Map forall k a. Map k a
Map.empty g
g 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 forall a. Ord a => a -> a -> Bool
> forall g k a. Map g k a -> g
generation Map g k a
m = forall g k a. (Ord g, Ord k) => Map g k a -> Map g k a
expire Map g k a
m { generation :: g
generation = g
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 forall a. Ord a => a -> a -> Bool
< 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)
aging :: Map g (Set k)
generation :: g
map :: Map k (Entry g a)
aging :: forall g k a. Map g k a -> Map g (Set k)
map :: forall g k a. Map g k a -> Map k (Entry g a)
generation :: forall g k a. Map g k a -> g
..} = case forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
Map.alterF (, forall a. a -> Maybe a
Just (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 :: Map k (Entry g a)
map=Map k (Entry g a)
m', aging :: Map g (Set k)
aging = forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith forall a. Semigroup a => a -> a -> a
(<>) g
g (forall a. a -> Set a
Set.singleton k
k) (forall g k.
(Ord g, Ord k) =>
g -> k -> Map g (Set k) -> Map g (Set k)
removeAging (forall g a. Entry g a -> g
gen Entry g a
old) k
k Map g (Set k)
aging)}
    (Maybe (Entry g a)
Nothing, Map k (Entry g a)
m')  -> Map g k a
m{map :: Map k (Entry g a)
map=Map k (Entry g a)
m', aging :: Map g (Set k)
aging = forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith forall a. Semigroup a => a -> a -> a
(<>) g
g (forall a. a -> Set a
Set.singleton k
k) Map g (Set 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 forall a. Ord a => a -> a -> Bool
< forall g k a. Map g k a -> g
generation Map g k a
m = (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)
aging :: Map g (Set k)
generation :: g
map :: Map k (Entry g a)
aging :: forall g k a. Map g k a -> Map g (Set k)
map :: forall g k a. Map g k a -> Map k (Entry g a)
generation :: forall g k a. Map g k a -> g
..} = case forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
Map.alterF 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')   -> (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')  -> (forall a. a -> Maybe a
Just (forall g a. Entry g a -> a
value Entry g a
old), Map g k a
m{map :: Map k (Entry g a)
map=Map k (Entry g a)
m', aging :: Map g (Set k)
aging = forall g k.
(Ord g, Ord k) =>
g -> k -> Map g (Set k) -> Map g (Set k)
removeAging (forall g a. Entry g a -> g
gen Entry g a
old) k
k Map g (Set k)
aging})
    ((Just Entry g a
old, Just Entry g a
new), Map k (Entry g a)
m') -> (forall a. a -> Maybe a
Just (forall g a. Entry g a -> a
value Entry g a
new), Map g k a
m{map :: Map k (Entry g a)
map=Map k (Entry g a)
m', aging :: Map g (Set k)
aging = forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith forall a. Semigroup a => a -> a -> a
(<>) g
g (forall a. a -> Set a
Set.singleton k
k) (forall g k.
(Ord g, Ord k) =>
g -> k -> Map g (Set k) -> Map g (Set k)
removeAging (forall g a. Entry g a -> g
gen Entry g a
old) k
k Map g (Set 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 = ((forall a. Maybe a
Nothing, forall a. Maybe a
Nothing), forall a. Maybe a
Nothing)
        f' (Just Entry g a
e) = case k -> a -> Maybe a
f k
k (forall g a. Entry g a -> a
value Entry g a
e) of
            Maybe a
Nothing -> ((forall a. a -> Maybe a
Just Entry g a
e, forall a. Maybe a
Nothing), forall a. Maybe a
Nothing)
            Just a
v  -> ((forall a. a -> Maybe a
Just Entry g a
e, forall a. a -> Maybe a
Just (forall g a. a -> g -> Entry g a
Entry a
v g
g)), forall a. a -> Maybe a
Just (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 = forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update (forall {a}. Set a -> Maybe (Set a)
nonNull forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 forall a. Set a -> Bool
Set.null Set a
s then forall a. Maybe a
Nothing else 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 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall g a. Entry g a -> a
value forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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)
aging :: Map g (Set k)
generation :: g
map :: Map k (Entry g a)
aging :: forall g k a. Map g k a -> Map g (Set k)
map :: forall g k a. Map g k a -> Map k (Entry g a)
generation :: forall g k a. Map g k a -> g
..} = case 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
gen :: g
value :: a
gen :: forall g a. Entry g a -> g
value :: forall g a. Entry g a -> a
..} -> Map g k a
m { map :: Map k (Entry g a)
map = forall k a. Ord k => k -> Map k a -> Map k a
Map.delete k
k Map k (Entry g a)
map, aging :: Map g (Set k)
aging = forall g k.
(Ord g, Ord k) =>
g -> k -> Map g (Set k) -> Map g (Set k)
removeAging g
gen k
k Map g (Set 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)
aging :: Map g (Set k)
generation :: g
map :: Map k (Entry g a)
aging :: forall g k a. Map g k a -> Map g (Set k)
map :: forall g k a. Map g k a -> Map k (Entry g a)
generation :: forall g k a. Map g k a -> g
..} = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall g a. Entry g a -> a
value forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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)
aging :: Map g (Set k)
generation :: g
map :: Map k (Entry g a)
aging :: forall g k a. Map g k a -> Map g (Set k)
map :: forall g k a. Map g k a -> Map k (Entry g a)
generation :: forall g k a. Map g k a -> g
..} = Map g k a
m{ map :: Map k (Entry g a)
map = Map k (Entry g a)
map', aging :: Map g (Set k)
aging = Map g (Set k)
aging'}
    where
        (Map g (Set k)
todo, Maybe (Set k)
exact, Map g (Set k)
later) = 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 forall a. Semigroup a => a -> a -> a
<> forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty (forall k a. k -> a -> Map k a
Map.singleton g
generation) Maybe (Set k)
exact
        map' :: Map k (Entry g a)
map' = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Map k (Entry g a)
map (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)
aging :: Map g (Set k)
generation :: g
map :: Map k (Entry g a)
aging :: forall g k a. Map g k a -> Map g (Set k)
map :: forall g k a. Map g k a -> Map k (Entry g a)
generation :: forall g k a. Map g k a -> g
..} = (forall k a. Map k a -> Int
Map.size Map k (Entry g a)
map, g
generation, forall (t :: * -> *) a. Foldable t => t a -> Int
length forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold Map g (Set k)
aging)