diff --git a/CHANGES.md b/CHANGES.md --- a/CHANGES.md +++ b/CHANGES.md @@ -1,3 +1,12 @@ +## 0.2.10.0 + + * Add `HashMap.alterF`. + + * Add `HashMap.keysSet`. + + * Make `HashMap.Strict.traverseWithKey` force the results before + installing them in the map. + ## 0.2.9.0 * Add `Ord/Ord1/Ord2` instances. (Thanks, Oleg Grenrus) diff --git a/Data/HashMap/Array.hs b/Data/HashMap/Array.hs --- a/Data/HashMap/Array.hs +++ b/Data/HashMap/Array.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE BangPatterns, CPP, MagicHash, Rank2Types, UnboxedTuples #-} +{-# LANGUAGE BangPatterns, CPP, MagicHash, Rank2Types, UnboxedTuples, ScopedTypeVariables #-} {-# OPTIONS_GHC -fno-full-laziness -funbox-strict-fields #-} -- | Zero based arrays. @@ -22,15 +22,19 @@ , write , index , indexM + , index# , update , updateWith' , unsafeUpdateM , insert , insertM , delete + , sameArray1 + , trim , unsafeFreeze , unsafeThaw + , unsafeSameArray , run , run2 , copy @@ -44,17 +48,19 @@ , map , map' , traverse - , filter + , traverse' , toList + , fromList ) where -import qualified Data.Traversable as Traversable -#if __GLASGOW_HASKELL__ < 709 -import Control.Applicative (Applicative) +#if !MIN_VERSION_base(4,8,0) +import Control.Applicative (Applicative (..), (<$>)) #endif +import Control.Applicative (liftA2) import Control.DeepSeq -import GHC.Exts(Int(..)) +import GHC.Exts(Int(..), Int#, reallyUnsafePtrEquality#, tagToEnum#, unsafeCoerce#, State#) import GHC.ST (ST(..)) +import Control.Monad.ST (stToIO) #if __GLASGOW_HASKELL__ >= 709 import Prelude hiding (filter, foldr, length, map, read, traverse) @@ -66,13 +72,13 @@ import GHC.Exts (SmallArray#, newSmallArray#, readSmallArray#, writeSmallArray#, indexSmallArray#, unsafeFreezeSmallArray#, unsafeThawSmallArray#, SmallMutableArray#, sizeofSmallArray#, copySmallArray#, thawSmallArray#, - sizeofSmallMutableArray#, copySmallMutableArray#) + sizeofSmallMutableArray#, copySmallMutableArray#, cloneSmallMutableArray#) #else import GHC.Exts (Array#, newArray#, readArray#, writeArray#, indexArray#, unsafeFreezeArray#, unsafeThawArray#, MutableArray#, sizeofArray#, copyArray#, thawArray#, - sizeofMutableArray#, copyMutableArray#) + sizeofMutableArray#, copyMutableArray#, cloneMutableArray#) #endif #if defined(ASSERTS) @@ -80,22 +86,71 @@ #endif import Data.HashMap.Unsafe (runST) +import Control.Monad ((>=>)) #if __GLASGOW_HASKELL__ >= 710 type Array# a = SmallArray# a type MutableArray# a = SmallMutableArray# a +newArray# :: Int# -> a -> State# d -> (# State# d, SmallMutableArray# d a #) newArray# = newSmallArray# + +unsafeFreezeArray# :: SmallMutableArray# d a + -> State# d -> (# State# d, SmallArray# a #) +unsafeFreezeArray# = unsafeFreezeSmallArray# + +readArray# :: SmallMutableArray# d a + -> Int# -> State# d -> (# State# d, a #) readArray# = readSmallArray# + +writeArray# :: SmallMutableArray# d a + -> Int# -> a -> State# d -> State# d writeArray# = writeSmallArray# + +indexArray# :: SmallArray# a -> Int# -> (# a #) indexArray# = indexSmallArray# -unsafeFreezeArray# = unsafeFreezeSmallArray# + +unsafeThawArray# :: SmallArray# a + -> State# d -> (# State# d, SmallMutableArray# d a #) unsafeThawArray# = unsafeThawSmallArray# + +sizeofArray# :: SmallArray# a -> Int# sizeofArray# = sizeofSmallArray# + +copyArray# :: SmallArray# a + -> Int# + -> SmallMutableArray# d a + -> Int# + -> Int# + -> State# d + -> State# d copyArray# = copySmallArray# + +cloneMutableArray# :: SmallMutableArray# s a + -> Int# + -> Int# + -> State# s + -> (# State# s, SmallMutableArray# s a #) +cloneMutableArray# = cloneSmallMutableArray# + +thawArray# :: SmallArray# a + -> Int# + -> Int# + -> State# d + -> (# State# d, SmallMutableArray# d a #) thawArray# = thawSmallArray# + +sizeofMutableArray# :: SmallMutableArray# s a -> Int# sizeofMutableArray# = sizeofSmallMutableArray# + +copyMutableArray# :: SmallMutableArray# d a + -> Int# + -> SmallMutableArray# d a + -> Int# + -> Int# + -> State# d + -> State# d copyMutableArray# = copySmallMutableArray# #endif @@ -126,6 +181,27 @@ instance Show a => Show (Array a) where show = show . toList +-- Determines whether two arrays have the same memory address. +-- This is more reliable than testing pointer equality on the +-- Array wrappers, but it's still slightly bogus. +unsafeSameArray :: Array a -> Array b -> Bool +unsafeSameArray (Array xs) (Array ys) = + tagToEnum# (unsafeCoerce# reallyUnsafePtrEquality# xs ys) + +sameArray1 :: (a -> b -> Bool) -> Array a -> Array b -> Bool +sameArray1 eq !xs0 !ys0 + | lenxs /= lenys = False + | otherwise = go 0 xs0 ys0 + where + go !k !xs !ys + | k == lenxs = True + | (# x #) <- index# xs k + , (# y #) <- index# ys k + = eq x y && go (k + 1) xs ys + + !lenxs = length xs0 + !lenys = length ys0 + length :: Array a -> Int length ary = I# (sizeofArray# (unArray ary)) {-# INLINE length #-} @@ -159,7 +235,10 @@ n0 = length ary0 go !ary !n !i | i >= n = () - | otherwise = rnf (index ary i) `seq` go ary n (i+1) + | (# x #) <- index# ary i + = rnf x `seq` go ary n (i+1) +-- We use index# just in case GHC can't see that the +-- relevant rnf is strict, or in case it actually isn't. {-# INLINE rnfArray #-} -- | Create a new mutable array of specified size, in the specified @@ -210,6 +289,12 @@ case indexArray# (unArray ary) i# of (# b #) -> b {-# INLINE index #-} +index# :: Array a -> Int -> (# a #) +index# ary _i@(I# i#) = + CHECK_BOUNDS("index#", length ary, _i) + indexArray# (unArray ary) i# +{-# INLINE index# #-} + indexM :: Array a -> Int -> ST s a indexM ary _i@(I# i#) = CHECK_BOUNDS("indexM", length ary, _i) @@ -256,6 +341,19 @@ case copyMutableArray# (unMArray src) sidx# (unMArray dst) didx# n# s# of s2 -> (# s2, () #) +cloneM :: MArray s a -> Int -> Int -> ST s (MArray s a) +cloneM _mary@(MArray mary#) _off@(I# off#) _len@(I# len#) = + CHECK_BOUNDS("cloneM_off", lengthM _mary, _off - 1) + CHECK_BOUNDS("cloneM_end", lengthM _mary, _off + _len - 1) + ST $ \ s -> + case cloneMutableArray# mary# off# len# s of + (# s', mary'# #) -> (# s', MArray mary'# #) + +-- | Create a new array of the @n@ first elements of @mary@. +trim :: MArray s a -> Int -> ST s (Array a) +trim mary n = cloneM mary 0 n >>= unsafeFreeze +{-# INLINE trim #-} + -- | /O(n)/ Insert an element at the given position in this array, -- increasing its size by one. insert :: Array e -> Int -> e -> Array e @@ -294,7 +392,9 @@ -- applying a function to it. Evaluates the element to WHNF before -- inserting it into the array. updateWith' :: Array e -> Int -> (e -> e) -> Array e -updateWith' ary idx f = update ary idx $! f (index ary idx) +updateWith' ary idx f + | (# x #) <- index# ary idx + = update ary idx $! f x {-# INLINE updateWith' #-} -- | /O(1)/ Update the element at the given position in this array, @@ -312,16 +412,20 @@ foldl' f = \ z0 ary0 -> go ary0 (length ary0) 0 z0 where go ary n i !z - | i >= n = z - | otherwise = go ary n (i+1) (f z (index ary i)) + | i >= n = z + | otherwise + = case index# ary i of + (# x #) -> go ary n (i+1) (f z x) {-# INLINE foldl' #-} foldr :: (a -> b -> b) -> b -> Array a -> b foldr f = \ z0 ary0 -> go ary0 (length ary0) 0 z0 where go ary n i z - | i >= n = z - | otherwise = f (index ary i) (go ary n (i+1) z) + | i >= n = z + | otherwise + = case index# ary i of + (# x #) -> f x (go ary n (i+1) z) {-# INLINE foldr #-} undefinedElem :: a @@ -363,7 +467,8 @@ go ary mary i n | i >= n = return mary | otherwise = do - write mary i $ f (index ary i) + x <- indexM ary i + write mary i $ f x go ary mary (i+1) n {-# INLINE map #-} @@ -378,7 +483,8 @@ go ary mary i n | i >= n = return mary | otherwise = do - write mary i $! f (index ary i) + x <- indexM ary i + write mary i $! f x go ary mary (i+1) n {-# INLINE map' #-} @@ -396,25 +502,82 @@ toList :: Array a -> [a] toList = foldr (:) [] +newtype STA a = STA {_runSTA :: forall s. MutableArray# s a -> ST s (Array a)} + +runSTA :: Int -> STA a -> Array a +runSTA !n (STA m) = runST $ new_ n >>= \ (MArray ar) -> m ar + traverse :: Applicative f => (a -> f b) -> Array a -> f (Array b) -traverse f = \ ary -> fromList (length ary) `fmap` - Traversable.traverse f (toList ary) -{-# INLINE traverse #-} +traverse f = \ !ary -> + let + !len = length ary + go !i + | i == len = pure $ STA $ \mary -> unsafeFreeze (MArray mary) + | (# x #) <- index# ary i + = liftA2 (\b (STA m) -> STA $ \mary -> + write (MArray mary) i b >> m mary) + (f x) (go (i + 1)) + in runSTA len <$> go 0 +{-# INLINE [1] traverse #-} -filter :: (a -> Bool) -> Array a -> Array a -filter p = \ ary -> - let !n = length ary - in run $ do - mary <- new_ n - go ary mary 0 0 n - where - go ary mary i j n - | i >= n = if i == j - then return mary - else do mary2 <- new_ j - copyM mary 0 mary2 0 j - return mary2 - | p el = write mary j el >> go ary mary (i+1) (j+1) n - | otherwise = go ary mary (i+1) j n - where el = index ary i -{-# INLINE filter #-} +-- TODO: Would it be better to just use a lazy traversal +-- and then force the elements of the result? My guess is +-- yes. +traverse' :: Applicative f => (a -> f b) -> Array a -> f (Array b) +traverse' f = \ !ary -> + let + !len = length ary + go !i + | i == len = pure $ STA $ \mary -> unsafeFreeze (MArray mary) + | (# x #) <- index# ary i + = liftA2 (\ !b (STA m) -> STA $ \mary -> + write (MArray mary) i b >> m mary) + (f x) (go (i + 1)) + in runSTA len <$> go 0 +{-# INLINE [1] traverse' #-} + +-- Traversing in ST, we don't need to get fancy; we +-- can just do it directly. +traverseST :: (a -> ST s b) -> Array a -> ST s (Array b) +traverseST f = \ ary0 -> + let + !len = length ary0 + go k !mary + | k == len = return mary + | otherwise = do + x <- indexM ary0 k + y <- f x + write mary k y + go (k + 1) mary + in new_ len >>= (go 0 >=> unsafeFreeze) +{-# INLINE traverseST #-} + +traverseIO :: (a -> IO b) -> Array a -> IO (Array b) +traverseIO f = \ ary0 -> + let + !len = length ary0 + go k !mary + | k == len = return mary + | otherwise = do + x <- stToIO $ indexM ary0 k + y <- f x + stToIO $ write mary k y + go (k + 1) mary + in stToIO (new_ len) >>= (go 0 >=> stToIO . unsafeFreeze) +{-# INLINE traverseIO #-} + + +-- Why don't we have similar RULES for traverse'? The efficient +-- way to traverse strictly in IO or ST is to force results as +-- they come in, which leads to different semantics. In particular, +-- we need to ensure that +-- +-- traverse' (\x -> print x *> pure undefined) xs +-- +-- will actually print all the values and then return undefined. +-- We could add a strict mapMWithIndex, operating in an arbitrary +-- Monad, that supported such rules, but we don't have that right now. +{-# RULES +"traverse/ST" forall f. traverse f = traverseST f +"traverse/IO" forall f. traverse f = traverseIO f + #-} diff --git a/Data/HashMap/Base.hs b/Data/HashMap/Base.hs --- a/Data/HashMap/Base.hs +++ b/Data/HashMap/Base.hs @@ -3,6 +3,12 @@ {-# LANGUAGE PatternGuards #-} {-# LANGUAGE RoleAnnotations #-} {-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE LambdaCase #-} +#if __GLASGOW_HASKELL__ >= 802 +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE UnboxedSums #-} +#endif {-# OPTIONS_GHC -fno-full-laziness -funbox-strict-fields #-} module Data.HashMap.Base @@ -28,6 +34,7 @@ , adjust , update , alter + , alterF -- * Combine -- ** Union @@ -89,6 +96,18 @@ , updateOrConcatWithKey , filterMapAux , equalKeys + , equalKeys1 + , lookupRecordCollision + , LookupRes(..) + , insert' + , delete' + , lookup' + , insertNewKey + , insertKeyExists + , deleteKeyExists + , insertModifying + , ptrEq + , adjust# ) where #if __GLASGOW_HASKELL__ < 710 @@ -129,6 +148,16 @@ import qualified Data.Hashable.Lifted as H #endif +#if __GLASGOW_HASKELL__ >= 802 +import GHC.Exts (TYPE, Int (..), Int#) +#endif + +#if MIN_VERSION_base(4,8,0) +import Data.Functor.Identity (Identity (..)) +#endif +import Control.Applicative (Const (..)) +import Data.Coerce (coerce) + -- | A set of values. A set cannot contain duplicate values. ------------------------------------------------------------------------ @@ -238,21 +267,41 @@ instance Traversable (HashMap k) where traverse f = traverseWithKey (const f) + {-# INLINABLE traverse #-} #if MIN_VERSION_base(4,9,0) instance Eq2 HashMap where - liftEq2 = equal + liftEq2 = equal2 instance Eq k => Eq1 (HashMap k) where - liftEq = equal (==) + liftEq = equal1 #endif instance (Eq k, Eq v) => Eq (HashMap k v) where - (==) = equal (==) (==) + (==) = equal1 (==) -equal :: (k -> k' -> Bool) -> (v -> v' -> Bool) +-- We rely on there being no Empty constructors in the tree! +-- This ensures that two equal HashMaps will have the same +-- shape, modulo the order of entries in Collisions. +equal1 :: Eq k + => (v -> v' -> Bool) + -> HashMap k v -> HashMap k v' -> Bool +equal1 eq = go + where + go Empty Empty = True + go (BitmapIndexed bm1 ary1) (BitmapIndexed bm2 ary2) + = bm1 == bm2 && A.sameArray1 go ary1 ary2 + go (Leaf h1 l1) (Leaf h2 l2) = h1 == h2 && leafEq l1 l2 + go (Full ary1) (Full ary2) = A.sameArray1 go ary1 ary2 + go (Collision h1 ary1) (Collision h2 ary2) + = h1 == h2 && isPermutationBy leafEq (A.toList ary1) (A.toList ary2) + go _ _ = False + + leafEq (L k1 v1) (L k2 v2) = k1 == k2 && eq v1 v2 + +equal2 :: (k -> k' -> Bool) -> (v -> v' -> Bool) -> HashMap k v -> HashMap k' v' -> Bool -equal eqk eqv t1 t2 = go (toList' t1 []) (toList' t2 []) +equal2 eqk eqv t1 t2 = go (toList' t1 []) (toList' t2 []) where -- If the two trees are the same, then their lists of 'Leaf's and -- 'Collision's read from left to right should be the same (modulo the @@ -311,8 +360,8 @@ leafCompare (L k v) (L k' v') = cmpk k k' `mappend` cmpv v v' -- Same as 'equal' but doesn't compare the values. -equalKeys :: (k -> k' -> Bool) -> HashMap k v -> HashMap k' v' -> Bool -equalKeys eq t1 t2 = go (toList' t1 []) (toList' t2 []) +equalKeys1 :: (k -> k' -> Bool) -> HashMap k v -> HashMap k' v' -> Bool +equalKeys1 eq t1 t2 = go (toList' t1 []) (toList' t2 []) where go (Leaf k1 l1 : tl1) (Leaf k2 l2 : tl2) | k1 == k2 && leafEq l1 l2 @@ -326,6 +375,22 @@ leafEq (L k _) (L k' _) = eq k k' +-- Same as 'equal1' but doesn't compare the values. +equalKeys :: Eq k => HashMap k v -> HashMap k v' -> Bool +equalKeys = go + where + go :: Eq k => HashMap k v -> HashMap k v' -> Bool + go Empty Empty = True + go (BitmapIndexed bm1 ary1) (BitmapIndexed bm2 ary2) + = bm1 == bm2 && A.sameArray1 go ary1 ary2 + go (Leaf h1 l1) (Leaf h2 l2) = h1 == h2 && leafEq l1 l2 + go (Full ary1) (Full ary2) = A.sameArray1 go ary1 ary2 + go (Collision h1 ary1) (Collision h2 ary2) + = h1 == h2 && isPermutationBy leafEq (A.toList ary1) (A.toList ary2) + go _ _ = False + + leafEq (L k1 _) (L k2 _) = k1 == k2 + #if MIN_VERSION_hashable(1,2,5) instance H.Hashable2 HashMap where liftHashWithSalt2 hk hv salt hm = go salt (toList' hm []) @@ -431,22 +496,122 @@ -- | /O(log n)/ Return the value to which the specified key is mapped, -- or 'Nothing' if this map contains no mapping for the key. lookup :: (Eq k, Hashable k) => k -> HashMap k v -> Maybe v -lookup k0 m0 = go h0 k0 0 m0 +#if __GLASGOW_HASKELL__ >= 802 +-- GHC does not yet perform a worker-wrapper transformation on +-- unboxed sums automatically. That seems likely to happen at some +-- point (possibly as early as GHC 8.6) but for now we do it manually. +lookup k m = case lookup# k m of + (# (# #) | #) -> Nothing + (# | a #) -> Just a +{-# INLINE lookup #-} + +lookup# :: (Eq k, Hashable k) => k -> HashMap k v -> (# (# #) | v #) +lookup# k m = lookupCont (\_ -> (# (# #) | #)) (\v _i -> (# | v #)) (hash k) k m +{-# INLINABLE lookup# #-} + +#else + +lookup k m = lookupCont (\_ -> Nothing) (\v _i -> Just v) (hash k) k m +{-# INLINABLE lookup #-} +#endif + +-- | lookup' is a version of lookup that takes the hash separately. +-- It is used to implement alterF. +lookup' :: Eq k => Hash -> k -> HashMap k v -> Maybe v +#if __GLASGOW_HASKELL__ >= 802 +-- GHC does not yet perform a worker-wrapper transformation on +-- unboxed sums automatically. That seems likely to happen at some +-- point (possibly as early as GHC 8.6) but for now we do it manually. +-- lookup' would probably prefer to be implemented in terms of its own +-- lookup'#, but it's not important enough and we don't want too much +-- code. +lookup' h k m = case lookupRecordCollision# h k m of + (# (# #) | #) -> Nothing + (# | (# a, _i #) #) -> Just a +{-# INLINE lookup' #-} +#else +lookup' h k m = lookupCont (\_ -> Nothing) (\v _i -> Just v) h k m +{-# INLINABLE lookup' #-} +#endif + +-- The result of a lookup, keeping track of if a hash collision occured. +-- If a collision did not occur then it will have the Int value (-1). +data LookupRes a = Absent | Present a !Int + +-- Internal helper for lookup. This version takes the precomputed hash so +-- that functions that make multiple calls to lookup and related functions +-- (insert, delete) only need to calculate the hash once. +-- +-- It is used by 'alterF' so that hash computation and key comparison only needs +-- to be performed once. With this information you can use the more optimized +-- versions of insert ('insertNewKey', 'insertKeyExists') and delete +-- ('deleteKeyExists') +-- +-- Outcomes: +-- Key not in map => Absent +-- Key in map, no collision => Present v (-1) +-- Key in map, collision => Present v position +lookupRecordCollision :: Eq k => Hash -> k -> HashMap k v -> LookupRes v +#if __GLASGOW_HASKELL__ >= 802 +lookupRecordCollision h k m = case lookupRecordCollision# h k m of + (# (# #) | #) -> Absent + (# | (# a, i #) #) -> Present a (I# i) -- GHC will eliminate the I# +{-# INLINE lookupRecordCollision #-} + +-- Why do we produce an Int# instead of an Int? Unfortunately, GHC is not +-- yet any good at unboxing things *inside* products, let alone sums. That +-- may be changing in GHC 8.6 or so (there is some work in progress), but +-- for now we use Int# explicitly here. We don't need to push the Int# +-- into lookupCont because inlining takes care of that. +lookupRecordCollision# :: Eq k => Hash -> k -> HashMap k v -> (# (# #) | (# v, Int# #) #) +lookupRecordCollision# h k m = + lookupCont (\_ -> (# (# #) | #)) (\v (I# i) -> (# | (# v, i #) #)) h k m +-- INLINABLE to specialize to the Eq instance. +{-# INLINABLE lookupRecordCollision# #-} + +#else /* GHC < 8.2 so there are no unboxed sums */ + +lookupRecordCollision h k m = lookupCont (\_ -> Absent) Present h k m +{-# INLINABLE lookupRecordCollision #-} +#endif + +-- A two-continuation version of lookupRecordCollision. This lets us +-- share source code between lookup and lookupRecordCollision without +-- risking any performance degradation. +-- +-- The absent continuation has type @((# #) -> r)@ instead of just @r@ +-- so we can be representation-polymorphic in the result type. Since +-- this whole thing is always inlined, we don't have to worry about +-- any extra CPS overhead. +lookupCont :: +#if __GLASGOW_HASKELL__ >= 802 + forall rep (r :: TYPE rep) k v. +#else + forall r k v. +#endif + Eq k + => ((# #) -> r) -- Absent continuation + -> (v -> Int -> r) -- Present continuation + -> Hash -- The hash of the key + -> k -> HashMap k v -> r +lookupCont absent present !h0 !k0 !m0 = go h0 k0 0 m0 where - h0 = hash k0 - go !_ !_ !_ Empty = Nothing + go :: Eq k => Hash -> k -> Int -> HashMap k v -> r + go !_ !_ !_ Empty = absent (# #) go h k _ (Leaf hx (L kx x)) - | h == hx && k == kx = Just x -- TODO: Split test in two - | otherwise = Nothing + | h == hx && k == kx = present x (-1) + | otherwise = absent (# #) go h k s (BitmapIndexed b v) - | b .&. m == 0 = Nothing - | otherwise = go h k (s+bitsPerSubkey) (A.index v (sparseIndex b m)) + | b .&. m == 0 = absent (# #) + | otherwise = + go h k (s+bitsPerSubkey) (A.index v (sparseIndex b m)) where m = mask h s - go h k s (Full v) = go h k (s+bitsPerSubkey) (A.index v (index h s)) + go h k 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 #-} + | h == hx = lookupInArrayCont absent present k v + | otherwise = absent (# #) +{-# INLINE lookupCont #-} -- | /O(log n)/ Return the value to which the specified key is mapped, -- or the default value if this map contains no mapping for the key. @@ -470,7 +635,7 @@ -- | Create a 'Collision' value with two 'Leaf' values. collision :: Hash -> Leaf k v -> Leaf k v -> HashMap k v -collision h e1 e2 = +collision h !e1 !e2 = let v = A.run $ do mary <- A.new 2 e1 A.write mary 1 e2 return mary @@ -488,9 +653,12 @@ -- key in this map. If this map previously contained a mapping for -- the key, the old value is replaced. insert :: (Eq k, Hashable k) => k -> v -> HashMap k v -> HashMap k v -insert k0 v0 m0 = go h0 k0 v0 0 m0 +insert k v m = insert' (hash k) k v m +{-# INLINABLE insert #-} + +insert' :: Eq k => Hash -> k -> v -> HashMap k v -> HashMap k v +insert' h0 k0 v0 m0 = go h0 k0 v0 0 m0 where - 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 @@ -521,8 +689,96 @@ 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 #-} +{-# INLINABLE insert' #-} +-- Insert optimized for the case when we know the key is not in the map. +-- +-- It is only valid to call this when the key does not exist in the map. +-- +-- We can skip: +-- - the key equality check on a Leaf +-- - check for its existence in the array for a hash collision +insertNewKey :: Hash -> k -> v -> HashMap k v -> HashMap k v +insertNewKey !h0 !k0 x0 !m0 = go h0 k0 x0 0 m0 + where + go !h !k x !_ Empty = Leaf h (L k x) + go h k x s (Leaf hy l@(L ky y)) + | hy == h = collision h l (L k x) + | otherwise = runST (two s h k x hy ky y) + go h k x s (BitmapIndexed b ary) + | b .&. m == 0 = + let !ary' = A.insert ary i $! Leaf h (L k x) + in bitmapIndexedOrFull (b .|. m) ary' + | otherwise = + let !st = A.index ary i + !st' = go h k x (s+bitsPerSubkey) st + in BitmapIndexed b (A.update ary i st') + where m = mask h s + i = sparseIndex b m + go h k x s (Full ary) = + let !st = A.index ary i + !st' = go h k x (s+bitsPerSubkey) st + in Full (update16 ary i st') + where i = index h s + go h k x s t@(Collision hy v) + | h == hy = Collision h (snocNewLeaf (L k x) v) + | otherwise = + go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) + where + snocNewLeaf :: Leaf k v -> A.Array (Leaf k v) -> A.Array (Leaf k v) + snocNewLeaf leaf ary = A.run $ do + let n = A.length ary + mary <- A.new_ (n + 1) + A.copy ary 0 mary 0 n + A.write mary n leaf + return mary +{-# NOINLINE insertNewKey #-} + + +-- Insert optimized for the case when we know the key is in the map. +-- +-- It is only valid to call this when the key exists in the map and you know the +-- hash collision position if there was one. This information can be obtained +-- from 'lookupRecordCollision'. If there is no collision pass (-1) as collPos +-- (first argument). +-- +-- We can skip the key equality check on a Leaf because we know the leaf must be +-- for this key. +insertKeyExists :: Int -> Hash -> k -> v -> HashMap k v -> HashMap k v +insertKeyExists !collPos0 !h0 !k0 x0 !m0 = go collPos0 h0 k0 x0 0 m0 + where + go !_collPos !h !k x !_s (Leaf _hy _kx) + = Leaf h (L k x) + go collPos h k x s (BitmapIndexed b ary) + | b .&. m == 0 = + let !ary' = A.insert ary i $ Leaf h (L k x) + in bitmapIndexedOrFull (b .|. m) ary' + | otherwise = + let !st = A.index ary i + !st' = go collPos h k x (s+bitsPerSubkey) st + in BitmapIndexed b (A.update ary i st') + where m = mask h s + i = sparseIndex b m + go collPos h k x s (Full ary) = + let !st = A.index ary i + !st' = go collPos h k x (s+bitsPerSubkey) st + in Full (update16 ary i st') + where i = index h s + go collPos h k x _s (Collision _hy v) + | collPos >= 0 = Collision h (setAtPosition collPos k x v) + | otherwise = Empty -- error "Internal error: go {collPos negative}" + go _ _ _ _ _ Empty = Empty -- error "Internal error: go Empty" + +{-# NOINLINE insertKeyExists #-} + +-- Replace the ith Leaf with Leaf k v. +-- +-- This does not check that @i@ is within bounds of the array. +setAtPosition :: Int -> k -> v -> A.Array (Leaf k v) -> A.Array (Leaf k v) +setAtPosition i k x ary = A.update ary i (L k x) +{-# INLINE setAtPosition #-} + + -- | In-place update version of 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) @@ -588,37 +844,79 @@ -- > where f new old = new + old 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 +-- We're not going to worry about allocating a function closure +-- to pass to insertModifying. See comments at 'adjust'. +insertWith f k new m = insertModifying new (\old -> (# f new old #)) k m +{-# INLINE insertWith #-} + +-- | @insertModifying@ is a lot like insertWith; we use it to implement alterF. +-- It takes a value to insert when the key is absent and a function +-- to apply to calculate a new value when the key is present. Thanks +-- to the unboxed unary tuple, we avoid introducing any unnecessary +-- thunks in the tree. +insertModifying :: (Eq k, Hashable k) => v -> (v -> (# v #)) -> k -> HashMap k v + -> HashMap k v +insertModifying x f k0 m0 = go h0 k0 0 m0 where - h0 = hash k0 - go !h !k x !_ Empty = Leaf h (L k x) - go h k x s (Leaf hy l@(L ky y)) + !h0 = hash k0 + go !h !k !_ Empty = Leaf h (L k x) + go h k s t@(Leaf hy l@(L ky y)) | hy == h = if ky == k - then Leaf h (L k (f x y)) + then case f y of + (# v' #) | ptrEq y v' -> t + | otherwise -> Leaf h (L k (v')) else collision h l (L k x) | otherwise = runST (two s h k x hy ky y) - go h k x s (BitmapIndexed b ary) + go h k s t@(BitmapIndexed b ary) | b .&. m == 0 = let ary' = A.insert ary i $! Leaf h (L k x) in bitmapIndexedOrFull (b .|. m) ary' | otherwise = - let st = A.index ary i - st' = go h k x (s+bitsPerSubkey) st - ary' = A.update ary i $! st' - in BitmapIndexed b ary' + let !st = A.index ary i + !st' = go h k (s+bitsPerSubkey) st + ary' = A.update ary i $! st' + in if ptrEq st st' + then t + else BitmapIndexed b ary' where m = mask h s i = sparseIndex b m - go h k x s (Full ary) = - let st = A.index ary i - st' = go h k x (s+bitsPerSubkey) st + go h k s t@(Full ary) = + let !st = A.index ary i + !st' = go h k (s+bitsPerSubkey) st ary' = update16 ary i $! st' - in Full ary' + in if ptrEq st st' + then t + else 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 #-} + go h k s t@(Collision hy v) + | h == hy = + let !v' = insertModifyingArr x f k v + in if A.unsafeSameArray v v' + then t + else Collision h v' + | otherwise = go h k s $ BitmapIndexed (mask hy s) (A.singleton t) +{-# INLINABLE insertModifying #-} +-- Like insertModifying for arrays; used to implement insertModifying +insertModifyingArr :: Eq k => v -> (v -> (# v #)) -> k -> A.Array (Leaf k v) + -> A.Array (Leaf k v) +insertModifyingArr x f k0 ary0 = go k0 ary0 0 (A.length ary0) + where + go !k !ary !i !n + | i >= n = A.run $ do + -- Not found, append to the end. + mary <- A.new_ (n + 1) + A.copy ary 0 mary 0 n + A.write mary n (L k x) + return mary + | otherwise = case A.index ary i of + (L kx y) | k == kx -> case f y of + (# y' #) -> if ptrEq y y' + then ary + else A.update ary i (L k y') + | otherwise -> go k ary (i+1) n +{-# INLINE insertModifyingArr #-} + -- | In-place update version of insertWith unsafeInsertWith :: forall k v. (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v @@ -658,9 +956,12 @@ -- | /O(log n)/ Remove the mapping for the specified key from this map -- if present. delete :: (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v -delete k0 m0 = go h0 k0 0 m0 +delete k m = delete' (hash k) k m +{-# INLINABLE delete #-} + +delete' :: Eq k => Hash -> k -> HashMap k v -> HashMap k v +delete' h0 k0 m0 = go h0 k0 0 m0 where - h0 = hash k0 go !_ !_ !_ Empty = Empty go h k _ t@(Leaf hy (L ky _)) | hy == h && ky == k = Empty @@ -708,36 +1009,106 @@ | otherwise -> Collision h (A.delete v i) Nothing -> t | otherwise = t -{-# INLINABLE delete #-} +{-# INLINABLE delete' #-} +-- | Delete optimized for the case when we know the key is in the map. +-- +-- It is only valid to call this when the key exists in the map and you know the +-- hash collision position if there was one. This information can be obtained +-- from 'lookupRecordCollision'. If there is no collision pass (-1) as collPos. +-- +-- We can skip: +-- - the key equality check on the leaf, if we reach a leaf it must be the key +deleteKeyExists :: Int -> Hash -> k -> HashMap k v -> HashMap k v +deleteKeyExists !collPos0 !h0 !k0 !m0 = go collPos0 h0 k0 0 m0 + where + go :: Int -> Hash -> k -> Int -> HashMap k v -> HashMap k v + go !_collPos !_h !_k !_s (Leaf _ _) = Empty + go collPos h k s (BitmapIndexed b ary) = + let !st = A.index ary i + !st' = go collPos h k (s+bitsPerSubkey) st + in case st' of + Empty | A.length ary == 1 -> Empty + | A.length ary == 2 -> + case (i, A.index ary 0, A.index ary 1) of + (0, _, l) | isLeafOrCollision l -> l + (1, l, _) | isLeafOrCollision l -> l + _ -> bIndexed + | otherwise -> bIndexed + where + bIndexed = BitmapIndexed (b .&. complement m) (A.delete ary i) + l | isLeafOrCollision l && A.length ary == 1 -> l + _ -> BitmapIndexed b (A.update ary i st') + where m = mask h s + i = sparseIndex b m + go collPos h k s (Full ary) = + let !st = A.index ary i + !st' = go collPos h k (s+bitsPerSubkey) st + in case st' of + Empty -> + let ary' = A.delete ary i + bm = fullNodeMask .&. complement (1 `unsafeShiftL` i) + in BitmapIndexed bm ary' + _ -> Full (A.update ary i st') + where i = index h s + go collPos h _ _ (Collision _hy v) + | A.length v == 2 + = if collPos == 0 + then Leaf h (A.index v 1) + else Leaf h (A.index v 0) + | otherwise = Collision h (A.delete v collPos) + go !_ !_ !_ !_ Empty = Empty -- error "Internal error: deleteKeyExists empty" +{-# NOINLINE deleteKeyExists #-} + -- | /O(log n)/ Adjust the value tied to a given key in this map only -- if it is present. Otherwise, leave the map alone. adjust :: (Eq k, Hashable k) => (v -> v) -> k -> HashMap k v -> HashMap k v -adjust f k0 m0 = go h0 k0 0 m0 +-- This operation really likes to leak memory, so using this +-- indirect implementation shouldn't hurt much. Furthermore, it allows +-- GHC to avoid a leak when the function is lazy. In particular, +-- +-- adjust (const x) k m +-- ==> adjust# (\v -> (# const x v #)) k m +-- ==> adjust# (\_ -> (# x #)) k m +adjust f k m = adjust# (\v -> (# f v #)) k m +{-# INLINE adjust #-} + +-- | Much like 'adjust', but not inherently leaky. +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)) + | hy == h && ky == k = case f y of + (# y' #) | ptrEq y y' -> t + | otherwise -> Leaf h (L k y') | otherwise = t go h k s t@(BitmapIndexed b ary) | b .&. m == 0 = t - | otherwise = let st = A.index ary i - st' = go h k (s+bitsPerSubkey) st + | otherwise = let !st = A.index ary i + !st' = go h k (s+bitsPerSubkey) st ary' = A.update ary i $! st' - in BitmapIndexed b ary' + in if ptrEq st st' + then t + else BitmapIndexed b ary' where m = mask h s i = sparseIndex b m - go h k s (Full ary) = + go h k s t@(Full ary) = let i = index h s - st = A.index ary i - st' = go h k (s+bitsPerSubkey) st + !st = A.index ary i + !st' = go h k (s+bitsPerSubkey) st ary' = update16 ary i $! st' - in Full ary' + in if ptrEq st st' + then t + else Full ary' go h k _ t@(Collision hy v) - | h == hy = Collision h (updateWith f k v) + | h == hy = let !v' = updateWith# f k v + in if A.unsafeSameArray v v' + then t + else Collision h v' | otherwise = t -{-# INLINABLE adjust #-} +{-# INLINABLE adjust# #-} -- | /O(log n)/ The expression (@'update' f k map@) updates the value @x@ at @k@, -- (if it is in the map). If (f k x) is @'Nothing', the element is deleted. @@ -751,12 +1122,162 @@ -- absence thereof. @alter@ can be used to insert, delete, or update a value in a -- map. In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. alter :: (Eq k, Hashable k) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v +-- TODO(m-renaud): Consider using specialized insert and delete for alter. alter f k m = case f (lookup k m) of Nothing -> delete k m Just v -> insert k v m {-# INLINABLE alter #-} +-- | /O(log n)/ The expression (@'alterF' f k map@) alters the value @x@ at +-- @k@, or absence thereof. @alterF@ can be used to insert, delete, or update +-- a value in a map. +-- +-- Note: 'alterF' is a flipped version of the 'at' combinator from +-- . +-- +-- @since 0.2.9 +alterF :: (Functor f, Eq k, Hashable k) + => (Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v) +-- We only calculate the hash once, but unless this is rewritten +-- by rules we may test for key equality multiple times. +-- We force the value of the map for consistency with the rewritten +-- version; otherwise someone could tell the difference using a lazy +-- @f@ and a functor that is similar to Const but not actually Const. +alterF f = \ !k !m -> + let + !h = hash k + mv = lookup' h k m + in (<$> f mv) $ \fres -> + case fres of + Nothing -> delete' h k m + Just v' -> insert' h k v' m + +-- We unconditionally rewrite alterF in RULES, but we expose an +-- unfolding just in case it's used in some way that prevents the +-- rule from firing. +{-# INLINABLE [0] alterF #-} + +#if MIN_VERSION_base(4,8,0) +-- This is just a bottom value. See the comment on the "alterFWeird" +-- rule. +test_bottom :: a +test_bottom = error "Data.HashMap.alterF internal error: hit test_bottom" + +-- We use this as an error result in RULES to ensure we don't get +-- any useless CallStack nonsense. +bogus# :: (# #) -> (# a #) +bogus# _ = error "Data.HashMap.alterF internal error: hit bogus#" + +{-# RULES +-- We probe the behavior of @f@ by applying it to Nothing and to +-- Just test_bottom. Based on the results, and how they relate to +-- each other, we choose the best implementation. + +"alterFWeird" forall f. alterF f = + alterFWeird (f Nothing) (f (Just test_bottom)) f + +-- This rule covers situations where alterF is used to simply insert or +-- delete in Identity (most likely via Control.Lens.At). We recognize here +-- (through the repeated @x@ on the LHS) that +-- +-- @f Nothing = f (Just bottom)@, +-- +-- which guarantees that @f@ doesn't care what its argument is, so +-- we don't have to either. +-- +-- Why only Identity? A variant of this rule is actually valid regardless of +-- the functor, but for some functors (e.g., []), it can lead to the +-- same keys being compared multiple times, which is bad if they're +-- ugly things like strings. This is unfortunate, since the rule is likely +-- a good idea for almost all realistic uses, but I don't like nasty +-- edge cases. +"alterFconstant" forall (f :: Maybe a -> Identity (Maybe a)) x. + alterFWeird x x f = \ !k !m -> + Identity (case runIdentity x of {Nothing -> delete k m; Just a -> insert k a m}) + +-- This rule handles the case where 'alterF' is used to do 'insertWith'-like +-- things. Whenever possible, GHC will get rid of the Maybe nonsense for us. +-- We delay this rule to stage 1 so alterFconstant has a chance to fire. +"alterFinsertWith" [1] forall (f :: Maybe a -> Identity (Maybe a)) x y. + alterFWeird (coerce (Just x)) (coerce (Just y)) f = + coerce (insertModifying x (\mold -> case runIdentity (f (Just mold)) of + Nothing -> bogus# (# #) + Just new -> (# new #))) + +-- Handle the case where someone uses 'alterF' instead of 'adjust'. This +-- rule is kind of picky; it will only work if the function doesn't +-- do anything between case matching on the Maybe and producing a result. +"alterFadjust" forall (f :: Maybe a -> Identity (Maybe a)) _y. + alterFWeird (coerce Nothing) (coerce (Just _y)) f = + coerce (adjust# (\x -> case runIdentity (f (Just x)) of + Just x' -> (# x' #) + Nothing -> bogus# (# #))) + +-- The simple specialization to Const; in this case we can look up +-- the key without caring what position it's in. This is only a tiny +-- optimization. +"alterFlookup" forall _ign1 _ign2 (f :: Maybe a -> Const r (Maybe a)). + alterFWeird _ign1 _ign2 f = \ !k !m -> Const (getConst (f (lookup k m))) + #-} + +-- This is a very unsafe version of alterF used for RULES. When calling +-- alterFWeird x y f, the following *must* hold: +-- +-- x = f Nothing +-- y = f (Just _|_) +-- +-- Failure to abide by these laws will make demons come out of your nose. +alterFWeird + :: (Functor f, Eq k, Hashable k) + => f (Maybe v) + -> f (Maybe v) + -> (Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v) +alterFWeird _ _ f = alterFEager f +{-# INLINE [0] alterFWeird #-} + +-- | This is the default version of alterF that we use in most non-trivial +-- cases. It's called "eager" because it looks up the given key in the map +-- eagerly, whether or not the given function requires that information. +alterFEager :: (Functor f, Eq k, Hashable k) + => (Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v) +alterFEager f !k m = (<$> f mv) $ \fres -> + case fres of + + ------------------------------ + -- Delete the key from the map. + Nothing -> case lookupRes of + + -- Key did not exist in the map to begin with, no-op + Absent -> m + + -- Key did exist + Present _ collPos -> deleteKeyExists collPos h k m + + ------------------------------ + -- Update value + Just v' -> case lookupRes of + + -- Key did not exist before, insert v' under a new key + Absent -> insertNewKey h k v' m + + -- Key existed before + Present v collPos -> + if v `ptrEq` v' + -- If the value is identical, no-op + then m + -- If the value changed, update the value. + else insertKeyExists collPos h k v' m + + where !h = hash k + !lookupRes = lookupRecordCollision h k m + !mv = case lookupRes of + Absent -> Nothing + Present v _ -> Just v +{-# INLINABLE alterFEager #-} +#endif + + ------------------------------------------------------------------------ -- * Combine @@ -874,7 +1395,9 @@ | 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) + x1 <- A.indexM ary1 i1 + x2 <- A.indexM ary2 i2 + A.write mary i $! f x1 x2 go (i+1) (i1+1) (i2+1) (m `unsafeShiftL` 1) | b1 .&. m /= 0 = do A.write mary i =<< A.indexM ary1 i1 @@ -905,8 +1428,10 @@ 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 (BitmapIndexed b ary) = BitmapIndexed b $ A.map go ary + go (Full ary) = Full $ A.map go ary + -- Why map strictly over collision arrays? Because there's no + -- point suspending the O(1) work this does for each leaf. go (Collision h ary) = Collision h $ A.map' (\ (L k v) -> L k (f k v)) ary {-# INLINE mapWithKey #-} @@ -919,10 +1444,17 @@ -- TODO: We should be able to use mutation to create the new -- 'HashMap'. --- | /O(n)/ Transform this map by accumulating an Applicative result --- from every value. -traverseWithKey :: Applicative f => (k -> v1 -> f v2) -> HashMap k v1 - -> f (HashMap k v2) +-- | /O(n)/ Perform an 'Applicative' action for each key-value pair +-- in a 'HashMap' and produce a 'HashMap' of all the results. +-- +-- Note: the order in which the actions occur is unspecified. In particular, +-- when the map contains hash collisions, the order in which the actions +-- associated with the keys involved will depend in an unspecified way on +-- their insertion order. +traverseWithKey + :: Applicative f + => (k -> v1 -> f v2) + -> HashMap k v1 -> f (HashMap k v2) traverseWithKey f = go where go Empty = pure Empty @@ -930,7 +1462,7 @@ 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 + Collision h <$> A.traverse' (\ (L k v) -> L k <$> f k v) ary {-# INLINE traverseWithKey #-} ------------------------------------------------------------------------ @@ -998,8 +1530,8 @@ -- | /O(n)/ Reduce this map by applying a binary operator to all -- elements, using the given starting value (typically the -- left-identity of the operator). Each application of the operator --- is evaluated before before using the result in the next --- application. This function is strict in the starting value. +-- is evaluated before using the result in the next application. +-- This function is strict in the starting value. foldl' :: (a -> v -> a) -> a -> HashMap k v -> a foldl' f = foldlWithKey' (\ z _ v -> f z v) {-# INLINE foldl' #-} @@ -1007,8 +1539,8 @@ -- | /O(n)/ Reduce this map by applying a binary operator to all -- elements, using the given starting value (typically the -- left-identity of the operator). Each application of the operator --- is evaluated before before using the result in the next --- application. This function is strict in the starting value. +-- is evaluated before using the result in the next application. +-- This function is strict in the starting value. foldlWithKey' :: (a -> k -> v -> a) -> a -> HashMap k v -> a foldlWithKey' f = go where @@ -1042,14 +1574,6 @@ ------------------------------------------------------------------------ -- * Filter --- | Create a new array of the @n@ first elements of @mary@. -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 #-} - -- | /O(n)/ Transform this map by applying a function to every value -- and retaining only some of them. mapMaybeWithKey :: (k -> v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2 @@ -1112,9 +1636,9 @@ ch <- A.read mary 0 case ch of t | isLeafOrCollision t -> return t - _ -> BitmapIndexed b <$> trim mary 1 + _ -> BitmapIndexed b <$> A.trim mary 1 _ -> do - ary2 <- trim mary j + ary2 <- A.trim mary j return $! if j == maxChildren then Full ary2 else BitmapIndexed b ary2 @@ -1141,9 +1665,9 @@ return $! Leaf h l _ | i == j -> do ary2 <- A.unsafeFreeze mary return $! Collision h ary2 - | otherwise -> do ary2 <- trim mary j + | otherwise -> do ary2 <- A.trim mary j return $! Collision h ary2 - | Just el <- onColl (A.index ary i) + | 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 #-} @@ -1196,18 +1720,25 @@ ------------------------------------------------------------------------ -- Array operations --- | /O(n)/ Lookup the value associated with the given key in this --- array. Returns 'Nothing' if the key wasn't found. -lookupInArray :: Eq k => k -> A.Array (Leaf k v) -> Maybe v -lookupInArray k0 ary0 = go k0 ary0 0 (A.length ary0) +-- | /O(n)/ Look up the value associated with the given key in an +-- array. +lookupInArrayCont :: +#if __GLASGOW_HASKELL__ >= 802 + forall rep (r :: TYPE rep) k v. +#else + forall r k v. +#endif + Eq k => ((# #) -> r) -> (v -> Int -> r) -> k -> A.Array (Leaf k v) -> r +lookupInArrayCont absent present k0 ary0 = go k0 ary0 0 (A.length ary0) where + go :: Eq k => k -> A.Array (Leaf k v) -> Int -> Int -> r go !k !ary !i !n - | i >= n = Nothing + | i >= n = absent (# #) | otherwise = case A.index ary i of (L kx v) - | k == kx -> Just v + | k == kx -> present v i | otherwise -> go k ary (i+1) n -{-# INLINABLE lookupInArray #-} +{-# INLINE lookupInArrayCont #-} -- | /O(n)/ Lookup the value associated with the given key in this -- array. Returns 'Nothing' if the key wasn't found. @@ -1222,15 +1753,18 @@ | 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) +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)) + (L kx y) | k == kx -> case f y of + (# y' #) + | ptrEq y y' -> ary + | otherwise -> A.update ary i (L k y') | otherwise -> go k ary (i+1) n -{-# INLINABLE updateWith #-} +{-# INLINABLE updateWith# #-} updateOrSnocWith :: Eq k => (v -> v -> v) -> k -> v -> A.Array (Leaf k v) -> A.Array (Leaf k v) @@ -1259,8 +1793,11 @@ 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 + -- TODO: instead of mapping and then folding, should we traverse? + -- We'll have to be careful to avoid allocating pairs or similar. + -- first: look up the position of each element of ary2 in ary1 - let indices = A.map (\(L k _) -> indexOf k ary1) ary2 + let indices = A.map' (\(L k _) -> indexOf k ary1) ary2 -- that tells us how large the overlap is: -- count number of Nothing constructors let nOnly2 = A.foldl' (\n -> maybe (n+1) (const n)) 0 indices @@ -1303,7 +1840,9 @@ -- | /O(n)/ Update the element at the given position in this array, by applying a function to it. update16With' :: A.Array e -> Int -> (e -> e) -> A.Array e -update16With' ary idx f = update16 ary idx $! f (A.index ary idx) +update16With' ary idx f + | (# x #) <- A.index# ary idx + = update16 ary idx $! f x {-# INLINE update16With' #-} -- | Unsafely clone an array of 16 elements. The length of the input diff --git a/Data/HashMap/Lazy.hs b/Data/HashMap/Lazy.hs --- a/Data/HashMap/Lazy.hs +++ b/Data/HashMap/Lazy.hs @@ -34,10 +34,10 @@ , singleton -- * Basic interface - , HM.null + , null , size , member - , HM.lookup + , lookup , lookupDefault , (!) , insert @@ -46,6 +46,7 @@ , adjust , update , alter + , alterF -- * Combine -- ** Union @@ -55,7 +56,7 @@ , unions -- * Transformations - , HM.map + , map , mapWithKey , traverseWithKey @@ -69,11 +70,11 @@ -- * Folds , foldl' , foldlWithKey' - , HM.foldr + , foldr , foldrWithKey -- * Filter - , HM.filter + , filter , filterWithKey , mapMaybe , mapMaybeWithKey @@ -86,9 +87,14 @@ , toList , fromList , fromListWith + + -- ** HashSets + , HS.keysSet ) where import Data.HashMap.Base as HM +import qualified Data.HashSet.Base as HS +import Prelude () -- $strictness -- diff --git a/Data/HashMap/Strict.hs b/Data/HashMap/Strict.hs --- a/Data/HashMap/Strict.hs +++ b/Data/HashMap/Strict.hs @@ -1,5 +1,4 @@ -{-# LANGUAGE BangPatterns, CPP, PatternGuards #-} -{-# LANGUAGE Trustworthy #-} +{-# LANGUAGE Safe #-} ------------------------------------------------------------------------ -- | @@ -34,10 +33,10 @@ , singleton -- * Basic interface - , HM.null + , null , size - , HM.member - , HM.lookup + , member + , lookup , lookupDefault , (!) , insert @@ -46,6 +45,7 @@ , adjust , update , alter + , alterF -- * Combine -- ** Union @@ -69,11 +69,11 @@ -- * Folds , foldl' , foldlWithKey' - , HM.foldr + , foldr , foldrWithKey -- * Filter - , HM.filter + , filter , filterWithKey , mapMaybe , mapMaybeWithKey @@ -86,426 +86,11 @@ , toList , fromList , fromListWith - ) where -import Data.Bits ((.&.), (.|.)) -import qualified Data.List as L -import Data.Hashable (Hashable) -import Prelude hiding (map) - -import qualified Data.HashMap.Array as A -import qualified Data.HashMap.Base as HM -import Data.HashMap.Base hiding ( - alter, adjust, fromList, fromListWith, insert, insertWith, differenceWith, - intersectionWith, intersectionWithKey, map, mapWithKey, mapMaybe, - mapMaybeWithKey, singleton, update, unionWith, unionWithKey) -import Data.HashMap.Unsafe (runST) - --- $strictness --- --- This module satisfies the following strictness properties: --- --- 1. Key arguments are evaluated to WHNF; --- --- 2. Keys and values are evaluated to WHNF before they are stored in --- the map. - ------------------------------------------------------------------------- --- * Construction - --- | /O(1)/ Construct a map with a single element. -singleton :: (Hashable k) => k -> v -> HashMap k v -singleton k !v = HM.singleton k v - ------------------------------------------------------------------------- --- * Basic interface - --- | /O(log n)/ Associate the specified value with the specified --- key in this map. If this map previously contained a mapping for --- the key, the old value is replaced. -insert :: (Eq k, Hashable k) => k -> v -> HashMap k v -> HashMap k v -insert k !v = HM.insert k v -{-# INLINABLE insert #-} - --- | /O(log n)/ Associate the value with the key in this map. If --- this map previously contained a mapping for the key, the old value --- is replaced by the result of applying the given function to the new --- and old value. Example: --- --- > insertWith f k v map --- > where f new old = new + old -insertWith :: (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v - -> HashMap k v -insertWith f k0 v0 m0 = go h0 k0 v0 0 m0 - where - h0 = hash k0 - go !h !k x !_ Empty = leaf h k x - go h k x s (Leaf hy l@(L ky y)) - | hy == h = if ky == k - then leaf h k (f x y) - else x `seq` (collision h l (L k x)) - | otherwise = x `seq` runST (two s h k x hy ky y) - go h k x s (BitmapIndexed b ary) - | b .&. m == 0 = - let ary' = A.insert ary i $! leaf h k x - in bitmapIndexedOrFull (b .|. m) ary' - | otherwise = - let st = A.index ary i - st' = go h k x (s+bitsPerSubkey) st - ary' = A.update ary i $! st' - in BitmapIndexed b ary' - where m = mask h s - i = sparseIndex b m - go h k x s (Full ary) = - let st = A.index ary i - st' = go h k x (s+bitsPerSubkey) st - ary' = update16 ary i $! st' - in Full ary' - where i = index h s - go h k x s t@(Collision hy v) - | h == hy = Collision h (updateOrSnocWith f k x v) - | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) -{-# INLINABLE insertWith #-} - --- | In-place update version of insertWith -unsafeInsertWith :: (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 !h !k x !_ Empty = return $! leaf h k x - go h k x s (Leaf hy l@(L ky y)) - | hy == h = if ky == k - then return $! leaf h k (f x y) - else do - let l' = x `seq` (L k x) - return $! collision h l l' - | otherwise = x `seq` two s h k x hy ky y - go h k x s t@(BitmapIndexed b ary) - | b .&. m == 0 = do - ary' <- A.insertM ary i $! leaf h k x - return $! bitmapIndexedOrFull (b .|. m) ary' - | otherwise = do - st <- A.indexM ary i - st' <- go h k x (s+bitsPerSubkey) st - A.unsafeUpdateM ary i st' - return t - where m = mask h s - i = sparseIndex b m - go h k x s t@(Full ary) = do - st <- A.indexM ary i - st' <- go h k x (s+bitsPerSubkey) st - A.unsafeUpdateM ary i st' - return t - where i = index h s - go h k x s t@(Collision hy v) - | h == hy = return $! Collision h (updateOrSnocWith f k x v) - | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) -{-# INLINABLE unsafeInsertWith #-} - --- | /O(log n)/ Adjust the value tied to a given key in this map only --- if it is present. Otherwise, leave the map alone. -adjust :: (Eq k, Hashable k) => (v -> v) -> k -> HashMap k v -> HashMap k v -adjust f k0 m0 = go h0 k0 0 m0 - where - h0 = hash k0 - go !_ !_ !_ Empty = Empty - go h k _ t@(Leaf hy (L ky y)) - | hy == h && ky == k = leaf h k (f y) - | otherwise = t - go h k s t@(BitmapIndexed b ary) - | b .&. m == 0 = t - | otherwise = let st = A.index ary i - st' = go h k (s+bitsPerSubkey) st - ary' = A.update ary i $! st' - in BitmapIndexed b ary' - where m = mask h s - i = sparseIndex b m - go h k s (Full ary) = - let i = index h s - st = A.index ary i - st' = go h k (s+bitsPerSubkey) st - ary' = update16 ary i $! st' - in Full ary' - go h k _ t@(Collision hy v) - | h == hy = Collision h (updateWith f k v) - | otherwise = t -{-# INLINABLE adjust #-} - --- | /O(log n)/ The expression (@'update' f k map@) updates the value @x@ at @k@, --- (if it is in the map). If (f k x) is @'Nothing', the element is deleted. --- If it is (@'Just' y), the key k is bound to the new value y. -update :: (Eq k, Hashable k) => (a -> Maybe a) -> k -> HashMap k a -> HashMap k a -update f = alter (>>= f) -{-# INLINABLE update #-} - --- | /O(log n)/ The expression (@'alter' f k map@) alters the value @x@ at @k@, or --- absence thereof. @alter@ can be used to insert, delete, or update a value in a --- map. In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. -alter :: (Eq k, Hashable k) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v -alter f k m = - case f (HM.lookup k m) of - Nothing -> delete k m - Just v -> insert k v m -{-# INLINABLE alter #-} - ------------------------------------------------------------------------- --- * Combine - --- | /O(n+m)/ The union of two maps. If a key occurs in both maps, --- the provided function (first argument) will be used to compute the result. -unionWith :: (Eq k, Hashable k) => (v -> v -> v) -> HashMap k v -> HashMap k v - -> HashMap k v -unionWith f = unionWithKey (const f) -{-# INLINE unionWith #-} - --- | /O(n+m)/ The union of two maps. If a key occurs in both maps, --- the provided function (first argument) will be used to compute the result. -unionWithKey :: (Eq k, Hashable k) => (k -> v -> v -> v) -> HashMap k v -> HashMap k v - -> HashMap k v -unionWithKey f = go 0 - where - -- empty vs. anything - go !_ t1 Empty = t1 - go _ Empty t2 = t2 - -- leaf vs. leaf - go s t1@(Leaf h1 l1@(L k1 v1)) t2@(Leaf h2 l2@(L k2 v2)) - | h1 == h2 = if k1 == k2 - then leaf h1 k1 (f k1 v1 v2) - else collision h1 l1 l2 - | otherwise = goDifferentHash s h1 h2 t1 t2 - go s t1@(Leaf h1 (L k1 v1)) t2@(Collision h2 ls2) - | h1 == h2 = Collision h1 (updateOrSnocWithKey f k1 v1 ls2) - | otherwise = goDifferentHash s h1 h2 t1 t2 - go s t1@(Collision h1 ls1) t2@(Leaf h2 (L k2 v2)) - | h1 == h2 = Collision h1 (updateOrSnocWithKey (flip . f) k2 v2 ls1) - | otherwise = goDifferentHash s h1 h2 t1 t2 - go s t1@(Collision h1 ls1) t2@(Collision h2 ls2) - | h1 == h2 = Collision h1 (updateOrConcatWithKey f ls1 ls2) - | otherwise = goDifferentHash s h1 h2 t1 t2 - -- branch vs. branch - 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' - -- leaf vs. branch - 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 #-} - ------------------------------------------------------------------------- --- * Transformations - --- | /O(n)/ Transform this map by applying a function to every value. -mapWithKey :: (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2 -mapWithKey f = go - where - go Empty = Empty - go (Leaf h (L k v)) = leaf h k (f k v) - go (BitmapIndexed b ary) = BitmapIndexed b $ A.map' go ary - go (Full ary) = Full $ A.map' go ary - go (Collision h ary) = - Collision h $ A.map' (\ (L k v) -> let !v' = f k v in L k v') ary -{-# INLINE mapWithKey #-} - --- | /O(n)/ Transform this map by applying a function to every value. -map :: (v1 -> v2) -> HashMap k v1 -> HashMap k v2 -map f = mapWithKey (const f) -{-# INLINE map #-} - - ------------------------------------------------------------------------- --- * Filter - --- | /O(n)/ Transform this map by applying a function to every value --- and retaining only some of them. -mapMaybeWithKey :: (k -> v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2 -mapMaybeWithKey f = filterMapAux onLeaf onColl - where onLeaf (Leaf h (L k v)) | Just v' <- f k v = Just (leaf h k v') - onLeaf _ = Nothing - - onColl (L k v) | Just v' <- f k v = Just (L k v') - | otherwise = Nothing -{-# INLINE mapMaybeWithKey #-} - --- | /O(n)/ Transform this map by applying a function to every value --- and retaining only some of them. -mapMaybe :: (v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2 -mapMaybe f = mapMaybeWithKey (const f) -{-# INLINE mapMaybe #-} - - --- TODO: Should we add a strict traverseWithKey? - ------------------------------------------------------------------------- --- * Difference and intersection - --- | /O(n*log m)/ Difference with a combining function. When two equal keys are --- encountered, the combining function is applied to the values of these keys. --- If it returns 'Nothing', the element is discarded (proper set difference). If --- it returns (@'Just' y@), the element is updated with a new value @y@. -differenceWith :: (Eq k, Hashable k) => (v -> w -> Maybe v) -> HashMap k v -> HashMap k w -> HashMap k v -differenceWith f a b = foldlWithKey' go empty a - where - go m k v = case HM.lookup k b of - Nothing -> insert k v m - Just w -> maybe m (\y -> insert k y m) (f v w) -{-# INLINABLE differenceWith #-} - --- | /O(n+m)/ Intersection of two maps. If a key occurs in both maps --- the provided function is used to combine the values from the two --- maps. -intersectionWith :: (Eq k, Hashable k) => (v1 -> v2 -> v3) -> HashMap k v1 - -> HashMap k v2 -> HashMap k v3 -intersectionWith f a b = foldlWithKey' go empty a - where - go m k v = case HM.lookup k b of - Just w -> insert k (f v w) m - _ -> m -{-# INLINABLE intersectionWith #-} - --- | /O(n+m)/ Intersection of two maps. If a key occurs in both maps --- the provided function is used to combine the values from the two --- maps. -intersectionWithKey :: (Eq k, Hashable k) => (k -> v1 -> v2 -> v3) - -> HashMap k v1 -> HashMap k v2 -> HashMap k v3 -intersectionWithKey f a b = foldlWithKey' go empty a - where - go m k v = case HM.lookup k b of - Just w -> insert k (f k v w) m - _ -> m -{-# INLINABLE intersectionWithKey #-} - ------------------------------------------------------------------------- --- ** Lists - --- | /O(n*log n)/ Construct a map with the supplied mappings. If the --- list contains duplicate mappings, the later mappings take --- precedence. -fromList :: (Eq k, Hashable k) => [(k, v)] -> HashMap k v -fromList = L.foldl' (\ m (k, !v) -> HM.unsafeInsert k v m) empty -{-# INLINABLE fromList #-} - --- | /O(n*log n)/ Construct a map from a list of elements. Uses --- the provided function f to merge duplicate entries (f newVal oldVal). --- --- For example: --- --- > fromListWith (+) [ (x, 1) | x <- xs ] --- --- will create a map with number of occurrences of each element in xs. --- --- > fromListWith (++) [ (k, [v]) | (k, v) <- xs ] --- --- will group all values by their keys in a list 'xs :: [(k, v)]' and --- return a 'HashMap k [v]'. -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 #-} - ------------------------------------------------------------------------- --- Array operations - -updateWith :: Eq k => (v -> v) -> k -> A.Array (Leaf k v) -> A.Array (Leaf k v) -updateWith f k0 ary0 = go k0 ary0 0 (A.length ary0) - where - go !k !ary !i !n - | i >= n = ary - | otherwise = case A.index ary i of - (L kx y) | k == kx -> let !v' = f y in A.update ary i (L k v') - | otherwise -> go k ary (i+1) n -{-# INLINABLE updateWith #-} - --- | Append the given key and value to the array. If the key is --- already present, instead update the value of the key by applying --- the given function to the new and old value (in that order). The --- value is always evaluated to WHNF before being inserted into the --- array. -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 #-} - --- | Append the given key and value to the array. If the key is --- already present, instead update the value of the key by applying --- the given function to the new and old value (in that order). The --- value is always evaluated to WHNF before being inserted into the --- array. -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 - -- Not found, append to the end. - mary <- A.new_ (n + 1) - A.copy ary 0 mary 0 n - let !l = v `seq` (L k v) - A.write mary n l - return mary - | otherwise = case A.index ary i of - (L kx y) | k == kx -> let !v' = f k v y in A.update ary i (L k v') - | otherwise -> go k v ary (i+1) n -{-# INLINABLE updateOrSnocWithKey #-} - ------------------------------------------------------------------------- --- Smart constructors --- --- These constructors make sure the value is in WHNF before it's --- inserted into the constructor. + -- ** HashSets + , HS.keysSet + ) where -leaf :: Hash -> k -> v -> HashMap k v -leaf h k !v = Leaf h (L k v) -{-# INLINE leaf #-} +import Data.HashMap.Strict.Base as HM +import qualified Data.HashSet.Base as HS +import Prelude () diff --git a/Data/HashMap/Strict/Base.hs b/Data/HashMap/Strict/Base.hs new file mode 100644 --- /dev/null +++ b/Data/HashMap/Strict/Base.hs @@ -0,0 +1,671 @@ +{-# LANGUAGE BangPatterns, CPP, PatternGuards, MagicHash, UnboxedTuples #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE Trustworthy #-} + +------------------------------------------------------------------------ +-- | +-- Module : Data.HashMap.Strict +-- Copyright : 2010-2012 Johan Tibell +-- License : BSD-style +-- Maintainer : johan.tibell@gmail.com +-- Stability : provisional +-- Portability : portable +-- +-- A map from /hashable/ keys to values. A map cannot contain +-- duplicate keys; each key can map to at most one value. A 'HashMap' +-- makes no guarantees as to the order of its elements. +-- +-- The implementation is based on /hash array mapped tries/. A +-- 'HashMap' is often faster than other tree-based set types, +-- especially when key comparison is expensive, as in the case of +-- strings. +-- +-- Many operations have a average-case complexity of /O(log n)/. The +-- implementation uses a large base (i.e. 16) so in practice these +-- operations are constant time. +module Data.HashMap.Strict.Base + ( + -- * Strictness properties + -- $strictness + + HashMap + + -- * Construction + , empty + , singleton + + -- * Basic interface + , HM.null + , size + , HM.member + , HM.lookup + , lookupDefault + , (!) + , insert + , insertWith + , delete + , adjust + , update + , alter + , alterF + + -- * Combine + -- ** Union + , union + , unionWith + , unionWithKey + , unions + + -- * Transformations + , map + , mapWithKey + , traverseWithKey + + -- * Difference and intersection + , difference + , differenceWith + , intersection + , intersectionWith + , intersectionWithKey + + -- * Folds + , foldl' + , foldlWithKey' + , HM.foldr + , foldrWithKey + + -- * Filter + , HM.filter + , filterWithKey + , mapMaybe + , mapMaybeWithKey + + -- * Conversions + , keys + , elems + + -- ** Lists + , toList + , fromList + , fromListWith + ) where + +import Data.Bits ((.&.), (.|.)) + +#if !MIN_VERSION_base(4,8,0) +import Control.Applicative (Applicative (..), (<$>)) +#endif +import qualified Data.List as L +import Data.Hashable (Hashable) +import Prelude hiding (map, lookup) + +import qualified Data.HashMap.Array as A +import qualified Data.HashMap.Base as HM +import Data.HashMap.Base hiding ( + alter, alterF, adjust, fromList, fromListWith, insert, insertWith, + differenceWith, intersectionWith, intersectionWithKey, map, mapWithKey, + mapMaybe, mapMaybeWithKey, singleton, update, unionWith, unionWithKey, + traverseWithKey) +import Data.HashMap.Unsafe (runST) +#if MIN_VERSION_base(4,8,0) +import Data.Functor.Identity +#endif +import Control.Applicative (Const (..)) +import Data.Coerce + +-- $strictness +-- +-- This module satisfies the following strictness properties: +-- +-- 1. Key arguments are evaluated to WHNF; +-- +-- 2. Keys and values are evaluated to WHNF before they are stored in +-- the map. + +------------------------------------------------------------------------ +-- * Construction + +-- | /O(1)/ Construct a map with a single element. +singleton :: (Hashable k) => k -> v -> HashMap k v +singleton k !v = HM.singleton k v + +------------------------------------------------------------------------ +-- * Basic interface + +-- | /O(log n)/ Associate the specified value with the specified +-- key in this map. If this map previously contained a mapping for +-- the key, the old value is replaced. +insert :: (Eq k, Hashable k) => k -> v -> HashMap k v -> HashMap k v +insert k !v = HM.insert k v +{-# INLINABLE insert #-} + +-- | /O(log n)/ Associate the value with the key in this map. If +-- this map previously contained a mapping for the key, the old value +-- is replaced by the result of applying the given function to the new +-- and old value. Example: +-- +-- > insertWith f k v map +-- > where f new old = new + old +insertWith :: (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v + -> HashMap k v +insertWith f k0 v0 m0 = go h0 k0 v0 0 m0 + where + h0 = hash k0 + go !h !k x !_ Empty = leaf h k x + go h k x s (Leaf hy l@(L ky y)) + | hy == h = if ky == k + then leaf h k (f x y) + else x `seq` (collision h l (L k x)) + | otherwise = x `seq` runST (two s h k x hy ky y) + go h k x s (BitmapIndexed b ary) + | b .&. m == 0 = + let ary' = A.insert ary i $! leaf h k x + in bitmapIndexedOrFull (b .|. m) ary' + | otherwise = + let st = A.index ary i + st' = go h k x (s+bitsPerSubkey) st + ary' = A.update ary i $! st' + in BitmapIndexed b ary' + where m = mask h s + i = sparseIndex b m + go h k x s (Full ary) = + let st = A.index ary i + st' = go h k x (s+bitsPerSubkey) st + ary' = update16 ary i $! st' + in Full ary' + where i = index h s + go h k x s t@(Collision hy v) + | h == hy = Collision h (updateOrSnocWith f k x v) + | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) +{-# INLINABLE insertWith #-} + +-- | In-place update version of insertWith +unsafeInsertWith :: (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 !h !k x !_ Empty = return $! leaf h k x + go h k x s (Leaf hy l@(L ky y)) + | hy == h = if ky == k + then return $! leaf h k (f x y) + else do + let l' = x `seq` (L k x) + return $! collision h l l' + | otherwise = x `seq` two s h k x hy ky y + go h k x s t@(BitmapIndexed b ary) + | b .&. m == 0 = do + ary' <- A.insertM ary i $! leaf h k x + return $! bitmapIndexedOrFull (b .|. m) ary' + | otherwise = do + st <- A.indexM ary i + st' <- go h k x (s+bitsPerSubkey) st + A.unsafeUpdateM ary i st' + return t + where m = mask h s + i = sparseIndex b m + go h k x s t@(Full ary) = do + st <- A.indexM ary i + st' <- go h k x (s+bitsPerSubkey) st + A.unsafeUpdateM ary i st' + return t + where i = index h s + go h k x s t@(Collision hy v) + | h == hy = return $! Collision h (updateOrSnocWith f k x v) + | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t) +{-# INLINABLE unsafeInsertWith #-} + +-- | /O(log n)/ Adjust the value tied to a given key in this map only +-- if it is present. Otherwise, leave the map alone. +adjust :: (Eq k, Hashable k) => (v -> v) -> k -> HashMap k v -> HashMap k v +adjust f k0 m0 = go h0 k0 0 m0 + where + h0 = hash k0 + go !_ !_ !_ Empty = Empty + go h k _ t@(Leaf hy (L ky y)) + | hy == h && ky == k = leaf h k (f y) + | otherwise = t + go h k s t@(BitmapIndexed b ary) + | b .&. m == 0 = t + | otherwise = let st = A.index ary i + st' = go h k (s+bitsPerSubkey) st + ary' = A.update ary i $! st' + in BitmapIndexed b ary' + where m = mask h s + i = sparseIndex b m + go h k s (Full ary) = + let i = index h s + st = A.index ary i + st' = go h k (s+bitsPerSubkey) st + ary' = update16 ary i $! st' + in Full ary' + go h k _ t@(Collision hy v) + | h == hy = Collision h (updateWith f k v) + | otherwise = t +{-# INLINABLE adjust #-} + +-- | /O(log n)/ The expression (@'update' f k map@) updates the value @x@ at @k@, +-- (if it is in the map). If (f k x) is @'Nothing', the element is deleted. +-- If it is (@'Just' y), the key k is bound to the new value y. +update :: (Eq k, Hashable k) => (a -> Maybe a) -> k -> HashMap k a -> HashMap k a +update f = alter (>>= f) +{-# INLINABLE update #-} + +-- | /O(log n)/ The expression (@'alter' f k map@) alters the value @x@ at @k@, or +-- absence thereof. @alter@ can be used to insert, delete, or update a value in a +-- map. In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. +alter :: (Eq k, Hashable k) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v +alter f k m = + case f (HM.lookup k m) of + Nothing -> delete k m + Just v -> insert k v m +{-# INLINABLE alter #-} + +-- | /O(log n)/ The expression (@'alterF' f k map@) alters the value @x@ at +-- @k@, or absence thereof. @alterF@ can be used to insert, delete, or update +-- a value in a map. +-- +-- Note: 'alterF' is a flipped version of the 'at' combinator from +-- . +-- +-- @since 0.2.9 +alterF :: (Functor f, Eq k, Hashable k) + => (Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v) +-- Special care is taken to only calculate the hash once. When we rewrite +-- with RULES, we also ensure that we only compare the key for equality +-- once. We force the value of the map for consistency with the rewritten +-- version; otherwise someone could tell the difference using a lazy +-- @f@ and a functor that is similar to Const but not actually Const. +alterF f = \ !k !m -> + let !h = hash k + mv = lookup' h k m + in (<$> f mv) $ \fres -> + case fres of + Nothing -> delete' h k m + Just !v' -> insert' h k v' m + +-- We rewrite this function unconditionally in RULES, but we expose +-- an unfolding just in case it's used in a context where the rules +-- don't fire. +{-# INLINABLE [0] alterF #-} + +#if MIN_VERSION_base(4,8,0) +-- See notes in Data.HashMap.Base +test_bottom :: a +test_bottom = error "Data.HashMap.alterF internal error: hit test_bottom" + +bogus# :: (# #) -> (# a #) +bogus# _ = error "Data.HashMap.alterF internal error: hit bogus#" + +impossibleAdjust :: a +impossibleAdjust = error "Data.HashMap.alterF internal error: impossible adjust" + +{-# RULES + +-- See detailed notes on alterF rules in Data.HashMap.Base. + +"alterFWeird" forall f. alterF f = + alterFWeird (f Nothing) (f (Just test_bottom)) f + +"alterFconstant" forall (f :: Maybe a -> Identity (Maybe a)) x. + alterFWeird x x f = \ !k !m -> + Identity (case runIdentity x of {Nothing -> delete k m; Just a -> insert k a m}) + +"alterFinsertWith" [1] forall (f :: Maybe a -> Identity (Maybe a)) x y. + alterFWeird (coerce (Just x)) (coerce (Just y)) f = + coerce (insertModifying x (\mold -> case runIdentity (f (Just mold)) of + Nothing -> bogus# (# #) + Just !new -> (# new #))) + +-- This rule is written a bit differently than the one for lazy +-- maps because the adjust here is strict. We could write it the +-- same general way anyway, but this seems simpler. +"alterFadjust" forall (f :: Maybe a -> Identity (Maybe a)) x. + alterFWeird (coerce Nothing) (coerce (Just x)) f = + coerce (adjust (\a -> case runIdentity (f (Just a)) of + Just a' -> a' + Nothing -> impossibleAdjust)) + +"alterFlookup" forall _ign1 _ign2 (f :: Maybe a -> Const r (Maybe a)) . + alterFWeird _ign1 _ign2 f = \ !k !m -> Const (getConst (f (lookup k m))) + #-} + +-- This is a very unsafe version of alterF used for RULES. When calling +-- alterFWeird x y f, the following *must* hold: +-- +-- x = f Nothing +-- y = f (Just _|_) +-- +-- Failure to abide by these laws will make demons come out of your nose. +alterFWeird + :: (Functor f, Eq k, Hashable k) + => f (Maybe v) + -> f (Maybe v) + -> (Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v) +alterFWeird _ _ f = alterFEager f +{-# INLINE [0] alterFWeird #-} + +-- | This is the default version of alterF that we use in most non-trivial +-- cases. It's called "eager" because it looks up the given key in the map +-- eagerly, whether or not the given function requires that information. +alterFEager :: (Functor f, Eq k, Hashable k) + => (Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v) +alterFEager f !k !m = (<$> f mv) $ \fres -> + case fres of + + ------------------------------ + -- Delete the key from the map. + Nothing -> case lookupRes of + + -- Key did not exist in the map to begin with, no-op + Absent -> m + + -- Key did exist, no collision + Present _ collPos -> deleteKeyExists collPos h k m + + ------------------------------ + -- Update value + Just v' -> case lookupRes of + + -- Key did not exist before, insert v' under a new key + Absent -> insertNewKey h k v' m + + -- Key existed before, no hash collision + Present v collPos -> v' `seq` + if v `ptrEq` v' + -- If the value is identical, no-op + then m + -- If the value changed, update the value. + else insertKeyExists collPos h k v' m + + where !h = hash k + !lookupRes = lookupRecordCollision h k m + !mv = case lookupRes of + Absent -> Nothing + Present v _ -> Just v +{-# INLINABLE alterFEager #-} +#endif + +------------------------------------------------------------------------ +-- * Combine + +-- | /O(n+m)/ The union of two maps. If a key occurs in both maps, +-- the provided function (first argument) will be used to compute the result. +unionWith :: (Eq k, Hashable k) => (v -> v -> v) -> HashMap k v -> HashMap k v + -> HashMap k v +unionWith f = unionWithKey (const f) +{-# INLINE unionWith #-} + +-- | /O(n+m)/ The union of two maps. If a key occurs in both maps, +-- the provided function (first argument) will be used to compute the result. +unionWithKey :: (Eq k, Hashable k) => (k -> v -> v -> v) -> HashMap k v -> HashMap k v + -> HashMap k v +unionWithKey f = go 0 + where + -- empty vs. anything + go !_ t1 Empty = t1 + go _ Empty t2 = t2 + -- leaf vs. leaf + go s t1@(Leaf h1 l1@(L k1 v1)) t2@(Leaf h2 l2@(L k2 v2)) + | h1 == h2 = if k1 == k2 + then leaf h1 k1 (f k1 v1 v2) + else collision h1 l1 l2 + | otherwise = goDifferentHash s h1 h2 t1 t2 + go s t1@(Leaf h1 (L k1 v1)) t2@(Collision h2 ls2) + | h1 == h2 = Collision h1 (updateOrSnocWithKey f k1 v1 ls2) + | otherwise = goDifferentHash s h1 h2 t1 t2 + go s t1@(Collision h1 ls1) t2@(Leaf h2 (L k2 v2)) + | h1 == h2 = Collision h1 (updateOrSnocWithKey (flip . f) k2 v2 ls1) + | otherwise = goDifferentHash s h1 h2 t1 t2 + go s t1@(Collision h1 ls1) t2@(Collision h2 ls2) + | h1 == h2 = Collision h1 (updateOrConcatWithKey f ls1 ls2) + | otherwise = goDifferentHash s h1 h2 t1 t2 + -- branch vs. branch + 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' + -- leaf vs. branch + 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 #-} + +------------------------------------------------------------------------ +-- * Transformations + +-- | /O(n)/ Transform this map by applying a function to every value. +mapWithKey :: (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2 +mapWithKey f = go + where + go Empty = Empty + go (Leaf h (L k v)) = leaf h k (f k v) + go (BitmapIndexed b ary) = BitmapIndexed b $ A.map' go ary + go (Full ary) = Full $ A.map' go ary + go (Collision h ary) = + Collision h $ A.map' (\ (L k v) -> let !v' = f k v in L k v') ary +{-# INLINE mapWithKey #-} + +-- | /O(n)/ Transform this map by applying a function to every value. +map :: (v1 -> v2) -> HashMap k v1 -> HashMap k v2 +map f = mapWithKey (const f) +{-# INLINE map #-} + + +------------------------------------------------------------------------ +-- * Filter + +-- | /O(n)/ Transform this map by applying a function to every value +-- and retaining only some of them. +mapMaybeWithKey :: (k -> v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2 +mapMaybeWithKey f = filterMapAux onLeaf onColl + where onLeaf (Leaf h (L k v)) | Just v' <- f k v = Just (leaf h k v') + onLeaf _ = Nothing + + onColl (L k v) | Just v' <- f k v = Just (L k v') + | otherwise = Nothing +{-# INLINE mapMaybeWithKey #-} + +-- | /O(n)/ Transform this map by applying a function to every value +-- and retaining only some of them. +mapMaybe :: (v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2 +mapMaybe f = mapMaybeWithKey (const f) +{-# INLINE mapMaybe #-} + +-- | /O(n)/ Perform an 'Applicative' action for each key-value pair +-- in a 'HashMap' and produce a 'HashMap' of all the results. Each 'HashMap' +-- will be strict in all its values. +-- +-- @ +-- traverseWithKey f = fmap ('map' id) . "Data.HashMap.Lazy".'Data.HashMap.Lazy.traverseWithKey' f +-- @ +-- +-- Note: the order in which the actions occur is unspecified. In particular, +-- when the map contains hash collisions, the order in which the actions +-- associated with the keys involved will depend in an unspecified way on +-- their insertion order. +traverseWithKey + :: Applicative f + => (k -> v1 -> f v2) + -> HashMap k v1 -> f (HashMap k v2) +traverseWithKey f = go + where + go Empty = pure Empty + go (Leaf h (L k v)) = leaf h k <$> f k v + go (BitmapIndexed b ary) = BitmapIndexed b <$> A.traverse' go ary + go (Full ary) = Full <$> A.traverse' go ary + go (Collision h ary) = + Collision h <$> A.traverse' (\ (L k v) -> (L k $!) <$> f k v) ary +{-# INLINE traverseWithKey #-} + +------------------------------------------------------------------------ +-- * Difference and intersection + +-- | /O(n*log m)/ Difference with a combining function. When two equal keys are +-- encountered, the combining function is applied to the values of these keys. +-- If it returns 'Nothing', the element is discarded (proper set difference). If +-- it returns (@'Just' y@), the element is updated with a new value @y@. +differenceWith :: (Eq k, Hashable k) => (v -> w -> Maybe v) -> HashMap k v -> HashMap k w -> HashMap k v +differenceWith f a b = foldlWithKey' go empty a + where + go m k v = case HM.lookup k b of + Nothing -> insert k v m + Just w -> maybe m (\y -> insert k y m) (f v w) +{-# INLINABLE differenceWith #-} + +-- | /O(n+m)/ Intersection of two maps. If a key occurs in both maps +-- the provided function is used to combine the values from the two +-- maps. +intersectionWith :: (Eq k, Hashable k) => (v1 -> v2 -> v3) -> HashMap k v1 + -> HashMap k v2 -> HashMap k v3 +intersectionWith f a b = foldlWithKey' go empty a + where + go m k v = case HM.lookup k b of + Just w -> insert k (f v w) m + _ -> m +{-# INLINABLE intersectionWith #-} + +-- | /O(n+m)/ Intersection of two maps. If a key occurs in both maps +-- the provided function is used to combine the values from the two +-- maps. +intersectionWithKey :: (Eq k, Hashable k) => (k -> v1 -> v2 -> v3) + -> HashMap k v1 -> HashMap k v2 -> HashMap k v3 +intersectionWithKey f a b = foldlWithKey' go empty a + where + go m k v = case HM.lookup k b of + Just w -> insert k (f k v w) m + _ -> m +{-# INLINABLE intersectionWithKey #-} + +------------------------------------------------------------------------ +-- ** Lists + +-- | /O(n*log n)/ Construct a map with the supplied mappings. If the +-- list contains duplicate mappings, the later mappings take +-- precedence. +fromList :: (Eq k, Hashable k) => [(k, v)] -> HashMap k v +fromList = L.foldl' (\ m (k, !v) -> HM.unsafeInsert k v m) empty +{-# INLINABLE fromList #-} + +-- | /O(n*log n)/ Construct a map from a list of elements. Uses +-- the provided function f to merge duplicate entries (f newVal oldVal). +-- +-- For example: +-- +-- > fromListWith (+) [ (x, 1) | x <- xs ] +-- +-- will create a map with number of occurrences of each element in xs. +-- +-- > fromListWith (++) [ (k, [v]) | (k, v) <- xs ] +-- +-- will group all values by their keys in a list 'xs :: [(k, v)]' and +-- return a 'HashMap k [v]'. +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 #-} + +------------------------------------------------------------------------ +-- Array operations + +updateWith :: Eq k => (v -> v) -> k -> A.Array (Leaf k v) -> A.Array (Leaf k v) +updateWith f k0 ary0 = go k0 ary0 0 (A.length ary0) + where + go !k !ary !i !n + | i >= n = ary + | otherwise = case A.index ary i of + (L kx y) | k == kx -> let !v' = f y in A.update ary i (L k v') + | otherwise -> go k ary (i+1) n +{-# INLINABLE updateWith #-} + +-- | Append the given key and value to the array. If the key is +-- already present, instead update the value of the key by applying +-- the given function to the new and old value (in that order). The +-- value is always evaluated to WHNF before being inserted into the +-- array. +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 #-} + +-- | Append the given key and value to the array. If the key is +-- already present, instead update the value of the key by applying +-- the given function to the new and old value (in that order). The +-- value is always evaluated to WHNF before being inserted into the +-- array. +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 + -- Not found, append to the end. + mary <- A.new_ (n + 1) + A.copy ary 0 mary 0 n + let !l = v `seq` (L k v) + A.write mary n l + return mary + | otherwise = case A.index ary i of + (L kx y) | k == kx -> let !v' = f k v y in A.update ary i (L k v') + | otherwise -> go k v ary (i+1) n +{-# INLINABLE updateOrSnocWithKey #-} + +------------------------------------------------------------------------ +-- Smart constructors +-- +-- These constructors make sure the value is in WHNF before it's +-- inserted into the constructor. + +leaf :: Hash -> k -> v -> HashMap k v +leaf h k = \ !v -> Leaf h (L k v) +{-# INLINE leaf #-} diff --git a/Data/HashSet.hs b/Data/HashSet.hs --- a/Data/HashSet.hs +++ b/Data/HashSet.hs @@ -1,10 +1,6 @@ -{-# LANGUAGE CPP, DeriveDataTypeable #-} -#if __GLASGOW_HASKELL__ >= 708 -{-# LANGUAGE RoleAnnotations #-} -{-# LANGUAGE TypeFamilies #-} -#endif +{-# LANGUAGE CPP #-} #if __GLASGOW_HASKELL__ >= 702 -{-# LANGUAGE Trustworthy #-} +{-# LANGUAGE Safe #-} #endif ------------------------------------------------------------------------ @@ -72,245 +68,5 @@ , fromMap ) where -import Control.DeepSeq (NFData(..)) -import Data.Data hiding (Typeable) -import Data.HashMap.Base (HashMap, foldrWithKey, equalKeys) -import Data.Hashable (Hashable(hashWithSalt)) -#if __GLASGOW_HASKELL__ >= 711 -import Data.Semigroup (Semigroup(..)) -#elif __GLASGOW_HASKELL__ < 709 -import Data.Monoid (Monoid(..)) -#endif -import GHC.Exts (build) -import Prelude hiding (filter, foldr, map, null) -import qualified Data.Foldable as Foldable -import qualified Data.HashMap.Lazy as H -import qualified Data.List as List -import Data.Typeable (Typeable) -import Text.Read - -#if __GLASGOW_HASKELL__ >= 708 -import qualified GHC.Exts as Exts -#endif - -#if MIN_VERSION_base(4,9,0) -import Data.Functor.Classes -#endif - -#if MIN_VERSION_hashable(1,2,5) -import qualified Data.Hashable.Lifted as H -#endif - --- | A set of values. A set cannot contain duplicate values. -newtype HashSet a = HashSet { - asMap :: HashMap a () - } deriving (Typeable) - -#if __GLASGOW_HASKELL__ >= 708 -type role HashSet nominal -#endif - -instance (NFData a) => NFData (HashSet a) where - rnf = rnf . asMap - {-# INLINE rnf #-} - -instance (Eq a) => Eq (HashSet a) where - HashSet a == HashSet b = equalKeys (==) a b - {-# INLINE (==) #-} - -#if MIN_VERSION_base(4,9,0) -instance Eq1 HashSet where - liftEq eq (HashSet a) (HashSet b) = equalKeys eq a b -#endif - -instance (Ord a) => Ord (HashSet a) where - compare (HashSet a) (HashSet b) = compare a b - {-# INLINE compare #-} - -#if MIN_VERSION_base(4,9,0) -instance Ord1 HashSet where - liftCompare c (HashSet a) (HashSet b) = liftCompare2 c compare a b -#endif - -instance Foldable.Foldable HashSet where - foldr = Data.HashSet.foldr - {-# INLINE foldr #-} - -#if __GLASGOW_HASKELL__ >= 711 -instance (Hashable a, Eq a) => Semigroup (HashSet a) where - (<>) = union - {-# INLINE (<>) #-} -#endif - -instance (Hashable a, Eq a) => Monoid (HashSet a) where - mempty = empty - {-# INLINE mempty #-} -#if __GLASGOW_HASKELL__ >= 711 - mappend = (<>) -#else - mappend = union -#endif - {-# INLINE mappend #-} - -instance (Eq a, Hashable a, Read a) => Read (HashSet a) where - readPrec = parens $ prec 10 $ do - Ident "fromList" <- lexP - xs <- readPrec - return (fromList xs) - - readListPrec = readListPrecDefault - -#if MIN_VERSION_base(4,9,0) -instance Show1 HashSet where - liftShowsPrec sp sl d m = - showsUnaryWith (liftShowsPrec sp sl) "fromList" d (toList m) -#endif - -instance (Show a) => Show (HashSet a) where - showsPrec d m = showParen (d > 10) $ - showString "fromList " . shows (toList m) - -instance (Data a, Eq a, Hashable a) => Data (HashSet a) 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 _ = hashSetDataType - dataCast1 f = gcast1 f - -#if MIN_VERSION_hashable(1,2,6) -instance H.Hashable1 HashSet where - liftHashWithSalt h s = H.liftHashWithSalt2 h hashWithSalt s . asMap -#endif - -instance (Hashable a) => Hashable (HashSet a) where - hashWithSalt salt = hashWithSalt salt . asMap - -fromListConstr :: Constr -fromListConstr = mkConstr hashSetDataType "fromList" [] Prefix - -hashSetDataType :: DataType -hashSetDataType = mkDataType "Data.HashSet" [fromListConstr] - --- | /O(1)/ Construct an empty set. -empty :: HashSet a -empty = HashSet H.empty - --- | /O(1)/ Construct a set with a single element. -singleton :: Hashable a => a -> HashSet a -singleton a = HashSet (H.singleton a ()) -{-# INLINABLE singleton #-} - --- | /O(1)/ Convert to the equivalent 'HashMap'. -toMap :: HashSet a -> HashMap a () -toMap = asMap - --- | /O(1)/ Convert from the equivalent 'HashMap'. -fromMap :: HashMap a () -> HashSet a -fromMap = HashSet - --- | /O(n+m)/ Construct a set containing all elements from both sets. --- --- To obtain good performance, the smaller set must be presented as --- the first argument. -union :: (Eq a, Hashable a) => HashSet a -> HashSet a -> HashSet a -union s1 s2 = HashSet $ H.union (asMap s1) (asMap s2) -{-# INLINE union #-} - --- TODO: Figure out the time complexity of 'unions'. - --- | Construct a set containing all elements from a list of sets. -unions :: (Eq a, Hashable a) => [HashSet a] -> HashSet a -unions = List.foldl' union empty -{-# INLINE unions #-} - --- | /O(1)/ Return 'True' if this set is empty, 'False' otherwise. -null :: HashSet a -> Bool -null = H.null . asMap -{-# INLINE null #-} - --- | /O(n)/ Return the number of elements in this set. -size :: HashSet a -> Int -size = H.size . asMap -{-# INLINE size #-} - --- | /O(log n)/ Return 'True' if the given value is present in this --- set, 'False' otherwise. -member :: (Eq a, Hashable a) => a -> HashSet a -> Bool -member a s = case H.lookup a (asMap s) of - Just _ -> True - _ -> False -{-# INLINABLE member #-} - --- | /O(log n)/ Add the specified value to this set. -insert :: (Eq a, Hashable a) => a -> HashSet a -> HashSet a -insert a = HashSet . H.insert a () . asMap -{-# INLINABLE insert #-} - --- | /O(log n)/ Remove the specified value from this set if --- present. -delete :: (Eq a, Hashable a) => a -> HashSet a -> HashSet a -delete a = HashSet . H.delete a . asMap -{-# INLINABLE delete #-} - --- | /O(n)/ Transform this set by applying a function to every value. --- The resulting set may be smaller than the source. -map :: (Hashable b, Eq b) => (a -> b) -> HashSet a -> HashSet b -map f = fromList . List.map f . toList -{-# INLINE map #-} - --- | /O(n)/ Difference of two sets. Return elements of the first set --- not existing in the second. -difference :: (Eq a, Hashable a) => HashSet a -> HashSet a -> HashSet a -difference (HashSet a) (HashSet b) = HashSet (H.difference a b) -{-# INLINABLE difference #-} - --- | /O(n)/ Intersection of two sets. Return elements present in both --- the first set and the second. -intersection :: (Eq a, Hashable a) => HashSet a -> HashSet a -> HashSet a -intersection (HashSet a) (HashSet b) = HashSet (H.intersection a b) -{-# INLINABLE intersection #-} - --- | /O(n)/ Reduce this set by applying a binary operator to all --- elements, using the given starting value (typically the --- left-identity of the operator). Each application of the operator --- is evaluated before before using the result in the next --- application. This function is strict in the starting value. -foldl' :: (a -> b -> a) -> a -> HashSet b -> a -foldl' f z0 = H.foldlWithKey' g z0 . asMap - where g z k _ = f z k -{-# INLINE foldl' #-} - --- | /O(n)/ Reduce this set by applying a binary operator to all --- elements, using the given starting value (typically the --- right-identity of the operator). -foldr :: (b -> a -> a) -> a -> HashSet b -> a -foldr f z0 = foldrWithKey g z0 . asMap - where g k _ z = f k z -{-# INLINE foldr #-} - --- | /O(n)/ Filter this set by retaining only elements satisfying a --- predicate. -filter :: (a -> Bool) -> HashSet a -> HashSet a -filter p = HashSet . H.filterWithKey q . asMap - where q k _ = p k -{-# INLINE filter #-} - --- | /O(n)/ Return a list of this set's elements. The list is --- produced lazily. -toList :: HashSet a -> [a] -toList t = build (\ c z -> foldrWithKey ((const .) c) z (asMap t)) -{-# INLINE toList #-} - --- | /O(n*min(W, n))/ Construct a set from a list of elements. -fromList :: (Eq a, Hashable a) => [a] -> HashSet a -fromList = HashSet . List.foldl' (\ m k -> H.insert k () m) H.empty -{-# INLINE fromList #-} - -#if __GLASGOW_HASKELL__ >= 708 -instance (Eq a, Hashable a) => Exts.IsList (HashSet a) where - type Item (HashSet a) = a - fromList = fromList - toList = toList -#endif +import Data.HashSet.Base +import Prelude () diff --git a/Data/HashSet/Base.hs b/Data/HashSet/Base.hs new file mode 100644 --- /dev/null +++ b/Data/HashSet/Base.hs @@ -0,0 +1,327 @@ +{-# LANGUAGE CPP, DeriveDataTypeable #-} +#if __GLASGOW_HASKELL__ >= 708 +{-# LANGUAGE RoleAnnotations #-} +{-# LANGUAGE TypeFamilies #-} +#endif +#if __GLASGOW_HASKELL__ >= 702 +{-# LANGUAGE Trustworthy #-} +#endif + +------------------------------------------------------------------------ +-- | +-- Module : Data.HashSet.Base +-- Copyright : 2011 Bryan O'Sullivan +-- License : BSD-style +-- Maintainer : johan.tibell@gmail.com +-- Stability : provisional +-- Portability : portable +-- +-- A set of /hashable/ values. A set cannot contain duplicate items. +-- A 'HashSet' makes no guarantees as to the order of its elements. +-- +-- The implementation is based on /hash array mapped trie/. A +-- 'HashSet' is often faster than other tree-based set types, +-- especially when value comparison is expensive, as in the case of +-- strings. +-- +-- Many operations have a average-case complexity of /O(log n)/. The +-- implementation uses a large base (i.e. 16) so in practice these +-- operations are constant time. + +module Data.HashSet.Base + ( + HashSet + + -- * Construction + , empty + , singleton + + -- * Combine + , union + , unions + + -- * Basic interface + , null + , size + , member + , insert + , delete + + -- * Transformations + , map + + -- * Difference and intersection + , difference + , intersection + + -- * Folds + , foldl' + , foldr + + -- * Filter + , filter + + -- * Conversions + + -- ** Lists + , toList + , fromList + + -- * HashMaps + , toMap + , fromMap + + -- Exported from Data.HashMap.{Strict, Lazy} + , keysSet + ) where + +import Control.DeepSeq (NFData(..)) +import Data.Data hiding (Typeable) +import Data.HashMap.Base (HashMap, foldrWithKey, equalKeys, equalKeys1) +import Data.Hashable (Hashable(hashWithSalt)) +#if __GLASGOW_HASKELL__ >= 711 +import Data.Semigroup (Semigroup(..)) +#elif __GLASGOW_HASKELL__ < 709 +import Data.Monoid (Monoid(..)) +#endif +import GHC.Exts (build) +import Prelude hiding (filter, foldr, map, null) +import qualified Data.Foldable as Foldable +import qualified Data.HashMap.Base as H +import qualified Data.List as List +import Data.Typeable (Typeable) +import Text.Read + +#if __GLASGOW_HASKELL__ >= 708 +import qualified GHC.Exts as Exts +#endif + +#if MIN_VERSION_base(4,9,0) +import Data.Functor.Classes +#endif + +#if MIN_VERSION_hashable(1,2,5) +import qualified Data.Hashable.Lifted as H +#endif + +import Data.Functor ((<$)) + +-- | A set of values. A set cannot contain duplicate values. +newtype HashSet a = HashSet { + asMap :: HashMap a () + } deriving (Typeable) + +#if __GLASGOW_HASKELL__ >= 708 +type role HashSet nominal +#endif + +instance (NFData a) => NFData (HashSet a) where + rnf = rnf . asMap + {-# INLINE rnf #-} + +instance (Eq a) => Eq (HashSet a) where + HashSet a == HashSet b = equalKeys a b + {-# INLINE (==) #-} + +#if MIN_VERSION_base(4,9,0) +instance Eq1 HashSet where + liftEq eq (HashSet a) (HashSet b) = equalKeys1 eq a b +#endif + +instance (Ord a) => Ord (HashSet a) where + compare (HashSet a) (HashSet b) = compare a b + {-# INLINE compare #-} + +#if MIN_VERSION_base(4,9,0) +instance Ord1 HashSet where + liftCompare c (HashSet a) (HashSet b) = liftCompare2 c compare a b +#endif + +instance Foldable.Foldable HashSet where + foldr = Data.HashSet.Base.foldr + {-# INLINE foldr #-} + +#if __GLASGOW_HASKELL__ >= 711 +instance (Hashable a, Eq a) => Semigroup (HashSet a) where + (<>) = union + {-# INLINE (<>) #-} +#endif + +instance (Hashable a, Eq a) => Monoid (HashSet a) where + mempty = empty + {-# INLINE mempty #-} +#if __GLASGOW_HASKELL__ >= 711 + mappend = (<>) +#else + mappend = union +#endif + {-# INLINE mappend #-} + +instance (Eq a, Hashable a, Read a) => Read (HashSet a) where + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + xs <- readPrec + return (fromList xs) + + readListPrec = readListPrecDefault + +#if MIN_VERSION_base(4,9,0) +instance Show1 HashSet where + liftShowsPrec sp sl d m = + showsUnaryWith (liftShowsPrec sp sl) "fromList" d (toList m) +#endif + +instance (Show a) => Show (HashSet a) where + showsPrec d m = showParen (d > 10) $ + showString "fromList " . shows (toList m) + +instance (Data a, Eq a, Hashable a) => Data (HashSet a) 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 _ = hashSetDataType + dataCast1 f = gcast1 f + +#if MIN_VERSION_hashable(1,2,6) +instance H.Hashable1 HashSet where + liftHashWithSalt h s = H.liftHashWithSalt2 h hashWithSalt s . asMap +#endif + +instance (Hashable a) => Hashable (HashSet a) where + hashWithSalt salt = hashWithSalt salt . asMap + +fromListConstr :: Constr +fromListConstr = mkConstr hashSetDataType "fromList" [] Prefix + +hashSetDataType :: DataType +hashSetDataType = mkDataType "Data.HashSet.Base.HashSet" [fromListConstr] + +-- | /O(1)/ Construct an empty set. +empty :: HashSet a +empty = HashSet H.empty + +-- | /O(1)/ Construct a set with a single element. +singleton :: Hashable a => a -> HashSet a +singleton a = HashSet (H.singleton a ()) +{-# INLINABLE singleton #-} + +-- | /O(1)/ Convert to the equivalent 'HashMap'. +toMap :: HashSet a -> HashMap a () +toMap = asMap + +-- | /O(1)/ Convert from the equivalent 'HashMap'. +fromMap :: HashMap a () -> HashSet a +fromMap = HashSet + +-- | /O(n)/ Produce a 'HashSet' of all the keys in the given 'HashMap'. +-- +-- @since 0.2.10.0 +keysSet :: HashMap k a -> HashSet k +keysSet m = fromMap (() <$ m) + +-- | /O(n+m)/ Construct a set containing all elements from both sets. +-- +-- To obtain good performance, the smaller set must be presented as +-- the first argument. +union :: (Eq a, Hashable a) => HashSet a -> HashSet a -> HashSet a +union s1 s2 = HashSet $ H.union (asMap s1) (asMap s2) +{-# INLINE union #-} + +-- TODO: Figure out the time complexity of 'unions'. + +-- | Construct a set containing all elements from a list of sets. +unions :: (Eq a, Hashable a) => [HashSet a] -> HashSet a +unions = List.foldl' union empty +{-# INLINE unions #-} + +-- | /O(1)/ Return 'True' if this set is empty, 'False' otherwise. +null :: HashSet a -> Bool +null = H.null . asMap +{-# INLINE null #-} + +-- | /O(n)/ Return the number of elements in this set. +size :: HashSet a -> Int +size = H.size . asMap +{-# INLINE size #-} + +-- | /O(log n)/ Return 'True' if the given value is present in this +-- set, 'False' otherwise. +member :: (Eq a, Hashable a) => a -> HashSet a -> Bool +member a s = case H.lookup a (asMap s) of + Just _ -> True + _ -> False +{-# INLINABLE member #-} + +-- | /O(log n)/ Add the specified value to this set. +insert :: (Eq a, Hashable a) => a -> HashSet a -> HashSet a +insert a = HashSet . H.insert a () . asMap +{-# INLINABLE insert #-} + +-- | /O(log n)/ Remove the specified value from this set if +-- present. +delete :: (Eq a, Hashable a) => a -> HashSet a -> HashSet a +delete a = HashSet . H.delete a . asMap +{-# INLINABLE delete #-} + +-- | /O(n)/ Transform this set by applying a function to every value. +-- The resulting set may be smaller than the source. +map :: (Hashable b, Eq b) => (a -> b) -> HashSet a -> HashSet b +map f = fromList . List.map f . toList +{-# INLINE map #-} + +-- | /O(n)/ Difference of two sets. Return elements of the first set +-- not existing in the second. +difference :: (Eq a, Hashable a) => HashSet a -> HashSet a -> HashSet a +difference (HashSet a) (HashSet b) = HashSet (H.difference a b) +{-# INLINABLE difference #-} + +-- | /O(n)/ Intersection of two sets. Return elements present in both +-- the first set and the second. +intersection :: (Eq a, Hashable a) => HashSet a -> HashSet a -> HashSet a +intersection (HashSet a) (HashSet b) = HashSet (H.intersection a b) +{-# INLINABLE intersection #-} + +-- | /O(n)/ Reduce this set by applying a binary operator to all +-- elements, using the given starting value (typically the +-- left-identity of the operator). Each application of the operator +-- is evaluated before before using the result in the next +-- application. This function is strict in the starting value. +foldl' :: (a -> b -> a) -> a -> HashSet b -> a +foldl' f z0 = H.foldlWithKey' g z0 . asMap + where g z k _ = f z k +{-# INLINE foldl' #-} + +-- | /O(n)/ Reduce this set by applying a binary operator to all +-- elements, using the given starting value (typically the +-- right-identity of the operator). +foldr :: (b -> a -> a) -> a -> HashSet b -> a +foldr f z0 = foldrWithKey g z0 . asMap + where g k _ z = f k z +{-# INLINE foldr #-} + +-- | /O(n)/ Filter this set by retaining only elements satisfying a +-- predicate. +filter :: (a -> Bool) -> HashSet a -> HashSet a +filter p = HashSet . H.filterWithKey q . asMap + where q k _ = p k +{-# INLINE filter #-} + +-- | /O(n)/ Return a list of this set's elements. The list is +-- produced lazily. +toList :: HashSet a -> [a] +toList t = build (\ c z -> foldrWithKey ((const .) c) z (asMap t)) +{-# INLINE toList #-} + +-- | /O(n*min(W, n))/ Construct a set from a list of elements. +fromList :: (Eq a, Hashable a) => [a] -> HashSet a +fromList = HashSet . List.foldl' (\ m k -> H.insert k () m) H.empty +{-# INLINE fromList #-} + +#if __GLASGOW_HASKELL__ >= 708 +instance (Eq a, Hashable a) => Exts.IsList (HashSet a) where + type Item (HashSet a) = a + fromList = fromList + toList = toList +#endif diff --git a/benchmarks/Benchmarks.hs b/benchmarks/Benchmarks.hs --- a/benchmarks/Benchmarks.hs +++ b/benchmarks/Benchmarks.hs @@ -6,6 +6,7 @@ import Control.DeepSeq.Generics (genericRnf) import Criterion.Main (bench, bgroup, defaultMain, env, nf, whnf) import Data.Bits ((.&.)) +import Data.Functor.Identity import Data.Hashable (Hashable) import qualified Data.ByteString as BS import qualified "hashmap" Data.HashMap as IHM @@ -227,6 +228,46 @@ , bench "ByteString" $ whnf (delete keysBS') hmbs , bench "Int" $ whnf (delete keysI') hmi ] + , bgroup "alterInsert" + [ bench "String" $ whnf (alterInsert elems) HM.empty + , bench "ByteString" $ whnf (alterInsert elemsBS) HM.empty + , bench "Int" $ whnf (alterInsert elemsI) HM.empty + ] + , bgroup "alterFInsert" + [ bench "String" $ whnf (alterFInsert elems) HM.empty + , bench "ByteString" $ whnf (alterFInsert elemsBS) HM.empty + , bench "Int" $ whnf (alterFInsert elemsI) HM.empty + ] + , bgroup "alterInsert-dup" + [ bench "String" $ whnf (alterInsert elems) hm + , bench "ByteString" $ whnf (alterInsert elemsBS) hmbs + , bench "Int" $ whnf (alterInsert elemsI) hmi + ] + , bgroup "alterFInsert-dup" + [ bench "String" $ whnf (alterFInsert elems) hm + , bench "ByteString" $ whnf (alterFInsert elemsBS) hmbs + , bench "Int" $ whnf (alterFInsert elemsI) hmi + ] + , bgroup "alterDelete" + [ bench "String" $ whnf (alterDelete keys) hm + , bench "ByteString" $ whnf (alterDelete keysBS) hmbs + , bench "Int" $ whnf (alterDelete keysI) hmi + ] + , bgroup "alterFDelete" + [ bench "String" $ whnf (alterFDelete keys) hm + , bench "ByteString" $ whnf (alterFDelete keysBS) hmbs + , bench "Int" $ whnf (alterFDelete keysI) hmi + ] + , bgroup "alterDelete-miss" + [ bench "String" $ whnf (alterDelete keys') hm + , bench "ByteString" $ whnf (alterDelete keysBS') hmbs + , bench "Int" $ whnf (alterDelete keysI') hmi + ] + , bgroup "alterFDelete-miss" + [ bench "String" $ whnf (alterFDelete keys') hm + , bench "ByteString" $ whnf (alterFDelete keysBS') hmbs + , bench "Int" $ whnf (alterFDelete keysI') hmi + ] -- Combine , bench "union" $ whnf (HM.union hmi) hmi2 @@ -309,6 +350,50 @@ -> HM.HashMap String Int #-} {-# SPECIALIZE delete :: [BS.ByteString] -> HM.HashMap BS.ByteString Int -> HM.HashMap BS.ByteString Int #-} + +alterInsert :: (Eq k, Hashable k) => [(k, Int)] -> HM.HashMap k Int + -> HM.HashMap k Int +alterInsert xs m0 = + foldl' (\m (k, v) -> HM.alter (const . Just $ v) k m) m0 xs +{-# SPECIALIZE alterInsert :: [(Int, Int)] -> HM.HashMap Int Int + -> HM.HashMap Int Int #-} +{-# SPECIALIZE alterInsert :: [(String, Int)] -> HM.HashMap String Int + -> HM.HashMap String Int #-} +{-# SPECIALIZE alterInsert :: [(BS.ByteString, Int)] -> HM.HashMap BS.ByteString Int + -> HM.HashMap BS.ByteString Int #-} + +alterDelete :: (Eq k, Hashable k) => [k] -> HM.HashMap k Int + -> HM.HashMap k Int +alterDelete xs m0 = + foldl' (\m k -> HM.alter (const Nothing) k m) m0 xs +{-# SPECIALIZE alterDelete :: [Int] -> HM.HashMap Int Int + -> HM.HashMap Int Int #-} +{-# SPECIALIZE alterDelete :: [String] -> HM.HashMap String Int + -> HM.HashMap String Int #-} +{-# SPECIALIZE alterDelete :: [BS.ByteString] -> HM.HashMap BS.ByteString Int + -> HM.HashMap BS.ByteString Int #-} + +alterFInsert :: (Eq k, Hashable k) => [(k, Int)] -> HM.HashMap k Int + -> HM.HashMap k Int +alterFInsert xs m0 = + foldl' (\m (k, v) -> runIdentity $ HM.alterF (const . Identity . Just $ v) k m) m0 xs +{-# SPECIALIZE alterFInsert :: [(Int, Int)] -> HM.HashMap Int Int + -> HM.HashMap Int Int #-} +{-# SPECIALIZE alterFInsert :: [(String, Int)] -> HM.HashMap String Int + -> HM.HashMap String Int #-} +{-# SPECIALIZE alterFInsert :: [(BS.ByteString, Int)] -> HM.HashMap BS.ByteString Int + -> HM.HashMap BS.ByteString Int #-} + +alterFDelete :: (Eq k, Hashable k) => [k] -> HM.HashMap k Int + -> HM.HashMap k Int +alterFDelete xs m0 = + foldl' (\m k -> runIdentity $ HM.alterF (const . Identity $ Nothing) k m) m0 xs +{-# SPECIALIZE alterFDelete :: [Int] -> HM.HashMap Int Int + -> HM.HashMap Int Int #-} +{-# SPECIALIZE alterFDelete :: [String] -> HM.HashMap String Int + -> HM.HashMap String Int #-} +{-# SPECIALIZE alterFDelete :: [BS.ByteString] -> HM.HashMap BS.ByteString Int + -> HM.HashMap BS.ByteString Int #-} ------------------------------------------------------------------------ -- * Map diff --git a/tests/HashMapProperties.hs b/tests/HashMapProperties.hs --- a/tests/HashMapProperties.hs +++ b/tests/HashMapProperties.hs @@ -13,13 +13,20 @@ import Data.Ord (comparing) #if defined(STRICT) import qualified Data.HashMap.Strict as HM +import qualified Data.Map.Strict as M #else import qualified Data.HashMap.Lazy as HM +import qualified Data.Map.Lazy as M #endif -import qualified Data.Map as M import Test.QuickCheck (Arbitrary, Property, (==>), (===)) import Test.Framework (Test, defaultMain, testGroup) import Test.Framework.Providers.QuickCheck2 (testProperty) +#if MIN_VERSION_base(4,8,0) +import Data.Functor.Identity (Identity (..)) +#endif +import Control.Applicative (Const (..)) +import Test.QuickCheck.Function (Fun, apply) +import Test.QuickCheck.Poly (A, B) -- Key type that generates more hash collisions. newtype Key = K { unK :: Int } @@ -168,6 +175,50 @@ pAlterDelete :: Key -> [(Key, Int)] -> Bool pAlterDelete k = M.alter (const Nothing) k `eq_` HM.alter (const Nothing) k + +-- We choose the list functor here because we don't fuss with +-- it in alterF rules and because it has a sufficiently interesting +-- structure to have a good chance of breaking if something is wrong. +pAlterF :: Key -> Fun (Maybe A) [Maybe A] -> [(Key, A)] -> Property +pAlterF k f xs = + fmap M.toAscList (M.alterF (apply f) k (M.fromList xs)) + === + fmap toAscList (HM.alterF (apply f) k (HM.fromList xs)) + +#if !MIN_VERSION_base(4,8,0) +newtype Identity a = Identity {runIdentity :: a} +instance Functor Identity where + fmap f (Identity x) = Identity (f x) +#endif + +pAlterFAdjust :: Key -> [(Key, Int)] -> Bool +pAlterFAdjust k = + runIdentity . M.alterF (Identity . fmap succ) k `eq_` + runIdentity . HM.alterF (Identity . fmap succ) k + +pAlterFInsert :: Key -> [(Key, Int)] -> Bool +pAlterFInsert k = + runIdentity . M.alterF (const . Identity . Just $ 3) k `eq_` + runIdentity . HM.alterF (const . Identity . Just $ 3) k + +pAlterFInsertWith :: Key -> Fun Int Int -> [(Key, Int)] -> Bool +pAlterFInsertWith k f = + runIdentity . M.alterF (Identity . Just . maybe 3 (apply f)) k `eq_` + runIdentity . HM.alterF (Identity . Just . maybe 3 (apply f)) k + +pAlterFDelete :: Key -> [(Key, Int)] -> Bool +pAlterFDelete k = + runIdentity . M.alterF (const (Identity Nothing)) k `eq_` + runIdentity . HM.alterF (const (Identity Nothing)) k + +pAlterFLookup :: Key + -> Fun (Maybe A) B + -> [(Key, A)] -> Bool +pAlterFLookup k f = + getConst . M.alterF (Const . apply f :: Maybe A -> Const B (Maybe A)) k + `eq` + getConst . HM.alterF (Const . apply f) k + ------------------------------------------------------------------------ -- ** Combine @@ -195,6 +246,11 @@ pMap :: [(Key, Int)] -> Bool pMap = M.map (+ 1) `eq_` HM.map (+ 1) +pTraverse :: [(Key, Int)] -> Bool +pTraverse xs = + L.sort (fmap (L.sort . M.toList) (M.traverseWithKey (\_ v -> [v + 1, v + 2]) (M.fromList (take 10 xs)))) + == L.sort (fmap (L.sort . HM.toList) (HM.traverseWithKey (\_ v -> [v + 1, v + 2]) (HM.fromList (take 10 xs)))) + ------------------------------------------------------------------------ -- ** Difference and intersection @@ -227,7 +283,7 @@ -- ** Folds pFoldr :: [(Int, Int)] -> Bool -pFoldr = (L.sort . M.fold (:) []) `eq` (L.sort . HM.foldr (:) []) +pFoldr = (L.sort . M.foldr (:) []) `eq` (L.sort . HM.foldr (:) []) pFoldrWithKey :: [(Int, Int)] -> Bool pFoldrWithKey = (sortByKey . M.foldrWithKey f []) `eq` @@ -317,6 +373,12 @@ , testProperty "alterAdjust" pAlterAdjust , testProperty "alterInsert" pAlterInsert , testProperty "alterDelete" pAlterDelete + , testProperty "alterF" pAlterF + , testProperty "alterFAdjust" pAlterFAdjust + , testProperty "alterFInsert" pAlterFInsert + , testProperty "alterFInsertWith" pAlterFInsertWith + , testProperty "alterFDelete" pAlterFDelete + , testProperty "alterFLookup" pAlterFLookup ] -- Combine , testProperty "union" pUnion @@ -325,6 +387,7 @@ , testProperty "unions" pUnions -- Transformations , testProperty "map" pMap + , testProperty "traverse" pTraverse -- Folds , testGroup "folds" [ testProperty "foldr" pFoldr @@ -370,6 +433,8 @@ -> Bool -- ^ True if the functions are equivalent eq f g xs = g (HM.fromList xs) == f (M.fromList xs) +infix 4 `eq` + eq_ :: (Eq k, Eq v, Hashable k, Ord k) => (Model k v -> Model k v) -- ^ Function that modifies a 'Model' -> (HM.HashMap k v -> HM.HashMap k v) -- ^ Function that modified a @@ -379,6 +444,8 @@ -> Bool -- ^ True if the functions are -- equivalent eq_ f g = (M.toAscList . f) `eq` (toAscList . g) + +infix 4 `eq_` ------------------------------------------------------------------------ -- * Test harness diff --git a/unordered-containers.cabal b/unordered-containers.cabal --- a/unordered-containers.cabal +++ b/unordered-containers.cabal @@ -1,5 +1,5 @@ name: unordered-containers -version: 0.2.9.0 +version: 0.2.10.0 synopsis: Efficient hashing-based container types description: Efficient hashing-based container types. The containers have been @@ -18,7 +18,7 @@ 2010 Edward Z. Yang category: Data build-type: Simple -cabal-version: >=1.8 +cabal-version: >=1.10 extra-source-files: CHANGES.md tested-with: GHC==8.4.1, GHC==8.2.2, GHC==8.0.2, GHC==7.10.3, GHC==7.8.4 @@ -34,15 +34,19 @@ other-modules: Data.HashMap.Array Data.HashMap.Base + Data.HashMap.Strict.Base Data.HashMap.List Data.HashMap.Unsafe Data.HashMap.UnsafeShift + Data.HashSet.Base build-depends: base >= 4.7 && < 5, deepseq >= 1.1, hashable >= 1.0.1.1 && < 1.3 + default-language: Haskell2010 + other-extensions: RoleAnnotations, UnboxedTuples, @@ -67,13 +71,14 @@ build-depends: base, - containers >= 0.4, + containers >= 0.5.8, hashable >= 1.0.1.1, QuickCheck >= 2.4.0.1, test-framework >= 0.3.3, test-framework-quickcheck2 >= 0.2.9, unordered-containers + default-language: Haskell2010 ghc-options: -Wall cpp-options: -DASSERTS @@ -84,13 +89,14 @@ build-depends: base, - containers >= 0.4, + containers >= 0.5.8, hashable >= 1.0.1.1, QuickCheck >= 2.4.0.1, test-framework >= 0.3.3, test-framework-quickcheck2 >= 0.2.9, unordered-containers + default-language: Haskell2010 ghc-options: -Wall cpp-options: -DASSERTS -DSTRICT @@ -108,6 +114,7 @@ test-framework-quickcheck2 >= 0.2.9, unordered-containers + default-language: Haskell2010 ghc-options: -Wall cpp-options: -DASSERTS @@ -125,6 +132,7 @@ test-framework >= 0.3.3, test-framework-quickcheck2 >= 0.2.9 + default-language: Haskell2010 ghc-options: -Wall cpp-options: -DASSERTS @@ -143,6 +151,7 @@ test-framework-quickcheck2, unordered-containers + default-language: Haskell2010 ghc-options: -Wall cpp-options: -DASSERTS @@ -161,6 +170,7 @@ test-framework-quickcheck2 >= 0.2.9, unordered-containers + default-language: Haskell2010 ghc-options: -Wall cpp-options: -DASSERTS @@ -177,15 +187,17 @@ Data.HashMap.Base Data.HashMap.Lazy Data.HashMap.Strict + Data.HashMap.Strict.Base Data.HashMap.Unsafe Data.HashMap.UnsafeShift Data.HashSet + Data.HashSet.Base Util.ByteString Util.Int Util.String build-depends: - base, + base >= 4.8.0, bytestring, containers, criterion >= 1.0 && < 1.3, @@ -196,6 +208,7 @@ mtl, random + default-language: Haskell2010 ghc-options: -Wall -O2 -rtsopts -fwarn-tabs -ferror-spans if flag(debug) cpp-options: -DASSERTS