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

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

-- |
-- A hash set, based on an STM-specialized hash array mapped trie.
newtype Set item
  = Set (A.SizedHamt item)
  deriving (Typeable)

-- |
-- Construct a new set.
{-# INLINEABLE new #-}
new :: STM (Set item)
new :: forall item. STM (Set item)
new =
  SizedHamt item -> Set item
forall item. SizedHamt item -> Set item
Set (SizedHamt item -> Set item)
-> STM (SizedHamt item) -> STM (Set item)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STM (SizedHamt item)
forall element. STM (SizedHamt element)
A.new

-- |
-- Construct a new set 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 (Set item)
newIO :: forall item. IO (Set item)
newIO =
  SizedHamt item -> Set item
forall item. SizedHamt item -> Set item
Set (SizedHamt item -> Set item)
-> IO (SizedHamt item) -> IO (Set item)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (SizedHamt item)
forall element. IO (SizedHamt element)
A.newIO

-- |
-- Check, whether the set is empty.
{-# INLINEABLE null #-}
null :: Set item -> STM Bool
null :: forall item. Set item -> STM Bool
null (Set SizedHamt item
hamt) =
  SizedHamt item -> STM Bool
forall element. SizedHamt element -> STM Bool
A.null SizedHamt item
hamt

-- |
-- Get the number of elements.
{-# INLINEABLE size #-}
size :: Set item -> STM Int
size :: forall item. Set item -> STM Int
size (Set SizedHamt item
hamt) =
  SizedHamt item -> STM Int
forall element. SizedHamt element -> STM Int
A.size SizedHamt item
hamt

-- |
-- Focus on an element with a strategy.
--
-- This function allows to perform simultaneous lookup and modification.
--
-- The strategy is over a unit since we already know,
-- which element we're focusing on and it doesn't make sense to replace it,
-- however we still can decide wether to keep or remove it.
{-# INLINEABLE focus #-}
focus :: (Hashable item) => B.Focus () STM result -> item -> Set item -> STM result
focus :: forall item result.
Hashable item =>
Focus () STM result -> item -> Set item -> STM result
focus Focus () STM result
unitFocus item
item (Set SizedHamt item
hamt) =
  Focus item STM result
-> (item -> item) -> item -> SizedHamt item -> STM result
forall key element result.
Hashable key =>
Focus element STM result
-> (element -> key) -> key -> SizedHamt element -> STM result
A.focus Focus item STM result
rowFocus item -> item
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id item
item SizedHamt item
hamt
  where
    rowFocus :: Focus item STM result
rowFocus =
      (() -> item)
-> (item -> ()) -> Focus () STM result -> Focus item STM result
forall (m :: * -> *) a b x.
Monad m =>
(a -> b) -> (b -> a) -> Focus a m x -> Focus b m x
B.mappingInput (item -> () -> item
forall a b. a -> b -> a
const item
item) (() -> item -> ()
forall a b. a -> b -> a
const ()) Focus () STM result
unitFocus

-- |
-- Lookup an element.
{-# INLINEABLE lookup #-}
lookup :: (Hashable item) => item -> Set item -> STM Bool
lookup :: forall item. Hashable item => item -> Set item -> STM Bool
lookup =
  Focus () STM Bool -> item -> Set item -> STM Bool
forall item result.
Hashable item =>
Focus () STM result -> item -> Set item -> STM result
focus ((Maybe () -> Bool) -> Focus () STM (Maybe ()) -> Focus () STM Bool
forall a b. (a -> b) -> Focus () STM a -> Focus () STM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Maybe () -> Bool
forall a. Maybe a -> Bool
isJust Focus () STM (Maybe ())
forall (m :: * -> *) a. Monad m => Focus a m (Maybe a)
B.lookup)

-- |
-- Insert a new element.
{-# INLINEABLE insert #-}
insert :: (Hashable item) => item -> Set item -> STM ()
insert :: forall item. Hashable item => item -> Set item -> STM ()
insert item
item (Set SizedHamt item
hamt) =
  (item -> item) -> item -> SizedHamt item -> STM ()
forall key element.
Hashable key =>
(element -> key) -> element -> SizedHamt element -> STM ()
A.insert item -> item
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id item
item SizedHamt item
hamt

-- |
-- Delete an element.
{-# INLINEABLE delete #-}
delete :: (Hashable item) => item -> Set item -> STM ()
delete :: forall item. Hashable item => item -> Set item -> STM ()
delete item
item (Set SizedHamt item
hamt) =
  Focus item STM ()
-> (item -> item) -> item -> SizedHamt item -> STM ()
forall key element result.
Hashable key =>
Focus element STM result
-> (element -> key) -> key -> SizedHamt element -> STM result
A.focus Focus item STM ()
forall (m :: * -> *) a. Monad m => Focus a m ()
B.delete item -> item
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id item
item SizedHamt item
hamt

-- |
-- Delete all the elements.
{-# INLINEABLE reset #-}
reset :: Set item -> STM ()
reset :: forall item. Set item -> STM ()
reset (Set SizedHamt item
hamt) =
  SizedHamt item -> STM ()
forall element. SizedHamt element -> STM ()
A.reset SizedHamt item
hamt

-- |
-- Stream the elements actively.
--
-- Amongst other features this function provides an interface to folding.
{-# INLINEABLE unfoldlM #-}
unfoldlM :: Set item -> UnfoldlM STM item
unfoldlM :: forall item. Set item -> UnfoldlM STM item
unfoldlM (Set SizedHamt item
hamt) =
  SizedHamt item -> UnfoldlM STM item
forall a. SizedHamt a -> UnfoldlM STM a
A.unfoldlM SizedHamt item
hamt

-- |
-- Stream the elements passively.
{-# INLINE listT #-}
listT :: Set item -> ListT STM item
listT :: forall item. Set item -> ListT STM item
listT (Set SizedHamt item
hamt) =
  SizedHamt item -> ListT STM item
forall a. SizedHamt a -> ListT STM a
A.listT SizedHamt item
hamt