-- |
-- 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 StmHamt.Prelude hiding (insert, lookup, delete, fold, null)
import StmHamt.Types
import qualified Focus as Focus
import qualified StmHamt.Hamt as Hamt


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

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

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

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

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

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

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

{-# INLINE lookup #-}
lookup :: (Eq key, Hashable key) => (element -> key) -> key -> SizedHamt element -> STM (Maybe element)
lookup :: (element -> key) -> key -> SizedHamt element -> STM (Maybe element)
lookup element -> key
elementToKey key
key (SizedHamt TVar Int
_ Hamt element
hamt) = (element -> key) -> key -> Hamt element -> STM (Maybe element)
forall key element.
(Eq key, 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 :: SizedHamt a -> UnfoldlM STM a
unfoldlM (SizedHamt TVar Int
_ Hamt a
hamt) = Hamt a -> UnfoldlM STM a
forall a. Hamt a -> UnfoldlM STM a
Hamt.unfoldlM Hamt a
hamt

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