Skip to content

Commit

Permalink
Reduce code size with unboxed unary tuples
Browse files Browse the repository at this point in the history
Experimentally make most strict and lazy functions share code
using unboxed unary tuples. I fear we may find this too expensive,
but it's an idea. Needs heavy benchmarking.

Addresses #64
  • Loading branch information
treeowl committed Feb 8, 2018
1 parent d672a11 commit 5b1218a
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 257 deletions.
88 changes: 48 additions & 40 deletions Data/HashMap/Base.hs
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,8 @@ module Data.HashMap.Base
, insertModifying
, ptrEq
, adjust#
, unionWithKey#
, unsafeInsertModifying
) where

#if __GLASGOW_HASKELL__ < 710
Expand Down Expand Up @@ -650,7 +652,7 @@ insert' h0 k0 v0 m0 = go h0 k0 v0 0 m0
else Full (update16 ary i st')
where i = index h s
go h k x s t@(Collision hy v)
| h == hy = Collision h (updateOrSnocWith const k x v)
| h == hy = Collision h (updateOrSnocWith (\v1 _ -> (# v1 #)) k x v)
| otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
{-# INLINABLE insert' #-}

Expand Down Expand Up @@ -773,7 +775,7 @@ unsafeInsert k0 v0 m0 = runST (go h0 k0 v0 0 m0)
return t
where i = index h s
go h k x s t@(Collision hy v)
| h == hy = return $! Collision h (updateOrSnocWith const k x v)
| h == hy = return $! Collision h (updateOrSnocWith (\v1 _ -> (# v1 #)) k x v)
| otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
{-# INLINABLE unsafeInsert #-}

Expand Down Expand Up @@ -809,30 +811,30 @@ insertWith :: (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v
-> HashMap k v
-- 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
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
insertModifying :: (Eq k, Hashable k) => ((# #) -> (# v #)) -> (v -> (# v #)) -> k -> HashMap k v
-> HashMap k v
insertModifying x f k0 m0 = go h0 k0 0 m0
where
!h0 = hash k0
go !h !k !_ Empty = Leaf h (L k x)
go !h !k !_ Empty = case x (# #) of (# new #) -> Leaf h (L k new)
go h k s t@(Leaf hy l@(L ky y))
| hy == h = if ky == k
then case f y of
(# v' #) | ptrEq y v' -> t
| otherwise -> Leaf h (L k (v'))
else collision h l (L k x)
| otherwise = runST (two s h k x hy ky y)
else case x (# #) of (# new #) -> collision h l (L k new)
| otherwise = case x (# #) of (# new #) -> runST (two s h k new hy ky y)
go h k s t@(BitmapIndexed b ary)
| b .&. m == 0 =
let ary' = A.insert ary i $! Leaf h (L k x)
let ary' = case x (# #) of (# new #) -> A.insert ary i $! Leaf h (L k new)
in bitmapIndexedOrFull (b .|. m) ary'
| otherwise =
let !st = A.index ary i
Expand Down Expand Up @@ -861,7 +863,7 @@ insertModifying x f k0 m0 = go h0 k0 0 m0
{-# INLINABLE insertModifying #-}

-- Like insertModifying for arrays; used to implement insertModifying
insertModifyingArr :: Eq k => v -> (v -> (# v #)) -> k -> A.Array (Leaf k v)
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
Expand All @@ -870,7 +872,7 @@ insertModifyingArr x f k0 ary0 = go k0 ary0 0 (A.length ary0)
-- 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)
case x (# #) of (# new #) -> A.write mary n (L k new)
return mary
| otherwise = case A.index ary i of
(L kx y) | k == kx -> case f y of
Expand All @@ -881,40 +883,40 @@ insertModifyingArr x f k0 ary0 = go k0 ary0 0 (A.length ary0)
{-# INLINE insertModifyingArr #-}

-- | In-place update version of insertWith
unsafeInsertWith :: forall k v. (Eq k, Hashable k)
=> (v -> v -> v) -> k -> v -> HashMap k v
unsafeInsertModifying :: forall k v. (Eq k, Hashable k)
=> ((# #) -> (# v #)) -> (v -> (# v #)) -> k -> HashMap k v
-> HashMap k v
unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
unsafeInsertModifying v0 f k0 m0 = runST (go h0 k0 0 m0)
where
h0 = hash k0
go :: Hash -> k -> v -> Shift -> HashMap k v -> ST s (HashMap k v)
go !h !k x !_ Empty = return $! Leaf h (L k x)
go h k x s (Leaf hy l@(L ky y))
go :: Hash -> k -> Shift -> HashMap k v -> ST s (HashMap k v)
go !h !k !_ Empty = case v0 (# #) of (# x #) -> return $! Leaf h (L k x)
go h k s (Leaf hy l@(L ky y))
| hy == h = if ky == k
then return $! Leaf h (L k (f x y))
else return $! collision h l (L k x)
| otherwise = two s h k x hy ky y
go h k x s t@(BitmapIndexed b ary)
then case f y of (# v #) -> return $! Leaf h (L k v)
else case v0 (# #) of (# x #) -> return $! collision h l (L k x)
| otherwise = case v0 (# #) of (# x #) -> two s h k x hy ky y
go h k s t@(BitmapIndexed b ary)
| b .&. m == 0 = do
ary' <- A.insertM ary i $! Leaf h (L k x)
ary' <- case v0 (# #) of (# x #) -> A.insertM ary i $! Leaf h (L k x)
return $! bitmapIndexedOrFull (b .|. m) ary'
| otherwise = do
st <- A.indexM ary i
st' <- go h k x (s+bitsPerSubkey) st
st' <- go h k (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
go h k s t@(Full ary) = do
st <- A.indexM ary i
st' <- go h k x (s+bitsPerSubkey) st
st' <- go h k (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 #-}
go h k s t@(Collision hy v)
| h == hy = return $! Collision h (insertModifyingArr v0 f k v)
| otherwise = go h k s $ BitmapIndexed (mask hy s) (A.singleton t)
{-# INLINABLE unsafeInsertModifying #-}

-- | /O(log n)/ Remove the mapping for the specified key from this map
-- if present.
Expand Down Expand Up @@ -1157,7 +1159,7 @@ bogus# _ = error "Data.HashMap.alterF internal error: hit bogus#"
-- 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
coerce (insertModifying (\_ -> (# x #)) (\mold -> case runIdentity (f (Just mold)) of
Nothing -> bogus# (# #)
Just new -> (# new #)))

Expand Down Expand Up @@ -1256,22 +1258,27 @@ unionWith f = unionWithKey (const f)
-- result.
unionWithKey :: (Eq k, Hashable k) => (k -> v -> v -> v) -> HashMap k v -> HashMap k v
-> HashMap k v
unionWithKey f = go 0
unionWithKey f m = unionWithKey# (\k v1 v2 -> (# f k v1 v2 #)) m
{-# INLINE unionWithKey #-}

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 (L k1 (f k1 v1 v2))
then case f k1 v1 v2 of (# v #) -> Leaf h1 (L k1 v)
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)
| h1 == h2 = Collision h1 (updateOrSnocWithKey (\q w x -> f q x w) 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)
Expand Down Expand Up @@ -1336,7 +1343,8 @@ unionWithKey f = go 0
where
m1 = mask h1 s
m2 = mask h2 s
{-# INLINE unionWithKey #-}
{-# INLINE unionWithKey# #-}


-- | Strict in the result of @f@.
unionArrayBy :: (a -> a -> a) -> Bitmap -> Bitmap -> A.Array a -> A.Array a
Expand Down Expand Up @@ -1667,7 +1675,7 @@ fromList = L.foldl' (\ m (k, v) -> unsafeInsert k v m) empty
-- | /O(n*log n)/ Construct a map from a list of elements. Uses
-- the provided function to merge duplicate entries.
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
fromListWith f = L.foldl' (\ m (k, v) -> unsafeInsertModifying (\_ -> (# v #)) (\y -> (# f v y #)) k m) empty
{-# INLINE fromListWith #-}

------------------------------------------------------------------------
Expand Down Expand Up @@ -1719,12 +1727,12 @@ updateWith# f k0 ary0 = go k0 ary0 0 (A.length ary0)
| otherwise -> go k ary (i+1) n
{-# INLINABLE updateWith# #-}

updateOrSnocWith :: Eq k => (v -> v -> v) -> k -> v -> A.Array (Leaf k v)
updateOrSnocWith :: Eq k => (v -> v -> (# v #)) -> k -> v -> A.Array (Leaf k v)
-> A.Array (Leaf k v)
updateOrSnocWith f = updateOrSnocWithKey (const f)
{-# INLINABLE updateOrSnocWith #-}

updateOrSnocWithKey :: Eq k => (k -> v -> v -> v) -> k -> v -> A.Array (Leaf k v)
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
Expand All @@ -1736,15 +1744,15 @@ updateOrSnocWithKey f k0 v0 ary0 = go k0 v0 ary0 0 (A.length ary0)
A.write mary n (L k v)
return mary
| otherwise = case A.index ary i of
(L kx y) | k == kx -> A.update ary i (L k (f k v y))
(L kx y) | k == kx -> case f k v y of (# y' #) -> A.update ary i (L k y')
| otherwise -> go k v ary (i+1) n
{-# INLINABLE updateOrSnocWithKey #-}

updateOrConcatWith :: Eq k => (v -> v -> v) -> A.Array (Leaf k v) -> A.Array (Leaf k v) -> A.Array (Leaf k v)
updateOrConcatWith :: Eq k => (v -> v -> (# v #)) -> A.Array (Leaf k v) -> A.Array (Leaf k v) -> A.Array (Leaf k v)
updateOrConcatWith f = updateOrConcatWithKey (const f)
{-# INLINABLE updateOrConcatWith #-}

updateOrConcatWithKey :: Eq k => (k -> v -> v -> v) -> A.Array (Leaf k v) -> A.Array (Leaf k v) -> A.Array (Leaf k v)
updateOrConcatWithKey :: 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
-- first: look up the position of each element of ary2 in ary1
let indices = A.map (\(L k _) -> indexOf k ary1) ary2
Expand All @@ -1763,7 +1771,7 @@ updateOrConcatWithKey f ary1 ary2 = A.run $ do
Just i1 -> do -- key occurs in both arrays, store combination in position i1
L k v1 <- A.indexM ary1 i1
L _ v2 <- A.indexM ary2 i2
A.write mary i1 (L k (f k v1 v2))
case f k v1 v2 of (# v' #) -> A.write mary i1 (L k v')
go iEnd (i2+1)
Nothing -> do -- key is only in ary2, append to end
A.write mary iEnd =<< A.indexM ary2 i2
Expand Down
Loading

0 comments on commit 5b1218a

Please sign in to comment.