module StmContainers.Map
  ( Map,
    new,
    newIO,
    null,
    size,
    focus,
    lookup,
    insert,
    delete,
    reset,
    unfoldlM,
    listT,
    listTNonAtomic,
  )
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 =
  Hamt (Product2 key value) -> Map key value
forall key value. Hamt (Product2 key value) -> Map key value
Map (Hamt (Product2 key value) -> Map key value)
-> STM (Hamt (Product2 key value)) -> STM (Map key value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STM (Hamt (Product2 key value))
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 =
  Hamt (Product2 key value) -> Map key value
forall key value. Hamt (Product2 key value) -> Map key value
Map (Hamt (Product2 key value) -> Map key value)
-> IO (Hamt (Product2 key value)) -> IO (Map key value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Hamt (Product2 key value))
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) =
  Hamt (Product2 key value) -> STM Bool
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 =
  (Int -> (key, value) -> STM Int)
-> Int -> UnfoldlM STM (key, value) -> STM Int
forall (m :: * -> *) output input.
Monad m =>
(output -> input -> m output)
-> output -> UnfoldlM m input -> m output
C.foldlM' (\Int
x (key, value)
_ -> Int -> STM Int
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Int
forall a. Enum a => a -> a
succ Int
x)) Int
0 (UnfoldlM STM (key, value) -> STM Int)
-> (Map key value -> UnfoldlM STM (key, value))
-> Map key value
-> STM Int
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
. Map key value -> UnfoldlM STM (key, value)
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) =
  Focus (Product2 key value) STM result
-> (Product2 key value -> key)
-> key
-> Hamt (Product2 key value)
-> STM result
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 =
      (value -> Product2 key value)
-> (Product2 key value -> value)
-> Focus value STM result
-> Focus (Product2 key value) STM result
forall (m :: * -> *) a b x.
Monad m =>
(a -> b) -> (b -> a) -> Focus a m x -> Focus b m x
B.mappingInput (\value
value -> key -> value -> Product2 key 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 =
  Focus value STM (Maybe value)
-> key -> Map key value -> STM (Maybe value)
forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
focus Focus value STM (Maybe value)
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) =
  STM Bool -> STM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void ((Product2 key value -> key)
-> Product2 key value -> Hamt (Product2 key value) -> STM Bool
forall key element.
Hashable key =>
(element -> key) -> element -> Hamt element -> STM Bool
A.insert (\(Product2 key
key value
_) -> key
key) (key -> value -> Product2 key value
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 =
  Focus value STM () -> key -> Map key value -> STM ()
forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
focus Focus value STM ()
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) =
  Hamt (Product2 key value) -> STM ()
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) =
  (Product2 key value -> (key, value))
-> UnfoldlM STM (Product2 key value) -> UnfoldlM STM (key, value)
forall a b. (a -> b) -> UnfoldlM STM a -> UnfoldlM STM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Product2 key
k value
v) -> (key
k, value
v)) (Hamt (Product2 key value) -> UnfoldlM STM (Product2 key value)
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) = (Product2 key value -> (key, value))
-> ListT STM (Product2 key value) -> ListT STM (key, value)
forall a b. (a -> b) -> ListT STM a -> ListT STM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Product2 key
k value
v) -> (key
k, value
v)) (Hamt (Product2 key value) -> ListT STM (Product2 key value)
forall a. Hamt a -> ListT STM a
A.listT Hamt (Product2 key value)
hamt)

-- |
-- Stream the associations passively.
-- Data may be inconsistent/out of date.
{-# INLINE listTNonAtomic #-}
listTNonAtomic :: Map key value -> ListT IO (key, value)
listTNonAtomic :: forall key value. Map key value -> ListT IO (key, value)
listTNonAtomic (Map Hamt (Product2 key value)
hamt) = (Product2 key value -> (key, value))
-> ListT IO (Product2 key value) -> ListT IO (key, value)
forall a b. (a -> b) -> ListT IO a -> ListT IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Product2 key
k value
v) -> (key
k, value
v)) (Hamt (Product2 key value) -> ListT IO (Product2 key value)
forall a. Hamt a -> ListT IO a
A.listTNonAtomic Hamt (Product2 key value)
hamt)