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

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


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

-- |
-- Construct a new map.
{-# INLINABLE new #-}
new :: STM (Map key value)
new :: 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.
{-# INLINABLE newIO #-}
newIO :: IO (Map key value)
newIO :: 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.
{-# INLINABLE null #-}
null :: Map key value -> STM Bool
null :: Map key value -> STM Bool
null (Map hamt :: 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.
{-# INLINABLE size #-}
size :: Map key value -> STM Int
size :: 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' (\ x :: Int
x _ -> Int -> STM Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Int
forall a. Enum a => a -> a
succ Int
x)) 0 (UnfoldlM STM (key, value) -> STM Int)
-> (Map key value -> UnfoldlM STM (key, value))
-> Map key value
-> STM Int
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 :: (Eq key, Hashable key) => B.Focus value STM result -> key -> Map key value -> STM result
focus :: Focus value STM result -> key -> Map key value -> STM result
focus valueFocus :: Focus value STM result
valueFocus key :: key
key (Map hamt :: 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.
(Eq key, 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
key _) -> 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
value -> key -> value -> Product2 key value
forall a b. a -> b -> Product2 a b
Product2 key
key value
value) (\(Product2 _ value :: value
value) -> value
value) Focus value STM result
valueFocus

-- |
-- Look up an item.
{-# INLINABLE lookup #-}
lookup :: (Eq key, Hashable key) => key -> Map key value -> STM (Maybe value)
lookup :: key -> Map key value -> STM (Maybe value)
lookup key :: key
key =
  Focus value STM (Maybe value)
-> key -> Map key value -> STM (Maybe value)
forall key value result.
(Eq key, 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 :: (Eq key, Hashable key) => value -> key -> Map key value -> STM ()
insert :: value -> key -> Map key value -> STM ()
insert value :: value
value key :: key
key (Map hamt :: 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.
(Eq key, Hashable key) =>
(element -> key) -> element -> Hamt element -> STM Bool
A.insert (\(Product2 key :: key
key _) -> 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.
{-# INLINABLE delete #-}
delete :: (Eq key, Hashable key) => key -> Map key value -> STM ()
delete :: key -> Map key value -> STM ()
delete key :: key
key =
  Focus value STM () -> key -> Map key value -> STM ()
forall key value result.
(Eq key, 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.
{-# INLINABLE reset #-}
reset :: Map key value -> STM ()
reset :: Map key value -> STM ()
reset (Map hamt :: 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.
{-# INLINABLE unfoldlM #-}
unfoldlM :: Map key value -> UnfoldlM STM (key, value)
unfoldlM :: Map key value -> UnfoldlM STM (key, value)
unfoldlM (Map hamt :: Hamt (Product2 key value)
hamt) =
  (Product2 key value -> (key, value))
-> UnfoldlM STM (Product2 key value) -> UnfoldlM STM (key, value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ (Product2 k :: key
k v :: 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 :: Map key value -> ListT STM (key, value)
listT (Map hamt :: Hamt (Product2 key value)
hamt) =
  (Product2 key value -> (key, value))
-> ListT STM (Product2 key value) -> ListT STM (key, value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ (Product2 k :: key
k v :: 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)