{-# LANGUAGE BangPatterns, CPP, DeriveDataTypeable, MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE LambdaCase #-}
#if __GLASGOW_HASKELL__ >= 802
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE UnboxedSums #-}
#endif
{-# OPTIONS_GHC -fno-full-laziness -funbox-strict-fields #-}
module Data.HashMap.Base
(
HashMap(..)
, Leaf(..)
, empty
, singleton
, null
, size
, member
, lookup
, lookupDefault
, (!)
, insert
, insertWith
, unsafeInsert
, delete
, adjust
, update
, alter
, alterF
, union
, unionWith
, unionWithKey
, unions
, map
, mapWithKey
, traverseWithKey
, difference
, differenceWith
, intersection
, intersectionWith
, intersectionWithKey
, foldl'
, foldlWithKey'
, foldr
, foldrWithKey
, mapMaybe
, mapMaybeWithKey
, filter
, filterWithKey
, keys
, elems
, toList
, fromList
, fromListWith
, Hash
, Bitmap
, bitmapIndexedOrFull
, collision
, hash
, mask
, index
, bitsPerSubkey
, fullNodeMask
, sparseIndex
, two
, unionArrayBy
, update16
, update16M
, update16With'
, updateOrConcatWith
, updateOrConcatWithKey
, filterMapAux
, equalKeys
, equalKeys1
, lookupRecordCollision
, LookupRes(..)
, insert'
, delete'
, lookup'
, insertNewKey
, insertKeyExists
, deleteKeyExists
, insertModifying
, ptrEq
, adjust#
) where
#if __GLASGOW_HASKELL__ < 710
import Control.Applicative ((<$>), Applicative(pure))
import Data.Monoid (Monoid(mempty, mappend))
import Data.Traversable (Traversable(..))
import Data.Word (Word)
#endif
#if __GLASGOW_HASKELL__ >= 711
import Data.Semigroup (Semigroup((<>)))
#endif
import Control.DeepSeq (NFData(rnf))
import Control.Monad.ST (ST)
import Data.Bits ((.&.), (.|.), complement, popCount)
import Data.Data hiding (Typeable)
import qualified Data.Foldable as Foldable
import qualified Data.List as L
import GHC.Exts ((==#), build, reallyUnsafePtrEquality#)
import Prelude hiding (filter, foldr, lookup, map, null, pred)
import Text.Read hiding (step)
import qualified Data.HashMap.Array as A
import qualified Data.Hashable as H
import Data.Hashable (Hashable)
import Data.HashMap.Unsafe (runST)
import Data.HashMap.UnsafeShift (unsafeShiftL, unsafeShiftR)
import Data.HashMap.List (isPermutationBy, unorderedCompare)
import Data.Typeable (Typeable)
import GHC.Exts (isTrue#)
import qualified GHC.Exts as Exts
#if MIN_VERSION_base(4,9,0)
import Data.Functor.Classes
#endif
#if MIN_VERSION_hashable(1,2,5)
import qualified Data.Hashable.Lifted as H
#endif
#if __GLASGOW_HASKELL__ >= 802
import GHC.Exts (TYPE, Int (..), Int#)
#endif
#if MIN_VERSION_base(4,8,0)
import Data.Functor.Identity (Identity (..))
#endif
import Control.Applicative (Const (..))
import Data.Coerce (coerce)
hash :: H.Hashable a => a -> Hash
hash = fromIntegral . H.hash
data Leaf k v = L !k v
deriving (Eq)
instance (NFData k, NFData v) => NFData (Leaf k v) where
rnf (L k v) = rnf k `seq` rnf v
data HashMap k v
= Empty
| BitmapIndexed !Bitmap !(A.Array (HashMap k v))
| Leaf !Hash !(Leaf k v)
| Full !(A.Array (HashMap k v))
| Collision !Hash !(A.Array (Leaf k v))
deriving (Typeable)
type role HashMap nominal representational
instance (NFData k, NFData v) => NFData (HashMap k v) where
rnf Empty = ()
rnf (BitmapIndexed _ ary) = rnf ary
rnf (Leaf _ l) = rnf l
rnf (Full ary) = rnf ary
rnf (Collision _ ary) = rnf ary
instance Functor (HashMap k) where
fmap = map
instance Foldable.Foldable (HashMap k) where
foldr f = foldrWithKey (const f)
#if __GLASGOW_HASKELL__ >= 711
instance (Eq k, Hashable k) => Semigroup (HashMap k v) where
(<>) = union
{-# INLINE (<>) #-}
#endif
instance (Eq k, Hashable k) => Monoid (HashMap k v) where
mempty = empty
{-# INLINE mempty #-}
#if __GLASGOW_HASKELL__ >= 711
mappend = (<>)
#else
mappend = union
#endif
{-# INLINE mappend #-}
instance (Data k, Data v, Eq k, Hashable k) => Data (HashMap k v) where
gfoldl f z m = z fromList `f` toList m
toConstr _ = fromListConstr
gunfold k z c = case constrIndex c of
1 -> k (z fromList)
_ -> error "gunfold"
dataTypeOf _ = hashMapDataType
dataCast2 f = gcast2 f
fromListConstr :: Constr
fromListConstr = mkConstr hashMapDataType "fromList" [] Prefix
hashMapDataType :: DataType
hashMapDataType = mkDataType "Data.HashMap.Base.HashMap" [fromListConstr]
type Hash = Word
type Bitmap = Word
type Shift = Int
#if MIN_VERSION_base(4,9,0)
instance Show2 HashMap where
liftShowsPrec2 spk slk spv slv d m =
showsUnaryWith (liftShowsPrec sp sl) "fromList" d (toList m)
where
sp = liftShowsPrec2 spk slk spv slv
sl = liftShowList2 spk slk spv slv
instance Show k => Show1 (HashMap k) where
liftShowsPrec = liftShowsPrec2 showsPrec showList
instance (Eq k, Hashable k, Read k) => Read1 (HashMap k) where
liftReadsPrec rp rl = readsData $
readsUnaryWith (liftReadsPrec rp' rl') "fromList" fromList
where
rp' = liftReadsPrec rp rl
rl' = liftReadList rp rl
#endif
instance (Eq k, Hashable k, Read k, Read e) => Read (HashMap k e) where
readPrec = parens $ prec 10 $ do
Ident "fromList" <- lexP
xs <- readPrec
return (fromList xs)
readListPrec = readListPrecDefault
instance (Show k, Show v) => Show (HashMap k v) where
showsPrec d m = showParen (d > 10) $
showString "fromList " . shows (toList m)
instance Traversable (HashMap k) where
traverse f = traverseWithKey (const f)
{-# INLINABLE traverse #-}
#if MIN_VERSION_base(4,9,0)
instance Eq2 HashMap where
liftEq2 = equal2
instance Eq k => Eq1 (HashMap k) where
liftEq = equal1
#endif
instance (Eq k, Eq v) => Eq (HashMap k v) where
(==) = equal1 (==)
equal1 :: Eq k
=> (v -> v' -> Bool)
-> HashMap k v -> HashMap k v' -> Bool
equal1 eq = go
where
go Empty Empty = True
go (BitmapIndexed bm1 ary1) (BitmapIndexed bm2 ary2)
= bm1 == bm2 && A.sameArray1 go ary1 ary2
go (Leaf h1 l1) (Leaf h2 l2) = h1 == h2 && leafEq l1 l2
go (Full ary1) (Full ary2) = A.sameArray1 go ary1 ary2
go (Collision h1 ary1) (Collision h2 ary2)
= h1 == h2 && isPermutationBy leafEq (A.toList ary1) (A.toList ary2)
go _ _ = False
leafEq (L k1 v1) (L k2 v2) = k1 == k2 && eq v1 v2
equal2 :: (k -> k' -> Bool) -> (v -> v' -> Bool)
-> HashMap k v -> HashMap k' v' -> Bool
equal2 eqk eqv t1 t2 = go (toList' t1 []) (toList' t2 [])
where
go (Leaf k1 l1 : tl1) (Leaf k2 l2 : tl2)
| k1 == k2 &&
leafEq l1 l2
= go tl1 tl2
go (Collision k1 ary1 : tl1) (Collision k2 ary2 : tl2)
| k1 == k2 &&
A.length ary1 == A.length ary2 &&
isPermutationBy leafEq (A.toList ary1) (A.toList ary2)
= go tl1 tl2
go [] [] = True
go _ _ = False
leafEq (L k v) (L k' v') = eqk k k' && eqv v v'
#if MIN_VERSION_base(4,9,0)
instance Ord2 HashMap where
liftCompare2 = cmp
instance Ord k => Ord1 (HashMap k) where
liftCompare = cmp compare
#endif
instance (Ord k, Ord v) => Ord (HashMap k v) where
compare = cmp compare compare
cmp :: (k -> k' -> Ordering) -> (v -> v' -> Ordering)
-> HashMap k v -> HashMap k' v' -> Ordering
cmp cmpk cmpv t1 t2 = go (toList' t1 []) (toList' t2 [])
where
go (Leaf k1 l1 : tl1) (Leaf k2 l2 : tl2)
= compare k1 k2 `mappend`
leafCompare l1 l2 `mappend`
go tl1 tl2
go (Collision k1 ary1 : tl1) (Collision k2 ary2 : tl2)
= compare k1 k2 `mappend`
compare (A.length ary1) (A.length ary2) `mappend`
unorderedCompare leafCompare (A.toList ary1) (A.toList ary2) `mappend`
go tl1 tl2
go (Leaf _ _ : _) (Collision _ _ : _) = LT
go (Collision _ _ : _) (Leaf _ _ : _) = GT
go [] [] = EQ
go [] _ = LT
go _ [] = GT
go _ _ = error "cmp: Should never happend, toList' includes non Leaf / Collision"
leafCompare (L k v) (L k' v') = cmpk k k' `mappend` cmpv v v'
equalKeys1 :: (k -> k' -> Bool) -> HashMap k v -> HashMap k' v' -> Bool
equalKeys1 eq t1 t2 = go (toList' t1 []) (toList' t2 [])
where
go (Leaf k1 l1 : tl1) (Leaf k2 l2 : tl2)
| k1 == k2 && leafEq l1 l2
= go tl1 tl2
go (Collision k1 ary1 : tl1) (Collision k2 ary2 : tl2)
| k1 == k2 && A.length ary1 == A.length ary2 &&
isPermutationBy leafEq (A.toList ary1) (A.toList ary2)
= go tl1 tl2
go [] [] = True
go _ _ = False
leafEq (L k _) (L k' _) = eq k k'
equalKeys :: Eq k => HashMap k v -> HashMap k v' -> Bool
equalKeys = go
where
go :: Eq k => HashMap k v -> HashMap k v' -> Bool
go Empty Empty = True
go (BitmapIndexed bm1 ary1) (BitmapIndexed bm2 ary2)
= bm1 == bm2 && A.sameArray1 go ary1 ary2
go (Leaf h1 l1) (Leaf h2 l2) = h1 == h2 && leafEq l1 l2
go (Full ary1) (Full ary2) = A.sameArray1 go ary1 ary2
go (Collision h1 ary1) (Collision h2 ary2)
= h1 == h2 && isPermutationBy leafEq (A.toList ary1) (A.toList ary2)
go _ _ = False
leafEq (L k1 _) (L k2 _) = k1 == k2
#if MIN_VERSION_hashable(1,2,5)
instance H.Hashable2 HashMap where
liftHashWithSalt2 hk hv salt hm = go salt (toList' hm [])
where
go s [] = s
go s (Leaf _ l : tl)
= s `hashLeafWithSalt` l `go` tl
go s (Collision h a : tl)
= (s `H.hashWithSalt` h) `hashCollisionWithSalt` a `go` tl
go s (_ : tl) = s `go` tl
hashLeafWithSalt s (L k v) = (s `hk` k) `hv` v
hashCollisionWithSalt s
= L.foldl' H.hashWithSalt s . arrayHashesSorted s
arrayHashesSorted s = L.sort . L.map (hashLeafWithSalt s) . A.toList
instance (Hashable k) => H.Hashable1 (HashMap k) where
liftHashWithSalt = H.liftHashWithSalt2 H.hashWithSalt
#endif
instance (Hashable k, Hashable v) => Hashable (HashMap k v) where
hashWithSalt salt hm = go salt (toList' hm [])
where
go :: Int -> [HashMap k v] -> Int
go s [] = s
go s (Leaf _ l : tl)
= s `hashLeafWithSalt` l `go` tl
go s (Collision h a : tl)
= (s `H.hashWithSalt` h) `hashCollisionWithSalt` a `go` tl
go s (_ : tl) = s `go` tl
hashLeafWithSalt :: Int -> Leaf k v -> Int
hashLeafWithSalt s (L k v) = s `H.hashWithSalt` k `H.hashWithSalt` v
hashCollisionWithSalt :: Int -> A.Array (Leaf k v) -> Int
hashCollisionWithSalt s
= L.foldl' H.hashWithSalt s . arrayHashesSorted s
arrayHashesSorted :: Int -> A.Array (Leaf k v) -> [Int]
arrayHashesSorted s = L.sort . L.map (hashLeafWithSalt s) . A.toList
toList' :: HashMap k v -> [HashMap k v] -> [HashMap k v]
toList' (BitmapIndexed _ ary) a = A.foldr toList' a ary
toList' (Full ary) a = A.foldr toList' a ary
toList' l@(Leaf _ _) a = l : a
toList' c@(Collision _ _) a = c : a
toList' Empty a = a
isLeafOrCollision :: HashMap k v -> Bool
isLeafOrCollision (Leaf _ _) = True
isLeafOrCollision (Collision _ _) = True
isLeafOrCollision _ = False
empty :: HashMap k v
empty = Empty
singleton :: (Hashable k) => k -> v -> HashMap k v
singleton k v = Leaf (hash k) (L k v)
null :: HashMap k v -> Bool
null Empty = True
null _ = False
size :: HashMap k v -> Int
size t = go t 0
where
go Empty !n = n
go (Leaf _ _) n = n + 1
go (BitmapIndexed _ ary) n = A.foldl' (flip go) n ary
go (Full ary) n = A.foldl' (flip go) n ary
go (Collision _ ary) n = n + A.length ary
member :: (Eq k, Hashable k) => k -> HashMap k a -> Bool
member k m = case lookup k m of
Nothing -> False
Just _ -> True
{-# INLINABLE member #-}
lookup :: (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
#if __GLASGOW_HASKELL__ >= 802
lookup k m = case lookup# k m of
(# (# #) | #) -> Nothing
(# | a #) -> Just a
{-# INLINE lookup #-}
lookup# :: (Eq k, Hashable k) => k -> HashMap k v -> (# (# #) | v #)
lookup# k m = lookupCont (\_ -> (# (# #) | #)) (\v _i -> (# | v #)) (hash k) k m
{-# INLINABLE lookup# #-}
#else
lookup k m = lookupCont (\_ -> Nothing) (\v _i -> Just v) (hash k) k m
{-# INLINABLE lookup #-}
#endif
lookup' :: Eq k => Hash -> k -> HashMap k v -> Maybe v
#if __GLASGOW_HASKELL__ >= 802
lookup' h k m = case lookupRecordCollision# h k m of
(# (# #) | #) -> Nothing
(# | (# a, _i #) #) -> Just a
{-# INLINE lookup' #-}
#else
lookup' h k m = lookupCont (\_ -> Nothing) (\v _i -> Just v) h k m
{-# INLINABLE lookup' #-}
#endif
data LookupRes a = Absent | Present a !Int
lookupRecordCollision :: Eq k => Hash -> k -> HashMap k v -> LookupRes v
#if __GLASGOW_HASKELL__ >= 802
lookupRecordCollision h k m = case lookupRecordCollision# h k m of
(# (# #) | #) -> Absent
(# | (# a, i #) #) -> Present a (I# i)
{-# INLINE lookupRecordCollision #-}
lookupRecordCollision# :: Eq k => Hash -> k -> HashMap k v -> (# (# #) | (# v, Int# #) #)
lookupRecordCollision# h k m =
lookupCont (\_ -> (# (# #) | #)) (\v (I# i) -> (# | (# v, i #) #)) h k m
{-# INLINABLE lookupRecordCollision# #-}
#else /* GHC < 8.2 so there are no unboxed sums */
lookupRecordCollision h k m = lookupCont (\_ -> Absent) Present h k m
{-# INLINABLE lookupRecordCollision #-}
#endif
lookupCont ::
#if __GLASGOW_HASKELL__ >= 802
forall rep (r :: TYPE rep) k v.
#else
forall r k v.
#endif
Eq k
=> ((# #) -> r)
-> (v -> Int -> r)
-> Hash
-> k -> HashMap k v -> r
lookupCont absent present !h0 !k0 !m0 = go h0 k0 0 m0
where
go :: Eq k => Hash -> k -> Int -> HashMap k v -> r
go !_ !_ !_ Empty = absent (# #)
go h k _ (Leaf hx (L kx x))
| h == hx && k == kx = present x (-1)
| otherwise = absent (# #)
go h k s (BitmapIndexed b v)
| b .&. m == 0 = absent (# #)
| otherwise =
go h k (s+bitsPerSubkey) (A.index v (sparseIndex b m))
where m = mask h s
go h k s (Full v) =
go h k (s+bitsPerSubkey) (A.index v (index h s))
go h k _ (Collision hx v)
| h == hx = lookupInArrayCont absent present k v
| otherwise = absent (# #)
{-# INLINE lookupCont #-}
lookupDefault :: (Eq k, Hashable k)
=> v
-> k -> HashMap k v -> v
lookupDefault def k t = case lookup k t of
Just v -> v
_ -> def
{-# INLINABLE lookupDefault #-}
(!) :: (Eq k, Hashable k) => HashMap k v -> k -> v
(!) m k = case lookup k m of
Just v -> v
Nothing -> error "Data.HashMap.Base.(!): key not found"
{-# INLINABLE (!) #-}
infixl 9 !
collision :: Hash -> Leaf k v -> Leaf k v -> HashMap k v
collision h !e1 !e2 =
let v = A.run $ do mary <- A.new 2 e1
A.write mary 1 e2
return mary
in Collision h v
{-# INLINE collision #-}
bitmapIndexedOrFull :: Bitmap -> A.Array (HashMap k v) -> HashMap k v
bitmapIndexedOrFull b ary
| b == fullNodeMask = Full ary
| otherwise = BitmapIndexed b ary
{-# INLINE bitmapIndexedOrFull #-}
insert :: (Eq k, Hashable k) => k -> v -> HashMap k v -> HashMap k v
insert k v m = insert' (hash k) k v m
{-# INLINABLE insert #-}
insert' :: Eq k => Hash -> k -> v -> HashMap k v -> HashMap k v
insert' h0 k0 v0 m0 = go h0 k0 v0 0 m0
where
go !h !k x !_ Empty = Leaf h (L k x)
go h k x s t@(Leaf hy l@(L ky y))
| hy == h = if ky == k
then if x `ptrEq` y
then t
else Leaf h (L k x)
else collision h l (L k x)
| otherwise = runST (two s h k x hy ky y)
go h k x s t@(BitmapIndexed b ary)
| b .&. m == 0 =
let !ary' = A.insert ary i $! Leaf h (L k x)
in bitmapIndexedOrFull (b .|. m) ary'
| otherwise =
let !st = A.index ary i
!st' = go h k x (s+bitsPerSubkey) st
in if st' `ptrEq` st
then t
else BitmapIndexed b (A.update ary i st')
where m = mask h s
i = sparseIndex b m
go h k x s t@(Full ary) =
let !st = A.index ary i
!st' = go h k x (s+bitsPerSubkey) st
in if st' `ptrEq` st
then t
else Full (update16 ary i st')
where i = index h s
go h k x s t@(Collision hy v)
| h == hy = Collision h (updateOrSnocWith const k x v)
| otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
{-# INLINABLE insert' #-}
insertNewKey :: Hash -> k -> v -> HashMap k v -> HashMap k v
insertNewKey !h0 !k0 x0 !m0 = go h0 k0 x0 0 m0
where
go !h !k x !_ Empty = Leaf h (L k x)
go h k x s (Leaf hy l@(L ky y))
| hy == h = collision h l (L k x)
| otherwise = runST (two s h k x hy ky y)
go h k x s (BitmapIndexed b ary)
| b .&. m == 0 =
let !ary' = A.insert ary i $! Leaf h (L k x)
in bitmapIndexedOrFull (b .|. m) ary'
| otherwise =
let !st = A.index ary i
!st' = go h k x (s+bitsPerSubkey) st
in BitmapIndexed b (A.update ary i st')
where m = mask h s
i = sparseIndex b m
go h k x s (Full ary) =
let !st = A.index ary i
!st' = go h k x (s+bitsPerSubkey) st
in Full (update16 ary i st')
where i = index h s
go h k x s t@(Collision hy v)
| h == hy = Collision h (snocNewLeaf (L k x) v)
| otherwise =
go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
where
snocNewLeaf :: Leaf k v -> A.Array (Leaf k v) -> A.Array (Leaf k v)
snocNewLeaf leaf ary = A.run $ do
let n = A.length ary
mary <- A.new_ (n + 1)
A.copy ary 0 mary 0 n
A.write mary n leaf
return mary
{-# NOINLINE insertNewKey #-}
insertKeyExists :: Int -> Hash -> k -> v -> HashMap k v -> HashMap k v
insertKeyExists !collPos0 !h0 !k0 x0 !m0 = go collPos0 h0 k0 x0 0 m0
where
go !_collPos !h !k x !_s (Leaf _hy _kx)
= Leaf h (L k x)
go collPos h k x s (BitmapIndexed b ary)
| b .&. m == 0 =
let !ary' = A.insert ary i $ Leaf h (L k x)
in bitmapIndexedOrFull (b .|. m) ary'
| otherwise =
let !st = A.index ary i
!st' = go collPos h k x (s+bitsPerSubkey) st
in BitmapIndexed b (A.update ary i st')
where m = mask h s
i = sparseIndex b m
go collPos h k x s (Full ary) =
let !st = A.index ary i
!st' = go collPos h k x (s+bitsPerSubkey) st
in Full (update16 ary i st')
where i = index h s
go collPos h k x _s (Collision _hy v)
| collPos >= 0 = Collision h (setAtPosition collPos k x v)
| otherwise = Empty
go _ _ _ _ _ Empty = Empty
{-# NOINLINE insertKeyExists #-}
setAtPosition :: Int -> k -> v -> A.Array (Leaf k v) -> A.Array (Leaf k v)
setAtPosition i k x ary = A.update ary i (L k x)
{-# INLINE setAtPosition #-}
unsafeInsert :: (Eq k, Hashable k) => k -> v -> HashMap k v -> HashMap k v
unsafeInsert k0 v0 m0 = runST (go h0 k0 v0 0 m0)
where
h0 = hash k0
go !h !k x !_ Empty = return $! Leaf h (L k x)
go h k x s t@(Leaf hy l@(L ky y))
| hy == h = if ky == k
then if x `ptrEq` y
then return t
else return $! Leaf h (L k x)
else return $! collision h l (L k x)
| otherwise = two s h k x hy ky y
go h k x s t@(BitmapIndexed b ary)
| b .&. m == 0 = do
ary' <- A.insertM ary i $! Leaf h (L k x)
return $! bitmapIndexedOrFull (b .|. m) ary'
| otherwise = do
st <- A.indexM ary i
st' <- go h k x (s+bitsPerSubkey) st
A.unsafeUpdateM ary i st'
return t
where m = mask h s
i = sparseIndex b m
go h k x s t@(Full ary) = do
st <- A.indexM ary i
st' <- go h k x (s+bitsPerSubkey) st
A.unsafeUpdateM ary i st'
return t
where i = index h s
go h k x s t@(Collision hy v)
| h == hy = return $! Collision h (updateOrSnocWith const k x v)
| otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
{-# INLINABLE unsafeInsert #-}
two :: Shift -> Hash -> k -> v -> Hash -> k -> v -> ST s (HashMap k v)
two = go
where
go s h1 k1 v1 h2 k2 v2
| bp1 == bp2 = do
st <- go (s+bitsPerSubkey) h1 k1 v1 h2 k2 v2
ary <- A.singletonM st
return $! BitmapIndexed bp1 ary
| otherwise = do
mary <- A.new 2 $ Leaf h1 (L k1 v1)
A.write mary idx2 $ Leaf h2 (L k2 v2)
ary <- A.unsafeFreeze mary
return $! BitmapIndexed (bp1 .|. bp2) ary
where
bp1 = mask h1 s
bp2 = mask h2 s
idx2 | index h1 s < index h2 s = 1
| otherwise = 0
{-# INLINE two #-}
insertWith :: (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v
-> HashMap k v
insertWith f k new m = insertModifying new (\old -> (# f new old #)) k m
{-# INLINE insertWith #-}
insertModifying :: (Eq k, Hashable k) => v -> (v -> (# v #)) -> k -> HashMap k v
-> HashMap k v
insertModifying x f k0 m0 = go h0 k0 0 m0
where
!h0 = hash k0
go !h !k !_ Empty = Leaf h (L k x)
go h k s t@(Leaf hy l@(L ky y))
| hy == h = if ky == k
then case f y of
(# v' #) | ptrEq y v' -> t
| otherwise -> Leaf h (L k (v'))
else collision h l (L k x)
| otherwise = runST (two s h k x hy ky y)
go h k s t@(BitmapIndexed b ary)
| b .&. m == 0 =
let ary' = A.insert ary i $! Leaf h (L k x)
in bitmapIndexedOrFull (b .|. m) ary'
| otherwise =
let !st = A.index ary i
!st' = go h k (s+bitsPerSubkey) st
ary' = A.update ary i $! st'
in if ptrEq st st'
then t
else BitmapIndexed b ary'
where m = mask h s
i = sparseIndex b m
go h k s t@(Full ary) =
let !st = A.index ary i
!st' = go h k (s+bitsPerSubkey) st
ary' = update16 ary i $! st'
in if ptrEq st st'
then t
else Full ary'
where i = index h s
go h k s t@(Collision hy v)
| h == hy =
let !v' = insertModifyingArr x f k v
in if A.unsafeSameArray v v'
then t
else Collision h v'
| otherwise = go h k s $ BitmapIndexed (mask hy s) (A.singleton t)
{-# INLINABLE insertModifying #-}
insertModifyingArr :: Eq k => v -> (v -> (# v #)) -> k -> A.Array (Leaf k v)
-> A.Array (Leaf k v)
insertModifyingArr x f k0 ary0 = go k0 ary0 0 (A.length ary0)
where
go !k !ary !i !n
| i >= n = A.run $ do
mary <- A.new_ (n + 1)
A.copy ary 0 mary 0 n
A.write mary n (L k x)
return mary
| otherwise = case A.index ary i of
(L kx y) | k == kx -> case f y of
(# y' #) -> if ptrEq y y'
then ary
else A.update ary i (L k y')
| otherwise -> go k ary (i+1) n
{-# INLINE insertModifyingArr #-}
unsafeInsertWith :: forall k v. (Eq k, Hashable k)
=> (v -> v -> v) -> k -> v -> HashMap k v
-> HashMap k v
unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
where
h0 = hash k0
go :: Hash -> k -> v -> Shift -> HashMap k v -> ST s (HashMap k v)
go !h !k x !_ Empty = return $! Leaf h (L k x)
go h k x s (Leaf hy l@(L ky y))
| hy == h = if ky == k
then return $! Leaf h (L k (f x y))
else return $! collision h l (L k x)
| otherwise = two s h k x hy ky y
go h k x s t@(BitmapIndexed b ary)
| b .&. m == 0 = do
ary' <- A.insertM ary i $! Leaf h (L k x)
return $! bitmapIndexedOrFull (b .|. m) ary'
| otherwise = do
st <- A.indexM ary i
st' <- go h k x (s+bitsPerSubkey) st
A.unsafeUpdateM ary i st'
return t
where m = mask h s
i = sparseIndex b m
go h k x s t@(Full ary) = do
st <- A.indexM ary i
st' <- go h k x (s+bitsPerSubkey) st
A.unsafeUpdateM ary i st'
return t
where i = index h s
go h k x s t@(Collision hy v)
| h == hy = return $! Collision h (updateOrSnocWith f k x v)
| otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
{-# INLINABLE unsafeInsertWith #-}
delete :: (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
delete k m = delete' (hash k) k m
{-# INLINABLE delete #-}
delete' :: Eq k => Hash -> k -> HashMap k v -> HashMap k v
delete' h0 k0 m0 = go h0 k0 0 m0
where
go !_ !_ !_ Empty = Empty
go h k _ t@(Leaf hy (L ky _))
| hy == h && ky == k = Empty
| otherwise = t
go h k s t@(BitmapIndexed b ary)
| b .&. m == 0 = t
| otherwise =
let !st = A.index ary i
!st' = go h k (s+bitsPerSubkey) st
in if st' `ptrEq` st
then t
else case st' of
Empty | A.length ary == 1 -> Empty
| A.length ary == 2 ->
case (i, A.index ary 0, A.index ary 1) of
(0, _, l) | isLeafOrCollision l -> l
(1, l, _) | isLeafOrCollision l -> l
_ -> bIndexed
| otherwise -> bIndexed
where
bIndexed = BitmapIndexed (b .&. complement m) (A.delete ary i)
l | isLeafOrCollision l && A.length ary == 1 -> l
_ -> BitmapIndexed b (A.update ary i st')
where m = mask h s
i = sparseIndex b m
go h k s t@(Full ary) =
let !st = A.index ary i
!st' = go h k (s+bitsPerSubkey) st
in if st' `ptrEq` st
then t
else case st' of
Empty ->
let ary' = A.delete ary i
bm = fullNodeMask .&. complement (1 `unsafeShiftL` i)
in BitmapIndexed bm ary'
_ -> Full (A.update ary i st')
where i = index h s
go h k _ t@(Collision hy v)
| h == hy = case indexOf k v of
Just i
| A.length v == 2 ->
if i == 0
then Leaf h (A.index v 1)
else Leaf h (A.index v 0)
| otherwise -> Collision h (A.delete v i)
Nothing -> t
| otherwise = t
{-# INLINABLE delete' #-}
deleteKeyExists :: Int -> Hash -> k -> HashMap k v -> HashMap k v
deleteKeyExists !collPos0 !h0 !k0 !m0 = go collPos0 h0 k0 0 m0
where
go :: Int -> Hash -> k -> Int -> HashMap k v -> HashMap k v
go !_collPos !_h !_k !_s (Leaf _ _) = Empty
go collPos h k s (BitmapIndexed b ary) =
let !st = A.index ary i
!st' = go collPos h k (s+bitsPerSubkey) st
in case st' of
Empty | A.length ary == 1 -> Empty
| A.length ary == 2 ->
case (i, A.index ary 0, A.index ary 1) of
(0, _, l) | isLeafOrCollision l -> l
(1, l, _) | isLeafOrCollision l -> l
_ -> bIndexed
| otherwise -> bIndexed
where
bIndexed = BitmapIndexed (b .&. complement m) (A.delete ary i)
l | isLeafOrCollision l && A.length ary == 1 -> l
_ -> BitmapIndexed b (A.update ary i st')
where m = mask h s
i = sparseIndex b m
go collPos h k s (Full ary) =
let !st = A.index ary i
!st' = go collPos h k (s+bitsPerSubkey) st
in case st' of
Empty ->
let ary' = A.delete ary i
bm = fullNodeMask .&. complement (1 `unsafeShiftL` i)
in BitmapIndexed bm ary'
_ -> Full (A.update ary i st')
where i = index h s
go collPos h _ _ (Collision _hy v)
| A.length v == 2
= if collPos == 0
then Leaf h (A.index v 1)
else Leaf h (A.index v 0)
| otherwise = Collision h (A.delete v collPos)
go !_ !_ !_ !_ Empty = Empty
{-# NOINLINE deleteKeyExists #-}
adjust :: (Eq k, Hashable k) => (v -> v) -> k -> HashMap k v -> HashMap k v
adjust f k m = adjust# (\v -> (# f v #)) k m
{-# INLINE adjust #-}
adjust# :: (Eq k, Hashable k) => (v -> (# v #)) -> k -> HashMap k v -> HashMap k v
adjust# f k0 m0 = go h0 k0 0 m0
where
h0 = hash k0
go !_ !_ !_ Empty = Empty
go h k _ t@(Leaf hy (L ky y))
| hy == h && ky == k = case f y of
(# y' #) | ptrEq y y' -> t
| otherwise -> Leaf h (L k y')
| otherwise = t
go h k s t@(BitmapIndexed b ary)
| b .&. m == 0 = t
| otherwise = let !st = A.index ary i
!st' = go h k (s+bitsPerSubkey) st
ary' = A.update ary i $! st'
in if ptrEq st st'
then t
else BitmapIndexed b ary'
where m = mask h s
i = sparseIndex b m
go h k s t@(Full ary) =
let i = index h s
!st = A.index ary i
!st' = go h k (s+bitsPerSubkey) st
ary' = update16 ary i $! st'
in if ptrEq st st'
then t
else Full ary'
go h k _ t@(Collision hy v)
| h == hy = let !v' = updateWith# f k v
in if A.unsafeSameArray v v'
then t
else Collision h v'
| otherwise = t
{-# INLINABLE adjust# #-}
update :: (Eq k, Hashable k) => (a -> Maybe a) -> k -> HashMap k a -> HashMap k a
update f = alter (>>= f)
{-# INLINABLE update #-}
alter :: (Eq k, Hashable k) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
alter f k m =
case f (lookup k m) of
Nothing -> delete k m
Just v -> insert k v m
{-# INLINABLE alter #-}
alterF :: (Functor f, Eq k, Hashable k)
=> (Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v)
alterF f = \ !k !m ->
let
!h = hash k
mv = lookup' h k m
in (<$> f mv) $ \fres ->
case fres of
Nothing -> delete' h k m
Just v' -> insert' h k v' m
{-# INLINABLE [0] alterF #-}
#if MIN_VERSION_base(4,8,0)
test_bottom :: a
test_bottom = error "Data.HashMap.alterF internal error: hit test_bottom"
bogus# :: (# #) -> (# a #)
bogus# _ = error "Data.HashMap.alterF internal error: hit bogus#"
{-# RULES
-- We probe the behavior of @f@ by applying it to Nothing and to
-- Just test_bottom. Based on the results, and how they relate to
-- each other, we choose the best implementation.
"alterFWeird" forall f. alterF f =
alterFWeird (f Nothing) (f (Just test_bottom)) f
-- This rule covers situations where alterF is used to simply insert or
-- delete in Identity (most likely via Control.Lens.At). We recognize here
-- (through the repeated @x@ on the LHS) that
--
-- @f Nothing = f (Just bottom)@,
--
-- which guarantees that @f@ doesn't care what its argument is, so
-- we don't have to either.
--
-- Why only Identity? A variant of this rule is actually valid regardless of
-- the functor, but for some functors (e.g., []), it can lead to the
-- same keys being compared multiple times, which is bad if they're
-- ugly things like strings. This is unfortunate, since the rule is likely
-- a good idea for almost all realistic uses, but I don't like nasty
-- edge cases.
"alterFconstant" forall (f :: Maybe a -> Identity (Maybe a)) x.
alterFWeird x x f = \ !k !m ->
Identity (case runIdentity x of {Nothing -> delete k m; Just a -> insert k a m})
-- This rule handles the case where 'alterF' is used to do 'insertWith'-like
-- things. Whenever possible, GHC will get rid of the Maybe nonsense for us.
-- We delay this rule to stage 1 so alterFconstant has a chance to fire.
"alterFinsertWith" [1] forall (f :: Maybe a -> Identity (Maybe a)) x y.
alterFWeird (coerce (Just x)) (coerce (Just y)) f =
coerce (insertModifying x (\mold -> case runIdentity (f (Just mold)) of
Nothing -> bogus# (# #)
Just new -> (# new #)))
-- Handle the case where someone uses 'alterF' instead of 'adjust'. This
-- rule is kind of picky; it will only work if the function doesn't
-- do anything between case matching on the Maybe and producing a result.
"alterFadjust" forall (f :: Maybe a -> Identity (Maybe a)) _y.
alterFWeird (coerce Nothing) (coerce (Just _y)) f =
coerce (adjust# (\x -> case runIdentity (f (Just x)) of
Just x' -> (# x' #)
Nothing -> bogus# (# #)))
-- The simple specialization to Const; in this case we can look up
-- the key without caring what position it's in. This is only a tiny
-- optimization.
"alterFlookup" forall _ign1 _ign2 (f :: Maybe a -> Const r (Maybe a)).
alterFWeird _ign1 _ign2 f = \ !k !m -> Const (getConst (f (lookup k m)))
#-}
alterFWeird
:: (Functor f, Eq k, Hashable k)
=> f (Maybe v)
-> f (Maybe v)
-> (Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v)
alterFWeird _ _ f = alterFEager f
{-# INLINE [0] alterFWeird #-}
alterFEager :: (Functor f, Eq k, Hashable k)
=> (Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v)
alterFEager f !k m = (<$> f mv) $ \fres ->
case fres of
Nothing -> case lookupRes of
Absent -> m
Present _ collPos -> deleteKeyExists collPos h k m
Just v' -> case lookupRes of
Absent -> insertNewKey h k v' m
Present v collPos ->
if v `ptrEq` v'
then m
else insertKeyExists collPos h k v' m
where !h = hash k
!lookupRes = lookupRecordCollision h k m
!mv = case lookupRes of
Absent -> Nothing
Present v _ -> Just v
{-# INLINABLE alterFEager #-}
#endif
union :: (Eq k, Hashable k) => HashMap k v -> HashMap k v -> HashMap k v
union = unionWith const
{-# INLINABLE union #-}
unionWith :: (Eq k, Hashable k) => (v -> v -> v) -> HashMap k v -> HashMap k v
-> HashMap k v
unionWith f = unionWithKey (const f)
{-# INLINE unionWith #-}
unionWithKey :: (Eq k, Hashable k) => (k -> v -> v -> v) -> HashMap k v -> HashMap k v
-> HashMap k v
unionWithKey f = go 0
where
go !_ t1 Empty = t1
go _ Empty t2 = t2
go s t1@(Leaf h1 l1@(L k1 v1)) t2@(Leaf h2 l2@(L k2 v2))
| h1 == h2 = if k1 == k2
then Leaf h1 (L k1 (f k1 v1 v2))
else collision h1 l1 l2
| otherwise = goDifferentHash s h1 h2 t1 t2
go s t1@(Leaf h1 (L k1 v1)) t2@(Collision h2 ls2)
| h1 == h2 = Collision h1 (updateOrSnocWithKey f k1 v1 ls2)
| otherwise = goDifferentHash s h1 h2 t1 t2
go s t1@(Collision h1 ls1) t2@(Leaf h2 (L k2 v2))
| h1 == h2 = Collision h1 (updateOrSnocWithKey (flip . f) k2 v2 ls1)
| otherwise = goDifferentHash s h1 h2 t1 t2
go s t1@(Collision h1 ls1) t2@(Collision h2 ls2)
| h1 == h2 = Collision h1 (updateOrConcatWithKey f ls1 ls2)
| otherwise = goDifferentHash s h1 h2 t1 t2
go s (BitmapIndexed b1 ary1) (BitmapIndexed b2 ary2) =
let b' = b1 .|. b2
ary' = unionArrayBy (go (s+bitsPerSubkey)) b1 b2 ary1 ary2
in bitmapIndexedOrFull b' ary'
go s (BitmapIndexed b1 ary1) (Full ary2) =
let ary' = unionArrayBy (go (s+bitsPerSubkey)) b1 fullNodeMask ary1 ary2
in Full ary'
go s (Full ary1) (BitmapIndexed b2 ary2) =
let ary' = unionArrayBy (go (s+bitsPerSubkey)) fullNodeMask b2 ary1 ary2
in Full ary'
go s (Full ary1) (Full ary2) =
let ary' = unionArrayBy (go (s+bitsPerSubkey)) fullNodeMask fullNodeMask
ary1 ary2
in Full ary'
go s (BitmapIndexed b1 ary1) t2
| b1 .&. m2 == 0 = let ary' = A.insert ary1 i t2
b' = b1 .|. m2
in bitmapIndexedOrFull b' ary'
| otherwise = let ary' = A.updateWith' ary1 i $ \st1 ->
go (s+bitsPerSubkey) st1 t2
in BitmapIndexed b1 ary'
where
h2 = leafHashCode t2
m2 = mask h2 s
i = sparseIndex b1 m2
go s t1 (BitmapIndexed b2 ary2)
| b2 .&. m1 == 0 = let ary' = A.insert ary2 i $! t1
b' = b2 .|. m1
in bitmapIndexedOrFull b' ary'
| otherwise = let ary' = A.updateWith' ary2 i $ \st2 ->
go (s+bitsPerSubkey) t1 st2
in BitmapIndexed b2 ary'
where
h1 = leafHashCode t1
m1 = mask h1 s
i = sparseIndex b2 m1
go s (Full ary1) t2 =
let h2 = leafHashCode t2
i = index h2 s
ary' = update16With' ary1 i $ \st1 -> go (s+bitsPerSubkey) st1 t2
in Full ary'
go s t1 (Full ary2) =
let h1 = leafHashCode t1
i = index h1 s
ary' = update16With' ary2 i $ \st2 -> go (s+bitsPerSubkey) t1 st2
in Full ary'
leafHashCode (Leaf h _) = h
leafHashCode (Collision h _) = h
leafHashCode _ = error "leafHashCode"
goDifferentHash s h1 h2 t1 t2
| m1 == m2 = BitmapIndexed m1 (A.singleton $! go (s+bitsPerSubkey) t1 t2)
| m1 < m2 = BitmapIndexed (m1 .|. m2) (A.pair t1 t2)
| otherwise = BitmapIndexed (m1 .|. m2) (A.pair t2 t1)
where
m1 = mask h1 s
m2 = mask h2 s
{-# INLINE unionWithKey #-}
unionArrayBy :: (a -> a -> a) -> Bitmap -> Bitmap -> A.Array a -> A.Array a
-> A.Array a
unionArrayBy f b1 b2 ary1 ary2 = A.run $ do
let b' = b1 .|. b2
mary <- A.new_ (popCount b')
let ba = b1 .&. b2
go !i !i1 !i2 !m
| m > b' = return ()
| b' .&. m == 0 = go i i1 i2 (m `unsafeShiftL` 1)
| ba .&. m /= 0 = do
x1 <- A.indexM ary1 i1
x2 <- A.indexM ary2 i2
A.write mary i $! f x1 x2
go (i+1) (i1+1) (i2+1) (m `unsafeShiftL` 1)
| b1 .&. m /= 0 = do
A.write mary i =<< A.indexM ary1 i1
go (i+1) (i1+1) (i2 ) (m `unsafeShiftL` 1)
| otherwise = do
A.write mary i =<< A.indexM ary2 i2
go (i+1) (i1 ) (i2+1) (m `unsafeShiftL` 1)
go 0 0 0 (b' .&. negate b')
return mary
{-# INLINE unionArrayBy #-}
unions :: (Eq k, Hashable k) => [HashMap k v] -> HashMap k v
unions = L.foldl' union empty
{-# INLINE unions #-}
mapWithKey :: (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2
mapWithKey f = go
where
go Empty = Empty
go (Leaf h (L k v)) = Leaf h $ L k (f k v)
go (BitmapIndexed b ary) = BitmapIndexed b $ A.map go ary
go (Full ary) = Full $ A.map go ary
go (Collision h ary) = Collision h $
A.map' (\ (L k v) -> L k (f k v)) ary
{-# INLINE mapWithKey #-}
map :: (v1 -> v2) -> HashMap k v1 -> HashMap k v2
map f = mapWithKey (const f)
{-# INLINE map #-}
traverseWithKey
:: Applicative f
=> (k -> v1 -> f v2)
-> HashMap k v1 -> f (HashMap k v2)
traverseWithKey f = go
where
go Empty = pure Empty
go (Leaf h (L k v)) = Leaf h . L k <$> f k v
go (BitmapIndexed b ary) = BitmapIndexed b <$> A.traverse go ary
go (Full ary) = Full <$> A.traverse go ary
go (Collision h ary) =
Collision h <$> A.traverse' (\ (L k v) -> L k <$> f k v) ary
{-# INLINE traverseWithKey #-}
difference :: (Eq k, Hashable k) => HashMap k v -> HashMap k w -> HashMap k v
difference a b = foldlWithKey' go empty a
where
go m k v = case lookup k b of
Nothing -> insert k v m
_ -> m
{-# INLINABLE difference #-}
differenceWith :: (Eq k, Hashable k) => (v -> w -> Maybe v) -> HashMap k v -> HashMap k w -> HashMap k v
differenceWith f a b = foldlWithKey' go empty a
where
go m k v = case lookup k b of
Nothing -> insert k v m
Just w -> maybe m (\y -> insert k y m) (f v w)
{-# INLINABLE differenceWith #-}
intersection :: (Eq k, Hashable k) => HashMap k v -> HashMap k w -> HashMap k v
intersection a b = foldlWithKey' go empty a
where
go m k v = case lookup k b of
Just _ -> insert k v m
_ -> m
{-# INLINABLE intersection #-}
intersectionWith :: (Eq k, Hashable k) => (v1 -> v2 -> v3) -> HashMap k v1
-> HashMap k v2 -> HashMap k v3
intersectionWith f a b = foldlWithKey' go empty a
where
go m k v = case lookup k b of
Just w -> insert k (f v w) m
_ -> m
{-# INLINABLE intersectionWith #-}
intersectionWithKey :: (Eq k, Hashable k) => (k -> v1 -> v2 -> v3)
-> HashMap k v1 -> HashMap k v2 -> HashMap k v3
intersectionWithKey f a b = foldlWithKey' go empty a
where
go m k v = case lookup k b of
Just w -> insert k (f k v w) m
_ -> m
{-# INLINABLE intersectionWithKey #-}
foldl' :: (a -> v -> a) -> a -> HashMap k v -> a
foldl' f = foldlWithKey' (\ z _ v -> f z v)
{-# INLINE foldl' #-}
foldlWithKey' :: (a -> k -> v -> a) -> a -> HashMap k v -> a
foldlWithKey' f = go
where
go !z Empty = z
go z (Leaf _ (L k v)) = f z k v
go z (BitmapIndexed _ ary) = A.foldl' go z ary
go z (Full ary) = A.foldl' go z ary
go z (Collision _ ary) = A.foldl' (\ z' (L k v) -> f z' k v) z ary
{-# INLINE foldlWithKey' #-}
foldr :: (v -> a -> a) -> a -> HashMap k v -> a
foldr f = foldrWithKey (const f)
{-# INLINE foldr #-}
foldrWithKey :: (k -> v -> a -> a) -> a -> HashMap k v -> a
foldrWithKey f = go
where
go z Empty = z
go z (Leaf _ (L k v)) = f k v z
go z (BitmapIndexed _ ary) = A.foldr (flip go) z ary
go z (Full ary) = A.foldr (flip go) z ary
go z (Collision _ ary) = A.foldr (\ (L k v) z' -> f k v z') z ary
{-# INLINE foldrWithKey #-}
mapMaybeWithKey :: (k -> v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
mapMaybeWithKey f = filterMapAux onLeaf onColl
where onLeaf (Leaf h (L k v)) | Just v' <- f k v = Just (Leaf h (L k v'))
onLeaf _ = Nothing
onColl (L k v) | Just v' <- f k v = Just (L k v')
| otherwise = Nothing
{-# INLINE mapMaybeWithKey #-}
mapMaybe :: (v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
mapMaybe f = mapMaybeWithKey (const f)
{-# INLINE mapMaybe #-}
filterWithKey :: forall k v. (k -> v -> Bool) -> HashMap k v -> HashMap k v
filterWithKey pred = filterMapAux onLeaf onColl
where onLeaf t@(Leaf _ (L k v)) | pred k v = Just t
onLeaf _ = Nothing
onColl el@(L k v) | pred k v = Just el
onColl _ = Nothing
{-# INLINE filterWithKey #-}
filterMapAux :: forall k v1 v2
. (HashMap k v1 -> Maybe (HashMap k v2))
-> (Leaf k v1 -> Maybe (Leaf k v2))
-> HashMap k v1
-> HashMap k v2
filterMapAux onLeaf onColl = go
where
go Empty = Empty
go t@Leaf{}
| Just t' <- onLeaf t = t'
| otherwise = Empty
go (BitmapIndexed b ary) = filterA ary b
go (Full ary) = filterA ary fullNodeMask
go (Collision h ary) = filterC ary h
filterA ary0 b0 =
let !n = A.length ary0
in runST $ do
mary <- A.new_ n
step ary0 mary b0 0 0 1 n
where
step :: A.Array (HashMap k v1) -> A.MArray s (HashMap k v2)
-> Bitmap -> Int -> Int -> Bitmap -> Int
-> ST s (HashMap k v2)
step !ary !mary !b i !j !bi n
| i >= n = case j of
0 -> return Empty
1 -> do
ch <- A.read mary 0
case ch of
t | isLeafOrCollision t -> return t
_ -> BitmapIndexed b <$> A.trim mary 1
_ -> do
ary2 <- A.trim mary j
return $! if j == maxChildren
then Full ary2
else BitmapIndexed b ary2
| bi .&. b == 0 = step ary mary b i j (bi `unsafeShiftL` 1) n
| otherwise = case go (A.index ary i) of
Empty -> step ary mary (b .&. complement bi) (i+1) j
(bi `unsafeShiftL` 1) n
t -> do A.write mary j t
step ary mary b (i+1) (j+1) (bi `unsafeShiftL` 1) n
filterC ary0 h =
let !n = A.length ary0
in runST $ do
mary <- A.new_ n
step ary0 mary 0 0 n
where
step :: A.Array (Leaf k v1) -> A.MArray s (Leaf k v2)
-> Int -> Int -> Int
-> ST s (HashMap k v2)
step !ary !mary i !j n
| i >= n = case j of
0 -> return Empty
1 -> do l <- A.read mary 0
return $! Leaf h l
_ | i == j -> do ary2 <- A.unsafeFreeze mary
return $! Collision h ary2
| otherwise -> do ary2 <- A.trim mary j
return $! Collision h ary2
| Just el <- onColl $! A.index ary i
= A.write mary j el >> step ary mary (i+1) (j+1) n
| otherwise = step ary mary (i+1) j n
{-# INLINE filterMapAux #-}
filter :: (v -> Bool) -> HashMap k v -> HashMap k v
filter p = filterWithKey (\_ v -> p v)
{-# INLINE filter #-}
keys :: HashMap k v -> [k]
keys = L.map fst . toList
{-# INLINE keys #-}
elems :: HashMap k v -> [v]
elems = L.map snd . toList
{-# INLINE elems #-}
toList :: HashMap k v -> [(k, v)]
toList t = build (\ c z -> foldrWithKey (curry c) z t)
{-# INLINE toList #-}
fromList :: (Eq k, Hashable k) => [(k, v)] -> HashMap k v
fromList = L.foldl' (\ m (k, v) -> unsafeInsert k v m) empty
{-# INLINABLE fromList #-}
fromListWith :: (Eq k, Hashable k) => (v -> v -> v) -> [(k, v)] -> HashMap k v
fromListWith f = L.foldl' (\ m (k, v) -> unsafeInsertWith f k v m) empty
{-# INLINE fromListWith #-}
lookupInArrayCont ::
#if __GLASGOW_HASKELL__ >= 802
forall rep (r :: TYPE rep) k v.
#else
forall r k v.
#endif
Eq k => ((# #) -> r) -> (v -> Int -> r) -> k -> A.Array (Leaf k v) -> r
lookupInArrayCont absent present k0 ary0 = go k0 ary0 0 (A.length ary0)
where
go :: Eq k => k -> A.Array (Leaf k v) -> Int -> Int -> r
go !k !ary !i !n
| i >= n = absent (# #)
| otherwise = case A.index ary i of
(L kx v)
| k == kx -> present v i
| otherwise -> go k ary (i+1) n
{-# INLINE lookupInArrayCont #-}
indexOf :: Eq k => k -> A.Array (Leaf k v) -> Maybe Int
indexOf k0 ary0 = go k0 ary0 0 (A.length ary0)
where
go !k !ary !i !n
| i >= n = Nothing
| otherwise = case A.index ary i of
(L kx _)
| k == kx -> Just i
| otherwise -> go k ary (i+1) n
{-# INLINABLE indexOf #-}
updateWith# :: Eq k => (v -> (# v #)) -> k -> A.Array (Leaf k v) -> A.Array (Leaf k v)
updateWith# f k0 ary0 = go k0 ary0 0 (A.length ary0)
where
go !k !ary !i !n
| i >= n = ary
| otherwise = case A.index ary i of
(L kx y) | k == kx -> case f y of
(# y' #)
| ptrEq y y' -> ary
| otherwise -> A.update ary i (L k y')
| otherwise -> go k ary (i+1) n
{-# INLINABLE updateWith# #-}
updateOrSnocWith :: Eq k => (v -> v -> v) -> k -> v -> A.Array (Leaf k v)
-> A.Array (Leaf k v)
updateOrSnocWith f = updateOrSnocWithKey (const f)
{-# INLINABLE updateOrSnocWith #-}
updateOrSnocWithKey :: Eq k => (k -> v -> v -> v) -> k -> v -> A.Array (Leaf k v)
-> A.Array (Leaf k v)
updateOrSnocWithKey f k0 v0 ary0 = go k0 v0 ary0 0 (A.length ary0)
where
go !k v !ary !i !n
| i >= n = A.run $ do
mary <- A.new_ (n + 1)
A.copy ary 0 mary 0 n
A.write mary n (L k v)
return mary
| otherwise = case A.index ary i of
(L kx y) | k == kx -> A.update ary i (L k (f k v y))
| otherwise -> go k v ary (i+1) n
{-# INLINABLE updateOrSnocWithKey #-}
updateOrConcatWith :: Eq k => (v -> v -> v) -> A.Array (Leaf k v) -> A.Array (Leaf k v) -> A.Array (Leaf k v)
updateOrConcatWith f = updateOrConcatWithKey (const f)
{-# INLINABLE updateOrConcatWith #-}
updateOrConcatWithKey :: Eq k => (k -> v -> v -> v) -> A.Array (Leaf k v) -> A.Array (Leaf k v) -> A.Array (Leaf k v)
updateOrConcatWithKey f ary1 ary2 = A.run $ do
let indices = A.map' (\(L k _) -> indexOf k ary1) ary2
let nOnly2 = A.foldl' (\n -> maybe (n+1) (const n)) 0 indices
let n1 = A.length ary1
let n2 = A.length ary2
mary <- A.new_ (n1 + nOnly2)
A.copy ary1 0 mary 0 n1
let go !iEnd !i2
| i2 >= n2 = return ()
| otherwise = case A.index indices i2 of
Just i1 -> do
L k v1 <- A.indexM ary1 i1
L _ v2 <- A.indexM ary2 i2
A.write mary i1 (L k (f k v1 v2))
go iEnd (i2+1)
Nothing -> do
A.write mary iEnd =<< A.indexM ary2 i2
go (iEnd+1) (i2+1)
go n1 0
return mary
{-# INLINABLE updateOrConcatWithKey #-}
update16 :: A.Array e -> Int -> e -> A.Array e
update16 ary idx b = runST (update16M ary idx b)
{-# INLINE update16 #-}
update16M :: A.Array e -> Int -> e -> ST s (A.Array e)
update16M ary idx b = do
mary <- clone16 ary
A.write mary idx b
A.unsafeFreeze mary
{-# INLINE update16M #-}
update16With' :: A.Array e -> Int -> (e -> e) -> A.Array e
update16With' ary idx f
| (# x #) <- A.index# ary idx
= update16 ary idx $! f x
{-# INLINE update16With' #-}
clone16 :: A.Array e -> ST s (A.MArray s e)
clone16 ary =
A.thaw ary 0 16
bitsPerSubkey :: Int
bitsPerSubkey = 4
maxChildren :: Int
maxChildren = fromIntegral $ 1 `unsafeShiftL` bitsPerSubkey
subkeyMask :: Bitmap
subkeyMask = 1 `unsafeShiftL` bitsPerSubkey - 1
sparseIndex :: Bitmap -> Bitmap -> Int
sparseIndex b m = popCount (b .&. (m - 1))
mask :: Word -> Shift -> Bitmap
mask w s = 1 `unsafeShiftL` index w s
{-# INLINE mask #-}
index :: Hash -> Shift -> Int
index w s = fromIntegral $ (unsafeShiftR w s) .&. subkeyMask
{-# INLINE index #-}
fullNodeMask :: Bitmap
fullNodeMask = complement (complement 0 `unsafeShiftL` maxChildren)
{-# INLINE fullNodeMask #-}
ptrEq :: a -> a -> Bool
ptrEq x y = isTrue# (reallyUnsafePtrEquality# x y ==# 1#)
{-# INLINE ptrEq #-}
instance (Eq k, Hashable k) => Exts.IsList (HashMap k v) where
type Item (HashMap k v) = (k, v)
fromList = fromList
toList = toList