module StmContainers.Bimap
  ( Bimap,
    new,
    newIO,
    null,
    size,
    focusLeft,
    focusRight,
    lookupLeft,
    lookupRight,
    insertLeft,
    insertRight,
    deleteLeft,
    deleteRight,
    reset,
    unfoldlM,
    listT,
  )
where

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

-- |
-- Bidirectional map.
-- Essentially, a bijection between subsets of its two argument types.
--
-- For one value of the left-hand type this map contains one value
-- of the right-hand type and vice versa.
data Bimap leftKey rightKey
  = Bimap !(A.Map leftKey rightKey) !(A.Map rightKey leftKey)
  deriving (Typeable)

-- |
-- Construct a new bimap.
{-# INLINE new #-}
new :: STM (Bimap leftKey rightKey)
new :: forall leftKey rightKey. STM (Bimap leftKey rightKey)
new =
  Map leftKey rightKey
-> Map rightKey leftKey -> Bimap leftKey rightKey
forall leftKey rightKey.
Map leftKey rightKey
-> Map rightKey leftKey -> Bimap leftKey rightKey
Bimap (Map leftKey rightKey
 -> Map rightKey leftKey -> Bimap leftKey rightKey)
-> STM (Map leftKey rightKey)
-> STM (Map rightKey leftKey -> Bimap leftKey rightKey)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STM (Map leftKey rightKey)
forall key value. STM (Map key value)
A.new STM (Map rightKey leftKey -> Bimap leftKey rightKey)
-> STM (Map rightKey leftKey) -> STM (Bimap leftKey rightKey)
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> STM (Map rightKey leftKey)
forall key value. STM (Map key value)
A.new

-- |
-- Construct a new bimap 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 (Bimap leftKey rightKey)
newIO :: forall leftKey rightKey. IO (Bimap leftKey rightKey)
newIO =
  Map leftKey rightKey
-> Map rightKey leftKey -> Bimap leftKey rightKey
forall leftKey rightKey.
Map leftKey rightKey
-> Map rightKey leftKey -> Bimap leftKey rightKey
Bimap (Map leftKey rightKey
 -> Map rightKey leftKey -> Bimap leftKey rightKey)
-> IO (Map leftKey rightKey)
-> IO (Map rightKey leftKey -> Bimap leftKey rightKey)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Map leftKey rightKey)
forall key value. IO (Map key value)
A.newIO IO (Map rightKey leftKey -> Bimap leftKey rightKey)
-> IO (Map rightKey leftKey) -> IO (Bimap leftKey rightKey)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IO (Map rightKey leftKey)
forall key value. IO (Map key value)
A.newIO

-- |
-- Check on being empty.
{-# INLINE null #-}
null :: Bimap leftKey rightKey -> STM Bool
null :: forall leftKey rightKey. Bimap leftKey rightKey -> STM Bool
null (Bimap Map leftKey rightKey
leftMap Map rightKey leftKey
_) =
  Map leftKey rightKey -> STM Bool
forall key value. Map key value -> STM Bool
A.null Map leftKey rightKey
leftMap

-- |
-- Get the number of elements.
{-# INLINE size #-}
size :: Bimap leftKey rightKey -> STM Int
size :: forall leftKey rightKey. Bimap leftKey rightKey -> STM Int
size (Bimap Map leftKey rightKey
leftMap Map rightKey leftKey
_) =
  Map leftKey rightKey -> STM Int
forall key value. Map key value -> STM Int
A.size Map leftKey rightKey
leftMap

-- |
-- Focus on a right value by the left value.
--
-- This function allows to perform composite operations in a single access
-- to a map item.
-- E.g., you can look up an item and delete it at the same time,
-- or update it and return the new value.
{-# INLINE focusLeft #-}
focusLeft :: (Hashable leftKey, Hashable rightKey) => B.Focus rightKey STM result -> leftKey -> Bimap leftKey rightKey -> STM result
focusLeft :: forall leftKey rightKey result.
(Hashable leftKey, Hashable rightKey) =>
Focus rightKey STM result
-> leftKey -> Bimap leftKey rightKey -> STM result
focusLeft Focus rightKey STM result
rightFocus leftKey
leftKey (Bimap Map leftKey rightKey
leftMap Map rightKey leftKey
rightMap) =
  do
    ((result
output, Change rightKey
change), Maybe rightKey
maybeRightKey) <- Focus rightKey STM ((result, Change rightKey), Maybe rightKey)
-> leftKey
-> Map leftKey rightKey
-> STM ((result, Change rightKey), Maybe rightKey)
forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
A.focus (Focus rightKey STM (result, Change rightKey)
-> Focus rightKey STM ((result, Change rightKey), Maybe rightKey)
forall (m :: * -> *) a b.
Monad m =>
Focus a m b -> Focus a m (b, Maybe a)
B.extractingInput (Focus rightKey STM result
-> Focus rightKey STM (result, Change rightKey)
forall (m :: * -> *) a b.
Monad m =>
Focus a m b -> Focus a m (b, Change a)
B.extractingChange Focus rightKey STM result
rightFocus)) leftKey
leftKey Map leftKey rightKey
leftMap
    case Change rightKey
change of
      Change rightKey
B.Leave ->
        () -> STM ()
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
      Change rightKey
B.Remove ->
        Maybe rightKey -> (rightKey -> STM ()) -> STM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe rightKey
maybeRightKey ((rightKey -> STM ()) -> STM ()) -> (rightKey -> STM ()) -> STM ()
forall a b. (a -> b) -> a -> b
$ \rightKey
oldRightKey -> rightKey -> Map rightKey leftKey -> STM ()
forall key value. Hashable key => key -> Map key value -> STM ()
A.delete rightKey
oldRightKey Map rightKey leftKey
rightMap
      B.Set rightKey
newRightKey ->
        do
          Maybe rightKey -> (rightKey -> STM ()) -> STM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe rightKey
maybeRightKey ((rightKey -> STM ()) -> STM ()) -> (rightKey -> STM ()) -> STM ()
forall a b. (a -> b) -> a -> b
$ \rightKey
rightKey -> rightKey -> Map rightKey leftKey -> STM ()
forall key value. Hashable key => key -> Map key value -> STM ()
A.delete rightKey
rightKey Map rightKey leftKey
rightMap
          Maybe leftKey
maybeReplacedLeftKey <- Focus leftKey STM (Maybe leftKey)
-> rightKey -> Map rightKey leftKey -> STM (Maybe leftKey)
forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
A.focus (Focus leftKey STM (Maybe leftKey)
forall (m :: * -> *) a. Monad m => Focus a m (Maybe a)
B.lookup Focus leftKey STM (Maybe leftKey)
-> Focus leftKey STM () -> Focus leftKey STM (Maybe leftKey)
forall a b.
Focus leftKey STM a -> Focus leftKey STM b -> Focus leftKey STM a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* leftKey -> Focus leftKey STM ()
forall (m :: * -> *) a. Monad m => a -> Focus a m ()
B.insert leftKey
leftKey) rightKey
newRightKey Map rightKey leftKey
rightMap
          Maybe leftKey -> (leftKey -> STM ()) -> STM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe leftKey
maybeReplacedLeftKey ((leftKey -> STM ()) -> STM ()) -> (leftKey -> STM ()) -> STM ()
forall a b. (a -> b) -> a -> b
$ \leftKey
replacedLeftKey -> leftKey -> Map leftKey rightKey -> STM ()
forall key value. Hashable key => key -> Map key value -> STM ()
A.delete leftKey
replacedLeftKey Map leftKey rightKey
leftMap
    result -> STM result
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return result
output

-- |
-- Focus on a left value by the right value.
--
-- This function allows to perform composite operations in a single access
-- to a map item.
-- E.g., you can look up an item and delete it at the same time,
-- or update it and return the new value.
{-# INLINE focusRight #-}
focusRight :: (Hashable leftKey, Hashable rightKey) => B.Focus leftKey STM result -> rightKey -> Bimap leftKey rightKey -> STM result
focusRight :: forall leftKey rightKey result.
(Hashable leftKey, Hashable rightKey) =>
Focus leftKey STM result
-> rightKey -> Bimap leftKey rightKey -> STM result
focusRight Focus leftKey STM result
valueFocus2 rightKey
rightKey (Bimap Map leftKey rightKey
leftMap Map rightKey leftKey
rightMap) =
  Focus leftKey STM result
-> rightKey -> Bimap rightKey leftKey -> STM result
forall leftKey rightKey result.
(Hashable leftKey, Hashable rightKey) =>
Focus rightKey STM result
-> leftKey -> Bimap leftKey rightKey -> STM result
focusLeft Focus leftKey STM result
valueFocus2 rightKey
rightKey (Map rightKey leftKey
-> Map leftKey rightKey -> Bimap rightKey leftKey
forall leftKey rightKey.
Map leftKey rightKey
-> Map rightKey leftKey -> Bimap leftKey rightKey
Bimap Map rightKey leftKey
rightMap Map leftKey rightKey
leftMap)

-- |
-- Look up a right value by the left value.
{-# INLINE lookupLeft #-}
lookupLeft :: (Hashable leftKey) => leftKey -> Bimap leftKey rightKey -> STM (Maybe rightKey)
lookupLeft :: forall leftKey rightKey.
Hashable leftKey =>
leftKey -> Bimap leftKey rightKey -> STM (Maybe rightKey)
lookupLeft leftKey
leftKey (Bimap Map leftKey rightKey
leftMap Map rightKey leftKey
_) =
  leftKey -> Map leftKey rightKey -> STM (Maybe rightKey)
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
A.lookup leftKey
leftKey Map leftKey rightKey
leftMap

-- |
-- Look up a left value by the right value.
{-# INLINE lookupRight #-}
lookupRight :: (Hashable rightKey) => rightKey -> Bimap leftKey rightKey -> STM (Maybe leftKey)
lookupRight :: forall rightKey leftKey.
Hashable rightKey =>
rightKey -> Bimap leftKey rightKey -> STM (Maybe leftKey)
lookupRight rightKey
rightKey (Bimap Map leftKey rightKey
_ Map rightKey leftKey
rightMap) =
  rightKey -> Map rightKey leftKey -> STM (Maybe leftKey)
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
A.lookup rightKey
rightKey Map rightKey leftKey
rightMap

-- |
-- Insert the association by the left value.
{-# INLINE insertLeft #-}
insertLeft :: (Hashable leftKey, Hashable rightKey) => rightKey -> leftKey -> Bimap leftKey rightKey -> STM ()
insertLeft :: forall leftKey rightKey.
(Hashable leftKey, Hashable rightKey) =>
rightKey -> leftKey -> Bimap leftKey rightKey -> STM ()
insertLeft rightKey
rightKey =
  Focus rightKey STM ()
-> leftKey -> Bimap leftKey rightKey -> STM ()
forall leftKey rightKey result.
(Hashable leftKey, Hashable rightKey) =>
Focus rightKey STM result
-> leftKey -> Bimap leftKey rightKey -> STM result
focusLeft (rightKey -> Focus rightKey STM ()
forall (m :: * -> *) a. Monad m => a -> Focus a m ()
B.insert rightKey
rightKey)

-- |
-- Insert the association by the right value.
{-# INLINE insertRight #-}
insertRight :: (Hashable leftKey, Hashable rightKey) => leftKey -> rightKey -> Bimap leftKey rightKey -> STM ()
insertRight :: forall leftKey rightKey.
(Hashable leftKey, Hashable rightKey) =>
leftKey -> rightKey -> Bimap leftKey rightKey -> STM ()
insertRight leftKey
leftKey rightKey
rightKey (Bimap Map leftKey rightKey
leftMap Map rightKey leftKey
rightMap) =
  leftKey -> rightKey -> Bimap rightKey leftKey -> STM ()
forall leftKey rightKey.
(Hashable leftKey, Hashable rightKey) =>
rightKey -> leftKey -> Bimap leftKey rightKey -> STM ()
insertLeft leftKey
leftKey rightKey
rightKey (Map rightKey leftKey
-> Map leftKey rightKey -> Bimap rightKey leftKey
forall leftKey rightKey.
Map leftKey rightKey
-> Map rightKey leftKey -> Bimap leftKey rightKey
Bimap Map rightKey leftKey
rightMap Map leftKey rightKey
leftMap)

-- |
-- Delete the association by the left value.
{-# INLINE deleteLeft #-}
deleteLeft :: (Hashable leftKey, Hashable rightKey) => leftKey -> Bimap leftKey rightKey -> STM ()
deleteLeft :: forall leftKey rightKey.
(Hashable leftKey, Hashable rightKey) =>
leftKey -> Bimap leftKey rightKey -> STM ()
deleteLeft leftKey
leftKey (Bimap Map leftKey rightKey
leftMap Map rightKey leftKey
rightMap) =
  Focus rightKey STM (Maybe rightKey)
-> leftKey -> Map leftKey rightKey -> STM (Maybe rightKey)
forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
A.focus Focus rightKey STM (Maybe rightKey)
forall (m :: * -> *) a. Monad m => Focus a m (Maybe a)
B.lookupAndDelete leftKey
leftKey Map leftKey rightKey
leftMap
    STM (Maybe rightKey) -> (Maybe rightKey -> STM ()) -> STM ()
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (rightKey -> STM ()) -> Maybe rightKey -> STM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\rightKey
rightKey -> rightKey -> Map rightKey leftKey -> STM ()
forall key value. Hashable key => key -> Map key value -> STM ()
A.delete rightKey
rightKey Map rightKey leftKey
rightMap)

-- |
-- Delete the association by the right value.
{-# INLINE deleteRight #-}
deleteRight :: (Hashable leftKey, Hashable rightKey) => rightKey -> Bimap leftKey rightKey -> STM ()
deleteRight :: forall leftKey rightKey.
(Hashable leftKey, Hashable rightKey) =>
rightKey -> Bimap leftKey rightKey -> STM ()
deleteRight rightKey
rightKey (Bimap Map leftKey rightKey
leftMap Map rightKey leftKey
rightMap) =
  rightKey -> Bimap rightKey leftKey -> STM ()
forall leftKey rightKey.
(Hashable leftKey, Hashable rightKey) =>
leftKey -> Bimap leftKey rightKey -> STM ()
deleteLeft rightKey
rightKey (Map rightKey leftKey
-> Map leftKey rightKey -> Bimap rightKey leftKey
forall leftKey rightKey.
Map leftKey rightKey
-> Map rightKey leftKey -> Bimap leftKey rightKey
Bimap Map rightKey leftKey
rightMap Map leftKey rightKey
leftMap)

-- |
-- Delete all the associations.
{-# INLINE reset #-}
reset :: Bimap leftKey rightKey -> STM ()
reset :: forall leftKey rightKey. Bimap leftKey rightKey -> STM ()
reset (Bimap Map leftKey rightKey
leftMap Map rightKey leftKey
rightMap) =
  do
    Map leftKey rightKey -> STM ()
forall key value. Map key value -> STM ()
A.reset Map leftKey rightKey
leftMap
    Map rightKey leftKey -> STM ()
forall key value. Map key value -> STM ()
A.reset Map rightKey leftKey
rightMap

-- |
-- Stream associations actively.
--
-- Amongst other features this function provides an interface to folding.
{-# INLINE unfoldlM #-}
unfoldlM :: Bimap leftKey rightKey -> UnfoldlM STM (leftKey, rightKey)
unfoldlM :: forall leftKey rightKey.
Bimap leftKey rightKey -> UnfoldlM STM (leftKey, rightKey)
unfoldlM (Bimap Map leftKey rightKey
leftMap Map rightKey leftKey
_) =
  Map leftKey rightKey -> UnfoldlM STM (leftKey, rightKey)
forall key value. Map key value -> UnfoldlM STM (key, value)
A.unfoldlM Map leftKey rightKey
leftMap

-- |
-- Stream the associations passively.
{-# INLINE listT #-}
listT :: Bimap key value -> ListT STM (key, value)
listT :: forall key value. Bimap key value -> ListT STM (key, value)
listT (Bimap Map key value
leftMap Map value key
_) =
  Map key value -> ListT STM (key, value)
forall key value. Map key value -> ListT STM (key, value)
A.listT Map key value
leftMap