1{-# LANGUAGE BangPatterns, CPP, DeriveDataTypeable, MagicHash #-} 2{-# LANGUAGE ScopedTypeVariables #-} 3{-# LANGUAGE PatternGuards #-} 4{-# LANGUAGE RoleAnnotations #-} 5{-# LANGUAGE TypeFamilies #-} 6{-# LANGUAGE UnboxedTuples #-} 7{-# LANGUAGE LambdaCase #-} 8#if __GLASGOW_HASKELL__ >= 802 9{-# LANGUAGE TypeInType #-} 10{-# LANGUAGE UnboxedSums #-} 11#endif 12{-# OPTIONS_GHC -fno-full-laziness -funbox-strict-fields #-} 13{-# OPTIONS_HADDOCK not-home #-} 14 15-- | = WARNING 16-- 17-- This module is considered __internal__. 18-- 19-- The Package Versioning Policy __does not apply__. 20-- 21-- The contents of this module may change __in any way whatsoever__ 22-- and __without any warning__ between minor versions of this package. 23-- 24-- Authors importing this module are expected to track development 25-- closely. 26 27module Data.HashMap.Internal 28 ( 29 HashMap(..) 30 , Leaf(..) 31 32 -- * Construction 33 , empty 34 , singleton 35 36 -- * Basic interface 37 , null 38 , size 39 , member 40 , lookup 41 , (!?) 42 , findWithDefault 43 , lookupDefault 44 , (!) 45 , insert 46 , insertWith 47 , unsafeInsert 48 , delete 49 , adjust 50 , update 51 , alter 52 , alterF 53 , isSubmapOf 54 , isSubmapOfBy 55 56 -- * Combine 57 -- ** Union 58 , union 59 , unionWith 60 , unionWithKey 61 , unions 62 63 -- ** Compose 64 , compose 65 66 -- * Transformations 67 , map 68 , mapWithKey 69 , traverseWithKey 70 71 -- * Difference and intersection 72 , difference 73 , differenceWith 74 , intersection 75 , intersectionWith 76 , intersectionWithKey 77 78 -- * Folds 79 , foldr' 80 , foldl' 81 , foldrWithKey' 82 , foldlWithKey' 83 , foldr 84 , foldl 85 , foldrWithKey 86 , foldlWithKey 87 , foldMapWithKey 88 89 -- * Filter 90 , mapMaybe 91 , mapMaybeWithKey 92 , filter 93 , filterWithKey 94 95 -- * Conversions 96 , keys 97 , elems 98 99 -- ** Lists 100 , toList 101 , fromList 102 , fromListWith 103 , fromListWithKey 104 105 -- Internals used by the strict version 106 , Hash 107 , Bitmap 108 , bitmapIndexedOrFull 109 , collision 110 , hash 111 , mask 112 , index 113 , bitsPerSubkey 114 , fullNodeMask 115 , sparseIndex 116 , two 117 , unionArrayBy 118 , update16 119 , update16M 120 , update16With' 121 , updateOrConcatWith 122 , updateOrConcatWithKey 123 , filterMapAux 124 , equalKeys 125 , equalKeys1 126 , lookupRecordCollision 127 , LookupRes(..) 128 , insert' 129 , delete' 130 , lookup' 131 , insertNewKey 132 , insertKeyExists 133 , deleteKeyExists 134 , insertModifying 135 , ptrEq 136 , adjust# 137 ) where 138 139#if __GLASGOW_HASKELL__ < 710 140import Control.Applicative ((<$>), Applicative(pure)) 141import Data.Monoid (Monoid(mempty, mappend)) 142import Data.Traversable (Traversable(..)) 143import Data.Word (Word) 144#endif 145#if __GLASGOW_HASKELL__ >= 711 146import Data.Semigroup (Semigroup((<>))) 147#endif 148import Control.DeepSeq (NFData(rnf)) 149import Control.Monad.ST (ST) 150import Data.Bits ((.&.), (.|.), complement, popCount, unsafeShiftL, unsafeShiftR) 151import Data.Data hiding (Typeable) 152import qualified Data.Foldable as Foldable 153#if MIN_VERSION_base(4,10,0) 154import Data.Bifoldable 155#endif 156import qualified Data.List as L 157import GHC.Exts ((==#), build, reallyUnsafePtrEquality#, inline) 158import Prelude hiding (filter, foldl, foldr, lookup, map, null, pred) 159import Text.Read hiding (step) 160 161import qualified Data.HashMap.Internal.Array as A 162import qualified Data.Hashable as H 163import Data.Hashable (Hashable) 164import Data.HashMap.Internal.Unsafe (runST) 165import Data.HashMap.Internal.List (isPermutationBy, unorderedCompare) 166import Data.Typeable (Typeable) 167 168import GHC.Exts (isTrue#) 169import qualified GHC.Exts as Exts 170 171#if MIN_VERSION_base(4,9,0) 172import Data.Functor.Classes 173import GHC.Stack 174#endif 175 176#if MIN_VERSION_hashable(1,2,5) 177import qualified Data.Hashable.Lifted as H 178#endif 179 180#if __GLASGOW_HASKELL__ >= 802 181import GHC.Exts (TYPE, Int (..), Int#) 182#endif 183 184#if MIN_VERSION_base(4,8,0) 185import Data.Functor.Identity (Identity (..)) 186#endif 187import Control.Applicative (Const (..)) 188import Data.Coerce (coerce) 189 190-- | A set of values. A set cannot contain duplicate values. 191------------------------------------------------------------------------ 192 193-- | Convenience function. Compute a hash value for the given value. 194hash :: H.Hashable a => a -> Hash 195hash = fromIntegral . H.hash 196 197data Leaf k v = L !k v 198 deriving (Eq) 199 200instance (NFData k, NFData v) => NFData (Leaf k v) where 201 rnf (L k v) = rnf k `seq` rnf v 202 203-- Invariant: The length of the 1st argument to 'Full' is 204-- 2^bitsPerSubkey 205 206-- | A map from keys to values. A map cannot contain duplicate keys; 207-- each key can map to at most one value. 208data HashMap k v 209 = Empty 210 | BitmapIndexed !Bitmap !(A.Array (HashMap k v)) 211 | Leaf !Hash !(Leaf k v) 212 | Full !(A.Array (HashMap k v)) 213 | Collision !Hash !(A.Array (Leaf k v)) 214 deriving (Typeable) 215 216type role HashMap nominal representational 217 218instance (NFData k, NFData v) => NFData (HashMap k v) where 219 rnf Empty = () 220 rnf (BitmapIndexed _ ary) = rnf ary 221 rnf (Leaf _ l) = rnf l 222 rnf (Full ary) = rnf ary 223 rnf (Collision _ ary) = rnf ary 224 225instance Functor (HashMap k) where 226 fmap = map 227 228instance Foldable.Foldable (HashMap k) where 229 foldMap f = foldMapWithKey (\ _k v -> f v) 230 {-# INLINE foldMap #-} 231 foldr = foldr 232 {-# INLINE foldr #-} 233 foldl = foldl 234 {-# INLINE foldl #-} 235 foldr' = foldr' 236 {-# INLINE foldr' #-} 237 foldl' = foldl' 238 {-# INLINE foldl' #-} 239#if MIN_VERSION_base(4,8,0) 240 null = null 241 {-# INLINE null #-} 242 length = size 243 {-# INLINE length #-} 244#endif 245 246#if MIN_VERSION_base(4,10,0) 247-- | @since 0.2.11 248instance Bifoldable HashMap where 249 bifoldMap f g = foldMapWithKey (\ k v -> f k `mappend` g v) 250 {-# INLINE bifoldMap #-} 251 bifoldr f g = foldrWithKey (\ k v acc -> k `f` (v `g` acc)) 252 {-# INLINE bifoldr #-} 253 bifoldl f g = foldlWithKey (\ acc k v -> (acc `f` k) `g` v) 254 {-# INLINE bifoldl #-} 255#endif 256 257#if __GLASGOW_HASKELL__ >= 711 258-- | '<>' = 'union' 259-- 260-- If a key occurs in both maps, the mapping from the first will be the mapping in the result. 261-- 262-- ==== __Examples__ 263-- 264-- >>> fromList [(1,'a'),(2,'b')] <> fromList [(2,'c'),(3,'d')] 265-- fromList [(1,'a'),(2,'b'),(3,'d')] 266instance (Eq k, Hashable k) => Semigroup (HashMap k v) where 267 (<>) = union 268 {-# INLINE (<>) #-} 269#endif 270 271-- | 'mempty' = 'empty' 272-- 273-- 'mappend' = 'union' 274-- 275-- If a key occurs in both maps, the mapping from the first will be the mapping in the result. 276-- 277-- ==== __Examples__ 278-- 279-- >>> mappend (fromList [(1,'a'),(2,'b')]) (fromList [(2,'c'),(3,'d')]) 280-- fromList [(1,'a'),(2,'b'),(3,'d')] 281instance (Eq k, Hashable k) => Monoid (HashMap k v) where 282 mempty = empty 283 {-# INLINE mempty #-} 284#if __GLASGOW_HASKELL__ >= 711 285 mappend = (<>) 286#else 287 mappend = union 288#endif 289 {-# INLINE mappend #-} 290 291instance (Data k, Data v, Eq k, Hashable k) => Data (HashMap k v) where 292 gfoldl f z m = z fromList `f` toList m 293 toConstr _ = fromListConstr 294 gunfold k z c = case constrIndex c of 295 1 -> k (z fromList) 296 _ -> error "gunfold" 297 dataTypeOf _ = hashMapDataType 298 dataCast2 f = gcast2 f 299 300fromListConstr :: Constr 301fromListConstr = mkConstr hashMapDataType "fromList" [] Prefix 302 303hashMapDataType :: DataType 304hashMapDataType = mkDataType "Data.HashMap.Internal.HashMap" [fromListConstr] 305 306type Hash = Word 307type Bitmap = Word 308type Shift = Int 309 310#if MIN_VERSION_base(4,9,0) 311instance Show2 HashMap where 312 liftShowsPrec2 spk slk spv slv d m = 313 showsUnaryWith (liftShowsPrec sp sl) "fromList" d (toList m) 314 where 315 sp = liftShowsPrec2 spk slk spv slv 316 sl = liftShowList2 spk slk spv slv 317 318instance Show k => Show1 (HashMap k) where 319 liftShowsPrec = liftShowsPrec2 showsPrec showList 320 321instance (Eq k, Hashable k, Read k) => Read1 (HashMap k) where 322 liftReadsPrec rp rl = readsData $ 323 readsUnaryWith (liftReadsPrec rp' rl') "fromList" fromList 324 where 325 rp' = liftReadsPrec rp rl 326 rl' = liftReadList rp rl 327#endif 328 329instance (Eq k, Hashable k, Read k, Read e) => Read (HashMap k e) where 330 readPrec = parens $ prec 10 $ do 331 Ident "fromList" <- lexP 332 xs <- readPrec 333 return (fromList xs) 334 335 readListPrec = readListPrecDefault 336 337instance (Show k, Show v) => Show (HashMap k v) where 338 showsPrec d m = showParen (d > 10) $ 339 showString "fromList " . shows (toList m) 340 341instance Traversable (HashMap k) where 342 traverse f = traverseWithKey (const f) 343 {-# INLINABLE traverse #-} 344 345#if MIN_VERSION_base(4,9,0) 346instance Eq2 HashMap where 347 liftEq2 = equal2 348 349instance Eq k => Eq1 (HashMap k) where 350 liftEq = equal1 351#endif 352 353-- | Note that, in the presence of hash collisions, equal @HashMap@s may 354-- behave differently, i.e. substitutivity may be violated: 355-- 356-- >>> data D = A | B deriving (Eq, Show) 357-- >>> instance Hashable D where hashWithSalt salt _d = salt 358-- 359-- >>> x = fromList [(A,1), (B,2)] 360-- >>> y = fromList [(B,2), (A,1)] 361-- 362-- >>> x == y 363-- True 364-- >>> toList x 365-- [(A,1),(B,2)] 366-- >>> toList y 367-- [(B,2),(A,1)] 368-- 369-- In general, the lack of substitutivity can be observed with any function 370-- that depends on the key ordering, such as folds and traversals. 371instance (Eq k, Eq v) => Eq (HashMap k v) where 372 (==) = equal1 (==) 373 374-- We rely on there being no Empty constructors in the tree! 375-- This ensures that two equal HashMaps will have the same 376-- shape, modulo the order of entries in Collisions. 377equal1 :: Eq k 378 => (v -> v' -> Bool) 379 -> HashMap k v -> HashMap k v' -> Bool 380equal1 eq = go 381 where 382 go Empty Empty = True 383 go (BitmapIndexed bm1 ary1) (BitmapIndexed bm2 ary2) 384 = bm1 == bm2 && A.sameArray1 go ary1 ary2 385 go (Leaf h1 l1) (Leaf h2 l2) = h1 == h2 && leafEq l1 l2 386 go (Full ary1) (Full ary2) = A.sameArray1 go ary1 ary2 387 go (Collision h1 ary1) (Collision h2 ary2) 388 = h1 == h2 && isPermutationBy leafEq (A.toList ary1) (A.toList ary2) 389 go _ _ = False 390 391 leafEq (L k1 v1) (L k2 v2) = k1 == k2 && eq v1 v2 392 393equal2 :: (k -> k' -> Bool) -> (v -> v' -> Bool) 394 -> HashMap k v -> HashMap k' v' -> Bool 395equal2 eqk eqv t1 t2 = go (toList' t1 []) (toList' t2 []) 396 where 397 -- If the two trees are the same, then their lists of 'Leaf's and 398 -- 'Collision's read from left to right should be the same (modulo the 399 -- order of elements in 'Collision'). 400 401 go (Leaf k1 l1 : tl1) (Leaf k2 l2 : tl2) 402 | k1 == k2 && 403 leafEq l1 l2 404 = go tl1 tl2 405 go (Collision k1 ary1 : tl1) (Collision k2 ary2 : tl2) 406 | k1 == k2 && 407 A.length ary1 == A.length ary2 && 408 isPermutationBy leafEq (A.toList ary1) (A.toList ary2) 409 = go tl1 tl2 410 go [] [] = True 411 go _ _ = False 412 413 leafEq (L k v) (L k' v') = eqk k k' && eqv v v' 414 415#if MIN_VERSION_base(4,9,0) 416instance Ord2 HashMap where 417 liftCompare2 = cmp 418 419instance Ord k => Ord1 (HashMap k) where 420 liftCompare = cmp compare 421#endif 422 423-- | The ordering is total and consistent with the `Eq` instance. However, 424-- nothing else about the ordering is specified, and it may change from 425-- version to version of either this package or of hashable. 426instance (Ord k, Ord v) => Ord (HashMap k v) where 427 compare = cmp compare compare 428 429cmp :: (k -> k' -> Ordering) -> (v -> v' -> Ordering) 430 -> HashMap k v -> HashMap k' v' -> Ordering 431cmp cmpk cmpv t1 t2 = go (toList' t1 []) (toList' t2 []) 432 where 433 go (Leaf k1 l1 : tl1) (Leaf k2 l2 : tl2) 434 = compare k1 k2 `mappend` 435 leafCompare l1 l2 `mappend` 436 go tl1 tl2 437 go (Collision k1 ary1 : tl1) (Collision k2 ary2 : tl2) 438 = compare k1 k2 `mappend` 439 compare (A.length ary1) (A.length ary2) `mappend` 440 unorderedCompare leafCompare (A.toList ary1) (A.toList ary2) `mappend` 441 go tl1 tl2 442 go (Leaf _ _ : _) (Collision _ _ : _) = LT 443 go (Collision _ _ : _) (Leaf _ _ : _) = GT 444 go [] [] = EQ 445 go [] _ = LT 446 go _ [] = GT 447 go _ _ = error "cmp: Should never happen, toList' includes non Leaf / Collision" 448 449 leafCompare (L k v) (L k' v') = cmpk k k' `mappend` cmpv v v' 450 451-- Same as 'equal' but doesn't compare the values. 452equalKeys1 :: (k -> k' -> Bool) -> HashMap k v -> HashMap k' v' -> Bool 453equalKeys1 eq t1 t2 = go (toList' t1 []) (toList' t2 []) 454 where 455 go (Leaf k1 l1 : tl1) (Leaf k2 l2 : tl2) 456 | k1 == k2 && leafEq l1 l2 457 = go tl1 tl2 458 go (Collision k1 ary1 : tl1) (Collision k2 ary2 : tl2) 459 | k1 == k2 && A.length ary1 == A.length ary2 && 460 isPermutationBy leafEq (A.toList ary1) (A.toList ary2) 461 = go tl1 tl2 462 go [] [] = True 463 go _ _ = False 464 465 leafEq (L k _) (L k' _) = eq k k' 466 467-- Same as 'equal1' but doesn't compare the values. 468equalKeys :: Eq k => HashMap k v -> HashMap k v' -> Bool 469equalKeys = go 470 where 471 go :: Eq k => HashMap k v -> HashMap k v' -> Bool 472 go Empty Empty = True 473 go (BitmapIndexed bm1 ary1) (BitmapIndexed bm2 ary2) 474 = bm1 == bm2 && A.sameArray1 go ary1 ary2 475 go (Leaf h1 l1) (Leaf h2 l2) = h1 == h2 && leafEq l1 l2 476 go (Full ary1) (Full ary2) = A.sameArray1 go ary1 ary2 477 go (Collision h1 ary1) (Collision h2 ary2) 478 = h1 == h2 && isPermutationBy leafEq (A.toList ary1) (A.toList ary2) 479 go _ _ = False 480 481 leafEq (L k1 _) (L k2 _) = k1 == k2 482 483#if MIN_VERSION_hashable(1,2,5) 484instance H.Hashable2 HashMap where 485 liftHashWithSalt2 hk hv salt hm = go salt (toList' hm []) 486 where 487 -- go :: Int -> [HashMap k v] -> Int 488 go s [] = s 489 go s (Leaf _ l : tl) 490 = s `hashLeafWithSalt` l `go` tl 491 -- For collisions we hashmix hash value 492 -- and then array of values' hashes sorted 493 go s (Collision h a : tl) 494 = (s `H.hashWithSalt` h) `hashCollisionWithSalt` a `go` tl 495 go s (_ : tl) = s `go` tl 496 497 -- hashLeafWithSalt :: Int -> Leaf k v -> Int 498 hashLeafWithSalt s (L k v) = (s `hk` k) `hv` v 499 500 -- hashCollisionWithSalt :: Int -> A.Array (Leaf k v) -> Int 501 hashCollisionWithSalt s 502 = L.foldl' H.hashWithSalt s . arrayHashesSorted s 503 504 -- arrayHashesSorted :: Int -> A.Array (Leaf k v) -> [Int] 505 arrayHashesSorted s = L.sort . L.map (hashLeafWithSalt s) . A.toList 506 507instance (Hashable k) => H.Hashable1 (HashMap k) where 508 liftHashWithSalt = H.liftHashWithSalt2 H.hashWithSalt 509#endif 510 511instance (Hashable k, Hashable v) => Hashable (HashMap k v) where 512 hashWithSalt salt hm = go salt hm 513 where 514 go :: Int -> HashMap k v -> Int 515 go s Empty = s 516 go s (BitmapIndexed _ a) = A.foldl' go s a 517 go s (Leaf h (L _ v)) 518 = s `H.hashWithSalt` h `H.hashWithSalt` v 519 -- For collisions we hashmix hash value 520 -- and then array of values' hashes sorted 521 go s (Full a) = A.foldl' go s a 522 go s (Collision h a) 523 = (s `H.hashWithSalt` h) `hashCollisionWithSalt` a 524 525 hashLeafWithSalt :: Int -> Leaf k v -> Int 526 hashLeafWithSalt s (L k v) = s `H.hashWithSalt` k `H.hashWithSalt` v 527 528 hashCollisionWithSalt :: Int -> A.Array (Leaf k v) -> Int 529 hashCollisionWithSalt s 530 = L.foldl' H.hashWithSalt s . arrayHashesSorted s 531 532 arrayHashesSorted :: Int -> A.Array (Leaf k v) -> [Int] 533 arrayHashesSorted s = L.sort . L.map (hashLeafWithSalt s) . A.toList 534 535 -- Helper to get 'Leaf's and 'Collision's as a list. 536toList' :: HashMap k v -> [HashMap k v] -> [HashMap k v] 537toList' (BitmapIndexed _ ary) a = A.foldr toList' a ary 538toList' (Full ary) a = A.foldr toList' a ary 539toList' l@(Leaf _ _) a = l : a 540toList' c@(Collision _ _) a = c : a 541toList' Empty a = a 542 543-- Helper function to detect 'Leaf's and 'Collision's. 544isLeafOrCollision :: HashMap k v -> Bool 545isLeafOrCollision (Leaf _ _) = True 546isLeafOrCollision (Collision _ _) = True 547isLeafOrCollision _ = False 548 549------------------------------------------------------------------------ 550-- * Construction 551 552-- | /O(1)/ Construct an empty map. 553empty :: HashMap k v 554empty = Empty 555 556-- | /O(1)/ Construct a map with a single element. 557singleton :: (Hashable k) => k -> v -> HashMap k v 558singleton k v = Leaf (hash k) (L k v) 559 560------------------------------------------------------------------------ 561-- * Basic interface 562 563-- | /O(1)/ Return 'True' if this map is empty, 'False' otherwise. 564null :: HashMap k v -> Bool 565null Empty = True 566null _ = False 567 568-- | /O(n)/ Return the number of key-value mappings in this map. 569size :: HashMap k v -> Int 570size t = go t 0 571 where 572 go Empty !n = n 573 go (Leaf _ _) n = n + 1 574 go (BitmapIndexed _ ary) n = A.foldl' (flip go) n ary 575 go (Full ary) n = A.foldl' (flip go) n ary 576 go (Collision _ ary) n = n + A.length ary 577 578-- | /O(log n)/ Return 'True' if the specified key is present in the 579-- map, 'False' otherwise. 580member :: (Eq k, Hashable k) => k -> HashMap k a -> Bool 581member k m = case lookup k m of 582 Nothing -> False 583 Just _ -> True 584{-# INLINABLE member #-} 585 586-- | /O(log n)/ Return the value to which the specified key is mapped, 587-- or 'Nothing' if this map contains no mapping for the key. 588lookup :: (Eq k, Hashable k) => k -> HashMap k v -> Maybe v 589#if __GLASGOW_HASKELL__ >= 802 590-- GHC does not yet perform a worker-wrapper transformation on 591-- unboxed sums automatically. That seems likely to happen at some 592-- point (possibly as early as GHC 8.6) but for now we do it manually. 593lookup k m = case lookup# k m of 594 (# (# #) | #) -> Nothing 595 (# | a #) -> Just a 596{-# INLINE lookup #-} 597 598lookup# :: (Eq k, Hashable k) => k -> HashMap k v -> (# (# #) | v #) 599lookup# k m = lookupCont (\_ -> (# (# #) | #)) (\v _i -> (# | v #)) (hash k) k 0 m 600{-# INLINABLE lookup# #-} 601 602#else 603 604lookup k m = lookupCont (\_ -> Nothing) (\v _i -> Just v) (hash k) k 0 m 605{-# INLINABLE lookup #-} 606#endif 607 608-- | lookup' is a version of lookup that takes the hash separately. 609-- It is used to implement alterF. 610lookup' :: Eq k => Hash -> k -> HashMap k v -> Maybe v 611#if __GLASGOW_HASKELL__ >= 802 612-- GHC does not yet perform a worker-wrapper transformation on 613-- unboxed sums automatically. That seems likely to happen at some 614-- point (possibly as early as GHC 8.6) but for now we do it manually. 615-- lookup' would probably prefer to be implemented in terms of its own 616-- lookup'#, but it's not important enough and we don't want too much 617-- code. 618lookup' h k m = case lookupRecordCollision# h k m of 619 (# (# #) | #) -> Nothing 620 (# | (# a, _i #) #) -> Just a 621{-# INLINE lookup' #-} 622#else 623lookup' h k m = lookupCont (\_ -> Nothing) (\v _i -> Just v) h k 0 m 624{-# INLINABLE lookup' #-} 625#endif 626 627-- The result of a lookup, keeping track of if a hash collision occured. 628-- If a collision did not occur then it will have the Int value (-1). 629data LookupRes a = Absent | Present a !Int 630 631-- Internal helper for lookup. This version takes the precomputed hash so 632-- that functions that make multiple calls to lookup and related functions 633-- (insert, delete) only need to calculate the hash once. 634-- 635-- It is used by 'alterF' so that hash computation and key comparison only needs 636-- to be performed once. With this information you can use the more optimized 637-- versions of insert ('insertNewKey', 'insertKeyExists') and delete 638-- ('deleteKeyExists') 639-- 640-- Outcomes: 641-- Key not in map => Absent 642-- Key in map, no collision => Present v (-1) 643-- Key in map, collision => Present v position 644lookupRecordCollision :: Eq k => Hash -> k -> HashMap k v -> LookupRes v 645#if __GLASGOW_HASKELL__ >= 802 646lookupRecordCollision h k m = case lookupRecordCollision# h k m of 647 (# (# #) | #) -> Absent 648 (# | (# a, i #) #) -> Present a (I# i) -- GHC will eliminate the I# 649{-# INLINE lookupRecordCollision #-} 650 651-- Why do we produce an Int# instead of an Int? Unfortunately, GHC is not 652-- yet any good at unboxing things *inside* products, let alone sums. That 653-- may be changing in GHC 8.6 or so (there is some work in progress), but 654-- for now we use Int# explicitly here. We don't need to push the Int# 655-- into lookupCont because inlining takes care of that. 656lookupRecordCollision# :: Eq k => Hash -> k -> HashMap k v -> (# (# #) | (# v, Int# #) #) 657lookupRecordCollision# h k m = 658 lookupCont (\_ -> (# (# #) | #)) (\v (I# i) -> (# | (# v, i #) #)) h k 0 m 659-- INLINABLE to specialize to the Eq instance. 660{-# INLINABLE lookupRecordCollision# #-} 661 662#else /* GHC < 8.2 so there are no unboxed sums */ 663 664lookupRecordCollision h k m = lookupCont (\_ -> Absent) Present h k 0 m 665{-# INLINABLE lookupRecordCollision #-} 666#endif 667 668-- A two-continuation version of lookupRecordCollision. This lets us 669-- share source code between lookup and lookupRecordCollision without 670-- risking any performance degradation. 671-- 672-- The absent continuation has type @((# #) -> r)@ instead of just @r@ 673-- so we can be representation-polymorphic in the result type. Since 674-- this whole thing is always inlined, we don't have to worry about 675-- any extra CPS overhead. 676-- 677-- The @Int@ argument is the offset of the subkey in the hash. When looking up 678-- keys at the top-level of a hashmap, the offset should be 0. When looking up 679-- keys at level @n@ of a hashmap, the offset should be @n * bitsPerSubkey@. 680lookupCont :: 681#if __GLASGOW_HASKELL__ >= 802 682 forall rep (r :: TYPE rep) k v. 683#else 684 forall r k v. 685#endif 686 Eq k 687 => ((# #) -> r) -- Absent continuation 688 -> (v -> Int -> r) -- Present continuation 689 -> Hash -- The hash of the key 690 -> k 691 -> Int -- The offset of the subkey in the hash. 692 -> HashMap k v -> r 693lookupCont absent present !h0 !k0 !s0 !m0 = go h0 k0 s0 m0 694 where 695 go :: Eq k => Hash -> k -> Int -> HashMap k v -> r 696 go !_ !_ !_ Empty = absent (# #) 697 go h k _ (Leaf hx (L kx x)) 698 | h == hx && k == kx = present x (-1) 699 | otherwise = absent (# #) 700 go h k s (BitmapIndexed b v) 701 | b .&. m == 0 = absent (# #) 702 | otherwise = 703 go h k (s+bitsPerSubkey) (A.index v (sparseIndex b m)) 704 where m = mask h s 705 go h k s (Full v) = 706 go h k (s+bitsPerSubkey) (A.index v (index h s)) 707 go h k _ (Collision hx v) 708 | h == hx = lookupInArrayCont absent present k v 709 | otherwise = absent (# #) 710{-# INLINE lookupCont #-} 711 712-- | /O(log n)/ Return the value to which the specified key is mapped, 713-- or 'Nothing' if this map contains no mapping for the key. 714-- 715-- This is a flipped version of 'lookup'. 716-- 717-- @since 0.2.11 718(!?) :: (Eq k, Hashable k) => HashMap k v -> k -> Maybe v 719(!?) m k = lookup k m 720{-# INLINE (!?) #-} 721 722 723-- | /O(log n)/ Return the value to which the specified key is mapped, 724-- or the default value if this map contains no mapping for the key. 725-- 726-- @since 0.2.11 727findWithDefault :: (Eq k, Hashable k) 728 => v -- ^ Default value to return. 729 -> k -> HashMap k v -> v 730findWithDefault def k t = case lookup k t of 731 Just v -> v 732 _ -> def 733{-# INLINABLE findWithDefault #-} 734 735 736-- | /O(log n)/ Return the value to which the specified key is mapped, 737-- or the default value if this map contains no mapping for the key. 738-- 739-- DEPRECATED: lookupDefault is deprecated as of version 0.2.11, replaced 740-- by 'findWithDefault'. 741lookupDefault :: (Eq k, Hashable k) 742 => v -- ^ Default value to return. 743 -> k -> HashMap k v -> v 744lookupDefault def k t = findWithDefault def k t 745{-# INLINE lookupDefault #-} 746 747-- | /O(log n)/ Return the value to which the specified key is mapped. 748-- Calls 'error' if this map contains no mapping for the key. 749#if MIN_VERSION_base(4,9,0) 750(!) :: (Eq k, Hashable k, HasCallStack) => HashMap k v -> k -> v 751#else 752(!) :: (Eq k, Hashable k) => HashMap k v -> k -> v 753#endif 754(!) m k = case lookup k m of 755 Just v -> v 756 Nothing -> error "Data.HashMap.Internal.(!): key not found" 757{-# INLINABLE (!) #-} 758 759infixl 9 ! 760 761-- | Create a 'Collision' value with two 'Leaf' values. 762collision :: Hash -> Leaf k v -> Leaf k v -> HashMap k v 763collision h !e1 !e2 = 764 let v = A.run $ do mary <- A.new 2 e1 765 A.write mary 1 e2 766 return mary 767 in Collision h v 768{-# INLINE collision #-} 769 770-- | Create a 'BitmapIndexed' or 'Full' node. 771bitmapIndexedOrFull :: Bitmap -> A.Array (HashMap k v) -> HashMap k v 772bitmapIndexedOrFull b ary 773 | b == fullNodeMask = Full ary 774 | otherwise = BitmapIndexed b ary 775{-# INLINE bitmapIndexedOrFull #-} 776 777-- | /O(log n)/ Associate the specified value with the specified 778-- key in this map. If this map previously contained a mapping for 779-- the key, the old value is replaced. 780insert :: (Eq k, Hashable k) => k -> v -> HashMap k v -> HashMap k v 781insert k v m = insert' (hash k) k v m 782{-# INLINABLE insert #-} 783 784insert' :: Eq k => Hash -> k -> v -> HashMap k v -> HashMap k v 785insert' h0 k0 v0 m0 = go h0 k0 v0 0 m0 786 where 787 go !h !k x !_ Empty = Leaf h (L k x) 788 go h k x s t@(Leaf hy l@(L ky y)) 789 | hy == h = if ky == k 790 then if x `ptrEq` y 791 then t 792 else Leaf h (L k x) 793 else collision h l (L k x) 794 | otherwise = runST (two s h k x hy t) 795 go h k x s t@(BitmapIndexed b ary) 796 | b .&. m == 0 = 797 let !ary' = A.insert ary i $! Leaf h (L k x) 798 in bitmapIndexedOrFull (b .|. m) ary' 799 | otherwise = 800 let !st = A.index ary i 801 !st' = go h k x (s+bitsPerSubkey) st 802 in if st' `ptrEq` st 803 then t 804 else BitmapIndexed b (A.update ary i st') 805 where m = mask h s 806 i = sparseIndex b m 807 go h k x s t@(Full ary) = 808 let !st = A.index ary i 809 !st' = go h k x (s+bitsPerSubkey) st 810 in if st' `ptrEq` st 811 then t 812 else Full (update16 ary i st') 813 where i = index h s 814 go h k x s t@(Collision hy v) 815 | h == hy = Collision h (updateOrSnocWith (\a _ -> (# a #)) k x v) 816 | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) 817{-# INLINABLE insert' #-} 818 819-- Insert optimized for the case when we know the key is not in the map. 820-- 821-- It is only valid to call this when the key does not exist in the map. 822-- 823-- We can skip: 824-- - the key equality check on a Leaf 825-- - check for its existence in the array for a hash collision 826insertNewKey :: Hash -> k -> v -> HashMap k v -> HashMap k v 827insertNewKey !h0 !k0 x0 !m0 = go h0 k0 x0 0 m0 828 where 829 go !h !k x !_ Empty = Leaf h (L k x) 830 go h k x s t@(Leaf hy l) 831 | hy == h = collision h l (L k x) 832 | otherwise = runST (two s h k x hy t) 833 go h k x s (BitmapIndexed b ary) 834 | b .&. m == 0 = 835 let !ary' = A.insert ary i $! Leaf h (L k x) 836 in bitmapIndexedOrFull (b .|. m) ary' 837 | otherwise = 838 let !st = A.index ary i 839 !st' = go h k x (s+bitsPerSubkey) st 840 in BitmapIndexed b (A.update ary i st') 841 where m = mask h s 842 i = sparseIndex b m 843 go h k x s (Full ary) = 844 let !st = A.index ary i 845 !st' = go h k x (s+bitsPerSubkey) st 846 in Full (update16 ary i st') 847 where i = index h s 848 go h k x s t@(Collision hy v) 849 | h == hy = Collision h (snocNewLeaf (L k x) v) 850 | otherwise = 851 go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) 852 where 853 snocNewLeaf :: Leaf k v -> A.Array (Leaf k v) -> A.Array (Leaf k v) 854 snocNewLeaf leaf ary = A.run $ do 855 let n = A.length ary 856 mary <- A.new_ (n + 1) 857 A.copy ary 0 mary 0 n 858 A.write mary n leaf 859 return mary 860{-# NOINLINE insertNewKey #-} 861 862 863-- Insert optimized for the case when we know the key is in the map. 864-- 865-- It is only valid to call this when the key exists in the map and you know the 866-- hash collision position if there was one. This information can be obtained 867-- from 'lookupRecordCollision'. If there is no collision pass (-1) as collPos 868-- (first argument). 869-- 870-- We can skip the key equality check on a Leaf because we know the leaf must be 871-- for this key. 872insertKeyExists :: Int -> Hash -> k -> v -> HashMap k v -> HashMap k v 873insertKeyExists !collPos0 !h0 !k0 x0 !m0 = go collPos0 h0 k0 x0 0 m0 874 where 875 go !_collPos !h !k x !_s (Leaf _hy _kx) 876 = Leaf h (L k x) 877 go collPos h k x s (BitmapIndexed b ary) 878 | b .&. m == 0 = 879 let !ary' = A.insert ary i $ Leaf h (L k x) 880 in bitmapIndexedOrFull (b .|. m) ary' 881 | otherwise = 882 let !st = A.index ary i 883 !st' = go collPos h k x (s+bitsPerSubkey) st 884 in BitmapIndexed b (A.update ary i st') 885 where m = mask h s 886 i = sparseIndex b m 887 go collPos h k x s (Full ary) = 888 let !st = A.index ary i 889 !st' = go collPos h k x (s+bitsPerSubkey) st 890 in Full (update16 ary i st') 891 where i = index h s 892 go collPos h k x _s (Collision _hy v) 893 | collPos >= 0 = Collision h (setAtPosition collPos k x v) 894 | otherwise = Empty -- error "Internal error: go {collPos negative}" 895 go _ _ _ _ _ Empty = Empty -- error "Internal error: go Empty" 896 897{-# NOINLINE insertKeyExists #-} 898 899-- Replace the ith Leaf with Leaf k v. 900-- 901-- This does not check that @i@ is within bounds of the array. 902setAtPosition :: Int -> k -> v -> A.Array (Leaf k v) -> A.Array (Leaf k v) 903setAtPosition i k x ary = A.update ary i (L k x) 904{-# INLINE setAtPosition #-} 905 906 907-- | In-place update version of insert 908unsafeInsert :: (Eq k, Hashable k) => k -> v -> HashMap k v -> HashMap k v 909unsafeInsert k0 v0 m0 = runST (go h0 k0 v0 0 m0) 910 where 911 h0 = hash k0 912 go !h !k x !_ Empty = return $! Leaf h (L k x) 913 go h k x s t@(Leaf hy l@(L ky y)) 914 | hy == h = if ky == k 915 then if x `ptrEq` y 916 then return t 917 else return $! Leaf h (L k x) 918 else return $! collision h l (L k x) 919 | otherwise = two s h k x hy t 920 go h k x s t@(BitmapIndexed b ary) 921 | b .&. m == 0 = do 922 ary' <- A.insertM ary i $! Leaf h (L k x) 923 return $! bitmapIndexedOrFull (b .|. m) ary' 924 | otherwise = do 925 st <- A.indexM ary i 926 st' <- go h k x (s+bitsPerSubkey) st 927 A.unsafeUpdateM ary i st' 928 return t 929 where m = mask h s 930 i = sparseIndex b m 931 go h k x s t@(Full ary) = do 932 st <- A.indexM ary i 933 st' <- go h k x (s+bitsPerSubkey) st 934 A.unsafeUpdateM ary i st' 935 return t 936 where i = index h s 937 go h k x s t@(Collision hy v) 938 | h == hy = return $! Collision h (updateOrSnocWith (\a _ -> (# a #)) k x v) 939 | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) 940{-# INLINABLE unsafeInsert #-} 941 942-- | Create a map from two key-value pairs which hashes don't collide. To 943-- enhance sharing, the second key-value pair is represented by the hash of its 944-- key and a singleton HashMap pairing its key with its value. 945-- 946-- Note: to avoid silly thunks, this function must be strict in the 947-- key. See issue #232. We don't need to force the HashMap argument 948-- because it's already in WHNF (having just been matched) and we 949-- just put it directly in an array. 950two :: Shift -> Hash -> k -> v -> Hash -> HashMap k v -> ST s (HashMap k v) 951two = go 952 where 953 go s h1 k1 v1 h2 t2 954 | bp1 == bp2 = do 955 st <- go (s+bitsPerSubkey) h1 k1 v1 h2 t2 956 ary <- A.singletonM st 957 return $ BitmapIndexed bp1 ary 958 | otherwise = do 959 mary <- A.new 2 $! Leaf h1 (L k1 v1) 960 A.write mary idx2 t2 961 ary <- A.unsafeFreeze mary 962 return $ BitmapIndexed (bp1 .|. bp2) ary 963 where 964 bp1 = mask h1 s 965 bp2 = mask h2 s 966 idx2 | index h1 s < index h2 s = 1 967 | otherwise = 0 968{-# INLINE two #-} 969 970-- | /O(log n)/ Associate the value with the key in this map. If 971-- this map previously contained a mapping for the key, the old value 972-- is replaced by the result of applying the given function to the new 973-- and old value. Example: 974-- 975-- > insertWith f k v map 976-- > where f new old = new + old 977insertWith :: (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v 978 -> HashMap k v 979-- We're not going to worry about allocating a function closure 980-- to pass to insertModifying. See comments at 'adjust'. 981insertWith f k new m = insertModifying new (\old -> (# f new old #)) k m 982{-# INLINE insertWith #-} 983 984-- | @insertModifying@ is a lot like insertWith; we use it to implement alterF. 985-- It takes a value to insert when the key is absent and a function 986-- to apply to calculate a new value when the key is present. Thanks 987-- to the unboxed unary tuple, we avoid introducing any unnecessary 988-- thunks in the tree. 989insertModifying :: (Eq k, Hashable k) => v -> (v -> (# v #)) -> k -> HashMap k v 990 -> HashMap k v 991insertModifying x f k0 m0 = go h0 k0 0 m0 992 where 993 !h0 = hash k0 994 go !h !k !_ Empty = Leaf h (L k x) 995 go h k s t@(Leaf hy l@(L ky y)) 996 | hy == h = if ky == k 997 then case f y of 998 (# v' #) | ptrEq y v' -> t 999 | otherwise -> Leaf h (L k (v')) 1000 else collision h l (L k x) 1001 | otherwise = runST (two s h k x hy t) 1002 go h k s t@(BitmapIndexed b ary) 1003 | b .&. m == 0 = 1004 let ary' = A.insert ary i $! Leaf h (L k x) 1005 in bitmapIndexedOrFull (b .|. m) ary' 1006 | otherwise = 1007 let !st = A.index ary i 1008 !st' = go h k (s+bitsPerSubkey) st 1009 ary' = A.update ary i $! st' 1010 in if ptrEq st st' 1011 then t 1012 else BitmapIndexed b ary' 1013 where m = mask h s 1014 i = sparseIndex b m 1015 go h k s t@(Full ary) = 1016 let !st = A.index ary i 1017 !st' = go h k (s+bitsPerSubkey) st 1018 ary' = update16 ary i $! st' 1019 in if ptrEq st st' 1020 then t 1021 else Full ary' 1022 where i = index h s 1023 go h k s t@(Collision hy v) 1024 | h == hy = 1025 let !v' = insertModifyingArr x f k v 1026 in if A.unsafeSameArray v v' 1027 then t 1028 else Collision h v' 1029 | otherwise = go h k s $ BitmapIndexed (mask hy s) (A.singleton t) 1030{-# INLINABLE insertModifying #-} 1031 1032-- Like insertModifying for arrays; used to implement insertModifying 1033insertModifyingArr :: Eq k => v -> (v -> (# v #)) -> k -> A.Array (Leaf k v) 1034 -> A.Array (Leaf k v) 1035insertModifyingArr x f k0 ary0 = go k0 ary0 0 (A.length ary0) 1036 where 1037 go !k !ary !i !n 1038 | i >= n = A.run $ do 1039 -- Not found, append to the end. 1040 mary <- A.new_ (n + 1) 1041 A.copy ary 0 mary 0 n 1042 A.write mary n (L k x) 1043 return mary 1044 | otherwise = case A.index ary i of 1045 (L kx y) | k == kx -> case f y of 1046 (# y' #) -> if ptrEq y y' 1047 then ary 1048 else A.update ary i (L k y') 1049 | otherwise -> go k ary (i+1) n 1050{-# INLINE insertModifyingArr #-} 1051 1052-- | In-place update version of insertWith 1053unsafeInsertWith :: forall k v. (Eq k, Hashable k) 1054 => (v -> v -> v) -> k -> v -> HashMap k v 1055 -> HashMap k v 1056unsafeInsertWith f k0 v0 m0 = unsafeInsertWithKey (const f) k0 v0 m0 1057{-# INLINABLE unsafeInsertWith #-} 1058 1059unsafeInsertWithKey :: forall k v. (Eq k, Hashable k) 1060 => (k -> v -> v -> v) -> k -> v -> HashMap k v 1061 -> HashMap k v 1062unsafeInsertWithKey f k0 v0 m0 = runST (go h0 k0 v0 0 m0) 1063 where 1064 h0 = hash k0 1065 go :: Hash -> k -> v -> Shift -> HashMap k v -> ST s (HashMap k v) 1066 go !h !k x !_ Empty = return $! Leaf h (L k x) 1067 go h k x s t@(Leaf hy l@(L ky y)) 1068 | hy == h = if ky == k 1069 then return $! Leaf h (L k (f k x y)) 1070 else return $! collision h l (L k x) 1071 | otherwise = two s h k x hy t 1072 go h k x s t@(BitmapIndexed b ary) 1073 | b .&. m == 0 = do 1074 ary' <- A.insertM ary i $! Leaf h (L k x) 1075 return $! bitmapIndexedOrFull (b .|. m) ary' 1076 | otherwise = do 1077 st <- A.indexM ary i 1078 st' <- go h k x (s+bitsPerSubkey) st 1079 A.unsafeUpdateM ary i st' 1080 return t 1081 where m = mask h s 1082 i = sparseIndex b m 1083 go h k x s t@(Full ary) = do 1084 st <- A.indexM ary i 1085 st' <- go h k x (s+bitsPerSubkey) st 1086 A.unsafeUpdateM ary i st' 1087 return t 1088 where i = index h s 1089 go h k x s t@(Collision hy v) 1090 | h == hy = return $! Collision h (updateOrSnocWithKey (\key a b -> (# f key a b #) ) k x v) 1091 | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) 1092{-# INLINABLE unsafeInsertWithKey #-} 1093 1094-- | /O(log n)/ Remove the mapping for the specified key from this map 1095-- if present. 1096delete :: (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v 1097delete k m = delete' (hash k) k m 1098{-# INLINABLE delete #-} 1099 1100delete' :: Eq k => Hash -> k -> HashMap k v -> HashMap k v 1101delete' h0 k0 m0 = go h0 k0 0 m0 1102 where 1103 go !_ !_ !_ Empty = Empty 1104 go h k _ t@(Leaf hy (L ky _)) 1105 | hy == h && ky == k = Empty 1106 | otherwise = t 1107 go h k s t@(BitmapIndexed b ary) 1108 | b .&. m == 0 = t 1109 | otherwise = 1110 let !st = A.index ary i 1111 !st' = go h k (s+bitsPerSubkey) st 1112 in if st' `ptrEq` st 1113 then t 1114 else case st' of 1115 Empty | A.length ary == 1 -> Empty 1116 | A.length ary == 2 -> 1117 case (i, A.index ary 0, A.index ary 1) of 1118 (0, _, l) | isLeafOrCollision l -> l 1119 (1, l, _) | isLeafOrCollision l -> l 1120 _ -> bIndexed 1121 | otherwise -> bIndexed 1122 where 1123 bIndexed = BitmapIndexed (b .&. complement m) (A.delete ary i) 1124 l | isLeafOrCollision l && A.length ary == 1 -> l 1125 _ -> BitmapIndexed b (A.update ary i st') 1126 where m = mask h s 1127 i = sparseIndex b m 1128 go h k s t@(Full ary) = 1129 let !st = A.index ary i 1130 !st' = go h k (s+bitsPerSubkey) st 1131 in if st' `ptrEq` st 1132 then t 1133 else case st' of 1134 Empty -> 1135 let ary' = A.delete ary i 1136 bm = fullNodeMask .&. complement (1 `unsafeShiftL` i) 1137 in BitmapIndexed bm ary' 1138 _ -> Full (A.update ary i st') 1139 where i = index h s 1140 go h k _ t@(Collision hy v) 1141 | h == hy = case indexOf k v of 1142 Just i 1143 | A.length v == 2 -> 1144 if i == 0 1145 then Leaf h (A.index v 1) 1146 else Leaf h (A.index v 0) 1147 | otherwise -> Collision h (A.delete v i) 1148 Nothing -> t 1149 | otherwise = t 1150{-# INLINABLE delete' #-} 1151 1152-- | Delete optimized for the case when we know the key is in the map. 1153-- 1154-- It is only valid to call this when the key exists in the map and you know the 1155-- hash collision position if there was one. This information can be obtained 1156-- from 'lookupRecordCollision'. If there is no collision pass (-1) as collPos. 1157-- 1158-- We can skip: 1159-- - the key equality check on the leaf, if we reach a leaf it must be the key 1160deleteKeyExists :: Int -> Hash -> k -> HashMap k v -> HashMap k v 1161deleteKeyExists !collPos0 !h0 !k0 !m0 = go collPos0 h0 k0 0 m0 1162 where 1163 go :: Int -> Hash -> k -> Int -> HashMap k v -> HashMap k v 1164 go !_collPos !_h !_k !_s (Leaf _ _) = Empty 1165 go collPos h k s (BitmapIndexed b ary) = 1166 let !st = A.index ary i 1167 !st' = go collPos h k (s+bitsPerSubkey) st 1168 in case st' of 1169 Empty | A.length ary == 1 -> Empty 1170 | A.length ary == 2 -> 1171 case (i, A.index ary 0, A.index ary 1) of 1172 (0, _, l) | isLeafOrCollision l -> l 1173 (1, l, _) | isLeafOrCollision l -> l 1174 _ -> bIndexed 1175 | otherwise -> bIndexed 1176 where 1177 bIndexed = BitmapIndexed (b .&. complement m) (A.delete ary i) 1178 l | isLeafOrCollision l && A.length ary == 1 -> l 1179 _ -> BitmapIndexed b (A.update ary i st') 1180 where m = mask h s 1181 i = sparseIndex b m 1182 go collPos h k s (Full ary) = 1183 let !st = A.index ary i 1184 !st' = go collPos h k (s+bitsPerSubkey) st 1185 in case st' of 1186 Empty -> 1187 let ary' = A.delete ary i 1188 bm = fullNodeMask .&. complement (1 `unsafeShiftL` i) 1189 in BitmapIndexed bm ary' 1190 _ -> Full (A.update ary i st') 1191 where i = index h s 1192 go collPos h _ _ (Collision _hy v) 1193 | A.length v == 2 1194 = if collPos == 0 1195 then Leaf h (A.index v 1) 1196 else Leaf h (A.index v 0) 1197 | otherwise = Collision h (A.delete v collPos) 1198 go !_ !_ !_ !_ Empty = Empty -- error "Internal error: deleteKeyExists empty" 1199{-# NOINLINE deleteKeyExists #-} 1200 1201-- | /O(log n)/ Adjust the value tied to a given key in this map only 1202-- if it is present. Otherwise, leave the map alone. 1203adjust :: (Eq k, Hashable k) => (v -> v) -> k -> HashMap k v -> HashMap k v 1204-- This operation really likes to leak memory, so using this 1205-- indirect implementation shouldn't hurt much. Furthermore, it allows 1206-- GHC to avoid a leak when the function is lazy. In particular, 1207-- 1208-- adjust (const x) k m 1209-- ==> adjust# (\v -> (# const x v #)) k m 1210-- ==> adjust# (\_ -> (# x #)) k m 1211adjust f k m = adjust# (\v -> (# f v #)) k m 1212{-# INLINE adjust #-} 1213 1214-- | Much like 'adjust', but not inherently leaky. 1215adjust# :: (Eq k, Hashable k) => (v -> (# v #)) -> k -> HashMap k v -> HashMap k v 1216adjust# f k0 m0 = go h0 k0 0 m0 1217 where 1218 h0 = hash k0 1219 go !_ !_ !_ Empty = Empty 1220 go h k _ t@(Leaf hy (L ky y)) 1221 | hy == h && ky == k = case f y of 1222 (# y' #) | ptrEq y y' -> t 1223 | otherwise -> Leaf h (L k y') 1224 | otherwise = t 1225 go h k s t@(BitmapIndexed b ary) 1226 | b .&. m == 0 = t 1227 | otherwise = let !st = A.index ary i 1228 !st' = go h k (s+bitsPerSubkey) st 1229 ary' = A.update ary i $! st' 1230 in if ptrEq st st' 1231 then t 1232 else BitmapIndexed b ary' 1233 where m = mask h s 1234 i = sparseIndex b m 1235 go h k s t@(Full ary) = 1236 let i = index h s 1237 !st = A.index ary i 1238 !st' = go h k (s+bitsPerSubkey) st 1239 ary' = update16 ary i $! st' 1240 in if ptrEq st st' 1241 then t 1242 else Full ary' 1243 go h k _ t@(Collision hy v) 1244 | h == hy = let !v' = updateWith# f k v 1245 in if A.unsafeSameArray v v' 1246 then t 1247 else Collision h v' 1248 | otherwise = t 1249{-# INLINABLE adjust# #-} 1250 1251-- | /O(log n)/ The expression @('update' f k map)@ updates the value @x@ at @k@ 1252-- (if it is in the map). If @(f x)@ is 'Nothing', the element is deleted. 1253-- If it is @('Just' y)@, the key @k@ is bound to the new value @y@. 1254update :: (Eq k, Hashable k) => (a -> Maybe a) -> k -> HashMap k a -> HashMap k a 1255update f = alter (>>= f) 1256{-# INLINABLE update #-} 1257 1258 1259-- | /O(log n)/ The expression @('alter' f k map)@ alters the value @x@ at @k@, or 1260-- absence thereof. 1261-- 1262-- 'alter' can be used to insert, delete, or update a value in a map. In short: 1263-- 1264-- @ 1265-- 'lookup' k ('alter' f k m) = f ('lookup' k m) 1266-- @ 1267alter :: (Eq k, Hashable k) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v 1268-- TODO(m-renaud): Consider using specialized insert and delete for alter. 1269alter f k m = 1270 case f (lookup k m) of 1271 Nothing -> delete k m 1272 Just v -> insert k v m 1273{-# INLINABLE alter #-} 1274 1275-- | /O(log n)/ The expression @('alterF' f k map)@ alters the value @x@ at 1276-- @k@, or absence thereof. 1277-- 1278-- 'alterF' can be used to insert, delete, or update a value in a map. 1279-- 1280-- Note: 'alterF' is a flipped version of the 'at' combinator from 1281-- <https://hackage.haskell.org/package/lens/docs/Control-Lens-At.html#v:at Control.Lens.At>. 1282-- 1283-- @since 0.2.10 1284alterF :: (Functor f, Eq k, Hashable k) 1285 => (Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v) 1286-- We only calculate the hash once, but unless this is rewritten 1287-- by rules we may test for key equality multiple times. 1288-- We force the value of the map for consistency with the rewritten 1289-- version; otherwise someone could tell the difference using a lazy 1290-- @f@ and a functor that is similar to Const but not actually Const. 1291alterF f = \ !k !m -> 1292 let 1293 !h = hash k 1294 mv = lookup' h k m 1295 in (<$> f mv) $ \fres -> 1296 case fres of 1297 Nothing -> maybe m (const (delete' h k m)) mv 1298 Just v' -> insert' h k v' m 1299 1300-- We unconditionally rewrite alterF in RULES, but we expose an 1301-- unfolding just in case it's used in some way that prevents the 1302-- rule from firing. 1303{-# INLINABLE [0] alterF #-} 1304 1305#if MIN_VERSION_base(4,8,0) 1306-- This is just a bottom value. See the comment on the "alterFWeird" 1307-- rule. 1308test_bottom :: a 1309test_bottom = error "Data.HashMap.alterF internal error: hit test_bottom" 1310 1311-- We use this as an error result in RULES to ensure we don't get 1312-- any useless CallStack nonsense. 1313bogus# :: (# #) -> (# a #) 1314bogus# _ = error "Data.HashMap.alterF internal error: hit bogus#" 1315 1316{-# RULES 1317-- We probe the behavior of @f@ by applying it to Nothing and to 1318-- Just test_bottom. Based on the results, and how they relate to 1319-- each other, we choose the best implementation. 1320 1321"alterFWeird" forall f. alterF f = 1322 alterFWeird (f Nothing) (f (Just test_bottom)) f 1323 1324-- This rule covers situations where alterF is used to simply insert or 1325-- delete in Identity (most likely via Control.Lens.At). We recognize here 1326-- (through the repeated @x@ on the LHS) that 1327-- 1328-- @f Nothing = f (Just bottom)@, 1329-- 1330-- which guarantees that @f@ doesn't care what its argument is, so 1331-- we don't have to either. 1332-- 1333-- Why only Identity? A variant of this rule is actually valid regardless of 1334-- the functor, but for some functors (e.g., []), it can lead to the 1335-- same keys being compared multiple times, which is bad if they're 1336-- ugly things like strings. This is unfortunate, since the rule is likely 1337-- a good idea for almost all realistic uses, but I don't like nasty 1338-- edge cases. 1339"alterFconstant" forall (f :: Maybe a -> Identity (Maybe a)) x. 1340 alterFWeird x x f = \ !k !m -> 1341 Identity (case runIdentity x of {Nothing -> delete k m; Just a -> insert k a m}) 1342 1343-- This rule handles the case where 'alterF' is used to do 'insertWith'-like 1344-- things. Whenever possible, GHC will get rid of the Maybe nonsense for us. 1345-- We delay this rule to stage 1 so alterFconstant has a chance to fire. 1346"alterFinsertWith" [1] forall (f :: Maybe a -> Identity (Maybe a)) x y. 1347 alterFWeird (coerce (Just x)) (coerce (Just y)) f = 1348 coerce (insertModifying x (\mold -> case runIdentity (f (Just mold)) of 1349 Nothing -> bogus# (# #) 1350 Just new -> (# new #))) 1351 1352-- Handle the case where someone uses 'alterF' instead of 'adjust'. This 1353-- rule is kind of picky; it will only work if the function doesn't 1354-- do anything between case matching on the Maybe and producing a result. 1355"alterFadjust" forall (f :: Maybe a -> Identity (Maybe a)) _y. 1356 alterFWeird (coerce Nothing) (coerce (Just _y)) f = 1357 coerce (adjust# (\x -> case runIdentity (f (Just x)) of 1358 Just x' -> (# x' #) 1359 Nothing -> bogus# (# #))) 1360 1361-- The simple specialization to Const; in this case we can look up 1362-- the key without caring what position it's in. This is only a tiny 1363-- optimization. 1364"alterFlookup" forall _ign1 _ign2 (f :: Maybe a -> Const r (Maybe a)). 1365 alterFWeird _ign1 _ign2 f = \ !k !m -> Const (getConst (f (lookup k m))) 1366 #-} 1367 1368-- This is a very unsafe version of alterF used for RULES. When calling 1369-- alterFWeird x y f, the following *must* hold: 1370-- 1371-- x = f Nothing 1372-- y = f (Just _|_) 1373-- 1374-- Failure to abide by these laws will make demons come out of your nose. 1375alterFWeird 1376 :: (Functor f, Eq k, Hashable k) 1377 => f (Maybe v) 1378 -> f (Maybe v) 1379 -> (Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v) 1380alterFWeird _ _ f = alterFEager f 1381{-# INLINE [0] alterFWeird #-} 1382 1383-- | This is the default version of alterF that we use in most non-trivial 1384-- cases. It's called "eager" because it looks up the given key in the map 1385-- eagerly, whether or not the given function requires that information. 1386alterFEager :: (Functor f, Eq k, Hashable k) 1387 => (Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v) 1388alterFEager f !k m = (<$> f mv) $ \fres -> 1389 case fres of 1390 1391 ------------------------------ 1392 -- Delete the key from the map. 1393 Nothing -> case lookupRes of 1394 1395 -- Key did not exist in the map to begin with, no-op 1396 Absent -> m 1397 1398 -- Key did exist 1399 Present _ collPos -> deleteKeyExists collPos h k m 1400 1401 ------------------------------ 1402 -- Update value 1403 Just v' -> case lookupRes of 1404 1405 -- Key did not exist before, insert v' under a new key 1406 Absent -> insertNewKey h k v' m 1407 1408 -- Key existed before 1409 Present v collPos -> 1410 if v `ptrEq` v' 1411 -- If the value is identical, no-op 1412 then m 1413 -- If the value changed, update the value. 1414 else insertKeyExists collPos h k v' m 1415 1416 where !h = hash k 1417 !lookupRes = lookupRecordCollision h k m 1418 !mv = case lookupRes of 1419 Absent -> Nothing 1420 Present v _ -> Just v 1421{-# INLINABLE alterFEager #-} 1422#endif 1423 1424-- | /O(n*log m)/ Inclusion of maps. A map is included in another map if the keys 1425-- are subsets and the corresponding values are equal: 1426-- 1427-- > isSubmapOf m1 m2 = keys m1 `isSubsetOf` keys m2 && 1428-- > and [ v1 == v2 | (k1,v1) <- toList m1; let v2 = m2 ! k1 ] 1429-- 1430-- ==== __Examples__ 1431-- 1432-- >>> fromList [(1,'a')] `isSubmapOf` fromList [(1,'a'),(2,'b')] 1433-- True 1434-- 1435-- >>> fromList [(1,'a'),(2,'b')] `isSubmapOf` fromList [(1,'a')] 1436-- False 1437-- 1438-- @since 0.2.12 1439isSubmapOf :: (Eq k, Hashable k, Eq v) => HashMap k v -> HashMap k v -> Bool 1440isSubmapOf = (inline isSubmapOfBy) (==) 1441{-# INLINABLE isSubmapOf #-} 1442 1443-- | /O(n*log m)/ Inclusion of maps with value comparison. A map is included in 1444-- another map if the keys are subsets and if the comparison function is true 1445-- for the corresponding values: 1446-- 1447-- > isSubmapOfBy cmpV m1 m2 = keys m1 `isSubsetOf` keys m2 && 1448-- > and [ v1 `cmpV` v2 | (k1,v1) <- toList m1; let v2 = m2 ! k1 ] 1449-- 1450-- ==== __Examples__ 1451-- 1452-- >>> isSubmapOfBy (<=) (fromList [(1,'a')]) (fromList [(1,'b'),(2,'c')]) 1453-- True 1454-- 1455-- >>> isSubmapOfBy (<=) (fromList [(1,'b')]) (fromList [(1,'a'),(2,'c')]) 1456-- False 1457-- 1458-- @since 0.2.12 1459isSubmapOfBy :: (Eq k, Hashable k) => (v1 -> v2 -> Bool) -> HashMap k v1 -> HashMap k v2 -> Bool 1460-- For maps without collisions the complexity is O(n*log m), where n is the size 1461-- of m1 and m the size of m2: the inclusion operation visits every leaf in m1 at least once. 1462-- For each leaf in m1, it looks up the key in m2. 1463-- 1464-- The worst case complexity is O(n*m). The worst case is when both hashmaps m1 1465-- and m2 are collision nodes for the same hash. Since collision nodes are 1466-- unsorted arrays, it requires for every key in m1 a linear search to to find a 1467-- matching key in m2, hence O(n*m). 1468isSubmapOfBy comp !m1 !m2 = go 0 m1 m2 1469 where 1470 -- An empty map is always a submap of any other map. 1471 go _ Empty _ = True 1472 1473 -- If the second map is empty and the first is not, it cannot be a submap. 1474 go _ _ Empty = False 1475 1476 -- If the first map contains only one entry, lookup the key in the second map. 1477 go s (Leaf h1 (L k1 v1)) t2 = lookupCont (\_ -> False) (\v2 _ -> comp v1 v2) h1 k1 s t2 1478 1479 -- In this case, we need to check that for each x in ls1, there is a y in 1480 -- ls2 such that x `comp` y. This is the worst case complexity-wise since it 1481 -- requires a O(m*n) check. 1482 go _ (Collision h1 ls1) (Collision h2 ls2) = 1483 h1 == h2 && subsetArray comp ls1 ls2 1484 1485 -- In this case, we only need to check the entries in ls2 with the hash h1. 1486 go s t1@(Collision h1 _) (BitmapIndexed b ls2) 1487 | b .&. m == 0 = False 1488 | otherwise = 1489 go (s+bitsPerSubkey) t1 (A.index ls2 (sparseIndex b m)) 1490 where m = mask h1 s 1491 1492 -- Similar to the previous case we need to traverse l2 at the index for the hash h1. 1493 go s t1@(Collision h1 _) (Full ls2) = 1494 go (s+bitsPerSubkey) t1 (A.index ls2 (index h1 s)) 1495 1496 -- In cases where the first and second map are BitmapIndexed or Full, 1497 -- traverse down the tree at the appropriate indices. 1498 go s (BitmapIndexed b1 ls1) (BitmapIndexed b2 ls2) = 1499 submapBitmapIndexed (go (s+bitsPerSubkey)) b1 ls1 b2 ls2 1500 go s (BitmapIndexed b1 ls1) (Full ls2) = 1501 submapBitmapIndexed (go (s+bitsPerSubkey)) b1 ls1 fullNodeMask ls2 1502 go s (Full ls1) (Full ls2) = 1503 submapBitmapIndexed (go (s+bitsPerSubkey)) fullNodeMask ls1 fullNodeMask ls2 1504 1505 -- Collision and Full nodes always contain at least two entries. Hence it 1506 -- cannot be a map of a leaf. 1507 go _ (Collision {}) (Leaf {}) = False 1508 go _ (BitmapIndexed {}) (Leaf {}) = False 1509 go _ (Full {}) (Leaf {}) = False 1510 go _ (BitmapIndexed {}) (Collision {}) = False 1511 go _ (Full {}) (Collision {}) = False 1512 go _ (Full {}) (BitmapIndexed {}) = False 1513{-# INLINABLE isSubmapOfBy #-} 1514 1515-- | /O(min n m))/ Checks if a bitmap indexed node is a submap of another. 1516submapBitmapIndexed :: (HashMap k v1 -> HashMap k v2 -> Bool) -> Bitmap -> A.Array (HashMap k v1) -> Bitmap -> A.Array (HashMap k v2) -> Bool 1517submapBitmapIndexed comp !b1 !ary1 !b2 !ary2 = subsetBitmaps && go 0 0 (b1Orb2 .&. negate b1Orb2) 1518 where 1519 go :: Int -> Int -> Bitmap -> Bool 1520 go !i !j !m 1521 | m > b1Orb2 = True 1522 1523 -- In case a key is both in ary1 and ary2, check ary1[i] <= ary2[j] and 1524 -- increment the indices i and j. 1525 | b1Andb2 .&. m /= 0 = comp (A.index ary1 i) (A.index ary2 j) && 1526 go (i+1) (j+1) (m `unsafeShiftL` 1) 1527 1528 -- In case a key occurs in ary1, but not ary2, only increment index j. 1529 | b2 .&. m /= 0 = go i (j+1) (m `unsafeShiftL` 1) 1530 1531 -- In case a key neither occurs in ary1 nor ary2, continue. 1532 | otherwise = go i j (m `unsafeShiftL` 1) 1533 1534 b1Andb2 = b1 .&. b2 1535 b1Orb2 = b1 .|. b2 1536 subsetBitmaps = b1Orb2 == b2 1537{-# INLINABLE submapBitmapIndexed #-} 1538 1539------------------------------------------------------------------------ 1540-- * Combine 1541 1542-- | /O(n+m)/ The union of two maps. If a key occurs in both maps, the 1543-- mapping from the first will be the mapping in the result. 1544-- 1545-- ==== __Examples__ 1546-- 1547-- >>> union (fromList [(1,'a'),(2,'b')]) (fromList [(2,'c'),(3,'d')]) 1548-- fromList [(1,'a'),(2,'b'),(3,'d')] 1549union :: (Eq k, Hashable k) => HashMap k v -> HashMap k v -> HashMap k v 1550union = unionWith const 1551{-# INLINABLE union #-} 1552 1553-- | /O(n+m)/ The union of two maps. If a key occurs in both maps, 1554-- the provided function (first argument) will be used to compute the 1555-- result. 1556unionWith :: (Eq k, Hashable k) => (v -> v -> v) -> HashMap k v -> HashMap k v 1557 -> HashMap k v 1558unionWith f = unionWithKey (const f) 1559{-# INLINE unionWith #-} 1560 1561-- | /O(n+m)/ The union of two maps. If a key occurs in both maps, 1562-- the provided function (first argument) will be used to compute the 1563-- result. 1564unionWithKey :: (Eq k, Hashable k) => (k -> v -> v -> v) -> HashMap k v -> HashMap k v 1565 -> HashMap k v 1566unionWithKey f = go 0 1567 where 1568 -- empty vs. anything 1569 go !_ t1 Empty = t1 1570 go _ Empty t2 = t2 1571 -- leaf vs. leaf 1572 go s t1@(Leaf h1 l1@(L k1 v1)) t2@(Leaf h2 l2@(L k2 v2)) 1573 | h1 == h2 = if k1 == k2 1574 then Leaf h1 (L k1 (f k1 v1 v2)) 1575 else collision h1 l1 l2 1576 | otherwise = goDifferentHash s h1 h2 t1 t2 1577 go s t1@(Leaf h1 (L k1 v1)) t2@(Collision h2 ls2) 1578 | h1 == h2 = Collision h1 (updateOrSnocWithKey (\k a b -> (# f k a b #)) k1 v1 ls2) 1579 | otherwise = goDifferentHash s h1 h2 t1 t2 1580 go s t1@(Collision h1 ls1) t2@(Leaf h2 (L k2 v2)) 1581 | h1 == h2 = Collision h1 (updateOrSnocWithKey (\k a b -> (# f k b a #)) k2 v2 ls1) 1582 | otherwise = goDifferentHash s h1 h2 t1 t2 1583 go s t1@(Collision h1 ls1) t2@(Collision h2 ls2) 1584 | h1 == h2 = Collision h1 (updateOrConcatWithKey f ls1 ls2) 1585 | otherwise = goDifferentHash s h1 h2 t1 t2 1586 -- branch vs. branch 1587 go s (BitmapIndexed b1 ary1) (BitmapIndexed b2 ary2) = 1588 let b' = b1 .|. b2 1589 ary' = unionArrayBy (go (s+bitsPerSubkey)) b1 b2 ary1 ary2 1590 in bitmapIndexedOrFull b' ary' 1591 go s (BitmapIndexed b1 ary1) (Full ary2) = 1592 let ary' = unionArrayBy (go (s+bitsPerSubkey)) b1 fullNodeMask ary1 ary2 1593 in Full ary' 1594 go s (Full ary1) (BitmapIndexed b2 ary2) = 1595 let ary' = unionArrayBy (go (s+bitsPerSubkey)) fullNodeMask b2 ary1 ary2 1596 in Full ary' 1597 go s (Full ary1) (Full ary2) = 1598 let ary' = unionArrayBy (go (s+bitsPerSubkey)) fullNodeMask fullNodeMask 1599 ary1 ary2 1600 in Full ary' 1601 -- leaf vs. branch 1602 go s (BitmapIndexed b1 ary1) t2 1603 | b1 .&. m2 == 0 = let ary' = A.insert ary1 i t2 1604 b' = b1 .|. m2 1605 in bitmapIndexedOrFull b' ary' 1606 | otherwise = let ary' = A.updateWith' ary1 i $ \st1 -> 1607 go (s+bitsPerSubkey) st1 t2 1608 in BitmapIndexed b1 ary' 1609 where 1610 h2 = leafHashCode t2 1611 m2 = mask h2 s 1612 i = sparseIndex b1 m2 1613 go s t1 (BitmapIndexed b2 ary2) 1614 | b2 .&. m1 == 0 = let ary' = A.insert ary2 i $! t1 1615 b' = b2 .|. m1 1616 in bitmapIndexedOrFull b' ary' 1617 | otherwise = let ary' = A.updateWith' ary2 i $ \st2 -> 1618 go (s+bitsPerSubkey) t1 st2 1619 in BitmapIndexed b2 ary' 1620 where 1621 h1 = leafHashCode t1 1622 m1 = mask h1 s 1623 i = sparseIndex b2 m1 1624 go s (Full ary1) t2 = 1625 let h2 = leafHashCode t2 1626 i = index h2 s 1627 ary' = update16With' ary1 i $ \st1 -> go (s+bitsPerSubkey) st1 t2 1628 in Full ary' 1629 go s t1 (Full ary2) = 1630 let h1 = leafHashCode t1 1631 i = index h1 s 1632 ary' = update16With' ary2 i $ \st2 -> go (s+bitsPerSubkey) t1 st2 1633 in Full ary' 1634 1635 leafHashCode (Leaf h _) = h 1636 leafHashCode (Collision h _) = h 1637 leafHashCode _ = error "leafHashCode" 1638 1639 goDifferentHash s h1 h2 t1 t2 1640 | m1 == m2 = BitmapIndexed m1 (A.singleton $! go (s+bitsPerSubkey) t1 t2) 1641 | m1 < m2 = BitmapIndexed (m1 .|. m2) (A.pair t1 t2) 1642 | otherwise = BitmapIndexed (m1 .|. m2) (A.pair t2 t1) 1643 where 1644 m1 = mask h1 s 1645 m2 = mask h2 s 1646{-# INLINE unionWithKey #-} 1647 1648-- | Strict in the result of @f@. 1649unionArrayBy :: (a -> a -> a) -> Bitmap -> Bitmap -> A.Array a -> A.Array a 1650 -> A.Array a 1651unionArrayBy f b1 b2 ary1 ary2 = A.run $ do 1652 let b' = b1 .|. b2 1653 mary <- A.new_ (popCount b') 1654 -- iterate over nonzero bits of b1 .|. b2 1655 -- it would be nice if we could shift m by more than 1 each time 1656 let ba = b1 .&. b2 1657 go !i !i1 !i2 !m 1658 | m > b' = return () 1659 | b' .&. m == 0 = go i i1 i2 (m `unsafeShiftL` 1) 1660 | ba .&. m /= 0 = do 1661 x1 <- A.indexM ary1 i1 1662 x2 <- A.indexM ary2 i2 1663 A.write mary i $! f x1 x2 1664 go (i+1) (i1+1) (i2+1) (m `unsafeShiftL` 1) 1665 | b1 .&. m /= 0 = do 1666 A.write mary i =<< A.indexM ary1 i1 1667 go (i+1) (i1+1) (i2 ) (m `unsafeShiftL` 1) 1668 | otherwise = do 1669 A.write mary i =<< A.indexM ary2 i2 1670 go (i+1) (i1 ) (i2+1) (m `unsafeShiftL` 1) 1671 go 0 0 0 (b' .&. negate b') -- XXX: b' must be non-zero 1672 return mary 1673 -- TODO: For the case where b1 .&. b2 == b1, i.e. when one is a 1674 -- subset of the other, we could use a slightly simpler algorithm, 1675 -- where we copy one array, and then update. 1676{-# INLINE unionArrayBy #-} 1677 1678-- TODO: Figure out the time complexity of 'unions'. 1679 1680-- | Construct a set containing all elements from a list of sets. 1681unions :: (Eq k, Hashable k) => [HashMap k v] -> HashMap k v 1682unions = L.foldl' union empty 1683{-# INLINE unions #-} 1684 1685 1686------------------------------------------------------------------------ 1687-- * Compose 1688 1689-- | Relate the keys of one map to the values of 1690-- the other, by using the values of the former as keys for lookups 1691-- in the latter. 1692-- 1693-- Complexity: \( O (n * \log(m)) \), where \(m\) is the size of the first argument 1694-- 1695-- >>> compose (fromList [('a', "A"), ('b', "B")]) (fromList [(1,'a'),(2,'b'),(3,'z')]) 1696-- fromList [(1,"A"),(2,"B")] 1697-- 1698-- @ 1699-- ('compose' bc ab '!?') = (bc '!?') <=< (ab '!?') 1700-- @ 1701-- 1702-- @since UNRELEASED 1703compose :: (Eq b, Hashable b) => HashMap b c -> HashMap a b -> HashMap a c 1704compose bc !ab 1705 | null bc = empty 1706 | otherwise = mapMaybe (bc !?) ab 1707 1708------------------------------------------------------------------------ 1709-- * Transformations 1710 1711-- | /O(n)/ Transform this map by applying a function to every value. 1712mapWithKey :: (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2 1713mapWithKey f = go 1714 where 1715 go Empty = Empty 1716 go (Leaf h (L k v)) = Leaf h $ L k (f k v) 1717 go (BitmapIndexed b ary) = BitmapIndexed b $ A.map go ary 1718 go (Full ary) = Full $ A.map go ary 1719 -- Why map strictly over collision arrays? Because there's no 1720 -- point suspending the O(1) work this does for each leaf. 1721 go (Collision h ary) = Collision h $ 1722 A.map' (\ (L k v) -> L k (f k v)) ary 1723{-# INLINE mapWithKey #-} 1724 1725-- | /O(n)/ Transform this map by applying a function to every value. 1726map :: (v1 -> v2) -> HashMap k v1 -> HashMap k v2 1727map f = mapWithKey (const f) 1728{-# INLINE map #-} 1729 1730-- TODO: We should be able to use mutation to create the new 1731-- 'HashMap'. 1732 1733-- | /O(n)/ Perform an 'Applicative' action for each key-value pair 1734-- in a 'HashMap' and produce a 'HashMap' of all the results. 1735-- 1736-- Note: the order in which the actions occur is unspecified. In particular, 1737-- when the map contains hash collisions, the order in which the actions 1738-- associated with the keys involved will depend in an unspecified way on 1739-- their insertion order. 1740traverseWithKey 1741 :: Applicative f 1742 => (k -> v1 -> f v2) 1743 -> HashMap k v1 -> f (HashMap k v2) 1744traverseWithKey f = go 1745 where 1746 go Empty = pure Empty 1747 go (Leaf h (L k v)) = Leaf h . L k <$> f k v 1748 go (BitmapIndexed b ary) = BitmapIndexed b <$> A.traverse go ary 1749 go (Full ary) = Full <$> A.traverse go ary 1750 go (Collision h ary) = 1751 Collision h <$> A.traverse' (\ (L k v) -> L k <$> f k v) ary 1752{-# INLINE traverseWithKey #-} 1753 1754------------------------------------------------------------------------ 1755-- * Difference and intersection 1756 1757-- | /O(n*log m)/ Difference of two maps. Return elements of the first map 1758-- not existing in the second. 1759difference :: (Eq k, Hashable k) => HashMap k v -> HashMap k w -> HashMap k v 1760difference a b = foldlWithKey' go empty a 1761 where 1762 go m k v = case lookup k b of 1763 Nothing -> insert k v m 1764 _ -> m 1765{-# INLINABLE difference #-} 1766 1767-- | /O(n*log m)/ Difference with a combining function. When two equal keys are 1768-- encountered, the combining function is applied to the values of these keys. 1769-- If it returns 'Nothing', the element is discarded (proper set difference). If 1770-- it returns (@'Just' y@), the element is updated with a new value @y@. 1771differenceWith :: (Eq k, Hashable k) => (v -> w -> Maybe v) -> HashMap k v -> HashMap k w -> HashMap k v 1772differenceWith f a b = foldlWithKey' go empty a 1773 where 1774 go m k v = case lookup k b of 1775 Nothing -> insert k v m 1776 Just w -> maybe m (\y -> insert k y m) (f v w) 1777{-# INLINABLE differenceWith #-} 1778 1779-- | /O(n*log m)/ Intersection of two maps. Return elements of the first 1780-- map for keys existing in the second. 1781intersection :: (Eq k, Hashable k) => HashMap k v -> HashMap k w -> HashMap k v 1782intersection a b = foldlWithKey' go empty a 1783 where 1784 go m k v = case lookup k b of 1785 Just _ -> insert k v m 1786 _ -> m 1787{-# INLINABLE intersection #-} 1788 1789-- | /O(n*log m)/ Intersection of two maps. If a key occurs in both maps 1790-- the provided function is used to combine the values from the two 1791-- maps. 1792intersectionWith :: (Eq k, Hashable k) => (v1 -> v2 -> v3) -> HashMap k v1 1793 -> HashMap k v2 -> HashMap k v3 1794intersectionWith f a b = foldlWithKey' go empty a 1795 where 1796 go m k v = case lookup k b of 1797 Just w -> insert k (f v w) m 1798 _ -> m 1799{-# INLINABLE intersectionWith #-} 1800 1801-- | /O(n*log m)/ Intersection of two maps. If a key occurs in both maps 1802-- the provided function is used to combine the values from the two 1803-- maps. 1804intersectionWithKey :: (Eq k, Hashable k) => (k -> v1 -> v2 -> v3) 1805 -> HashMap k v1 -> HashMap k v2 -> HashMap k v3 1806intersectionWithKey f a b = foldlWithKey' go empty a 1807 where 1808 go m k v = case lookup k b of 1809 Just w -> insert k (f k v w) m 1810 _ -> m 1811{-# INLINABLE intersectionWithKey #-} 1812 1813------------------------------------------------------------------------ 1814-- * Folds 1815 1816-- | /O(n)/ Reduce this map by applying a binary operator to all 1817-- elements, using the given starting value (typically the 1818-- left-identity of the operator). Each application of the operator 1819-- is evaluated before using the result in the next application. 1820-- This function is strict in the starting value. 1821foldl' :: (a -> v -> a) -> a -> HashMap k v -> a 1822foldl' f = foldlWithKey' (\ z _ v -> f z v) 1823{-# INLINE foldl' #-} 1824 1825-- | /O(n)/ Reduce this map by applying a binary operator to all 1826-- elements, using the given starting value (typically the 1827-- right-identity of the operator). Each application of the operator 1828-- is evaluated before using the result in the next application. 1829-- This function is strict in the starting value. 1830foldr' :: (v -> a -> a) -> a -> HashMap k v -> a 1831foldr' f = foldrWithKey' (\ _ v z -> f v z) 1832{-# INLINE foldr' #-} 1833 1834-- | /O(n)/ Reduce this map by applying a binary operator to all 1835-- elements, using the given starting value (typically the 1836-- left-identity of the operator). Each application of the operator 1837-- is evaluated before using the result in the next application. 1838-- This function is strict in the starting value. 1839foldlWithKey' :: (a -> k -> v -> a) -> a -> HashMap k v -> a 1840foldlWithKey' f = go 1841 where 1842 go !z Empty = z 1843 go z (Leaf _ (L k v)) = f z k v 1844 go z (BitmapIndexed _ ary) = A.foldl' go z ary 1845 go z (Full ary) = A.foldl' go z ary 1846 go z (Collision _ ary) = A.foldl' (\ z' (L k v) -> f z' k v) z ary 1847{-# INLINE foldlWithKey' #-} 1848 1849-- | /O(n)/ Reduce this map by applying a binary operator to all 1850-- elements, using the given starting value (typically the 1851-- right-identity of the operator). Each application of the operator 1852-- is evaluated before using the result in the next application. 1853-- This function is strict in the starting value. 1854foldrWithKey' :: (k -> v -> a -> a) -> a -> HashMap k v -> a 1855foldrWithKey' f = flip go 1856 where 1857 go Empty z = z 1858 go (Leaf _ (L k v)) !z = f k v z 1859 go (BitmapIndexed _ ary) !z = A.foldr' go z ary 1860 go (Full ary) !z = A.foldr' go z ary 1861 go (Collision _ ary) !z = A.foldr' (\ (L k v) z' -> f k v z') z ary 1862{-# INLINE foldrWithKey' #-} 1863 1864-- | /O(n)/ Reduce this map by applying a binary operator to all 1865-- elements, using the given starting value (typically the 1866-- right-identity of the operator). 1867foldr :: (v -> a -> a) -> a -> HashMap k v -> a 1868foldr f = foldrWithKey (const f) 1869{-# INLINE foldr #-} 1870 1871-- | /O(n)/ Reduce this map by applying a binary operator to all 1872-- elements, using the given starting value (typically the 1873-- left-identity of the operator). 1874foldl :: (a -> v -> a) -> a -> HashMap k v -> a 1875foldl f = foldlWithKey (\a _k v -> f a v) 1876{-# INLINE foldl #-} 1877 1878-- | /O(n)/ Reduce this map by applying a binary operator to all 1879-- elements, using the given starting value (typically the 1880-- right-identity of the operator). 1881foldrWithKey :: (k -> v -> a -> a) -> a -> HashMap k v -> a 1882foldrWithKey f = flip go 1883 where 1884 go Empty z = z 1885 go (Leaf _ (L k v)) z = f k v z 1886 go (BitmapIndexed _ ary) z = A.foldr go z ary 1887 go (Full ary) z = A.foldr go z ary 1888 go (Collision _ ary) z = A.foldr (\ (L k v) z' -> f k v z') z ary 1889{-# INLINE foldrWithKey #-} 1890 1891-- | /O(n)/ Reduce this map by applying a binary operator to all 1892-- elements, using the given starting value (typically the 1893-- left-identity of the operator). 1894foldlWithKey :: (a -> k -> v -> a) -> a -> HashMap k v -> a 1895foldlWithKey f = go 1896 where 1897 go z Empty = z 1898 go z (Leaf _ (L k v)) = f z k v 1899 go z (BitmapIndexed _ ary) = A.foldl go z ary 1900 go z (Full ary) = A.foldl go z ary 1901 go z (Collision _ ary) = A.foldl (\ z' (L k v) -> f z' k v) z ary 1902{-# INLINE foldlWithKey #-} 1903 1904-- | /O(n)/ Reduce the map by applying a function to each element 1905-- and combining the results with a monoid operation. 1906foldMapWithKey :: Monoid m => (k -> v -> m) -> HashMap k v -> m 1907foldMapWithKey f = go 1908 where 1909 go Empty = mempty 1910 go (Leaf _ (L k v)) = f k v 1911 go (BitmapIndexed _ ary) = A.foldMap go ary 1912 go (Full ary) = A.foldMap go ary 1913 go (Collision _ ary) = A.foldMap (\ (L k v) -> f k v) ary 1914{-# INLINE foldMapWithKey #-} 1915 1916------------------------------------------------------------------------ 1917-- * Filter 1918 1919-- | /O(n)/ Transform this map by applying a function to every value 1920-- and retaining only some of them. 1921mapMaybeWithKey :: (k -> v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2 1922mapMaybeWithKey f = filterMapAux onLeaf onColl 1923 where onLeaf (Leaf h (L k v)) | Just v' <- f k v = Just (Leaf h (L k v')) 1924 onLeaf _ = Nothing 1925 1926 onColl (L k v) | Just v' <- f k v = Just (L k v') 1927 | otherwise = Nothing 1928{-# INLINE mapMaybeWithKey #-} 1929 1930-- | /O(n)/ Transform this map by applying a function to every value 1931-- and retaining only some of them. 1932mapMaybe :: (v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2 1933mapMaybe f = mapMaybeWithKey (const f) 1934{-# INLINE mapMaybe #-} 1935 1936-- | /O(n)/ Filter this map by retaining only elements satisfying a 1937-- predicate. 1938filterWithKey :: forall k v. (k -> v -> Bool) -> HashMap k v -> HashMap k v 1939filterWithKey pred = filterMapAux onLeaf onColl 1940 where onLeaf t@(Leaf _ (L k v)) | pred k v = Just t 1941 onLeaf _ = Nothing 1942 1943 onColl el@(L k v) | pred k v = Just el 1944 onColl _ = Nothing 1945{-# INLINE filterWithKey #-} 1946 1947 1948-- | Common implementation for 'filterWithKey' and 'mapMaybeWithKey', 1949-- allowing the former to former to reuse terms. 1950filterMapAux :: forall k v1 v2 1951 . (HashMap k v1 -> Maybe (HashMap k v2)) 1952 -> (Leaf k v1 -> Maybe (Leaf k v2)) 1953 -> HashMap k v1 1954 -> HashMap k v2 1955filterMapAux onLeaf onColl = go 1956 where 1957 go Empty = Empty 1958 go t@Leaf{} 1959 | Just t' <- onLeaf t = t' 1960 | otherwise = Empty 1961 go (BitmapIndexed b ary) = filterA ary b 1962 go (Full ary) = filterA ary fullNodeMask 1963 go (Collision h ary) = filterC ary h 1964 1965 filterA ary0 b0 = 1966 let !n = A.length ary0 1967 in runST $ do 1968 mary <- A.new_ n 1969 step ary0 mary b0 0 0 1 n 1970 where 1971 step :: A.Array (HashMap k v1) -> A.MArray s (HashMap k v2) 1972 -> Bitmap -> Int -> Int -> Bitmap -> Int 1973 -> ST s (HashMap k v2) 1974 step !ary !mary !b i !j !bi n 1975 | i >= n = case j of 1976 0 -> return Empty 1977 1 -> do 1978 ch <- A.read mary 0 1979 case ch of 1980 t | isLeafOrCollision t -> return t 1981 _ -> BitmapIndexed b <$> A.trim mary 1 1982 _ -> do 1983 ary2 <- A.trim mary j 1984 return $! if j == maxChildren 1985 then Full ary2 1986 else BitmapIndexed b ary2 1987 | bi .&. b == 0 = step ary mary b i j (bi `unsafeShiftL` 1) n 1988 | otherwise = case go (A.index ary i) of 1989 Empty -> step ary mary (b .&. complement bi) (i+1) j 1990 (bi `unsafeShiftL` 1) n 1991 t -> do A.write mary j t 1992 step ary mary b (i+1) (j+1) (bi `unsafeShiftL` 1) n 1993 1994 filterC ary0 h = 1995 let !n = A.length ary0 1996 in runST $ do 1997 mary <- A.new_ n 1998 step ary0 mary 0 0 n 1999 where 2000 step :: A.Array (Leaf k v1) -> A.MArray s (Leaf k v2) 2001 -> Int -> Int -> Int 2002 -> ST s (HashMap k v2) 2003 step !ary !mary i !j n 2004 | i >= n = case j of 2005 0 -> return Empty 2006 1 -> do l <- A.read mary 0 2007 return $! Leaf h l 2008 _ | i == j -> do ary2 <- A.unsafeFreeze mary 2009 return $! Collision h ary2 2010 | otherwise -> do ary2 <- A.trim mary j 2011 return $! Collision h ary2 2012 | Just el <- onColl $! A.index ary i 2013 = A.write mary j el >> step ary mary (i+1) (j+1) n 2014 | otherwise = step ary mary (i+1) j n 2015{-# INLINE filterMapAux #-} 2016 2017-- | /O(n)/ Filter this map by retaining only elements which values 2018-- satisfy a predicate. 2019filter :: (v -> Bool) -> HashMap k v -> HashMap k v 2020filter p = filterWithKey (\_ v -> p v) 2021{-# INLINE filter #-} 2022 2023------------------------------------------------------------------------ 2024-- * Conversions 2025 2026-- TODO: Improve fusion rules by modelled them after the Prelude ones 2027-- on lists. 2028 2029-- | /O(n)/ Return a list of this map's keys. The list is produced 2030-- lazily. 2031keys :: HashMap k v -> [k] 2032keys = L.map fst . toList 2033{-# INLINE keys #-} 2034 2035-- | /O(n)/ Return a list of this map's values. The list is produced 2036-- lazily. 2037elems :: HashMap k v -> [v] 2038elems = L.map snd . toList 2039{-# INLINE elems #-} 2040 2041------------------------------------------------------------------------ 2042-- ** Lists 2043 2044-- | /O(n)/ Return a list of this map's elements. The list is 2045-- produced lazily. The order of its elements is unspecified. 2046toList :: HashMap k v -> [(k, v)] 2047toList t = build (\ c z -> foldrWithKey (curry c) z t) 2048{-# INLINE toList #-} 2049 2050-- | /O(n)/ Construct a map with the supplied mappings. If the list 2051-- contains duplicate mappings, the later mappings take precedence. 2052fromList :: (Eq k, Hashable k) => [(k, v)] -> HashMap k v 2053fromList = L.foldl' (\ m (k, v) -> unsafeInsert k v m) empty 2054{-# INLINABLE fromList #-} 2055 2056-- | /O(n*log n)/ Construct a map from a list of elements. Uses 2057-- the provided function @f@ to merge duplicate entries with 2058-- @(f newVal oldVal)@. 2059-- 2060-- === Examples 2061-- 2062-- Given a list @xs@, create a map with the number of occurrences of each 2063-- element in @xs@: 2064-- 2065-- > let xs = ['a', 'b', 'a'] 2066-- > in fromListWith (+) [ (x, 1) | x <- xs ] 2067-- > 2068-- > = fromList [('a', 2), ('b', 1)] 2069-- 2070-- Given a list of key-value pairs @xs :: [(k, v)]@, group all values by their 2071-- keys and return a @HashMap k [v]@. 2072-- 2073-- > let xs = [('a', 1), ('b', 2), ('a', 3)] 2074-- > in fromListWith (++) [ (k, [v]) | (k, v) <- xs ] 2075-- > 2076-- > = fromList [('a', [3, 1]), ('b', [2])] 2077-- 2078-- Note that the lists in the resulting map contain elements in reverse order 2079-- from their occurences in the original list. 2080-- 2081-- More generally, duplicate entries are accumulated as follows; 2082-- this matters when @f@ is not commutative or not associative. 2083-- 2084-- > fromListWith f [(k, a), (k, b), (k, c), (k, d)] 2085-- > = fromList [(k, f d (f c (f b a)))] 2086fromListWith :: (Eq k, Hashable k) => (v -> v -> v) -> [(k, v)] -> HashMap k v 2087fromListWith f = L.foldl' (\ m (k, v) -> unsafeInsertWith f k v m) empty 2088{-# INLINE fromListWith #-} 2089 2090-- | /O(n*log n)/ Construct a map from a list of elements. Uses 2091-- the provided function to merge duplicate entries. 2092-- 2093-- === Examples 2094-- 2095-- Given a list of key-value pairs where the keys are of different flavours, e.g: 2096-- 2097-- > data Key = Div | Sub 2098-- 2099-- and the values need to be combined differently when there are duplicates, 2100-- depending on the key: 2101-- 2102-- > combine Div = div 2103-- > combine Sub = (-) 2104-- 2105-- then @fromListWithKey@ can be used as follows: 2106-- 2107-- > fromListWithKey combine [(Div, 2), (Div, 6), (Sub, 2), (Sub, 3)] 2108-- > = fromList [(Div, 3), (Sub, 1)] 2109-- 2110-- More generally, duplicate entries are accumulated as follows; 2111-- 2112-- > fromListWith f [(k, a), (k, b), (k, c), (k, d)] 2113-- > = fromList [(k, f k d (f k c (f k b a)))] 2114-- 2115-- @since 0.2.11 2116fromListWithKey :: (Eq k, Hashable k) => (k -> v -> v -> v) -> [(k, v)] -> HashMap k v 2117fromListWithKey f = L.foldl' (\ m (k, v) -> unsafeInsertWithKey f k v m) empty 2118{-# INLINE fromListWithKey #-} 2119 2120------------------------------------------------------------------------ 2121-- Array operations 2122 2123-- | /O(n)/ Look up the value associated with the given key in an 2124-- array. 2125lookupInArrayCont :: 2126#if __GLASGOW_HASKELL__ >= 802 2127 forall rep (r :: TYPE rep) k v. 2128#else 2129 forall r k v. 2130#endif 2131 Eq k => ((# #) -> r) -> (v -> Int -> r) -> k -> A.Array (Leaf k v) -> r 2132lookupInArrayCont absent present k0 ary0 = go k0 ary0 0 (A.length ary0) 2133 where 2134 go :: Eq k => k -> A.Array (Leaf k v) -> Int -> Int -> r 2135 go !k !ary !i !n 2136 | i >= n = absent (# #) 2137 | otherwise = case A.index ary i of 2138 (L kx v) 2139 | k == kx -> present v i 2140 | otherwise -> go k ary (i+1) n 2141{-# INLINE lookupInArrayCont #-} 2142 2143-- | /O(n)/ Lookup the value associated with the given key in this 2144-- array. Returns 'Nothing' if the key wasn't found. 2145indexOf :: Eq k => k -> A.Array (Leaf k v) -> Maybe Int 2146indexOf k0 ary0 = go k0 ary0 0 (A.length ary0) 2147 where 2148 go !k !ary !i !n 2149 | i >= n = Nothing 2150 | otherwise = case A.index ary i of 2151 (L kx _) 2152 | k == kx -> Just i 2153 | otherwise -> go k ary (i+1) n 2154{-# INLINABLE indexOf #-} 2155 2156updateWith# :: Eq k => (v -> (# v #)) -> k -> A.Array (Leaf k v) -> A.Array (Leaf k v) 2157updateWith# f k0 ary0 = go k0 ary0 0 (A.length ary0) 2158 where 2159 go !k !ary !i !n 2160 | i >= n = ary 2161 | otherwise = case A.index ary i of 2162 (L kx y) | k == kx -> case f y of 2163 (# y' #) 2164 | ptrEq y y' -> ary 2165 | otherwise -> A.update ary i (L k y') 2166 | otherwise -> go k ary (i+1) n 2167{-# INLINABLE updateWith# #-} 2168 2169updateOrSnocWith :: Eq k => (v -> v -> (# v #)) -> k -> v -> A.Array (Leaf k v) 2170 -> A.Array (Leaf k v) 2171updateOrSnocWith f = updateOrSnocWithKey (const f) 2172{-# INLINABLE updateOrSnocWith #-} 2173 2174updateOrSnocWithKey :: Eq k => (k -> v -> v -> (# v #)) -> k -> v -> A.Array (Leaf k v) 2175 -> A.Array (Leaf k v) 2176updateOrSnocWithKey f k0 v0 ary0 = go k0 v0 ary0 0 (A.length ary0) 2177 where 2178 go !k v !ary !i !n 2179 | i >= n = A.run $ do 2180 -- Not found, append to the end. 2181 mary <- A.new_ (n + 1) 2182 A.copy ary 0 mary 0 n 2183 A.write mary n (L k v) 2184 return mary 2185 | L kx y <- A.index ary i 2186 , k == kx 2187 , (# v2 #) <- f k v y 2188 = A.update ary i (L k v2) 2189 | otherwise 2190 = go k v ary (i+1) n 2191{-# INLINABLE updateOrSnocWithKey #-} 2192 2193updateOrConcatWith :: Eq k => (v -> v -> v) -> A.Array (Leaf k v) -> A.Array (Leaf k v) -> A.Array (Leaf k v) 2194updateOrConcatWith f = updateOrConcatWithKey (const f) 2195{-# INLINABLE updateOrConcatWith #-} 2196 2197updateOrConcatWithKey :: Eq k => (k -> v -> v -> v) -> A.Array (Leaf k v) -> A.Array (Leaf k v) -> A.Array (Leaf k v) 2198updateOrConcatWithKey f ary1 ary2 = A.run $ do 2199 -- TODO: instead of mapping and then folding, should we traverse? 2200 -- We'll have to be careful to avoid allocating pairs or similar. 2201 2202 -- first: look up the position of each element of ary2 in ary1 2203 let indices = A.map' (\(L k _) -> indexOf k ary1) ary2 2204 -- that tells us how large the overlap is: 2205 -- count number of Nothing constructors 2206 let nOnly2 = A.foldl' (\n -> maybe (n+1) (const n)) 0 indices 2207 let n1 = A.length ary1 2208 let n2 = A.length ary2 2209 -- copy over all elements from ary1 2210 mary <- A.new_ (n1 + nOnly2) 2211 A.copy ary1 0 mary 0 n1 2212 -- append or update all elements from ary2 2213 let go !iEnd !i2 2214 | i2 >= n2 = return () 2215 | otherwise = case A.index indices i2 of 2216 Just i1 -> do -- key occurs in both arrays, store combination in position i1 2217 L k v1 <- A.indexM ary1 i1 2218 L _ v2 <- A.indexM ary2 i2 2219 A.write mary i1 (L k (f k v1 v2)) 2220 go iEnd (i2+1) 2221 Nothing -> do -- key is only in ary2, append to end 2222 A.write mary iEnd =<< A.indexM ary2 i2 2223 go (iEnd+1) (i2+1) 2224 go n1 0 2225 return mary 2226{-# INLINABLE updateOrConcatWithKey #-} 2227 2228-- | /O(n*m)/ Check if the first array is a subset of the second array. 2229subsetArray :: Eq k => (v1 -> v2 -> Bool) -> A.Array (Leaf k v1) -> A.Array (Leaf k v2) -> Bool 2230subsetArray cmpV ary1 ary2 = A.length ary1 <= A.length ary2 && A.all inAry2 ary1 2231 where 2232 inAry2 (L k1 v1) = lookupInArrayCont (\_ -> False) (\v2 _ -> cmpV v1 v2) k1 ary2 2233 {-# INLINE inAry2 #-} 2234 2235------------------------------------------------------------------------ 2236-- Manually unrolled loops 2237 2238-- | /O(n)/ Update the element at the given position in this array. 2239update16 :: A.Array e -> Int -> e -> A.Array e 2240update16 ary idx b = runST (update16M ary idx b) 2241{-# INLINE update16 #-} 2242 2243-- | /O(n)/ Update the element at the given position in this array. 2244update16M :: A.Array e -> Int -> e -> ST s (A.Array e) 2245update16M ary idx b = do 2246 mary <- clone16 ary 2247 A.write mary idx b 2248 A.unsafeFreeze mary 2249{-# INLINE update16M #-} 2250 2251-- | /O(n)/ Update the element at the given position in this array, by applying a function to it. 2252update16With' :: A.Array e -> Int -> (e -> e) -> A.Array e 2253update16With' ary idx f 2254 | (# x #) <- A.index# ary idx 2255 = update16 ary idx $! f x 2256{-# INLINE update16With' #-} 2257 2258-- | Unsafely clone an array of 16 elements. The length of the input 2259-- array is not checked. 2260clone16 :: A.Array e -> ST s (A.MArray s e) 2261clone16 ary = 2262 A.thaw ary 0 16 2263 2264------------------------------------------------------------------------ 2265-- Bit twiddling 2266 2267bitsPerSubkey :: Int 2268bitsPerSubkey = 4 2269 2270maxChildren :: Int 2271maxChildren = 1 `unsafeShiftL` bitsPerSubkey 2272 2273subkeyMask :: Bitmap 2274subkeyMask = 1 `unsafeShiftL` bitsPerSubkey - 1 2275 2276sparseIndex :: Bitmap -> Bitmap -> Int 2277sparseIndex b m = popCount (b .&. (m - 1)) 2278 2279mask :: Word -> Shift -> Bitmap 2280mask w s = 1 `unsafeShiftL` index w s 2281{-# INLINE mask #-} 2282 2283-- | Mask out the 'bitsPerSubkey' bits used for indexing at this level 2284-- of the tree. 2285index :: Hash -> Shift -> Int 2286index w s = fromIntegral $ (unsafeShiftR w s) .&. subkeyMask 2287{-# INLINE index #-} 2288 2289-- | A bitmask with the 'bitsPerSubkey' least significant bits set. 2290fullNodeMask :: Bitmap 2291fullNodeMask = complement (complement 0 `unsafeShiftL` maxChildren) 2292{-# INLINE fullNodeMask #-} 2293 2294-- | Check if two the two arguments are the same value. N.B. This 2295-- function might give false negatives (due to GC moving objects.) 2296ptrEq :: a -> a -> Bool 2297ptrEq x y = isTrue# (reallyUnsafePtrEquality# x y ==# 1#) 2298{-# INLINE ptrEq #-} 2299 2300------------------------------------------------------------------------ 2301-- IsList instance 2302instance (Eq k, Hashable k) => Exts.IsList (HashMap k v) where 2303 type Item (HashMap k v) = (k, v) 2304 fromList = fromList 2305 toList = toList 2306