{-# 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.
--
-- NOTE: The entries (key - value pairs) 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)
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 decrementRestCapacity #-}
decrementRestCapacity :: (HasCallStack, PrimMonad m) => VUM.MVector (PrimState m) Int -> String -> m ()
decrementRestCapacity :: forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MVector (PrimState m) Int -> String -> m ()
decrementRestCapacity MVector (PrimState m) Int
restCap String
funcName = do
  Int
rest <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead MVector (PrimState m) 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 m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite MVector (PrimState m) Int
restCap Int
0 (Int
rest Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  () -> m ()
forall a. a -> m 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 = 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
  let !maxCapHM :: Int
maxCapHM = Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
  MVector (PrimState m) Int
restCapHM <- Int -> Int -> m (MVector (PrimState m) 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 (PrimState m) Int
keyHM <- Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
k
  MVector (PrimState m) a
valHM <- Int -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
k
  MVector (PrimState m) Bit
usedHM <- Int -> Bit -> m (MVector (PrimState m) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
k (Bit -> m (MVector (PrimState m) Bit))
-> Bit -> m (MVector (PrimState m) Bit)
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
False
  HashMap (PrimState m) a -> m (HashMap (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HashMap {Int
MVector (PrimState m) a
MVector (PrimState m) Int
MVector (PrimState m) 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
maxCapHM :: Int
restCapHM :: MVector (PrimState m) Int
maskHM :: Int
keyHM :: MVector (PrimState m) Int
valHM :: MVector (PrimState m) a
usedHM :: MVector (PrimState m) Bit
..}

-- | \(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 = do
  HashMap (PrimState m) a
hm <- Int -> m (HashMap (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (HashMap (PrimState m) a)
new Int
n
  Vector (Int, a) -> ((Int, a) -> m ()) -> m ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector (Int, a)
xs (((Int, a) -> m ()) -> m ()) -> ((Int, a) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(!Int
i, !a
x) -> do
    HashMap (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
HashMap (PrimState m) a -> Int -> a -> m ()
insert HashMap (PrimState m) a
hm Int
i a
x
  HashMap (PrimState m) a -> m (HashMap (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HashMap (PrimState m) a
hm

-- | \(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)\) (Internal) Hash value calculation.
--
-- @since 1.1.0.0
{-# INLINE 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

-- | \(O(1)\) (Internal) Hashed slot search.
--
-- @since 1.1.0.0
{-# INLINE index #-}
index :: (PrimMonad m) => HashMap (PrimState m) a -> Int -> m Int
index :: forall (m :: * -> *) a.
PrimMonad m =>
HashMap (PrimState m) a -> Int -> m Int
index 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
..} Int
k = Int -> m Int
inner (HashMap (PrimState m) a -> Int -> Int
forall a s. HashMap a s -> Int -> Int
hash HashMap (PrimState m) 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 (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 (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

-- | \(O(1)\) Return the value to which the specified key is mapped, or `Nothing` if this map
-- contains no mapping for the key.
--
-- @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 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
..} Int
k = do
  Int
i <- HashMap (PrimState m) a -> Int -> m Int
forall (m :: * -> *) a.
PrimMonad m =>
HashMap (PrimState m) a -> Int -> m Int
index HashMap (PrimState m) a
hm Int
k
  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 (PrimState m) Bit
usedHM Int
i
  if Bool
b
    then a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> m a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
valHM Int
i
    else Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing

-- | \(O(1)\) Checks whether the hash map contains the element.
--
-- @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 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
..} Int
k = do
  Int
i <- HashMap (PrimState m) a -> Int -> m Int
forall (m :: * -> *) a.
PrimMonad m =>
HashMap (PrimState m) a -> Int -> m Int
index HashMap (PrimState m) a
hm Int
k
  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 (PrimState m) Bit
usedHM Int
i
  -- TODO: is this key check necessary
  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 (PrimState m) Int
keyHM Int
i
  Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> m Bool) -> Bool -> m 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

-- | \(O(1)\) Checks whether the hash map does not contain the element.
--
-- @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 = Bool -> Bool
not (Bool -> Bool) -> m Bool -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap (PrimState m) a -> Int -> m Bool
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m) =>
HashMap (PrimState m) a -> Int -> m Bool
member HashMap (PrimState m) a
hm Int
k

-- | \(O(1)\) Inserts a \((k, v)\) pair.
--
-- @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 ()) -> m (Maybe a) -> m ()
forall a b. (a -> b) -> a -> b
$ HashMap (PrimState m) a -> Int -> a -> m (Maybe a)
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

-- | \(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}}))\).
--
-- @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 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 -> a -> a
f Int
k a
v = do
  Int
i <- HashMap (PrimState m) a -> Int -> m Int
forall (m :: * -> *) a.
PrimMonad m =>
HashMap (PrimState m) a -> Int -> m Int
index HashMap (PrimState m) a
hm Int
k
  Bit Bool
b <- MVector (PrimState m) Bit -> Int -> Bit -> m Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector (PrimState m) Bit
usedHM Int
i (Bit -> m Bit) -> Bit -> m Bit
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
True
  if Bool
b
    then do
      -- modify the existing entry
      MVector (PrimState m) a -> (a -> a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState m) a
valHM (a -> a -> a
f a
v) Int
i
    else do
      -- insert the new \((k, v)\) pair
      MVector (PrimState m) Int -> String -> m ()
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MVector (PrimState m) Int -> String -> m ()
decrementRestCapacity MVector (PrimState m) Int
restCapHM String
"insertWith"
      MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) Int
keyHM Int
i Int
k
      MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
valHM Int
i a
v

-- | \(O(1)\) Inserts a \((k, v)\) pair and returns the old value, or `Nothing` if no such entry
-- exists.
--
-- @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 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
..} Int
k a
v = do
  Int
i <- HashMap (PrimState m) a -> Int -> m Int
forall (m :: * -> *) a.
PrimMonad m =>
HashMap (PrimState m) a -> Int -> m Int
index HashMap (PrimState m) a
hm Int
k
  Bit Bool
b <- MVector (PrimState m) Bit -> Int -> Bit -> m Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector (PrimState m) Bit
usedHM Int
i (Bit -> m Bit) -> Bit -> m 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) -> m a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) a -> Int -> a -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
VGM.exchange MVector (PrimState m) a
valHM Int
i a
v
    else do
      -- insert the new (key, value) pair
      MVector (PrimState m) Int -> String -> m ()
forall (m :: * -> *).
(HasCallStack, PrimMonad m) =>
MVector (PrimState m) Int -> String -> m ()
decrementRestCapacity MVector (PrimState m) Int
restCapHM String
"exchange"
      MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) Int
keyHM Int
i Int
k
      MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
valHM Int
i a
v
      Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing

-- | \(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 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 -> a
f Int
k = do
  Int
i <- HashMap (PrimState m) a -> Int -> m Int
forall (m :: * -> *) a.
PrimMonad m =>
HashMap (PrimState m) a -> Int -> m Int
index HashMap (PrimState m) a
hm Int
k
  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 (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 -> a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState m) a
valHM a -> a
f Int
i

-- | \(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 <- HashMap (PrimState m) a -> Int -> m Int
forall (m :: * -> *) a.
PrimMonad m =>
HashMap (PrimState m) a -> Int -> m Int
index HashMap (PrimState m) a
hm Int
k
  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 (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 {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
  MVector (PrimState m) Bit -> Bit -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
VGM.set MVector (PrimState m) Bit
usedHM (Bit -> m ()) -> Bit -> m ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
False
  MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector (PrimState m) Int
restCapHM Int
0 Int
maxCapHM

-- | \(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 {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
  Vector Bit
used <- MVector (PrimState m) Bit -> m (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) Bit
usedHM
  Vector Int
keys_ <- MVector (PrimState m) Int -> m (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) Int
keyHM
  Vector Int -> m (Vector Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector Int -> m (Vector Int)) -> Vector Int -> m (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_

-- | \(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 {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
  Vector Bit
used <- MVector (PrimState m) Bit -> m (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) Bit
usedHM
  Vector a
vals <- MVector (PrimState m) a -> m (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) a
valHM
  Vector a -> m (Vector a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector a -> m (Vector a)) -> Vector a -> m (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

-- | \(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 {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
  Vector Bit
used <- MVector (PrimState m) Bit -> m (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) Bit
usedHM
  Vector Int
keys_ <- MVector (PrimState m) Int -> m (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) Int
keyHM
  Vector a
vals <- MVector (PrimState m) a -> m (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) a
valHM
  Vector (Int, a) -> m (Vector (Int, a))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector (Int, a) -> m (Vector (Int, a)))
-> Vector (Int, a) -> m (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