-- |
-- HAMT API,
-- optimized for a fast 'size' operation.
-- That however comes at the cost of a small overhead in the other operations.
module StmHamt.SizedHamt
  ( SizedHamt,
    new,
    newIO,
    null,
    size,
    focus,
    insert,
    lookup,
    reset,
    unfoldlM,
    listT,
  )
where

import qualified Focus as Focus
import qualified StmHamt.Hamt as Hamt
import StmHamt.Prelude hiding (delete, fold, insert, lookup, null)
import StmHamt.Types

{-# INLINE new #-}
new :: STM (SizedHamt element)
new :: forall element. STM (SizedHamt element)
new = forall element. TVar Int -> Hamt element -> SizedHamt element
SizedHamt forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> STM (TVar a)
newTVar Int
0 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. STM (Hamt a)
Hamt.new

{-# INLINE newIO #-}
newIO :: IO (SizedHamt element)
newIO :: forall element. IO (SizedHamt element)
newIO = forall element. TVar Int -> Hamt element -> SizedHamt element
SizedHamt forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> IO (TVar a)
newTVarIO Int
0 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. IO (Hamt a)
Hamt.newIO

-- |
-- /O(1)/.
{-# INLINE null #-}
null :: SizedHamt element -> STM Bool
null :: forall element. SizedHamt element -> STM Bool
null (SizedHamt TVar Int
sizeVar Hamt element
_) = (forall a. Eq a => a -> a -> Bool
== Int
0) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. TVar a -> STM a
readTVar TVar Int
sizeVar

-- |
-- /O(1)/.
{-# INLINE size #-}
size :: SizedHamt element -> STM Int
size :: forall element. SizedHamt element -> STM Int
size (SizedHamt TVar Int
sizeVar Hamt element
_) = forall a. TVar a -> STM a
readTVar TVar Int
sizeVar

{-# INLINE reset #-}
reset :: SizedHamt element -> STM ()
reset :: forall element. SizedHamt element -> STM ()
reset (SizedHamt TVar Int
sizeVar Hamt element
hamt) =
  do
    forall a. Hamt a -> STM ()
Hamt.reset Hamt element
hamt
    forall a. TVar a -> a -> STM ()
writeTVar TVar Int
sizeVar Int
0

{-# INLINE focus #-}
focus :: (Hashable key) => Focus element STM result -> (element -> key) -> key -> SizedHamt element -> STM result
focus :: forall key element result.
Hashable key =>
Focus element STM result
-> (element -> key) -> key -> SizedHamt element -> STM result
focus Focus element STM result
focus element -> key
elementToKey key
key (SizedHamt TVar Int
sizeVar Hamt element
hamt) =
  do
    (result
result, Maybe (Int -> Int)
sizeModifier) <- forall key element result.
Hashable key =>
Focus element STM result
-> (element -> key) -> key -> Hamt element -> STM result
Hamt.focus Focus element STM (result, Maybe (Int -> Int))
newFocus element -> key
elementToKey key
key Hamt element
hamt
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe (Int -> Int)
sizeModifier (forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
sizeVar)
    forall (m :: * -> *) a. Monad m => a -> m a
return result
result
  where
    newFocus :: Focus element STM (result, Maybe (Int -> Int))
newFocus = forall (m :: * -> *) sizeChange a b.
Monad m =>
sizeChange
-> sizeChange
-> sizeChange
-> Focus a m b
-> Focus a m (b, sizeChange)
Focus.testingSizeChange (forall a. a -> Maybe a
Just forall a. Enum a => a -> a
pred) forall a. Maybe a
Nothing (forall a. a -> Maybe a
Just forall a. Enum a => a -> a
succ) Focus element STM result
focus

{-# INLINE insert #-}
insert :: (Hashable key) => (element -> key) -> element -> SizedHamt element -> STM ()
insert :: forall key element.
Hashable key =>
(element -> key) -> element -> SizedHamt element -> STM ()
insert element -> key
elementToKey element
element (SizedHamt TVar Int
sizeVar Hamt element
hamt) =
  do
    Bool
inserted <- forall key element.
Hashable key =>
(element -> key) -> element -> Hamt element -> STM Bool
Hamt.insert element -> key
elementToKey element
element Hamt element
hamt
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
inserted (forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
sizeVar forall a. Enum a => a -> a
succ)

{-# INLINE lookup #-}
lookup :: (Hashable key) => (element -> key) -> key -> SizedHamt element -> STM (Maybe element)
lookup :: forall key element.
Hashable key =>
(element -> key) -> key -> SizedHamt element -> STM (Maybe element)
lookup element -> key
elementToKey key
key (SizedHamt TVar Int
_ Hamt element
hamt) = forall key element.
Hashable key =>
(element -> key) -> key -> Hamt element -> STM (Maybe element)
Hamt.lookup element -> key
elementToKey key
key Hamt element
hamt

{-# INLINE unfoldlM #-}
unfoldlM :: SizedHamt a -> UnfoldlM STM a
unfoldlM :: forall a. SizedHamt a -> UnfoldlM STM a
unfoldlM (SizedHamt TVar Int
_ Hamt a
hamt) = forall a. Hamt a -> UnfoldlM STM a
Hamt.unfoldlM Hamt a
hamt

{-# INLINE listT #-}
listT :: SizedHamt a -> ListT STM a
listT :: forall a. SizedHamt a -> ListT STM a
listT (SizedHamt TVar Int
_ Hamt a
hamt) = forall a. Hamt a -> ListT STM a
Hamt.listT Hamt a
hamt