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 =
  forall key value. Map key (Set value) -> Multimap key value
Multimap forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 =
  forall key value. Map key (Set value) -> Multimap key value
Multimap forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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) =
  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) = 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 = 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 <- forall item. STM (Set item)
B.new
                forall item. Hashable item => item -> Set item -> STM ()
B.insert value
value Set value
set
                forall (m :: * -> *) a. Monad m => a -> m a
return (result
output, forall a. a -> Change a
C.Set Set value
set)
            Change ()
_ ->
              forall (m :: * -> *) a. Monad m => a -> m a
return (result
output, forall a. Change a
C.Leave)
    reveal :: Set value -> STM (result, Change (Set value))
reveal Set value
set = do
      result
output <- 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 <- forall a. a -> a -> Bool -> a
bool forall a. Change a
C.Leave forall a. Change a
C.Remove forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall item. Set item -> STM Bool
B.null Set value
set
      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) =
  forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False) (forall item. Hashable item => item -> Set item -> STM Bool
B.lookup value
value) forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< 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) =
  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) = 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 = 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 <- forall item. STM (Set item)
B.new
          forall item. Hashable item => item -> Set item -> STM ()
B.insert value
value Set value
set
          forall (m :: * -> *) a. Monad m => a -> m a
return ((), forall a. a -> Change a
C.Set Set value
set)
        reveal :: Set value -> STM ((), Change (Set value))
reveal Set value
set = do
          forall item. Hashable item => item -> Set item -> STM ()
B.insert value
value Set value
set
          forall (m :: * -> *) a. Monad m => a -> m a
return ((), 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) = 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 = forall element (m :: * -> *) result.
m (result, Change element)
-> (element -> m (result, Change element))
-> Focus element m result
Focus forall {m :: * -> *} {a}. Monad m => m ((), Change a)
conceal Set value -> STM ((), Change (Set value))
reveal
      where
        conceal :: m ((), Change a)
conceal = forall {m :: * -> *} {b}. Monad m => b -> m ((), b)
returnChange forall a. Change a
C.Leave
        reveal :: Set value -> STM ((), Change (Set value))
reveal Set value
set = do
          forall item. Hashable item => item -> Set item -> STM ()
B.delete value
value Set value
set
          forall item. Set item -> STM Bool
B.null Set value
set forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall {m :: * -> *} {b}. Monad m => b -> m ((), b)
returnChange forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. a -> a -> Bool -> a
bool forall a. Change a
C.Leave forall a. Change a
C.Remove
        returnChange :: b -> m ((), b)
returnChange b
c = 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) =
  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) =
  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) =
  forall key value. Map key value -> UnfoldlM STM (key, value)
A.unfoldlM Map key (Set value)
m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(key
key, Set value
s) -> (key
key,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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) =
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst (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) =
  forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
A.lookup key
key Map key (Set value)
m) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty 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) =
  forall key value. Map key value -> ListT STM (key, value)
A.listT Map key (Set value)
m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(key
key, Set value
s) -> (key
key,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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) =
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst (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) =
  forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
A.lookup key
key Map key (Set value)
m) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty forall item. Set item -> ListT STM item
B.listT