{-# LANGUAGE BangPatterns, CPP, DeriveDataTypeable, MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PatternGuards #-}
#if __GLASGOW_HASKELL__ >= 708
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE TypeFamilies #-}
#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
, union
, unionWith
, unionWithKey
, unions
, map
, mapWithKey
, traverseWithKey
, difference
, 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
) 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)
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.PopCount (popCount)
import Data.HashMap.Unsafe (runST)
import Data.HashMap.UnsafeShift (unsafeShiftL, unsafeShiftR)
import Data.Typeable (Typeable)
#if __GLASGOW_HASKELL__ >= 707
import GHC.Exts (isTrue#)
#endif
#if __GLASGOW_HASKELL__ >= 708
import qualified GHC.Exts as Exts
#endif
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)
#if __GLASGOW_HASKELL__ >= 708
type role HashMap nominal representational
#endif
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
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)
instance (Eq k, Eq v) => Eq (HashMap k v) where
(==) = equal
equal :: (Eq k, Eq v) => HashMap k v -> HashMap k v -> Bool
equal t1 t2 = go (toList' t1 []) (toList' t2 [])
where
go (Leaf k1 l1 : tl1) (Leaf k2 l2 : tl2)
| k1 == k2 && l1 == l2
= go tl1 tl2
go (Collision k1 ary1 : tl1) (Collision k2 ary2 : tl2)
| k1 == k2 && A.length ary1 == A.length ary2 &&
L.null (A.toList ary1 L.\\ A.toList ary2)
= go tl1 tl2
go [] [] = True
go _ _ = False
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
arrayHashesSorted :: A.Array (Leaf k v) -> [Int]
arrayHashesSorted = L.sort . L.map leafValueHash . A.toList
leafValueHash :: Leaf k v -> Int
leafValueHash (L _ v) = H.hash v
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
lookup k0 m0 = go h0 k0 0 m0
where
h0 = hash k0
go !_ !_ !_ Empty = Nothing
go h k _ (Leaf hx (L kx x))
| h == hx && k == kx = Just x
| otherwise = Nothing
go h k s (BitmapIndexed b v)
| b .&. m == 0 = Nothing
| 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 = lookupInArray k v
| otherwise = Nothing
{-# INLINABLE lookup #-}
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 k0 v0 m0 = go h0 k0 v0 0 m0
where
h0 = hash k0
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 #-}
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 k0 v0 m0 = go h0 k0 v0 0 m0
where
h0 = hash k0
go !h !k x !_ Empty = Leaf h (L k x)
go h k x s (Leaf hy l@(L ky y))
| hy == h = if ky == k
then Leaf h (L k (f x y))
else 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
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 :: 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 k0 m0 = go h0 k0 0 m0
where
h0 = hash k0
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 #-}
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 (L 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 (lookup k m) of
Nothing -> delete k m
Just v -> insert k v m
{-# INLINABLE alter #-}
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
A.write mary i $! f (A.index ary1 i1) (A.index ary2 i2)
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 #-}
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 #-}
trim :: A.MArray s a -> Int -> ST s (A.Array a)
trim mary n = do
mary2 <- A.new_ n
A.copyM mary 0 mary2 0 n
A.unsafeFreeze mary2
{-# INLINE trim #-}
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 <$> trim mary 1
_ -> do
ary2 <- 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 <- 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 #-}
lookupInArray :: Eq k => k -> A.Array (Leaf k v) -> Maybe v
lookupInArray 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 v)
| k == kx -> Just v
| otherwise -> go k ary (i+1) n
{-# INLINABLE lookupInArray #-}
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 -> A.update ary i (L k (f 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 = update16 ary idx $! f (A.index ary idx)
{-# INLINE update16With' #-}
clone16 :: A.Array e -> ST s (A.MArray s e)
clone16 ary =
#if __GLASGOW_HASKELL__ >= 702
A.thaw ary 0 16
#else
do mary <- A.new_ 16
A.indexM ary 0 >>= A.write mary 0
A.indexM ary 1 >>= A.write mary 1
A.indexM ary 2 >>= A.write mary 2
A.indexM ary 3 >>= A.write mary 3
A.indexM ary 4 >>= A.write mary 4
A.indexM ary 5 >>= A.write mary 5
A.indexM ary 6 >>= A.write mary 6
A.indexM ary 7 >>= A.write mary 7
A.indexM ary 8 >>= A.write mary 8
A.indexM ary 9 >>= A.write mary 9
A.indexM ary 10 >>= A.write mary 10
A.indexM ary 11 >>= A.write mary 11
A.indexM ary 12 >>= A.write mary 12
A.indexM ary 13 >>= A.write mary 13
A.indexM ary 14 >>= A.write mary 14
A.indexM ary 15 >>= A.write mary 15
return mary
#endif
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
#if __GLASGOW_HASKELL__ < 707
ptrEq x y = reallyUnsafePtrEquality# x y ==# 1#
#else
ptrEq x y = isTrue# (reallyUnsafePtrEquality# x y ==# 1#)
#endif
{-# INLINE ptrEq #-}
#if __GLASGOW_HASKELL__ >= 708
instance (Eq k, Hashable k) => Exts.IsList (HashMap k v) where
type Item (HashMap k v) = (k, v)
fromList = fromList
toList = toList
#endif