{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UnliftedNewtypes #-}
{-# OPTIONS_GHC -Wno-incomplete-patterns #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}
{-# OPTIONS_HADDOCK hide #-}

module Data.HashMap.Mutable.Linear.Internal where

import qualified Control.Functor.Linear as Control
import Data.Array.Mutable.Linear (Array)
import qualified Data.Array.Mutable.Linear as Array
import qualified Data.Function as NonLinear
import Data.Functor.Identity hiding (runIdentity)
import qualified Data.Functor.Linear as Data
import Data.Hashable
import qualified Data.Maybe as NonLinear
import Data.Unrestricted.Linear
import Prelude.Linear hiding (filter, insert, lookup, mapMaybe, read, (+))
import Unsafe.Coerce (unsafeCoerce)
import qualified Unsafe.Linear as Unsafe
import Prelude ((+))
import qualified Prelude

-- # Implementation Notes
-- This is a simple implementatation of robin hood hashing.
--
-- See these links:
--

-- * https://programming.guide/robin-hood-hashing.html

-- * https://andre.arko.net/2017/08/24/robin-hood-hashing/

-- * https://cs.uwaterloo.ca/research/tr/1986/CS-86-14.pdf

--

-- # Constants
--------------------------------------------------

-- | When to trigger a resize.
--
-- A high load factor usually is not desirable because it makes operations
-- do more probes. A very low one is also not desirable since there're some
-- operations which take time relative to the 'capacity'.
--
-- This should be between (0, 1)
--
-- The value 0.75 is what Java uses:
-- https://docs.oracle.com/javase/10/docs/api/java/util/HashMap.html
constMaxLoadFactor :: Float
constMaxLoadFactor :: Float
constMaxLoadFactor = Float
0.75

-- | When resizing, the capacity will be multiplied by this amount.
--
-- This should be greater than one.
constGrowthFactor :: Int
constGrowthFactor :: Int
constGrowthFactor = Int
2

-- # Core Data Types
--------------------------------------------------

-- | A mutable hashmap with a linear interface.
data HashMap k v where
  -- |
  -- @loadFactor m = size m / cap m@
  --
  -- Invariants:
  -- - array is non-empty
  -- - (count / capacity) <= constMaxLoadFactor.
  HashMap ::
    -- | The number of stored (key, value) pairs.
    !Int ->
    -- | Capacity of the underlying array (cached here)
    !Int ->
    -- | Underlying array.
    !(RobinArr k v) %1 ->
    HashMap k v

-- | An array of Robin values
--
-- Each cell is Nothing if empty and is a RobinVal with the correct
-- PSL otherwise.
type RobinArr k v = Array (Maybe (RobinVal k v))

-- | Robin values are triples of the key, value and PSL
-- (the probe sequence length).
data RobinVal k v = RobinVal !PSL !k v
  deriving (Int -> RobinVal k v -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> RobinVal k v -> ShowS
forall k v. (Show k, Show v) => [RobinVal k v] -> ShowS
forall k v. (Show k, Show v) => RobinVal k v -> String
showList :: [RobinVal k v] -> ShowS
$cshowList :: forall k v. (Show k, Show v) => [RobinVal k v] -> ShowS
show :: RobinVal k v -> String
$cshow :: forall k v. (Show k, Show v) => RobinVal k v -> String
showsPrec :: Int -> RobinVal k v -> ShowS
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> RobinVal k v -> ShowS
Show)

incRobinValPSL :: RobinVal k v -> RobinVal k v
incRobinValPSL :: forall k v. RobinVal k v -> RobinVal k v
incRobinValPSL (RobinVal (PSL Int
p) k
k v
v) = forall k v. PSL -> k -> v -> RobinVal k v
RobinVal (Int -> PSL
PSL (Int
p forall a. Num a => a -> a -> a
+ Int
1)) k
k v
v

decRobinValPSL :: RobinVal k v -> RobinVal k v
decRobinValPSL :: forall k v. RobinVal k v -> RobinVal k v
decRobinValPSL (RobinVal (PSL Int
p) k
k v
v) = forall k v. PSL -> k -> v -> RobinVal k v
RobinVal (Int -> PSL
PSL (Int
p forall a. AdditiveGroup a => a %1 -> a %1 -> a
- Int
1)) k
k v
v

-- | A probe sequence length
newtype PSL = PSL Int
  deriving (PSL -> PSL -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PSL -> PSL -> Bool
$c/= :: PSL -> PSL -> Bool
== :: PSL -> PSL -> Bool
$c== :: PSL -> PSL -> Bool
Prelude.Eq, Eq PSL
PSL -> PSL -> Bool
PSL -> PSL -> Ordering
PSL -> PSL -> PSL
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: PSL -> PSL -> PSL
$cmin :: PSL -> PSL -> PSL
max :: PSL -> PSL -> PSL
$cmax :: PSL -> PSL -> PSL
>= :: PSL -> PSL -> Bool
$c>= :: PSL -> PSL -> Bool
> :: PSL -> PSL -> Bool
$c> :: PSL -> PSL -> Bool
<= :: PSL -> PSL -> Bool
$c<= :: PSL -> PSL -> Bool
< :: PSL -> PSL -> Bool
$c< :: PSL -> PSL -> Bool
compare :: PSL -> PSL -> Ordering
$ccompare :: PSL -> PSL -> Ordering
Prelude.Ord, Integer -> PSL
PSL -> PSL
PSL -> PSL -> PSL
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
fromInteger :: Integer -> PSL
$cfromInteger :: Integer -> PSL
signum :: PSL -> PSL
$csignum :: PSL -> PSL
abs :: PSL -> PSL
$cabs :: PSL -> PSL
negate :: PSL -> PSL
$cnegate :: PSL -> PSL
* :: PSL -> PSL -> PSL
$c* :: PSL -> PSL -> PSL
- :: PSL -> PSL -> PSL
$c- :: PSL -> PSL -> PSL
+ :: PSL -> PSL -> PSL
$c+ :: PSL -> PSL -> PSL
Prelude.Num, Int -> PSL -> ShowS
[PSL] -> ShowS
PSL -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PSL] -> ShowS
$cshowList :: [PSL] -> ShowS
show :: PSL -> String
$cshow :: PSL -> String
showsPrec :: Int -> PSL -> ShowS
$cshowsPrec :: Int -> PSL -> ShowS
Prelude.Show)

-- | At minimum, we need to store hashable
-- and identifiable keys
type Keyed k = (Prelude.Eq k, Hashable k)

-- | The results of searching for where to insert a key.
--
-- PSL's on the constructors are the probes spent from the query, this
-- might be different than PSL's of the cell at the returned index
-- (in case of `IndexToSwap` constructor).
data ProbeResult k v where
  -- | An empty cell at index to insert a new element with PSL.
  IndexToInsert :: !PSL -> !Int -> ProbeResult k v
  -- | A matching cell at index with a PSL and a value to update.
  IndexToUpdate :: v -> !PSL -> !Int -> ProbeResult k v
  -- | An occupied, richer, cell which should be evicted when inserting
  -- the new element. The swapped-out cell will then need to be inserted
  -- with a higher PSL.
  IndexToSwap :: RobinVal k v -> !PSL -> !Int -> ProbeResult k v

-- # Construction and Modification
--------------------------------------------------

-- | Run a computation with an empty 'HashMap' with given capacity.
empty ::
  forall k v b.
  Keyed k =>
  Int ->
  (HashMap k v %1 -> Ur b) %1 ->
  Ur b
empty :: forall k v b. Keyed k => Int -> (HashMap k v %1 -> Ur b) %1 -> Ur b
empty Int
size HashMap k v %1 -> Ur b
scope =
  let cap :: Int
cap = forall a. (Dupable a, Ord a) => a %1 -> a %1 -> a
max Int
1 Int
size
   in forall a b.
HasCallStack =>
Int -> a -> (Array a %1 -> Ur b) %1 -> Ur b
Array.alloc Int
cap forall a. Maybe a
Nothing (\Array (Maybe (RobinVal k v))
arr -> HashMap k v %1 -> Ur b
scope (forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
0 Int
cap Array (Maybe (RobinVal k v))
arr))

-- | Create an empty HashMap, using another as a uniqueness proof.
allocBeside :: Keyed k => Int -> HashMap k' v' %1 -> (HashMap k v, HashMap k' v')
allocBeside :: forall k k' v' v.
Keyed k =>
Int -> HashMap k' v' %1 -> (HashMap k v, HashMap k' v')
allocBeside Int
size (HashMap Int
s' Int
c' RobinArr k' v'
arr) =
  let cap :: Int
cap = forall a. (Dupable a, Ord a) => a %1 -> a %1 -> a
max Int
1 Int
size
   in forall a b. Int -> a -> Array b %1 -> (Array a, Array b)
Array.allocBeside Int
cap forall a. Maybe a
Nothing RobinArr k' v'
arr forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Array (Maybe (RobinVal k v))
arr', RobinArr k' v'
arr'') ->
        (forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
size Int
cap Array (Maybe (RobinVal k v))
arr', forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
s' Int
c' RobinArr k' v'
arr'')

-- | Run a computation with an 'HashMap' containing given key-value pairs.
fromList ::
  forall k v b.
  Keyed k =>
  [(k, v)] ->
  (HashMap k v %1 -> Ur b) %1 ->
  Ur b
fromList :: forall k v b.
Keyed k =>
[(k, v)] -> (HashMap k v %1 -> Ur b) %1 -> Ur b
fromList [(k, v)]
xs HashMap k v %1 -> Ur b
scope =
  let cap :: Int
cap =
        forall a. (Dupable a, Ord a) => a %1 -> a %1 -> a
max
          Int
1
          (forall a b. (RealFrac a, Integral b) => a -> b
ceiling @Float @Int (forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (t :: * -> *) a. Foldable t => t a -> Int
Prelude.length [(k, v)]
xs) forall a. Fractional a => a -> a -> a
/ Float
constMaxLoadFactor))
   in forall a b.
HasCallStack =>
Int -> a -> (Array a %1 -> Ur b) %1 -> Ur b
Array.alloc
        Int
cap
        forall a. Maybe a
Nothing
        (\Array (Maybe (RobinVal k v))
arr -> HashMap k v %1 -> Ur b
scope (forall k v. Keyed k => [(k, v)] -> HashMap k v %1 -> HashMap k v
insertAll [(k, v)]
xs (forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
0 Int
cap Array (Maybe (RobinVal k v))
arr)))

-- | The most general modification function; which can insert, update or delete
-- a value of the key, while collecting an effect in the form of an arbitrary
-- 'Control.Functor'.
alterF :: (Keyed k, Control.Functor f) => (Maybe v -> f (Ur (Maybe v))) -> k -> HashMap k v %1 -> f (HashMap k v)
alterF :: forall k (f :: * -> *) v.
(Keyed k, Functor f) =>
(Maybe v -> f (Ur (Maybe v)))
-> k -> HashMap k v %1 -> f (HashMap k v)
alterF Maybe v -> f (Ur (Maybe v))
f k
key HashMap k v
hm =
  forall k v. Keyed k => k -> HashMap k v %1 -> (Ur Int, HashMap k v)
idealIndexForKey k
key HashMap k v
hm forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur Int
idx, HashMap k v
hm') ->
    forall k v.
Keyed k =>
k
-> PSL
-> Int
-> HashMap k v
%1 -> (# HashMap k v, ProbeResult k v #)
probeFrom k
key PSL
0 Int
idx HashMap k v
hm' forall a b c. (# a, b #) %1 -> ((# a, b #) %1 -> c) %1 -> c
`chainU` \case
      -- The key does not exist, and there is an empty cell to insert.
      (# HashMap Int
count Int
cap RobinArr k v
arr, IndexToInsert PSL
psl Int
ix #) ->
        Maybe v -> f (Ur (Maybe v))
f forall a. Maybe a
Nothing forall (f :: * -> *) a b.
Functor f =>
f a %1 -> (a %1 -> b) %1 -> f b
Control.<&> \case
          -- We don't need to insert anything.
          Ur Maybe v
Nothing -> forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
count Int
cap RobinArr k v
arr
          -- We need to insert a new key.
          Ur (Just v
v) ->
            forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap
              (Int
count forall a. Num a => a -> a -> a
+ Int
1)
              Int
cap
              (forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr Int
ix (forall a. a -> Maybe a
Just (forall k v. PSL -> k -> v -> RobinVal k v
RobinVal PSL
psl k
key v
v)))
              forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& forall k v. Keyed k => HashMap k v %1 -> HashMap k v
growMapIfNecessary
      -- The key exists.
      (# HashMap Int
count Int
cap RobinArr k v
arr, IndexToUpdate v
v PSL
psl Int
ix #) ->
        Maybe v -> f (Ur (Maybe v))
f (forall a. a -> Maybe a
Just v
v) forall (f :: * -> *) a b.
Functor f =>
f a %1 -> (a %1 -> b) %1 -> f b
Control.<&> \case
          -- We need to delete it.
          Ur Maybe v
Nothing ->
            forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr Int
ix forall a. Maybe a
Nothing forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \RobinArr k v
arr' ->
              forall k v.
Keyed k =>
Int -> Int -> RobinArr k v %1 -> Int -> RobinArr k v
shiftSegmentBackward Int
1 Int
cap RobinArr k v
arr' ((Int
ix forall a. Num a => a -> a -> a
+ Int
1) forall a. Integral a => a -> a -> a
`mod` Int
cap) forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \RobinArr k v
arr'' ->
                forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap
                  (Int
count forall a. AdditiveGroup a => a %1 -> a %1 -> a
- Int
1)
                  Int
cap
                  RobinArr k v
arr''
          -- We need to update it.
          Ur (Just v
new) ->
            forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap
              Int
count
              Int
cap
              (forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr Int
ix (forall a. a -> Maybe a
Just (forall k v. PSL -> k -> v -> RobinVal k v
RobinVal PSL
psl k
key v
new)))
      -- The key does not exist, but there is a key to evict.
      (# HashMap Int
count Int
cap RobinArr k v
arr, IndexToSwap RobinVal k v
evicted PSL
psl Int
ix #) ->
        Maybe v -> f (Ur (Maybe v))
f forall a. Maybe a
Nothing forall (f :: * -> *) a b.
Functor f =>
f a %1 -> (a %1 -> b) %1 -> f b
Control.<&> \case
          -- We don't need to insert anything.
          Ur Maybe v
Nothing -> forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
count Int
cap RobinArr k v
arr
          -- We need to insert a new key.
          Ur (Just v
v) ->
            forall k v.
Keyed k =>
HashMap k v %1 -> Int -> RobinVal k v -> HashMap k v
tryInsertAtIndex
              ( forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap
                  Int
count
                  Int
cap
                  (forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr Int
ix (forall a. a -> Maybe a
Just (forall k v. PSL -> k -> v -> RobinVal k v
RobinVal PSL
psl k
key v
v)))
              )
              ((Int
ix forall a. Num a => a -> a -> a
+ Int
1) forall a. Integral a => a -> a -> a
`mod` Int
cap)
              (forall k v. RobinVal k v -> RobinVal k v
incRobinValPSL RobinVal k v
evicted)
              forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& forall k v. Keyed k => HashMap k v %1 -> HashMap k v
growMapIfNecessary
{-# INLINE alterF #-}

-- aspiwack: I'm implementing `alter` in terms of `alterF`, because, at this
-- point, we may have some bug fixes and so on and so forth. And maintaining two
-- functions this size is quite a bit unpleasant. Nevertheless, the extra boxing
-- required by the intermediate `Ur` call, there, makes it so that the
-- specialisation of `alterF` to `Identity` doesn't quite yield the code that we
-- would like, it's a bit costlier than it should. So in an ideal word, we would
-- implement both manually. In the future probably.

-- | A general modification function; which can insert, update or delete
-- a value of the key. See 'alterF', for an even more general function.
alter :: Keyed k => (Maybe v -> Maybe v) -> k -> HashMap k v %1 -> HashMap k v
alter :: forall k v.
Keyed k =>
(Maybe v -> Maybe v) -> k -> HashMap k v %1 -> HashMap k v
alter Maybe v -> Maybe v
f k
key HashMap k v
hm = forall a. Identity a %1 -> a
runIdentity forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall k (f :: * -> *) v.
(Keyed k, Functor f) =>
(Maybe v -> f (Ur (Maybe v)))
-> k -> HashMap k v %1 -> f (HashMap k v)
alterF (\Maybe v
v -> forall a. a -> Identity a
Identity (forall a. a -> Ur a
Ur (Maybe v -> Maybe v
f Maybe v
v))) k
key HashMap k v
hm
  where
    runIdentity :: Identity a %1 -> a
    runIdentity :: forall a. Identity a %1 -> a
runIdentity (Identity a
x) = a
x
{-# INLINE alter #-}

-- | Insert a key value pair to a 'HashMap'. It overwrites the previous
-- value if it exists.
insert :: Keyed k => k -> v -> HashMap k v %1 -> HashMap k v
insert :: forall k v. Keyed k => k -> v -> HashMap k v %1 -> HashMap k v
insert k
k v
v = forall k v.
Keyed k =>
(Maybe v -> Maybe v) -> k -> HashMap k v %1 -> HashMap k v
alter (\Maybe v
_ -> forall a. a -> Maybe a
Just v
v) k
k

-- | Delete a key from a 'HashMap'. Does nothing if the key does not
-- exist.
delete :: Keyed k => k -> HashMap k v %1 -> HashMap k v
delete :: forall k v. Keyed k => k -> HashMap k v %1 -> HashMap k v
delete = forall k v.
Keyed k =>
(Maybe v -> Maybe v) -> k -> HashMap k v %1 -> HashMap k v
alter (\Maybe v
_ -> forall a. Maybe a
Nothing)

-- | 'insert' (in the provided order) the given key-value pairs to
-- the hashmap.
insertAll :: Keyed k => [(k, v)] -> HashMap k v %1 -> HashMap k v
insertAll :: forall k v. Keyed k => [(k, v)] -> HashMap k v %1 -> HashMap k v
insertAll [] HashMap k v
hmap = HashMap k v
hmap
insertAll ((k
k, v
v) : [(k, v)]
xs) HashMap k v
hmap = forall k v. Keyed k => [(k, v)] -> HashMap k v %1 -> HashMap k v
insertAll [(k, v)]
xs (forall k v. Keyed k => k -> v -> HashMap k v %1 -> HashMap k v
insert k
k v
v HashMap k v
hmap)

-- TODO: Do a resize first on the length of the input.

-- | A version of 'fmap' which can throw out the elements.
--
-- Complexity: O(capacity hm)
mapMaybe :: Keyed k => (v -> Maybe v') -> HashMap k v %1 -> HashMap k v'
mapMaybe :: forall k v v'.
Keyed k =>
(v -> Maybe v') -> HashMap k v %1 -> HashMap k v'
mapMaybe v -> Maybe v'
f = forall k v v'.
Keyed k =>
(k -> v -> Maybe v') -> HashMap k v %1 -> HashMap k v'
mapMaybeWithKey (\k
_k v
v -> v -> Maybe v'
f v
v)

-- | Same as 'mapMaybe', but also has access to the keys.
mapMaybeWithKey ::
  forall k v v'.
  Keyed k =>
  (k -> v -> Maybe v') ->
  HashMap k v %1 ->
  HashMap k v'
mapMaybeWithKey :: forall k v v'.
Keyed k =>
(k -> v -> Maybe v') -> HashMap k v %1 -> HashMap k v'
mapMaybeWithKey k -> v -> Maybe v'
_ (HashMap Int
0 Int
cap RobinArr k v
arr) = forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
0 Int
cap (forall a b. a %1 -> b
Unsafe.coerce RobinArr k v
arr)
mapMaybeWithKey k -> v -> Maybe v'
f (HashMap Int
_ Int
cap RobinArr k v
arr) =
  forall a. Array a %1 -> (Ur Int, Array a)
Array.size RobinArr k v
arr forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur Int
size, RobinArr k v
arr1) ->
    Int
-> Int
-> (Bool, Int)
-> Int
-> RobinArr k v
%1 -> (Ur Int, RobinArr k v)
mapAndPushBack Int
0 (Int
size forall a. AdditiveGroup a => a %1 -> a %1 -> a
- Int
1) (Bool
False, Int
0) Int
0 RobinArr k v
arr1 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur Int
c, RobinArr k v
arr2) ->
      forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
c Int
cap (forall a b. a %1 -> b
Unsafe.coerce RobinArr k v
arr2)
  where
    f' :: k -> v -> Maybe v
    f' :: k -> v -> Maybe v
f' k
k v
v = forall a b. a -> b
unsafeCoerce (k -> v -> Maybe v'
f k
k v
v)

    -- Going from arr[0] to arr[size-1] map each element while
    -- simultaneously pushing elements back if some earlier element(s)
    -- were deleted in a contiguous segment and if the current
    -- element has PSL > 0. Maintain a counter of how
    -- far to push elements back. At arr[size-1] if needed, call
    -- shiftSegmentBackward with the counter at arr[0].
    mapAndPushBack ::
      Int -> -- Current index
      Int -> -- Last index of array which is (size-1)
      (Bool, Int) -> -- (b,n) s.t. b iff open space n cells before current cell
      Int -> -- Count of present key-value pairs
      RobinArr k v %1 ->
      (Ur Int, RobinArr k v) -- The new count and fully mapped array
    mapAndPushBack :: Int
-> Int
-> (Bool, Int)
-> Int
-> RobinArr k v
%1 -> (Ur Int, RobinArr k v)
mapAndPushBack Int
ix Int
end (Bool
shift, Int
dec) Int
count RobinArr k v
arr
      | (Int
ix forall a. Ord a => a %1 -> a %1 -> Bool
> Int
end) =
          if Bool
shift
            then (forall a. a -> Ur a
Ur Int
count, forall k v.
Keyed k =>
Int -> Int -> RobinArr k v %1 -> Int -> RobinArr k v
shiftSegmentBackward Int
dec (Int
end forall a. Num a => a -> a -> a
+ Int
1) RobinArr k v
arr Int
0)
            else (forall a. a -> Ur a
Ur Int
count, RobinArr k v
arr)
      | Bool
otherwise =
          forall a. Array a %1 -> Int -> (Ur a, Array a)
Array.unsafeRead RobinArr k v
arr Int
ix forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
            (Ur Maybe (RobinVal k v)
Nothing, RobinArr k v
arr1) ->
              Int
-> Int
-> (Bool, Int)
-> Int
-> RobinArr k v
%1 -> (Ur Int, RobinArr k v)
mapAndPushBack (Int
ix forall a. Num a => a -> a -> a
+ Int
1) Int
end (Bool
False, Int
0) Int
count RobinArr k v
arr1
            (Ur (Just (RobinVal (PSL Int
p) k
k v
v)), RobinArr k v
arr1) -> case k -> v -> Maybe v
f' k
k v
v of
              Maybe v
Nothing ->
                forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr1 Int
ix forall a. Maybe a
Nothing
                  forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \RobinArr k v
arr2 -> Int
-> Int
-> (Bool, Int)
-> Int
-> RobinArr k v
%1 -> (Ur Int, RobinArr k v)
mapAndPushBack (Int
ix forall a. Num a => a -> a -> a
+ Int
1) Int
end (Bool
True, Int
dec forall a. Num a => a -> a -> a
+ Int
1) Int
count RobinArr k v
arr2
              Just v
v' -> case Bool
shift of
                Bool
False ->
                  forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr1 Int
ix (forall a. a -> Maybe a
Just (forall k v. PSL -> k -> v -> RobinVal k v
RobinVal (Int -> PSL
PSL Int
p) k
k v
v'))
                    forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \RobinArr k v
arr2 -> Int
-> Int
-> (Bool, Int)
-> Int
-> RobinArr k v
%1 -> (Ur Int, RobinArr k v)
mapAndPushBack (Int
ix forall a. Num a => a -> a -> a
+ Int
1) Int
end (Bool
False, Int
0) (Int
count forall a. Num a => a -> a -> a
+ Int
1) RobinArr k v
arr2
                Bool
True -> case Int
dec forall a. Ord a => a %1 -> a %1 -> Bool
<= Int
p of
                  Bool
False ->
                    forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr1 (Int
ix forall a. AdditiveGroup a => a %1 -> a %1 -> a
- Int
p) (forall a. a -> Maybe a
Just (forall k v. PSL -> k -> v -> RobinVal k v
RobinVal PSL
0 k
k v
v'))
                      forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \RobinArr k v
arr2 -> case Int
p forall a. Eq a => a %1 -> a %1 -> Bool
== Int
0 of
                        Bool
False ->
                          forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr2 Int
ix forall a. Maybe a
Nothing
                            forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \RobinArr k v
arr3 -> Int
-> Int
-> (Bool, Int)
-> Int
-> RobinArr k v
%1 -> (Ur Int, RobinArr k v)
mapAndPushBack (Int
ix forall a. Num a => a -> a -> a
+ Int
1) Int
end (Bool
True, Int
p) (Int
count forall a. Num a => a -> a -> a
+ Int
1) RobinArr k v
arr3
                        Bool
True -> Int
-> Int
-> (Bool, Int)
-> Int
-> RobinArr k v
%1 -> (Ur Int, RobinArr k v)
mapAndPushBack (Int
ix forall a. Num a => a -> a -> a
+ Int
1) Int
end (Bool
False, Int
0) (Int
count forall a. Num a => a -> a -> a
+ Int
1) RobinArr k v
arr2
                  Bool
True ->
                    forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr1 (Int
ix forall a. AdditiveGroup a => a %1 -> a %1 -> a
- Int
dec) (forall a. a -> Maybe a
Just (forall k v. PSL -> k -> v -> RobinVal k v
RobinVal (Int -> PSL
PSL (Int
p forall a. AdditiveGroup a => a %1 -> a %1 -> a
- Int
dec)) k
k v
v'))
                      forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \RobinArr k v
arr2 ->
                        forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr2 Int
ix forall a. Maybe a
Nothing
                          forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \RobinArr k v
arr3 -> Int
-> Int
-> (Bool, Int)
-> Int
-> RobinArr k v
%1 -> (Ur Int, RobinArr k v)
mapAndPushBack (Int
ix forall a. Num a => a -> a -> a
+ Int
1) Int
end (Bool
True, Int
dec) (Int
count forall a. Num a => a -> a -> a
+ Int
1) RobinArr k v
arr3

-- | Complexity: O(capacity hm)
filterWithKey :: Keyed k => (k -> v -> Bool) -> HashMap k v %1 -> HashMap k v
filterWithKey :: forall k v.
Keyed k =>
(k -> v -> Bool) -> HashMap k v %1 -> HashMap k v
filterWithKey k -> v -> Bool
f =
  forall k v v'.
Keyed k =>
(k -> v -> Maybe v') -> HashMap k v %1 -> HashMap k v'
mapMaybeWithKey
    (\k
k v
v -> if k -> v -> Bool
f k
k v
v then forall a. a -> Maybe a
Just v
v else forall a. Maybe a
Nothing)

-- | Complexity: O(capacity hm)
filter :: Keyed k => (v -> Bool) -> HashMap k v %1 -> HashMap k v
filter :: forall k v. Keyed k => (v -> Bool) -> HashMap k v %1 -> HashMap k v
filter v -> Bool
f = forall k v.
Keyed k =>
(k -> v -> Bool) -> HashMap k v %1 -> HashMap k v
filterWithKey (\k
_k v
v -> v -> Bool
f v
v)

-- | Union of two maps using the provided function on conflicts.
--
-- Complexity: O(min(capacity hm1, capacity hm2)
unionWith ::
  Keyed k =>
  (v -> v -> v) ->
  HashMap k v %1 ->
  HashMap k v %1 ->
  HashMap k v
unionWith :: forall k v.
Keyed k =>
(v -> v -> v) -> HashMap k v %1 -> HashMap k v %1 -> HashMap k v
unionWith v -> v -> v
onConflict (HashMap k v
hm1 :: HashMap k v) HashMap k v
hm2 =
  -- To insert the elements in smaller map to the larger map, we
  -- compare their capacities, and flip the arguments if necessary.
  forall k v. HashMap k v %1 -> (Ur Int, HashMap k v)
capacity HashMap k v
hm1 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur Int
cap1, HashMap k v
hm1') ->
    forall k v. HashMap k v %1 -> (Ur Int, HashMap k v)
capacity HashMap k v
hm2 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur Int
cap2, HashMap k v
hm2') ->
      if Int
cap1 forall a. Ord a => a %1 -> a %1 -> Bool
> Int
cap2
        then (v -> v -> v) -> HashMap k v %1 -> Ur [(k, v)] %1 -> HashMap k v
go v -> v -> v
onConflict HashMap k v
hm1' (forall k v. HashMap k v %1 -> Ur [(k, v)]
toList HashMap k v
hm2')
        else (v -> v -> v) -> HashMap k v %1 -> Ur [(k, v)] %1 -> HashMap k v
go (\v
v2 v
v1 -> v -> v -> v
onConflict v
v1 v
v2) HashMap k v
hm2' (forall k v. HashMap k v %1 -> Ur [(k, v)]
toList HashMap k v
hm1')
  where
    go ::
      (v -> v -> v) ->
      HashMap k v %1 -> -- larger map
      Ur [(k, v)] %1 -> -- contents of the smaller map
      HashMap k v
    go :: (v -> v -> v) -> HashMap k v %1 -> Ur [(k, v)] %1 -> HashMap k v
go v -> v -> v
_ HashMap k v
hm (Ur []) = HashMap k v
hm
    go v -> v -> v
f HashMap k v
hm (Ur ((k
k, v
vr) : [(k, v)]
xs)) =
      forall k v.
Keyed k =>
(Maybe v -> Maybe v) -> k -> HashMap k v %1 -> HashMap k v
alter
        ( \case
            Maybe v
Nothing -> forall a. a -> Maybe a
Just v
vr
            Just v
vl -> forall a. a -> Maybe a
Just (v -> v -> v
f v
vl v
vr)
        )
        k
k
        HashMap k v
hm
        forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \HashMap k v
hm -> (v -> v -> v) -> HashMap k v %1 -> Ur [(k, v)] %1 -> HashMap k v
go v -> v -> v
f HashMap k v
hm (forall a. a -> Ur a
Ur [(k, v)]
xs)

-- | A right-biased union.
--
-- Complexity: O(min(capacity hm1, capacity hm2)
union :: Keyed k => HashMap k v %1 -> HashMap k v %1 -> HashMap k v
union :: forall k v.
Keyed k =>
HashMap k v %1 -> HashMap k v %1 -> HashMap k v
union HashMap k v
hm1 HashMap k v
hm2 = forall k v.
Keyed k =>
(v -> v -> v) -> HashMap k v %1 -> HashMap k v %1 -> HashMap k v
unionWith (\v
_v1 v
v2 -> v
v2) HashMap k v
hm1 HashMap k v
hm2

-- | Intersection of two maps with the provided combine function.
--
-- Complexity: O(min(capacity hm1, capacity hm2)
intersectionWith ::
  Keyed k =>
  (a -> b -> c) ->
  HashMap k a %1 ->
  HashMap k b %1 ->
  HashMap k c
intersectionWith :: forall k a b c.
Keyed k =>
(a -> b -> c) -> HashMap k a %1 -> HashMap k b %1 -> HashMap k c
intersectionWith a -> b -> c
combine (HashMap k a
hm1 :: HashMap k a') HashMap k b
hm2 =
  forall k k' v' v.
Keyed k =>
Int -> HashMap k' v' %1 -> (HashMap k v, HashMap k' v')
allocBeside Int
0 HashMap k a
hm1 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(HashMap k c
hmNew, HashMap k a
hm1') ->
    forall k v. HashMap k v %1 -> (Ur Int, HashMap k v)
capacity HashMap k a
hm1' forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur Int
cap1, HashMap k a
hm1'') ->
      forall k v. HashMap k v %1 -> (Ur Int, HashMap k v)
capacity HashMap k b
hm2 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur Int
cap2, HashMap k b
hm2') ->
        if Int
cap1 forall a. Ord a => a %1 -> a %1 -> Bool
> Int
cap2
          then forall a b c.
(a -> b -> c)
-> HashMap k a
%1 -> Ur [(k, b)]
%1 -> HashMap k c
%1 -> HashMap k c
go a -> b -> c
combine HashMap k a
hm1'' (forall k v. HashMap k v %1 -> Ur [(k, v)]
toList HashMap k b
hm2') HashMap k c
hmNew
          else forall a b c.
(a -> b -> c)
-> HashMap k a
%1 -> Ur [(k, b)]
%1 -> HashMap k c
%1 -> HashMap k c
go (\b
v2 a
v1 -> a -> b -> c
combine a
v1 b
v2) HashMap k b
hm2' (forall k v. HashMap k v %1 -> Ur [(k, v)]
toList HashMap k a
hm1'') HashMap k c
hmNew
  where
    -- Iterate over the smaller map, while checking for the matches
    -- on the bigger map; and accumulate results on a third map.
    go ::
      (a -> b -> c) ->
      HashMap k a %1 ->
      Ur [(k, b)] %1 ->
      HashMap k c %1 ->
      HashMap k c
    go :: forall a b c.
(a -> b -> c)
-> HashMap k a
%1 -> Ur [(k, b)]
%1 -> HashMap k c
%1 -> HashMap k c
go a -> b -> c
_ HashMap k a
hm (Ur []) HashMap k c
acc = HashMap k a
hm forall a b. Consumable a => a %1 -> b %1 -> b
`lseq` HashMap k c
acc
    go a -> b -> c
f HashMap k a
hm (Ur ((k
k, b
b) : [(k, b)]
xs)) HashMap k c
acc =
      forall k v.
Keyed k =>
k -> HashMap k v %1 -> (Ur (Maybe v), HashMap k v)
lookup k
k HashMap k a
hm forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
        (Ur Maybe a
Nothing, HashMap k a
hm') -> forall a b c.
(a -> b -> c)
-> HashMap k a
%1 -> Ur [(k, b)]
%1 -> HashMap k c
%1 -> HashMap k c
go a -> b -> c
f HashMap k a
hm' (forall a. a -> Ur a
Ur [(k, b)]
xs) HashMap k c
acc
        (Ur (Just a
a), HashMap k a
hm') -> forall a b c.
(a -> b -> c)
-> HashMap k a
%1 -> Ur [(k, b)]
%1 -> HashMap k c
%1 -> HashMap k c
go a -> b -> c
f HashMap k a
hm' (forall a. a -> Ur a
Ur [(k, b)]
xs) (forall k v. Keyed k => k -> v -> HashMap k v %1 -> HashMap k v
insert k
k (a -> b -> c
f a
a b
b) HashMap k c
acc)

-- |
-- Reduce the 'HashMap' 'capacity' to decrease wasted memory. Returns
-- a semantically identical 'HashMap'.
--
-- This is only useful after a lot of deletes.
--
-- Complexity: O(capacity hm)
shrinkToFit :: Keyed k => HashMap k a %1 -> HashMap k a
shrinkToFit :: forall k v. Keyed k => HashMap k v %1 -> HashMap k v
shrinkToFit HashMap k a
hm =
  forall k v. HashMap k v %1 -> (Ur Int, HashMap k v)
size HashMap k a
hm forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur Int
size, HashMap k a
hm') ->
    let targetSize :: Int
targetSize =
          forall a b. (RealFrac a, Integral b) => a -> b
ceiling
            (forall a. Ord a => a -> a -> a
Prelude.max Float
1 (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
size forall a. Fractional a => a -> a -> a
Prelude./ Float
constMaxLoadFactor))
     in forall k v. Keyed k => Int -> HashMap k v %1 -> HashMap k v
resize Int
targetSize HashMap k a
hm'

-- # Querying
--------------------------------------------------

-- | Number of key-value pairs inside the 'HashMap'
size :: HashMap k v %1 -> (Ur Int, HashMap k v)
size :: forall k v. HashMap k v %1 -> (Ur Int, HashMap k v)
size (HashMap Int
ct Int
cap RobinArr k v
arr) = (forall a. a -> Ur a
Ur Int
ct, forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
ct Int
cap RobinArr k v
arr)

-- | Maximum number of elements the HashMap can store without
-- resizing. However, for performance reasons, the 'HashMap' might be
-- before full.
--
-- Use 'shrinkToFit' to reduce the wasted space.
capacity :: HashMap k v %1 -> (Ur Int, HashMap k v)
capacity :: forall k v. HashMap k v %1 -> (Ur Int, HashMap k v)
capacity (HashMap Int
ct Int
cap RobinArr k v
arr) = (forall a. a -> Ur a
Ur Int
cap, forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
ct Int
cap RobinArr k v
arr)

-- | Look up a value from a 'HashMap'.
lookup :: Keyed k => k -> HashMap k v %1 -> (Ur (Maybe v), HashMap k v)
lookup :: forall k v.
Keyed k =>
k -> HashMap k v %1 -> (Ur (Maybe v), HashMap k v)
lookup k
k HashMap k v
hm =
  forall k v. Keyed k => k -> HashMap k v %1 -> (Ur Int, HashMap k v)
idealIndexForKey k
k HashMap k v
hm forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur Int
idx, HashMap k v
hm') ->
    forall k v.
Keyed k =>
k
-> PSL
-> Int
-> HashMap k v
%1 -> (# HashMap k v, ProbeResult k v #)
probeFrom k
k PSL
0 Int
idx HashMap k v
hm' forall a b c. (# a, b #) %1 -> ((# a, b #) %1 -> c) %1 -> c
`chainU` \case
      (# HashMap k v
h, IndexToUpdate v
v PSL
_ Int
_ #) ->
        (forall a. a -> Ur a
Ur (forall a. a -> Maybe a
Just v
v), HashMap k v
h)
      (# HashMap k v
h, IndexToInsert PSL
_ Int
_ #) ->
        (forall a. a -> Ur a
Ur forall a. Maybe a
Nothing, HashMap k v
h)
      (# HashMap k v
h, IndexToSwap RobinVal k v
_ PSL
_ Int
_ #) ->
        (forall a. a -> Ur a
Ur forall a. Maybe a
Nothing, HashMap k v
h)

-- | Check if the given key exists.
member :: Keyed k => k -> HashMap k v %1 -> (Ur Bool, HashMap k v)
member :: forall k v.
Keyed k =>
k -> HashMap k v %1 -> (Ur Bool, HashMap k v)
member k
k HashMap k v
hm =
  forall k v.
Keyed k =>
k -> HashMap k v %1 -> (Ur (Maybe v), HashMap k v)
lookup k
k HashMap k v
hm forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
    (Ur Maybe v
Nothing, HashMap k v
hm') -> (forall a. a -> Ur a
Ur Bool
False, HashMap k v
hm')
    (Ur (Just v
_), HashMap k v
hm') -> (forall a. a -> Ur a
Ur Bool
True, HashMap k v
hm')

-- | Converts a HashMap to a lazy list.
toList :: HashMap k v %1 -> Ur [(k, v)]
toList :: forall k v. HashMap k v %1 -> Ur [(k, v)]
toList (HashMap Int
_ Int
_ RobinArr k v
arr) =
  forall a. Array a %1 -> Ur [a]
Array.toList RobinArr k v
arr forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur [Maybe (RobinVal k v)]
elems) ->
    [Maybe (RobinVal k v)]
elems
      forall a b. a -> (a -> b) -> b
NonLinear.& forall a. [Maybe a] -> [a]
NonLinear.catMaybes
      forall a b. a -> (a -> b) -> b
NonLinear.& forall a b. (a -> b) -> [a] -> [b]
Prelude.map (\(RobinVal PSL
_ k
k v
v) -> (k
k, v
v))
      forall a b. a -> (a -> b) -> b
NonLinear.& forall a. a -> Ur a
Ur

-- # Instances
--------------------------------------------------

instance Consumable (HashMap k v) where
  consume :: HashMap k v %1 -> ()
  consume :: HashMap k v %1 -> ()
consume (HashMap Int
_ Int
_ RobinArr k v
arr) = forall a. Consumable a => a %1 -> ()
consume RobinArr k v
arr

instance Dupable (HashMap k v) where
  dup2 :: HashMap k v %1 -> (HashMap k v, HashMap k v)
dup2 (HashMap Int
i Int
c RobinArr k v
arr) =
    forall a. Dupable a => a %1 -> (a, a)
dup2 RobinArr k v
arr forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(RobinArr k v
a1, RobinArr k v
a2) ->
      (forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
i Int
c RobinArr k v
a1, forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
i Int
c RobinArr k v
a2)

instance Data.Functor (HashMap k) where
  fmap :: forall a b. (a %1 -> b) -> HashMap k a %1 -> HashMap k b
fmap a %1 -> b
f (HashMap Int
s Int
c RobinArr k a
arr) =
    forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
s Int
c forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$
      forall (f :: * -> *) a b. Functor f => (a %1 -> b) -> f a %1 -> f b
Data.fmap
        ( \case
            Maybe (RobinVal k a)
Nothing -> forall a. Maybe a
Nothing
            Just (RobinVal PSL
p k
k a
v) -> forall a. a -> Maybe a
Just (forall k v. PSL -> k -> v -> RobinVal k v
RobinVal PSL
p k
k (a %1 -> b
f a
v))
        )
        RobinArr k a
arr

instance Prelude.Semigroup (HashMap k v) where
  <> :: HashMap k v -> HashMap k v -> HashMap k v
(<>) = forall a. HasCallStack => String -> a
error String
"Prelude.<>: invariant violation, unrestricted HashMap"

instance Keyed k => Semigroup (HashMap k v) where
  <> :: HashMap k v %1 -> HashMap k v %1 -> HashMap k v
(<>) = forall k v.
Keyed k =>
HashMap k v %1 -> HashMap k v %1 -> HashMap k v
union

-- # Internal library
--------------------------------------------------

_debugShow :: (Show k, Show v) => HashMap k v %1 -> String
_debugShow :: forall k v. (Show k, Show v) => HashMap k v %1 -> String
_debugShow (HashMap Int
_ Int
_ RobinArr k v
robinArr) =
  forall a. Array a %1 -> Ur [a]
Array.toList RobinArr k v
robinArr forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur [Maybe (RobinVal k v)]
xs) -> forall a. Show a => a -> String
show [Maybe (RobinVal k v)]
xs

idealIndexForKey ::
  Keyed k =>
  k ->
  HashMap k v %1 ->
  (Ur Int, HashMap k v)
idealIndexForKey :: forall k v. Keyed k => k -> HashMap k v %1 -> (Ur Int, HashMap k v)
idealIndexForKey k
k (HashMap Int
sz Int
cap RobinArr k v
arr) =
  (forall a. a -> Ur a
Ur (forall a. Integral a => a -> a -> a
mod (forall a. Hashable a => a -> Int
hash k
k) Int
cap), forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
sz Int
cap RobinArr k v
arr)

-- | Given a key, psl of the probe so far, current unread index, and
-- a full hashmap, return a probe result: the place the key already
-- exists, a place to swap from, or an unfilled cell to write over.
probeFrom ::
  Keyed k =>
  k ->
  PSL ->
  Int ->
  HashMap k v %1 ->
  (# HashMap k v, ProbeResult k v #)
probeFrom :: forall k v.
Keyed k =>
k
-> PSL
-> Int
-> HashMap k v
%1 -> (# HashMap k v, ProbeResult k v #)
probeFrom k
k PSL
p Int
ix (HashMap Int
ct Int
cap RobinArr k v
arr) =
  forall a. Array a %1 -> Int -> (Ur a, Array a)
Array.unsafeRead RobinArr k v
arr Int
ix forall a b c. a %1 -> (a %1 -> (# b, c #)) %1 -> (# b, c #)
`chainU'` \case
    (Ur Maybe (RobinVal k v)
Nothing, RobinArr k v
arr') ->
      (# forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
ct Int
cap RobinArr k v
arr', forall k v. PSL -> Int -> ProbeResult k v
IndexToInsert PSL
p Int
ix #)
    (Ur (Just robinVal' :: RobinVal k v
robinVal'@(RobinVal PSL
psl k
k' v
v')), RobinArr k v
arr') ->
      case k
k forall a. Eq a => a -> a -> Bool
Prelude.== k
k' of
        -- Note: in the True case, we must have p == psl
        Bool
True -> (# forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
ct Int
cap RobinArr k v
arr', forall v k. v -> PSL -> Int -> ProbeResult k v
IndexToUpdate v
v' PSL
psl Int
ix #)
        Bool
False -> case PSL
psl forall a. Ord a => a -> a -> Bool
Prelude.< PSL
p of
          Bool
True -> (# forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
ct Int
cap RobinArr k v
arr', forall k v. RobinVal k v -> PSL -> Int -> ProbeResult k v
IndexToSwap RobinVal k v
robinVal' PSL
p Int
ix #)
          Bool
False ->
            forall k v.
Keyed k =>
k
-> PSL
-> Int
-> HashMap k v
%1 -> (# HashMap k v, ProbeResult k v #)
probeFrom k
k (PSL
p forall a. Num a => a -> a -> a
+ PSL
1) ((Int
ix forall a. Num a => a -> a -> a
+ Int
1) forall a. Integral a => a -> a -> a
`mod` Int
cap) (forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
ct Int
cap RobinArr k v
arr')

-- | Try to insert at a given index with a given PSL. So the
-- probing starts from the given index (with the given PSL).
tryInsertAtIndex ::
  Keyed k =>
  HashMap k v %1 ->
  Int ->
  RobinVal k v ->
  HashMap k v
tryInsertAtIndex :: forall k v.
Keyed k =>
HashMap k v %1 -> Int -> RobinVal k v -> HashMap k v
tryInsertAtIndex HashMap k v
hmap Int
ix (RobinVal PSL
psl k
key v
val) =
  forall k v.
Keyed k =>
k
-> PSL
-> Int
-> HashMap k v
%1 -> (# HashMap k v, ProbeResult k v #)
probeFrom k
key PSL
psl Int
ix HashMap k v
hmap forall a b c. (# a, b #) %1 -> ((# a, b #) %1 -> c) %1 -> c
`chainU` \case
    (# HashMap Int
ct Int
cap RobinArr k v
arr, IndexToUpdate v
_ PSL
psl' Int
ix' #) ->
      forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr Int
ix' (forall a. a -> Maybe a
Just forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall k v. PSL -> k -> v -> RobinVal k v
RobinVal PSL
psl' k
key v
val)
        forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
ct Int
cap
    (# HashMap Int
ct Int
cap RobinArr k v
arr, IndexToInsert PSL
psl' Int
ix' #) ->
      forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr Int
ix' (forall a. a -> Maybe a
Just forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall k v. PSL -> k -> v -> RobinVal k v
RobinVal PSL
psl' k
key v
val)
        forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap (Int
ct forall a. Num a => a -> a -> a
+ Int
1) Int
cap
    (# HashMap Int
ct Int
cap RobinArr k v
arr, IndexToSwap RobinVal k v
oldVal PSL
psl' Int
ix' #) ->
      forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr Int
ix' (forall a. a -> Maybe a
Just forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall k v. PSL -> k -> v -> RobinVal k v
RobinVal PSL
psl' k
key v
val)
        forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
ct Int
cap
        forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \HashMap k v
hm -> forall k v.
Keyed k =>
HashMap k v %1 -> Int -> RobinVal k v -> HashMap k v
tryInsertAtIndex HashMap k v
hm ((Int
ix' forall a. Num a => a -> a -> a
+ Int
1) forall a. Integral a => a -> a -> a
`mod` Int
cap) (forall k v. RobinVal k v -> RobinVal k v
incRobinValPSL RobinVal k v
oldVal)

-- | Shift all cells with PSLs > 0 in a continuous segment
-- following the deleted cell, backwards by one and decrement
-- their PSLs.
shiftSegmentBackward ::
  Keyed k =>
  Int ->
  Int ->
  RobinArr k v %1 ->
  Int ->
  RobinArr k v
shiftSegmentBackward :: forall k v.
Keyed k =>
Int -> Int -> RobinArr k v %1 -> Int -> RobinArr k v
shiftSegmentBackward Int
dec Int
s RobinArr k v
arr Int
ix =
  forall a. Array a %1 -> Int -> (Ur a, Array a)
Array.unsafeRead RobinArr k v
arr Int
ix forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
    (Ur Maybe (RobinVal k v)
Nothing, RobinArr k v
arr') -> RobinArr k v
arr'
    (Ur (Just (RobinVal PSL
0 k
_ v
_)), RobinArr k v
arr') -> RobinArr k v
arr'
    (Ur (Just RobinVal k v
val), RobinArr k v
arr') ->
      forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr' Int
ix forall a. Maybe a
Nothing forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \RobinArr k v
arr'' ->
        forall k v.
Keyed k =>
Int -> Int -> RobinArr k v %1 -> Int -> RobinArr k v
shiftSegmentBackward
          Int
dec
          Int
s
          (forall a. Array a %1 -> Int -> a -> Array a
Array.unsafeWrite RobinArr k v
arr'' ((Int
ix forall a. AdditiveGroup a => a %1 -> a %1 -> a
- Int
dec forall a. Num a => a -> a -> a
+ Int
s) forall a. Integral a => a -> a -> a
`mod` Int
s) (forall a. a -> Maybe a
Just forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall k v. RobinVal k v -> RobinVal k v
decRobinValPSL RobinVal k v
val))
          ((Int
ix forall a. Num a => a -> a -> a
+ Int
1) forall a. Integral a => a -> a -> a
`mod` Int
s)

-- TODO: This does twice as much writes than necessary, it first empties
-- the cell, just to update it again at the next call. We can save some
-- writes by only emptying the last cell.

-- | Makes sure that the map is not exceeding its utilization threshold
-- (constMaxUtilization), resizes (constGrowthFactor) if necessary.
growMapIfNecessary :: Keyed k => HashMap k v %1 -> HashMap k v
growMapIfNecessary :: forall k v. Keyed k => HashMap k v %1 -> HashMap k v
growMapIfNecessary (HashMap Int
sz Int
cap RobinArr k v
arr) =
  let load :: Float
load = forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
sz forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
cap
   in if Float
load forall a. Ord a => a -> a -> Bool
Prelude.< Float
constMaxLoadFactor
        then forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
sz Int
cap RobinArr k v
arr
        else
          let newCap :: Int
newCap = forall a. (Dupable a, Ord a) => a %1 -> a %1 -> a
max Int
1 (Int
cap forall a. Multiplicative a => a %1 -> a %1 -> a
* Int
constGrowthFactor)
           in forall k v. Keyed k => Int -> HashMap k v %1 -> HashMap k v
resize Int
newCap (forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
sz Int
cap RobinArr k v
arr)

-- | Resizes the HashMap to given capacity.
--
-- Invariant: Given capacity should be greater than the size, this is not
-- checked.
resize :: Keyed k => Int -> HashMap k v %1 -> HashMap k v
resize :: forall k v. Keyed k => Int -> HashMap k v %1 -> HashMap k v
resize Int
targetSize (HashMap Int
_ Int
_ RobinArr k v
arr) =
  forall a b. Int -> a -> Array b %1 -> (Array a, Array b)
Array.allocBeside Int
targetSize forall a. Maybe a
Nothing RobinArr k v
arr forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(RobinArr k v
newArr, RobinArr k v
oldArr) ->
    forall a. Array a %1 -> Ur [a]
Array.toList RobinArr k v
oldArr forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \(Ur [Maybe (RobinVal k v)]
elems) ->
      let xs :: [(k, v)]
xs =
            [Maybe (RobinVal k v)]
elems
              forall a b. a -> (a -> b) -> b
NonLinear.& forall a. [Maybe a] -> [a]
NonLinear.catMaybes
              forall a b. a -> (a -> b) -> b
NonLinear.& forall a b. (a -> b) -> [a] -> [b]
Prelude.map (\(RobinVal PSL
_ k
k v
v) -> (k
k, v
v))
       in forall k v. Keyed k => [(k, v)] -> HashMap k v %1 -> HashMap k v
insertAll [(k, v)]
xs (forall k v. Int -> Int -> RobinArr k v -> HashMap k v
HashMap Int
0 Int
targetSize RobinArr k v
newArr)

-- TODO: 'insertAll' keeps checking capacity on each insert. We should
-- replace it with a faster unsafe variant.

-- TODO: Remove the below workarounds once we are on GHC 9.2.
--
-- We have to use these functions below because:
--

-- * GHC <9.2 does not allow linear `case` statements.

-- * LambdaCase workaround does not work, because (&) does not work with

--   unlifted types.
chainU :: (# a, b #) %1 -> ((# a, b #) %1 -> c) %1 -> c
chainU :: forall a b c. (# a, b #) %1 -> ((# a, b #) %1 -> c) %1 -> c
chainU (# a, b #)
x (# a, b #) %1 -> c
f = (# a, b #) %1 -> c
f (# a, b #)
x

chainU' :: a %1 -> (a %1 -> (# b, c #)) %1 -> (# b, c #)
chainU' :: forall a b c. a %1 -> (a %1 -> (# b, c #)) %1 -> (# b, c #)
chainU' a
x a %1 -> (# b, c #)
f = a %1 -> (# b, c #)
f a
x