module StmContainers.Multimap
  ( Multimap,
    new,
    newIO,
    null,
    focus,
    lookup,
    lookupByKey,
    insert,
    delete,
    deleteByKey,
    reset,
    unfoldlM,
    unfoldlMKeys,
    unfoldlMByKey,
    listT,
    listTKeys,
    listTByKey,
  )
where

import qualified Focus as C
import qualified StmContainers.Map as A
import StmContainers.Prelude hiding (delete, empty, foldM, insert, lookup, null, toList)
import qualified StmContainers.Set as B

-- |
-- A multimap, based on an STM-specialized hash array mapped trie.
--
-- Basically it's just a wrapper API around @'A.Map' key ('B.Set' value)@.
newtype Multimap key value
  = Multimap (A.Map key (B.Set value))
  deriving (Typeable)

-- |
-- Construct a new multimap.
{-# INLINE new #-}
new :: STM (Multimap key value)
new :: forall key value. STM (Multimap key value)
new =
  Map key (Set value) -> Multimap key value
forall key value. Map key (Set value) -> Multimap key value
Multimap (Map key (Set value) -> Multimap key value)
-> STM (Map key (Set value)) -> STM (Multimap key value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STM (Map key (Set value))
forall key value. STM (Map key value)
A.new

-- |
-- Construct a new multimap in IO.
--
-- This is useful for creating it on a top-level using 'unsafePerformIO',
-- because using 'atomically' inside 'unsafePerformIO' isn't possible.
{-# INLINE newIO #-}
newIO :: IO (Multimap key value)
newIO :: forall key value. IO (Multimap key value)
newIO =
  Map key (Set value) -> Multimap key value
forall key value. Map key (Set value) -> Multimap key value
Multimap (Map key (Set value) -> Multimap key value)
-> IO (Map key (Set value)) -> IO (Multimap key value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Map key (Set value))
forall key value. IO (Map key value)
A.newIO

-- |
-- Check on being empty.
{-# INLINE null #-}
null :: Multimap key value -> STM Bool
null :: forall key value. Multimap key value -> STM Bool
null (Multimap Map key (Set value)
map) =
  Map key (Set value) -> STM Bool
forall key value. Map key value -> STM Bool
A.null Map key (Set value)
map

-- |
-- Focus on an item by the value and the key.
--
-- This function allows to perform simultaneous lookup and modification.
--
-- The focus is over a unit since we already know,
-- which value we're focusing on and it doesn't make sense to replace it,
-- however we still can decide wether to keep or remove it.
{-# INLINE focus #-}
focus :: (Hashable key, Hashable value) => C.Focus () STM result -> value -> key -> Multimap key value -> STM result
focus :: forall key value result.
(Hashable key, Hashable value) =>
Focus () STM result
-> value -> key -> Multimap key value -> STM result
focus unitFocus :: Focus () STM result
unitFocus@(Focus STM (result, Change ())
concealUnit () -> STM (result, Change ())
_) value
value key
key (Multimap Map key (Set value)
map) = Focus (Set value) STM result
-> key -> Map key (Set value) -> STM result
forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
A.focus Focus (Set value) STM result
setFocus key
key Map key (Set value)
map
  where
    setFocus :: Focus (Set value) STM result
setFocus = STM (result, Change (Set value))
-> (Set value -> STM (result, Change (Set value)))
-> Focus (Set value) STM result
forall element (m :: * -> *) result.
m (result, Change element)
-> (element -> m (result, Change element))
-> Focus element m result
C.Focus STM (result, Change (Set value))
conceal Set value -> STM (result, Change (Set value))
reveal
      where
        conceal :: STM (result, Change (Set value))
conceal = do
          (result
output, Change ()
change) <- STM (result, Change ())
concealUnit
          case Change ()
change of
            C.Set () ->
              do
                Set value
set <- STM (Set value)
forall item. STM (Set item)
B.new
                value -> Set value -> STM ()
forall item. Hashable item => item -> Set item -> STM ()
B.insert value
value Set value
set
                (result, Change (Set value)) -> STM (result, Change (Set value))
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return (result
output, Set value -> Change (Set value)
forall a. a -> Change a
C.Set Set value
set)
            Change ()
_ ->
              (result, Change (Set value)) -> STM (result, Change (Set value))
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return (result
output, Change (Set value)
forall a. Change a
C.Leave)
    reveal :: Set value -> STM (result, Change (Set value))
reveal Set value
set = do
      result
output <- Focus () STM result -> value -> Set value -> STM result
forall item result.
Hashable item =>
Focus () STM result -> item -> Set item -> STM result
B.focus Focus () STM result
unitFocus value
value Set value
set
      Change (Set value)
change <- Change (Set value)
-> Change (Set value) -> Bool -> Change (Set value)
forall a. a -> a -> Bool -> a
bool Change (Set value)
forall a. Change a
C.Leave Change (Set value)
forall a. Change a
C.Remove (Bool -> Change (Set value))
-> STM Bool -> STM (Change (Set value))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set value -> STM Bool
forall item. Set item -> STM Bool
B.null Set value
set
      (result, Change (Set value)) -> STM (result, Change (Set value))
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return (result
output, Change (Set value)
change)

-- |
-- Look up an item by a value and a key.
{-# INLINE lookup #-}
lookup :: (Hashable key, Hashable value) => value -> key -> Multimap key value -> STM Bool
lookup :: forall key value.
(Hashable key, Hashable value) =>
value -> key -> Multimap key value -> STM Bool
lookup value
value key
key (Multimap Map key (Set value)
m) =
  STM Bool
-> (Set value -> STM Bool) -> Maybe (Set value) -> STM Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Bool -> STM Bool
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False) (value -> Set value -> STM Bool
forall item. Hashable item => item -> Set item -> STM Bool
B.lookup value
value) (Maybe (Set value) -> STM Bool)
-> STM (Maybe (Set value)) -> STM Bool
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< key -> Map key (Set value) -> STM (Maybe (Set value))
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
A.lookup key
key Map key (Set value)
m

-- |
-- Look up all values by key.
{-# INLINE lookupByKey #-}
lookupByKey :: (Hashable key) => key -> Multimap key value -> STM (Maybe (B.Set value))
lookupByKey :: forall key value.
Hashable key =>
key -> Multimap key value -> STM (Maybe (Set value))
lookupByKey key
key (Multimap Map key (Set value)
m) =
  key -> Map key (Set value) -> STM (Maybe (Set value))
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
A.lookup key
key Map key (Set value)
m

-- |
-- Insert an item.
{-# INLINEABLE insert #-}
insert :: (Hashable key, Hashable value) => value -> key -> Multimap key value -> STM ()
insert :: forall key value.
(Hashable key, Hashable value) =>
value -> key -> Multimap key value -> STM ()
insert value
value key
key (Multimap Map key (Set value)
map) = Focus (Set value) STM () -> key -> Map key (Set value) -> STM ()
forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
A.focus Focus (Set value) STM ()
setFocus key
key Map key (Set value)
map
  where
    setFocus :: Focus (Set value) STM ()
setFocus = STM ((), Change (Set value))
-> (Set value -> STM ((), Change (Set value)))
-> Focus (Set value) STM ()
forall element (m :: * -> *) result.
m (result, Change element)
-> (element -> m (result, Change element))
-> Focus element m result
Focus STM ((), Change (Set value))
conceal Set value -> STM ((), Change (Set value))
reveal
      where
        conceal :: STM ((), Change (Set value))
conceal = do
          Set value
set <- STM (Set value)
forall item. STM (Set item)
B.new
          value -> Set value -> STM ()
forall item. Hashable item => item -> Set item -> STM ()
B.insert value
value Set value
set
          ((), Change (Set value)) -> STM ((), Change (Set value))
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return ((), Set value -> Change (Set value)
forall a. a -> Change a
C.Set Set value
set)
        reveal :: Set value -> STM ((), Change (Set value))
reveal Set value
set = do
          value -> Set value -> STM ()
forall item. Hashable item => item -> Set item -> STM ()
B.insert value
value Set value
set
          ((), Change (Set value)) -> STM ((), Change (Set value))
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return ((), Change (Set value)
forall a. Change a
C.Leave)

-- |
-- Delete an item by a value and a key.
{-# INLINEABLE delete #-}
delete :: (Hashable key, Hashable value) => value -> key -> Multimap key value -> STM ()
delete :: forall key value.
(Hashable key, Hashable value) =>
value -> key -> Multimap key value -> STM ()
delete value
value key
key (Multimap Map key (Set value)
map) = Focus (Set value) STM () -> key -> Map key (Set value) -> STM ()
forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
A.focus Focus (Set value) STM ()
setFocus key
key Map key (Set value)
map
  where
    setFocus :: Focus (Set value) STM ()
setFocus = STM ((), Change (Set value))
-> (Set value -> STM ((), Change (Set value)))
-> Focus (Set value) STM ()
forall element (m :: * -> *) result.
m (result, Change element)
-> (element -> m (result, Change element))
-> Focus element m result
Focus STM ((), Change (Set value))
forall {m :: * -> *} {a}. Monad m => m ((), Change a)
conceal Set value -> STM ((), Change (Set value))
reveal
      where
        conceal :: m ((), Change a)
conceal = Change a -> m ((), Change a)
forall {m :: * -> *} {b}. Monad m => b -> m ((), b)
returnChange Change a
forall a. Change a
C.Leave
        reveal :: Set value -> STM ((), Change (Set value))
reveal Set value
set = do
          value -> Set value -> STM ()
forall item. Hashable item => item -> Set item -> STM ()
B.delete value
value Set value
set
          Set value -> STM Bool
forall item. Set item -> STM Bool
B.null Set value
set STM Bool
-> (Bool -> STM ((), Change (Set value)))
-> STM ((), Change (Set value))
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Change (Set value) -> STM ((), Change (Set value))
forall {m :: * -> *} {b}. Monad m => b -> m ((), b)
returnChange (Change (Set value) -> STM ((), Change (Set value)))
-> (Bool -> Change (Set value))
-> Bool
-> STM ((), Change (Set value))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Change (Set value)
-> Change (Set value) -> Bool -> Change (Set value)
forall a. a -> a -> Bool -> a
bool Change (Set value)
forall a. Change a
C.Leave Change (Set value)
forall a. Change a
C.Remove
        returnChange :: b -> m ((), b)
returnChange b
c = ((), b) -> m ((), b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((), b
c)

-- |
-- Delete all values associated with the key.
{-# INLINEABLE deleteByKey #-}
deleteByKey :: (Hashable key) => key -> Multimap key value -> STM ()
deleteByKey :: forall key value.
Hashable key =>
key -> Multimap key value -> STM ()
deleteByKey key
key (Multimap Map key (Set value)
map) =
  key -> Map key (Set value) -> STM ()
forall key value. Hashable key => key -> Map key value -> STM ()
A.delete key
key Map key (Set value)
map

-- |
-- Delete all the associations.
{-# INLINE reset #-}
reset :: Multimap key value -> STM ()
reset :: forall key value. Multimap key value -> STM ()
reset (Multimap Map key (Set value)
map) =
  Map key (Set value) -> STM ()
forall key value. Map key value -> STM ()
A.reset Map key (Set value)
map

-- |
-- Stream associations actively.
--
-- Amongst other features this function provides an interface to folding.
unfoldlM :: Multimap key value -> UnfoldlM STM (key, value)
unfoldlM :: forall key value. Multimap key value -> UnfoldlM STM (key, value)
unfoldlM (Multimap Map key (Set value)
m) =
  Map key (Set value) -> UnfoldlM STM (key, Set value)
forall key value. Map key value -> UnfoldlM STM (key, value)
A.unfoldlM Map key (Set value)
m UnfoldlM STM (key, Set value)
-> ((key, Set value) -> UnfoldlM STM (key, value))
-> UnfoldlM STM (key, value)
forall a b.
UnfoldlM STM a -> (a -> UnfoldlM STM b) -> UnfoldlM STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(key
key, Set value
s) -> (key
key,) (value -> (key, value))
-> UnfoldlM STM value -> UnfoldlM STM (key, value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set value -> UnfoldlM STM value
forall item. Set item -> UnfoldlM STM item
B.unfoldlM Set value
s

-- |
-- Stream keys actively.
unfoldlMKeys :: Multimap key value -> UnfoldlM STM key
unfoldlMKeys :: forall key value. Multimap key value -> UnfoldlM STM key
unfoldlMKeys (Multimap Map key (Set value)
m) =
  ((key, Set value) -> key)
-> UnfoldlM STM (key, Set value) -> UnfoldlM STM key
forall a b. (a -> b) -> UnfoldlM STM a -> UnfoldlM STM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (key, Set value) -> key
forall a b. (a, b) -> a
fst (Map key (Set value) -> UnfoldlM STM (key, Set value)
forall key value. Map key value -> UnfoldlM STM (key, value)
A.unfoldlM Map key (Set value)
m)

-- |
-- Stream values by a key actively.
unfoldlMByKey :: (Hashable key) => key -> Multimap key value -> UnfoldlM STM value
unfoldlMByKey :: forall key value.
Hashable key =>
key -> Multimap key value -> UnfoldlM STM value
unfoldlMByKey key
key (Multimap Map key (Set value)
m) =
  STM (Maybe (Set value)) -> UnfoldlM STM (Maybe (Set value))
forall (m :: * -> *) a. Monad m => m a -> UnfoldlM m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (key -> Map key (Set value) -> STM (Maybe (Set value))
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
A.lookup key
key Map key (Set value)
m) UnfoldlM STM (Maybe (Set value))
-> (Maybe (Set value) -> UnfoldlM STM value) -> UnfoldlM STM value
forall a b.
UnfoldlM STM a -> (a -> UnfoldlM STM b) -> UnfoldlM STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= UnfoldlM STM value
-> (Set value -> UnfoldlM STM value)
-> Maybe (Set value)
-> UnfoldlM STM value
forall b a. b -> (a -> b) -> Maybe a -> b
maybe UnfoldlM STM value
forall a. Monoid a => a
mempty Set value -> UnfoldlM STM value
forall item. Set item -> UnfoldlM STM item
B.unfoldlM

-- |
-- Stream associations passively.
listT :: Multimap key value -> ListT STM (key, value)
listT :: forall key value. Multimap key value -> ListT STM (key, value)
listT (Multimap Map key (Set value)
m) =
  Map key (Set value) -> ListT STM (key, Set value)
forall key value. Map key value -> ListT STM (key, value)
A.listT Map key (Set value)
m ListT STM (key, Set value)
-> ((key, Set value) -> ListT STM (key, value))
-> ListT STM (key, value)
forall a b. ListT STM a -> (a -> ListT STM b) -> ListT STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(key
key, Set value
s) -> (key
key,) (value -> (key, value))
-> ListT STM value -> ListT STM (key, value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set value -> ListT STM value
forall item. Set item -> ListT STM item
B.listT Set value
s

-- |
-- Stream keys passively.
listTKeys :: Multimap key value -> ListT STM key
listTKeys :: forall key value. Multimap key value -> ListT STM key
listTKeys (Multimap Map key (Set value)
m) =
  ((key, Set value) -> key)
-> ListT STM (key, Set value) -> ListT STM key
forall a b. (a -> b) -> ListT STM a -> ListT STM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (key, Set value) -> key
forall a b. (a, b) -> a
fst (Map key (Set value) -> ListT STM (key, Set value)
forall key value. Map key value -> ListT STM (key, value)
A.listT Map key (Set value)
m)

-- |
-- Stream values by a key passively.
listTByKey :: (Hashable key) => key -> Multimap key value -> ListT STM value
listTByKey :: forall key value.
Hashable key =>
key -> Multimap key value -> ListT STM value
listTByKey key
key (Multimap Map key (Set value)
m) =
  STM (Maybe (Set value)) -> ListT STM (Maybe (Set value))
forall (m :: * -> *) a. Monad m => m a -> ListT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (key -> Map key (Set value) -> STM (Maybe (Set value))
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
A.lookup key
key Map key (Set value)
m) ListT STM (Maybe (Set value))
-> (Maybe (Set value) -> ListT STM value) -> ListT STM value
forall a b. ListT STM a -> (a -> ListT STM b) -> ListT STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ListT STM value
-> (Set value -> ListT STM value)
-> Maybe (Set value)
-> ListT STM value
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ListT STM value
forall a. Monoid a => a
mempty Set value -> ListT STM value
forall item. Set item -> ListT STM item
B.listT