{-# LANGUAGE BangPatterns, CPP, PatternGuards, MagicHash, UnboxedTuples #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE Trustworthy #-}
{-# OPTIONS_HADDOCK not-home #-}
module Data.HashMap.Internal.Strict
(
HashMap
, empty
, singleton
, HM.null
, size
, HM.member
, HM.lookup
, (HM.!?)
, HM.findWithDefault
, lookupDefault
, (!)
, insert
, insertWith
, delete
, adjust
, update
, alter
, alterF
, isSubmapOf
, isSubmapOfBy
, union
, unionWith
, unionWithKey
, unions
, map
, mapWithKey
, traverseWithKey
, difference
, differenceWith
, intersection
, intersectionWith
, intersectionWithKey
, foldMapWithKey
, foldr'
, foldl'
, foldrWithKey'
, foldlWithKey'
, HM.foldr
, HM.foldl
, foldrWithKey
, foldlWithKey
, HM.filter
, filterWithKey
, mapMaybe
, mapMaybeWithKey
, keys
, elems
, toList
, fromList
, fromListWith
, fromListWithKey
) where
import Data.Bits ((.&.), (.|.))
#if !MIN_VERSION_base(4,8,0)
import Control.Applicative (Applicative (..), (<$>))
#endif
import qualified Data.List as L
import Data.Hashable (Hashable)
import Prelude hiding (map, lookup)
import qualified Data.HashMap.Internal.Array as A
import qualified Data.HashMap.Internal as HM
import Data.HashMap.Internal hiding (
alter, alterF, adjust, fromList, fromListWith, fromListWithKey,
insert, insertWith,
differenceWith, intersectionWith, intersectionWithKey, map, mapWithKey,
mapMaybe, mapMaybeWithKey, singleton, update, unionWith, unionWithKey,
traverseWithKey)
import Data.HashMap.Internal.Unsafe (runST)
#if MIN_VERSION_base(4,8,0)
import Data.Functor.Identity
#endif
import Control.Applicative (Const (..))
import Data.Coerce
singleton :: (Hashable k) => k -> v -> HashMap k v
singleton k !v = HM.singleton k v
insert :: (Eq k, Hashable k) => k -> v -> HashMap k v -> HashMap k v
insert k !v = HM.insert k v
{-# INLINABLE insert #-}
insertWith :: (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v
-> HashMap k v
insertWith f k0 v0 m0 = go h0 k0 v0 0 m0
where
h0 = hash k0
go !h !k x !_ Empty = leaf h k x
go h k x s t@(Leaf hy l@(L ky y))
| hy == h = if ky == k
then leaf h k (f x y)
else x `seq` (collision h l (L k x))
| otherwise = x `seq` runST (two s h k x hy t)
go h k x s (BitmapIndexed b ary)
| b .&. m == 0 =
let ary' = A.insert ary i $! leaf h k x
in bitmapIndexedOrFull (b .|. m) ary'
| otherwise =
let st = A.index ary i
st' = go h k x (s+bitsPerSubkey) st
ary' = A.update ary i $! st'
in BitmapIndexed b ary'
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
ary' = update16 ary i $! st'
in Full ary'
where i = index h s
go h k x s t@(Collision hy v)
| h == hy = Collision h (updateOrSnocWith f k x v)
| otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
{-# INLINABLE insertWith #-}
unsafeInsertWith :: (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v
-> HashMap k v
unsafeInsertWith f k0 v0 m0 = unsafeInsertWithKey (const f) k0 v0 m0
{-# INLINABLE unsafeInsertWith #-}
unsafeInsertWithKey :: (Eq k, Hashable k) => (k -> v -> v -> v) -> k -> v -> HashMap k v
-> HashMap k v
unsafeInsertWithKey f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
where
h0 = hash k0
go !h !k x !_ Empty = return $! leaf h k x
go h k x s t@(Leaf hy l@(L ky y))
| hy == h = if ky == k
then return $! leaf h k (f k x y)
else do
let l' = x `seq` (L k x)
return $! collision h l l'
| otherwise = x `seq` two s h k x hy t
go h k x s t@(BitmapIndexed b ary)
| b .&. m == 0 = do
ary' <- A.insertM ary i $! leaf h 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 (updateOrSnocWithKey f k x v)
| otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
{-# INLINABLE unsafeInsertWithKey #-}
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 = leaf h k (f 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 BitmapIndexed b ary'
where m = mask h s
i = sparseIndex b m
go h k s (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 Full ary'
go h k _ t@(Collision hy v)
| h == hy = Collision h (updateWith f k 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 (HM.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 -> maybe m (const (delete' h k m)) mv
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#"
impossibleAdjust :: a
impossibleAdjust = error "Data.HashMap.alterF internal error: impossible adjust"
{-# RULES
-- See detailed notes on alterF rules in Data.HashMap.Internal.
"alterFWeird" forall f. alterF f =
alterFWeird (f Nothing) (f (Just test_bottom)) f
"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})
"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 #)))
-- This rule is written a bit differently than the one for lazy
-- maps because the adjust here is strict. We could write it the
-- same general way anyway, but this seems simpler.
"alterFadjust" forall (f :: Maybe a -> Identity (Maybe a)) x.
alterFWeird (coerce Nothing) (coerce (Just x)) f =
coerce (adjust (\a -> case runIdentity (f (Just a)) of
Just a' -> a'
Nothing -> impossibleAdjust))
"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 -> v' `seq`
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
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 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 #-}
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 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) -> let !v' = f k v in L k v') ary
{-# INLINE mapWithKey #-}
map :: (v1 -> v2) -> HashMap k v1 -> HashMap k v2
map f = mapWithKey (const f)
{-# INLINE map #-}
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 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 #-}
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 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 #-}
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 HM.lookup k b of
Nothing -> insert k v m
Just w -> maybe m (\y -> insert k y m) (f v w)
{-# INLINABLE differenceWith #-}
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 HM.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 HM.lookup k b of
Just w -> insert k (f k v w) m
_ -> m
{-# INLINABLE intersectionWithKey #-}
fromList :: (Eq k, Hashable k) => [(k, v)] -> HashMap k v
fromList = L.foldl' (\ m (k, !v) -> HM.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 #-}
fromListWithKey :: (Eq k, Hashable k) => (k -> v -> v -> v) -> [(k, v)] -> HashMap k v
fromListWithKey f = L.foldl' (\ m (k, v) -> unsafeInsertWithKey f k v m) empty
{-# INLINE fromListWithKey #-}
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 -> let !v' = f y in A.update ary i (L k v')
| 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
let !l = v `seq` (L k v)
A.write mary n l
return mary
| otherwise = case A.index ary i of
(L kx y) | k == kx -> let !v' = f k v y in A.update ary i (L k v')
| otherwise -> go k v ary (i+1) n
{-# INLINABLE updateOrSnocWithKey #-}
leaf :: Hash -> k -> v -> HashMap k v
leaf h k = \ !v -> Leaf h (L k v)
{-# INLINE leaf #-}