module StmContainers.Map
  ( Map,
    new,
    newIO,
    null,
    size,
    focus,
    lookup,
    insert,
    delete,
    reset,
    unfoldlM,
    listT,
  )
where

import qualified DeferredFolds.UnfoldlM as C
import qualified Focus as B
import StmContainers.Prelude hiding (delete, empty, foldM, insert, lookup, null, toList)
import qualified StmHamt.Hamt as A

-- |
-- Hash-table, based on STM-specialized Hash Array Mapped Trie.
newtype Map key value
  = Map (A.Hamt (Product2 key value))

-- |
-- Construct a new map.
{-# INLINEABLE new #-}
new :: STM (Map key value)
new :: forall key value. STM (Map key value)
new =
  forall key value. Hamt (Product2 key value) -> Map key value
Map forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. STM (Hamt a)
A.new

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

-- |
-- Check, whether the map is empty.
{-# INLINEABLE null #-}
null :: Map key value -> STM Bool
null :: forall key value. Map key value -> STM Bool
null (Map Hamt (Product2 key value)
hamt) =
  forall a. Hamt a -> STM Bool
A.null Hamt (Product2 key value)
hamt

-- |
-- Get the number of elements.
{-# INLINEABLE size #-}
size :: Map key value -> STM Int
size :: forall key value. Map key value -> STM Int
size =
  forall (m :: * -> *) output input.
Monad m =>
(output -> input -> m output)
-> output -> UnfoldlM m input -> m output
C.foldlM' (\Int
x (key, value)
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. Enum a => a -> a
succ Int
x)) Int
0 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall key value. Map key value -> UnfoldlM STM (key, value)
unfoldlM

-- |
-- Focus on a value by the key.
--
-- This function allows to perform composite operations in a single access
-- to the map's row.
-- E.g., you can look up a value and delete it at the same time,
-- or update it and return the new value.
{-# INLINE focus #-}
focus :: (Hashable key) => B.Focus value STM result -> key -> Map key value -> STM result
focus :: forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
focus Focus value STM result
valueFocus key
key (Map Hamt (Product2 key value)
hamt) =
  forall key element result.
Hashable key =>
Focus element STM result
-> (element -> key) -> key -> Hamt element -> STM result
A.focus Focus (Product2 key value) STM result
rowFocus (\(Product2 key
key value
_) -> key
key) key
key Hamt (Product2 key value)
hamt
  where
    rowFocus :: Focus (Product2 key value) STM result
rowFocus =
      forall (m :: * -> *) a b x.
Monad m =>
(a -> b) -> (b -> a) -> Focus a m x -> Focus b m x
B.mappingInput (\value
value -> forall a b. a -> b -> Product2 a b
Product2 key
key value
value) (\(Product2 key
_ value
value) -> value
value) Focus value STM result
valueFocus

-- |
-- Look up an item.
{-# INLINEABLE lookup #-}
lookup :: (Hashable key) => key -> Map key value -> STM (Maybe value)
lookup :: forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
lookup key
key =
  forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
focus forall (m :: * -> *) a. Monad m => Focus a m (Maybe a)
B.lookup key
key

-- |
-- Insert a value at a key.
{-# INLINE insert #-}
insert :: (Hashable key) => value -> key -> Map key value -> STM ()
insert :: forall key value.
Hashable key =>
value -> key -> Map key value -> STM ()
insert value
value key
key (Map Hamt (Product2 key value)
hamt) =
  forall (f :: * -> *) a. Functor f => f a -> f ()
void (forall key element.
Hashable key =>
(element -> key) -> element -> Hamt element -> STM Bool
A.insert (\(Product2 key
key value
_) -> key
key) (forall a b. a -> b -> Product2 a b
Product2 key
key value
value) Hamt (Product2 key value)
hamt)

-- |
-- Delete an item by a key.
{-# INLINEABLE delete #-}
delete :: (Hashable key) => key -> Map key value -> STM ()
delete :: forall key value. Hashable key => key -> Map key value -> STM ()
delete key
key =
  forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
focus forall (m :: * -> *) a. Monad m => Focus a m ()
B.delete key
key

-- |
-- Delete all the associations.
{-# INLINEABLE reset #-}
reset :: Map key value -> STM ()
reset :: forall key value. Map key value -> STM ()
reset (Map Hamt (Product2 key value)
hamt) =
  forall a. Hamt a -> STM ()
A.reset Hamt (Product2 key value)
hamt

-- |
-- Stream the associations actively.
--
-- Amongst other features this function provides an interface to folding.
{-# INLINEABLE unfoldlM #-}
unfoldlM :: Map key value -> UnfoldlM STM (key, value)
unfoldlM :: forall key value. Map key value -> UnfoldlM STM (key, value)
unfoldlM (Map Hamt (Product2 key value)
hamt) =
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Product2 key
k value
v) -> (key
k, value
v)) (forall a. Hamt a -> UnfoldlM STM a
A.unfoldlM Hamt (Product2 key value)
hamt)

-- |
-- Stream the associations passively.
{-# INLINE listT #-}
listT :: Map key value -> ListT STM (key, value)
listT :: forall key value. Map key value -> ListT STM (key, value)
listT (Map Hamt (Product2 key value)
hamt) =
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Product2 key
k value
v) -> (key
k, value
v)) (forall a. Hamt a -> ListT STM a
A.listT Hamt (Product2 key value)
hamt)