{-# LANGUAGE RecordWildCards #-}

-- original implementaion:
-- <https://github.com/maspypy/library/blob/main/ds/hashmap.hpp>

-- | A dense, fast `Int` hash map with a fixed-sized `capacity` of \(n\). Most operations are
-- performed in \(O(1)\) time, but in average.
--
-- ==== Capacity limitation
-- Access to each key creates a new entry. Note that entries cannot be invalidated due to the
-- internal implementation (called /open addressing/).
--
-- ==== __Example__
-- Create a `HashMap` with `capacity` \(10\):
--
-- >>> import AtCoder.Extra.HashMap qualified as HM
-- >>> hm <- HM.new @_ @Int 10
--
-- `insert`, `lookup` and other functions are available in \(O(1)\) averaged time:
--
-- >>> HM.insert hm 0 100
-- >>> HM.insert hm 10 101
-- >>> HM.size hm
-- 2
--
-- >>> HM.lookup hm 0
-- Just 100
--
-- >>> HM.lookup hm 10
-- Just 101
--
-- @since 1.1.0.0
module AtCoder.Extra.HashMap
  ( -- * HashMap
    HashMap,

    -- * Constructors
    new,
    build,

    -- * Metadata
    capacity,
    size,

    -- * Lookups
    lookup,
    member,
    notMember,

    -- * Modifications

    -- ** Insertions
    insert,
    insertWith,
    exchange,

    -- ** Updates
    modify,
    modifyM,

    -- ** Reset
    clear,

    -- * Conversions

    -- ** Safe conversions
    keys,
    elems,
    assocs,

    -- ** Unsafe conversions
    unsafeKeys,
    unsafeElems,
    unsafeAssocs,
  )
where

import AtCoder.Internal.Assert qualified as ACIA
import Control.Monad (void, when)
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST)
import Data.Bit (Bit (..))
import Data.Bits (Bits (xor, (.&.)), (.>>.))
import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import Data.Word (Word64)
import GHC.Stack (HasCallStack)
import Prelude hiding (lookup)

-- | A dense, fast `Int` hash map with a fixed-sized capacity of \(n\).
--
-- @since 1.1.0.0
data HashMap s a = HashMap
  { -- | Maximum number of elements.
    forall s a. HashMap s a -> Int
maxCapHM :: {-# UNPACK #-} !Int,
    -- | The number of elements that can be added.
    forall s a. HashMap s a -> MVector s Int
restCapHM :: !(VUM.MVector s Int),
    -- | Bit mask for powerset iteration on indexing.
    forall s a. HashMap s a -> Int
maskHM :: {-# UNPACK #-} !Int,
    -- | Original key to the hash index.
    forall s a. HashMap s a -> MVector s Int
keyHM :: !(VUM.MVector s Int),
    -- | Values to the hash index.
    forall s a. HashMap s a -> MVector s a
valHM :: !(VUM.MVector s a),
    -- | Whether the slot is used or not.
    forall s a. HashMap s a -> MVector s Bit
usedHM :: !(VUM.MVector s Bit)
  }

{-# INLINE decrementRestCapacityST #-}
decrementRestCapacityST :: (HasCallStack) => VUM.MVector s Int -> String -> ST s ()
decrementRestCapacityST :: forall s. HasCallStack => MVector s Int -> String -> ST s ()
decrementRestCapacityST MVector s Int
restCap String
funcName = do
  Int
rest <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector s Int
MVector (PrimState (ST s)) Int
restCap Int
0
  let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
rest Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (String -> ()) -> String -> ()
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.HashMap." String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
funcName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": out of capacity"
  MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
restCap Int
0 (Int
rest Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | \(O(n)\) Creates a `HashMap` of capacity \(n\).
--
-- @since 1.1.0.0
{-# INLINE new #-}
new :: (PrimMonad m, VU.Unbox a) => Int -> m (HashMap (PrimState m) a)
new :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (HashMap (PrimState m) a)
new Int
n = ST (PrimState m) (HashMap (PrimState m) a)
-> m (HashMap (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (HashMap (PrimState m) a)
 -> m (HashMap (PrimState m) a))
-> ST (PrimState m) (HashMap (PrimState m) a)
-> m (HashMap (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Int -> ST (PrimState m) (HashMap (PrimState m) a)
forall a s. Unbox a => Int -> ST s (HashMap s a)
newST Int
n

-- | \(O(n)\) Creates a `HashMap` of capacity \(n\) with initial entries.
--
-- @since 1.1.0.0
{-# INLINE build #-}
build :: (PrimMonad m, VU.Unbox a) => Int -> VU.Vector (Int, a) -> m (HashMap (PrimState m) a)
build :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> Vector (Int, a) -> m (HashMap (PrimState m) a)
build Int
n Vector (Int, a)
xs = ST (PrimState m) (HashMap (PrimState m) a)
-> m (HashMap (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (HashMap (PrimState m) a)
 -> m (HashMap (PrimState m) a))
-> ST (PrimState m) (HashMap (PrimState m) a)
-> m (HashMap (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Int
-> Vector (Int, a) -> ST (PrimState m) (HashMap (PrimState m) a)
forall a s. Unbox a => Int -> Vector (Int, a) -> ST s (HashMap s a)
buildST Int
n Vector (Int, a)
xs

-- | \(O(1)\) Returns the maximum number of elements the hash map can store.
--
-- @since 1.1.0.0
{-# INLINE capacity #-}
capacity :: HashMap s a -> Int
capacity :: forall s a. HashMap s a -> Int
capacity = HashMap s a -> Int
forall s a. HashMap s a -> Int
maxCapHM

-- | \(O(1)\) Returns the number of elements in the hash map.
--
-- @since 1.1.0.0
{-# INLINE size #-}
size :: (PrimMonad m) => HashMap (PrimState m) a -> m Int
size :: forall (m :: * -> *) a.
PrimMonad m =>
HashMap (PrimState m) a -> m Int
size HashMap {Int
MVector (PrimState m) a
MVector (PrimState m) Int
MVector (PrimState m) Bit
maxCapHM :: forall s a. HashMap s a -> Int
restCapHM :: forall s a. HashMap s a -> MVector s Int
maskHM :: forall s a. HashMap s a -> Int
keyHM :: forall s a. HashMap s a -> MVector s Int
valHM :: forall s a. HashMap s a -> MVector s a
usedHM :: forall s a. HashMap s a -> MVector s Bit
maxCapHM :: Int
restCapHM :: MVector (PrimState m) Int
maskHM :: Int
keyHM :: MVector (PrimState m) Int
valHM :: MVector (PrimState m) a
usedHM :: MVector (PrimState m) Bit
..} = do
  !Int
rest <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector (PrimState m) Int
restCapHM Int
0
  Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Int
maxCapHM Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
rest

-- | \(O(1)\) Return the value to which the specified key is mapped, or `Nothing` if this map
-- contains no mapping for the key.
--
-- ==== Constraint
-- - The rest capacity must be non-zero. Otherwise it loops forever.
--
-- @since 1.1.0.0
{-# INLINE lookup #-}
lookup :: (HasCallStack, VU.Unbox a, PrimMonad m) => HashMap (PrimState m) a -> Int -> m (Maybe a)
lookup :: forall a (m :: * -> *).
(HasCallStack, Unbox a, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m (Maybe a)
lookup HashMap (PrimState m) a
hm Int
k = ST (PrimState m) (Maybe a) -> m (Maybe a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe a) -> m (Maybe a))
-> ST (PrimState m) (Maybe a) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ HashMap (PrimState m) a -> Int -> ST (PrimState m) (Maybe a)
forall a s.
(HasCallStack, Unbox a) =>
HashMap s a -> Int -> ST s (Maybe a)
lookupST HashMap (PrimState m) a
hm Int
k

-- | \(O(1)\) Checks whether the hash map contains the element.
--
-- ==== Constraint
-- - The rest capacity must be non-zero. Otherwise it loops forever.
--
-- @since 1.1.0.0
{-# INLINE member #-}
member :: (HasCallStack, PrimMonad m) => HashMap (PrimState m) a -> Int -> m Bool
member :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m Bool
member HashMap (PrimState m) a
hm Int
k = ST (PrimState m) Bool -> m Bool
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Bool -> m Bool)
-> ST (PrimState m) Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ HashMap (PrimState m) a -> Int -> ST (PrimState m) Bool
forall s a. HasCallStack => HashMap s a -> Int -> ST s Bool
memberST HashMap (PrimState m) a
hm Int
k

-- | \(O(1)\) Checks whether the hash map does not contain the element.
--
-- ==== Constraint
-- - The rest capacity must be non-zero. Otherwise it loops forever.
--
-- @since 1.1.0.0
{-# INLINE notMember #-}
notMember :: (HasCallStack, PrimMonad m) => HashMap (PrimState m) a -> Int -> m Bool
notMember :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m Bool
notMember HashMap (PrimState m) a
hm Int
k = ST (PrimState m) Bool -> m Bool
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Bool -> m Bool)
-> ST (PrimState m) Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> ST (PrimState m) Bool -> ST (PrimState m) Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap (PrimState m) a -> Int -> ST (PrimState m) Bool
forall s a. HasCallStack => HashMap s a -> Int -> ST s Bool
memberST HashMap (PrimState m) a
hm Int
k

-- | \(O(1)\) Inserts a \((k, v)\) pair.
--
-- ==== Constraint
-- - The rest capacity must be non-zero. Otherwise it loops forever.
--
-- @since 1.1.0.0
{-# INLINE insert #-}
insert :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> Int -> a -> m ()
insert :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
insert HashMap (PrimState m) a
hm Int
k a
v = m (Maybe a) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m (Maybe a) -> m ())
-> (ST (PrimState m) (Maybe a) -> m (Maybe a))
-> ST (PrimState m) (Maybe a)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ST (PrimState m) (Maybe a) -> m (Maybe a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe a) -> m ())
-> ST (PrimState m) (Maybe a) -> m ()
forall a b. (a -> b) -> a -> b
$ HashMap (PrimState m) a -> Int -> a -> ST (PrimState m) (Maybe a)
forall a s.
(HasCallStack, Unbox a) =>
HashMap s a -> Int -> a -> ST s (Maybe a)
exchangeST HashMap (PrimState m) a
hm Int
k a
v

-- | \(O(1)\) Inserts a \((k, v)\) pair. If the key exists, the function will insert the pair
-- \((k, f(v_{\mathrm{new}}, v_{\mathrm{old}}))\).
--
-- ==== Constraint
-- - The rest capacity must be non-zero. Otherwise it loops forever.
--
-- @since 1.1.0.0
{-# INLINE insertWith #-}
insertWith :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> (a -> a -> a) -> Int -> a -> m ()
insertWith :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> (a -> a -> a) -> Int -> a -> m ()
insertWith HashMap (PrimState m) a
hm a -> a -> a
f Int
k a
v = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ HashMap (PrimState m) a
-> (a -> a -> a) -> Int -> a -> ST (PrimState m) ()
forall a s.
(HasCallStack, Unbox a) =>
HashMap s a -> (a -> a -> a) -> Int -> a -> ST s ()
insertWithST HashMap (PrimState m) a
hm a -> a -> a
f Int
k a
v

-- | \(O(1)\) Inserts a \((k, v)\) pair and returns the old value, or `Nothing` if no such entry
-- exists.
--
-- ==== Constraint
-- - The rest capacity must be non-zero. Otherwise it loops forever.
--
-- @since 1.1.0.0
{-# INLINE exchange #-}
exchange :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> Int -> a -> m (Maybe a)
exchange :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m (Maybe a)
exchange HashMap (PrimState m) a
hm Int
k a
v = ST (PrimState m) (Maybe a) -> m (Maybe a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe a) -> m (Maybe a))
-> ST (PrimState m) (Maybe a) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ HashMap (PrimState m) a -> Int -> a -> ST (PrimState m) (Maybe a)
forall a s.
(HasCallStack, Unbox a) =>
HashMap s a -> Int -> a -> ST s (Maybe a)
exchangeST HashMap (PrimState m) a
hm Int
k a
v

-- | \(O(1)\) Modifies the element at the given key. Does nothing if no such entry exists.
--
-- @since 1.1.0.0
{-# INLINE modify #-}
modify :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> (a -> a) -> Int -> m ()
modify :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> (a -> a) -> Int -> m ()
modify HashMap (PrimState m) a
hm a -> a
f Int
k = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ HashMap (PrimState m) a -> (a -> a) -> Int -> ST (PrimState m) ()
forall a s.
(HasCallStack, Unbox a) =>
HashMap s a -> (a -> a) -> Int -> ST s ()
modifyST HashMap (PrimState m) a
hm a -> a
f Int
k

-- | \(O(1)\) Modifies the element at the given key. Does nothing if no such entry exists.
--
-- @since 1.1.0.0
{-# INLINE modifyM #-}
modifyM :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> (a -> m a) -> Int -> m ()
modifyM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> (a -> m a) -> Int -> m ()
modifyM hm :: HashMap (PrimState m) a
hm@HashMap {Int
MVector (PrimState m) a
MVector (PrimState m) Int
MVector (PrimState m) Bit
maxCapHM :: forall s a. HashMap s a -> Int
restCapHM :: forall s a. HashMap s a -> MVector s Int
maskHM :: forall s a. HashMap s a -> Int
keyHM :: forall s a. HashMap s a -> MVector s Int
valHM :: forall s a. HashMap s a -> MVector s a
usedHM :: forall s a. HashMap s a -> MVector s Bit
maxCapHM :: Int
restCapHM :: MVector (PrimState m) Int
maskHM :: Int
keyHM :: MVector (PrimState m) Int
valHM :: MVector (PrimState m) a
usedHM :: MVector (PrimState m) Bit
..} a -> m a
f Int
k = do
  Int
i <- ST (PrimState m) Int -> m Int
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Int -> m Int) -> ST (PrimState m) Int -> m Int
forall a b. (a -> b) -> a -> b
$ HashMap (PrimState m) a -> Int -> ST (PrimState m) Int
forall s a. HasCallStack => HashMap s a -> Int -> ST s Int
indexST HashMap (PrimState m) a
hm Int
k
  Bit Bool
b <- ST (PrimState m) Bit -> m Bit
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Bit -> m Bit) -> ST (PrimState m) Bit -> m Bit
forall a b. (a -> b) -> a -> b
$ MVector (PrimState (ST (PrimState m))) Bit
-> Int -> ST (PrimState m) Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Bit
MVector (PrimState (ST (PrimState m))) Bit
usedHM Int
i
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    MVector (PrimState m) a -> (a -> m a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Int -> m ()
VGM.modifyM MVector (PrimState m) a
valHM a -> m a
f Int
i

-- | \(O(n)\) Clears all the elements.
--
-- @since 1.1.0.0
{-# INLINE clear #-}
clear :: (PrimMonad m) => HashMap (PrimState m) a -> m ()
clear :: forall (m :: * -> *) a.
PrimMonad m =>
HashMap (PrimState m) a -> m ()
clear HashMap (PrimState m) a
hm = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ HashMap (PrimState m) a -> ST (PrimState m) ()
forall s a. HashMap s a -> ST s ()
clearST HashMap (PrimState m) a
hm

-- | \(O(n)\) Enumerates the keys in the hash map.
--
-- @since 1.1.0.0
{-# INLINE keys #-}
keys :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector Int)
keys :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> m (Vector Int)
keys HashMap (PrimState m) a
hm = Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector a
VU.force (Vector Int -> Vector Int) -> m (Vector Int) -> m (Vector Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap (PrimState m) a -> m (Vector Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> m (Vector Int)
unsafeKeys HashMap (PrimState m) a
hm

-- | \(O(n)\) Enumerates the elements (values) in the hash map.
--
-- @since 1.1.0.0
{-# INLINE elems #-}
elems :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector a)
elems :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> m (Vector a)
elems HashMap (PrimState m) a
hm = Vector a -> Vector a
forall a. Unbox a => Vector a -> Vector a
VU.force (Vector a -> Vector a) -> m (Vector a) -> m (Vector a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap (PrimState m) a -> m (Vector a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> m (Vector a)
unsafeElems HashMap (PrimState m) a
hm

-- | \(O(n)\) Enumerates the key-value pairs in the hash map.
--
-- @since 1.1.0.0
{-# INLINE assocs #-}
assocs :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector (Int, a))
assocs :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> m (Vector (Int, a))
assocs HashMap (PrimState m) a
hm = Vector (Int, a) -> Vector (Int, a)
forall a. Unbox a => Vector a -> Vector a
VU.force (Vector (Int, a) -> Vector (Int, a))
-> m (Vector (Int, a)) -> m (Vector (Int, a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap (PrimState m) a -> m (Vector (Int, a))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> m (Vector (Int, a))
unsafeAssocs HashMap (PrimState m) a
hm

-- | \(O(n)\) Enumerates the keys in the hash map.
--
-- @since 1.1.0.0
{-# INLINE unsafeKeys #-}
unsafeKeys :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector Int)
unsafeKeys :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> m (Vector Int)
unsafeKeys HashMap (PrimState m) a
hm = ST (PrimState m) (Vector Int) -> m (Vector Int)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Vector Int) -> m (Vector Int))
-> ST (PrimState m) (Vector Int) -> m (Vector Int)
forall a b. (a -> b) -> a -> b
$ HashMap (PrimState m) a -> ST (PrimState m) (Vector Int)
forall a s. Unbox a => HashMap s a -> ST s (Vector Int)
unsafeKeysST HashMap (PrimState m) a
hm

-- | \(O(n)\) Enumerates the elements (values) in the hash map.
--
-- @since 1.1.0.0
{-# INLINE unsafeElems #-}
unsafeElems :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector a)
unsafeElems :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> m (Vector a)
unsafeElems HashMap (PrimState m) a
hm = ST (PrimState m) (Vector a) -> m (Vector a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Vector a) -> m (Vector a))
-> ST (PrimState m) (Vector a) -> m (Vector a)
forall a b. (a -> b) -> a -> b
$ HashMap (PrimState m) a -> ST (PrimState m) (Vector a)
forall a s. Unbox a => HashMap s a -> ST s (Vector a)
unsafeElemsST HashMap (PrimState m) a
hm

-- | \(O(n)\) Enumerates the key-value pairs in the hash map.
--
-- @since 1.1.0.0
{-# INLINE unsafeAssocs #-}
unsafeAssocs :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector (Int, a))
unsafeAssocs :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> m (Vector (Int, a))
unsafeAssocs HashMap (PrimState m) a
hm = ST (PrimState m) (Vector (Int, a)) -> m (Vector (Int, a))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Vector (Int, a)) -> m (Vector (Int, a)))
-> ST (PrimState m) (Vector (Int, a)) -> m (Vector (Int, a))
forall a b. (a -> b) -> a -> b
$ HashMap (PrimState m) a -> ST (PrimState m) (Vector (Int, a))
forall a s. Unbox a => HashMap s a -> ST s (Vector (Int, a))
unsafeAssocsST HashMap (PrimState m) a
hm

-- -------------------------------------------------------------------------------------------------
-- Internal
-- -------------------------------------------------------------------------------------------------

{-# INLINEABLE newST #-}
newST :: (VU.Unbox a) => Int -> ST s (HashMap s a)
newST :: forall a s. Unbox a => Int -> ST s (HashMap s a)
newST Int
n = do
  let !k0 :: Int
k0 = Int
1
  let !k :: Int
k = (Int -> Bool) -> (Int -> Int) -> Int -> Int
forall a. (a -> Bool) -> (a -> a) -> a -> a
until (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2) Int
k0
  -- we need extra space
  let !maxCapHM :: Int
maxCapHM = Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
  MVector s Int
restCapHM <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
1 Int
maxCapHM
  let !maskHM :: Int
maskHM = Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
  MVector s Int
keyHM <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
k
  MVector s a
valHM <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
k
  MVector s Bit
usedHM <- Int -> Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
k (Bit -> ST s (MVector (PrimState (ST s)) Bit))
-> Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
False
  HashMap s a -> ST s (HashMap s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HashMap {Int
MVector s a
MVector s Int
MVector s Bit
maxCapHM :: Int
restCapHM :: MVector s Int
maskHM :: Int
keyHM :: MVector s Int
valHM :: MVector s a
usedHM :: MVector s Bit
maxCapHM :: Int
restCapHM :: MVector s Int
maskHM :: Int
keyHM :: MVector s Int
valHM :: MVector s a
usedHM :: MVector s Bit
..}

{-# INLINEABLE buildST #-}
buildST :: (VU.Unbox a) => Int -> VU.Vector (Int, a) -> ST s (HashMap s a)
buildST :: forall a s. Unbox a => Int -> Vector (Int, a) -> ST s (HashMap s a)
buildST Int
n Vector (Int, a)
xs = do
  HashMap s a
hm <- Int -> ST s (HashMap s a)
forall a s. Unbox a => Int -> ST s (HashMap s a)
newST Int
n
  Vector (Int, a) -> ((Int, a) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector (Int, a)
xs (((Int, a) -> ST s ()) -> ST s ())
-> ((Int, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
i, !a
x) -> do
    ST s (Maybe a) -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s (Maybe a) -> ST s ()) -> ST s (Maybe a) -> ST s ()
forall a b. (a -> b) -> a -> b
$ HashMap s a -> Int -> a -> ST s (Maybe a)
forall a s.
(HasCallStack, Unbox a) =>
HashMap s a -> Int -> a -> ST s (Maybe a)
exchangeST HashMap s a
hm Int
i a
x
  HashMap s a -> ST s (HashMap s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HashMap s a
hm

-- TODO: no need of INLINE?

-- | \(O(1)\) (Internal) Hash value calculation.
{-# INLINEABLE hash #-}
hash :: HashMap a s -> Int -> Int
hash :: forall a s. HashMap a s -> Int -> Int
hash HashMap a s
hm Int
x = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Int) -> Word64 -> Int
forall a b. (a -> b) -> a -> b
$ (Word64
x3 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
x3 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
.>>. Int
31)) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (HashMap a s -> Int
forall s a. HashMap s a -> Int
maskHM HashMap a s
hm)
  where
    fixedRandom, x1, x2, x3 :: Word64
    fixedRandom :: Word64
fixedRandom = Word64
321896566547
    x1 :: Word64
x1 = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
fixedRandom
    x2 :: Word64
x2 = (Word64
x1 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
x1 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
.>>. Int
30)) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
0xbf58476d1ce4e5b9
    x3 :: Word64
x3 = (Word64
x2 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (Word64
x2 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
.>>. Int
27)) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
0x94d049bb133111eb

-- TODO: INLINE indexST?

-- | \(O(1)\) (Internal) Hashed slot search.
--
-- ==== Constraint
-- - The rest capacity must be non-zero. Otherwise it loops forever.
{-# INLINEABLE indexST #-}
indexST :: (HasCallStack) => HashMap s a -> Int -> ST s Int
indexST :: forall s a. HasCallStack => HashMap s a -> Int -> ST s Int
indexST hm :: HashMap s a
hm@HashMap {Int
MVector s a
MVector s Int
MVector s Bit
maxCapHM :: forall s a. HashMap s a -> Int
restCapHM :: forall s a. HashMap s a -> MVector s Int
maskHM :: forall s a. HashMap s a -> Int
keyHM :: forall s a. HashMap s a -> MVector s Int
valHM :: forall s a. HashMap s a -> MVector s a
usedHM :: forall s a. HashMap s a -> MVector s Bit
maxCapHM :: Int
restCapHM :: MVector s Int
maskHM :: Int
keyHM :: MVector s Int
valHM :: MVector s a
usedHM :: MVector s Bit
..} Int
k = do
  Int -> ST s Int
forall {m :: * -> *}.
(PrimState m ~ s, PrimMonad m) =>
Int -> m Int
inner (HashMap s a -> Int -> Int
forall a s. HashMap a s -> Int -> Int
hash HashMap s a
hm Int
k)
  where
    inner :: Int -> m Int
inner !Int
h = do
      Bit Bool
b <- MVector (PrimState m) Bit -> Int -> m Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Bit
MVector (PrimState m) Bit
usedHM Int
h
      -- already there?
      Int
k' <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState m) Int
keyHM Int
h
      if Bool
b Bool -> Bool -> Bool
&& Int
k' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
k
        then Int -> m Int
inner (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
maskHM
        else Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
h

{-# INLINEABLE lookupST #-}
lookupST :: (HasCallStack, VU.Unbox a) => HashMap s a -> Int -> ST s (Maybe a)
lookupST :: forall a s.
(HasCallStack, Unbox a) =>
HashMap s a -> Int -> ST s (Maybe a)
lookupST hm :: HashMap s a
hm@HashMap {Int
MVector s a
MVector s Int
MVector s Bit
maxCapHM :: forall s a. HashMap s a -> Int
restCapHM :: forall s a. HashMap s a -> MVector s Int
maskHM :: forall s a. HashMap s a -> Int
keyHM :: forall s a. HashMap s a -> MVector s Int
valHM :: forall s a. HashMap s a -> MVector s a
usedHM :: forall s a. HashMap s a -> MVector s Bit
maxCapHM :: Int
restCapHM :: MVector s Int
maskHM :: Int
keyHM :: MVector s Int
valHM :: MVector s a
usedHM :: MVector s Bit
..} Int
k = do
  Int
i <- HashMap s a -> Int -> ST s Int
forall s a. HasCallStack => HashMap s a -> Int -> ST s Int
indexST HashMap s a
hm Int
k
  Bit Bool
b <- MVector (PrimState (ST s)) Bit -> Int -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Bit
MVector (PrimState (ST s)) Bit
usedHM Int
i
  if Bool
b
    then a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> ST s a -> ST s (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
valHM Int
i
    else Maybe a -> ST s (Maybe a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing

{-# INLINEABLE memberST #-}
memberST :: (HasCallStack) => HashMap s a -> Int -> ST s Bool
memberST :: forall s a. HasCallStack => HashMap s a -> Int -> ST s Bool
memberST hm :: HashMap s a
hm@HashMap {Int
MVector s a
MVector s Int
MVector s Bit
maxCapHM :: forall s a. HashMap s a -> Int
restCapHM :: forall s a. HashMap s a -> MVector s Int
maskHM :: forall s a. HashMap s a -> Int
keyHM :: forall s a. HashMap s a -> MVector s Int
valHM :: forall s a. HashMap s a -> MVector s a
usedHM :: forall s a. HashMap s a -> MVector s Bit
maxCapHM :: Int
restCapHM :: MVector s Int
maskHM :: Int
keyHM :: MVector s Int
valHM :: MVector s a
usedHM :: MVector s Bit
..} Int
k = do
  Int
i <- HashMap s a -> Int -> ST s Int
forall s a. HasCallStack => HashMap s a -> Int -> ST s Int
indexST HashMap s a
hm Int
k
  Bit Bool
b <- MVector (PrimState (ST s)) Bit -> Int -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Bit
MVector (PrimState (ST s)) Bit
usedHM Int
i
  -- TODO: is this key check necessary
  Int
k' <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
keyHM Int
i
  Bool -> ST s Bool
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> ST s Bool) -> Bool -> ST s Bool
forall a b. (a -> b) -> a -> b
$ Bool
b Bool -> Bool -> Bool
&& Int
k' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k

{-# INLINEABLE insertWithST #-}
insertWithST :: (HasCallStack, VU.Unbox a) => HashMap s a -> (a -> a -> a) -> Int -> a -> ST s ()
insertWithST :: forall a s.
(HasCallStack, Unbox a) =>
HashMap s a -> (a -> a -> a) -> Int -> a -> ST s ()
insertWithST hm :: HashMap s a
hm@HashMap {Int
MVector s a
MVector s Int
MVector s Bit
maxCapHM :: forall s a. HashMap s a -> Int
restCapHM :: forall s a. HashMap s a -> MVector s Int
maskHM :: forall s a. HashMap s a -> Int
keyHM :: forall s a. HashMap s a -> MVector s Int
valHM :: forall s a. HashMap s a -> MVector s a
usedHM :: forall s a. HashMap s a -> MVector s Bit
maxCapHM :: Int
restCapHM :: MVector s Int
maskHM :: Int
keyHM :: MVector s Int
valHM :: MVector s a
usedHM :: MVector s Bit
..} a -> a -> a
f Int
k a
v = do
  Int
i <- HashMap s a -> Int -> ST s Int
forall s a. HasCallStack => HashMap s a -> Int -> ST s Int
indexST HashMap s a
hm Int
k
  Bit Bool
b <- MVector (PrimState (ST s)) Bit -> Int -> Bit -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector s Bit
MVector (PrimState (ST s)) Bit
usedHM Int
i (Bit -> ST s Bit) -> Bit -> ST s Bit
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
True
  if Bool
b
    then do
      -- modify the existing entry
      MVector (PrimState (ST s)) a -> (a -> a) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s a
MVector (PrimState (ST s)) a
valHM (a -> a -> a
f a
v) Int
i
    else do
      -- insert the new \((k, v)\) pair
      MVector s Int -> String -> ST s ()
forall s. HasCallStack => MVector s Int -> String -> ST s ()
decrementRestCapacityST MVector s Int
restCapHM String
"insertWith"
      MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
keyHM Int
i Int
k
      MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
valHM Int
i a
v

{-# INLINEABLE exchangeST #-}
exchangeST :: (HasCallStack, VU.Unbox a) => HashMap s a -> Int -> a -> ST s (Maybe a)
exchangeST :: forall a s.
(HasCallStack, Unbox a) =>
HashMap s a -> Int -> a -> ST s (Maybe a)
exchangeST hm :: HashMap s a
hm@HashMap {Int
MVector s a
MVector s Int
MVector s Bit
maxCapHM :: forall s a. HashMap s a -> Int
restCapHM :: forall s a. HashMap s a -> MVector s Int
maskHM :: forall s a. HashMap s a -> Int
keyHM :: forall s a. HashMap s a -> MVector s Int
valHM :: forall s a. HashMap s a -> MVector s a
usedHM :: forall s a. HashMap s a -> MVector s Bit
maxCapHM :: Int
restCapHM :: MVector s Int
maskHM :: Int
keyHM :: MVector s Int
valHM :: MVector s a
usedHM :: MVector s Bit
..} Int
k a
v = do
  Int
i <- HashMap s a -> Int -> ST s Int
forall s a. HasCallStack => HashMap s a -> Int -> ST s Int
indexST HashMap s a
hm Int
k
  Bit Bool
b <- MVector (PrimState (ST s)) Bit -> Int -> Bit -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector s Bit
MVector (PrimState (ST s)) Bit
usedHM Int
i (Bit -> ST s Bit) -> Bit -> ST s Bit
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
True
  if Bool
b
    then do
      -- overwrite the existing entry
      a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> ST s a -> ST s (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) a -> Int -> a -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector s a
MVector (PrimState (ST s)) a
valHM Int
i a
v
    else do
      -- insert the new (key, value) pair
      MVector s Int -> String -> ST s ()
forall s. HasCallStack => MVector s Int -> String -> ST s ()
decrementRestCapacityST MVector s Int
restCapHM String
"exchange"
      MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
keyHM Int
i Int
k
      MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
valHM Int
i a
v
      Maybe a -> ST s (Maybe a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing

{-# INLINEABLE modifyST #-}
modifyST :: (HasCallStack, VU.Unbox a) => HashMap s a -> (a -> a) -> Int -> ST s ()
modifyST :: forall a s.
(HasCallStack, Unbox a) =>
HashMap s a -> (a -> a) -> Int -> ST s ()
modifyST hm :: HashMap s a
hm@HashMap {Int
MVector s a
MVector s Int
MVector s Bit
maxCapHM :: forall s a. HashMap s a -> Int
restCapHM :: forall s a. HashMap s a -> MVector s Int
maskHM :: forall s a. HashMap s a -> Int
keyHM :: forall s a. HashMap s a -> MVector s Int
valHM :: forall s a. HashMap s a -> MVector s a
usedHM :: forall s a. HashMap s a -> MVector s Bit
maxCapHM :: Int
restCapHM :: MVector s Int
maskHM :: Int
keyHM :: MVector s Int
valHM :: MVector s a
usedHM :: MVector s Bit
..} a -> a
f Int
k = do
  Int
i <- HashMap s a -> Int -> ST s Int
forall s a. HasCallStack => HashMap s a -> Int -> ST s Int
indexST HashMap s a
hm Int
k
  Bit Bool
b <- MVector (PrimState (ST s)) Bit -> Int -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Bit
MVector (PrimState (ST s)) Bit
usedHM Int
i
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    MVector (PrimState (ST s)) a -> (a -> a) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s a
MVector (PrimState (ST s)) a
valHM a -> a
f Int
i

{-# INLINEABLE clearST #-}
clearST :: HashMap s a -> ST s ()
clearST :: forall s a. HashMap s a -> ST s ()
clearST HashMap {Int
MVector s a
MVector s Int
MVector s Bit
maxCapHM :: forall s a. HashMap s a -> Int
restCapHM :: forall s a. HashMap s a -> MVector s Int
maskHM :: forall s a. HashMap s a -> Int
keyHM :: forall s a. HashMap s a -> MVector s Int
valHM :: forall s a. HashMap s a -> MVector s a
usedHM :: forall s a. HashMap s a -> MVector s Bit
maxCapHM :: Int
restCapHM :: MVector s Int
maskHM :: Int
keyHM :: MVector s Int
valHM :: MVector s a
usedHM :: MVector s Bit
..} = do
  MVector (PrimState (ST s)) Bit -> Bit -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
VGM.set MVector s Bit
MVector (PrimState (ST s)) Bit
usedHM (Bit -> ST s ()) -> Bit -> ST s ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
False
  MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
restCapHM Int
0 Int
maxCapHM

{-# INLINEABLE unsafeKeysST #-}
unsafeKeysST :: (VU.Unbox a) => HashMap s a -> ST s (VU.Vector Int)
unsafeKeysST :: forall a s. Unbox a => HashMap s a -> ST s (Vector Int)
unsafeKeysST HashMap {Int
MVector s a
MVector s Int
MVector s Bit
maxCapHM :: forall s a. HashMap s a -> Int
restCapHM :: forall s a. HashMap s a -> MVector s Int
maskHM :: forall s a. HashMap s a -> Int
keyHM :: forall s a. HashMap s a -> MVector s Int
valHM :: forall s a. HashMap s a -> MVector s a
usedHM :: forall s a. HashMap s a -> MVector s Bit
maxCapHM :: Int
restCapHM :: MVector s Int
maskHM :: Int
keyHM :: MVector s Int
valHM :: MVector s a
usedHM :: MVector s Bit
..} = do
  Vector Bit
used <- MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Bit
MVector (PrimState (ST s)) Bit
usedHM
  Vector Int
keys_ <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
keyHM
  Vector Int -> ST s (Vector Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector Int -> ST s (Vector Int))
-> Vector Int -> ST s (Vector Int)
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool) -> Vector Int -> Vector Int
forall a. Unbox a => (Int -> a -> Bool) -> Vector a -> Vector a
VU.ifilter (Bool -> Int -> Bool
forall a b. a -> b -> a
const (Bool -> Int -> Bool) -> (Int -> Bool) -> Int -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bit -> Bool
unBit (Bit -> Bool) -> (Int -> Bit) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector Bit
used VG.!)) Vector Int
keys_

{-# INLINEABLE unsafeElemsST #-}
unsafeElemsST :: (VU.Unbox a) => HashMap s a -> ST s (VU.Vector a)
unsafeElemsST :: forall a s. Unbox a => HashMap s a -> ST s (Vector a)
unsafeElemsST HashMap {Int
MVector s a
MVector s Int
MVector s Bit
maxCapHM :: forall s a. HashMap s a -> Int
restCapHM :: forall s a. HashMap s a -> MVector s Int
maskHM :: forall s a. HashMap s a -> Int
keyHM :: forall s a. HashMap s a -> MVector s Int
valHM :: forall s a. HashMap s a -> MVector s a
usedHM :: forall s a. HashMap s a -> MVector s Bit
maxCapHM :: Int
restCapHM :: MVector s Int
maskHM :: Int
keyHM :: MVector s Int
valHM :: MVector s a
usedHM :: MVector s Bit
..} = do
  Vector Bit
used <- MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Bit
MVector (PrimState (ST s)) Bit
usedHM
  Vector a
vals <- MVector (PrimState (ST s)) a -> ST s (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s a
MVector (PrimState (ST s)) a
valHM
  Vector a -> ST s (Vector a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector a -> ST s (Vector a)) -> Vector a -> ST s (Vector a)
forall a b. (a -> b) -> a -> b
$ (Int -> a -> Bool) -> Vector a -> Vector a
forall a. Unbox a => (Int -> a -> Bool) -> Vector a -> Vector a
VU.ifilter (Bool -> a -> Bool
forall a b. a -> b -> a
const (Bool -> a -> Bool) -> (Int -> Bool) -> Int -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bit -> Bool
unBit (Bit -> Bool) -> (Int -> Bit) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector Bit
used VG.!)) Vector a
vals

{-# INLINEABLE unsafeAssocsST #-}
unsafeAssocsST :: (VU.Unbox a) => HashMap s a -> ST s (VU.Vector (Int, a))
unsafeAssocsST :: forall a s. Unbox a => HashMap s a -> ST s (Vector (Int, a))
unsafeAssocsST HashMap {Int
MVector s a
MVector s Int
MVector s Bit
maxCapHM :: forall s a. HashMap s a -> Int
restCapHM :: forall s a. HashMap s a -> MVector s Int
maskHM :: forall s a. HashMap s a -> Int
keyHM :: forall s a. HashMap s a -> MVector s Int
valHM :: forall s a. HashMap s a -> MVector s a
usedHM :: forall s a. HashMap s a -> MVector s Bit
maxCapHM :: Int
restCapHM :: MVector s Int
maskHM :: Int
keyHM :: MVector s Int
valHM :: MVector s a
usedHM :: MVector s Bit
..} = do
  Vector Bit
used <- MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Bit
MVector (PrimState (ST s)) Bit
usedHM
  Vector Int
keys_ <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
keyHM
  Vector a
vals <- MVector (PrimState (ST s)) a -> ST s (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s a
MVector (PrimState (ST s)) a
valHM
  Vector (Int, a) -> ST s (Vector (Int, a))
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector (Int, a) -> ST s (Vector (Int, a)))
-> Vector (Int, a) -> ST s (Vector (Int, a))
forall a b. (a -> b) -> a -> b
$ (Int -> (Int, a) -> Bool) -> Vector (Int, a) -> Vector (Int, a)
forall a. Unbox a => (Int -> a -> Bool) -> Vector a -> Vector a
VU.ifilter (Bool -> (Int, a) -> Bool
forall a b. a -> b -> a
const (Bool -> (Int, a) -> Bool)
-> (Int -> Bool) -> Int -> (Int, a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bit -> Bool
unBit (Bit -> Bool) -> (Int -> Bit) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector Bit
used VG.!)) (Vector (Int, a) -> Vector (Int, a))
-> Vector (Int, a) -> Vector (Int, a)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector a -> Vector (Int, a)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector Int
keys_ Vector a
vals