Skip to content

Commit ba708e3

Browse files
committed
Make sure we perform all array lookups eagerly
Some of these changes probably don't make a difference in context, but they at least make it easier to see that we're doing the right thing.
1 parent 8201efe commit ba708e3

File tree

2 files changed

+16
-5
lines changed

2 files changed

+16
-5
lines changed

Data/HashMap/Array.hs

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ module Data.HashMap.Array
2222
, write
2323
, index
2424
, indexM
25+
, index#
2526
, update
2627
, updateWith'
2728
, unsafeUpdateM
@@ -234,7 +235,10 @@ rnfArray ary0 = go ary0 n0 0
234235
n0 = length ary0
235236
go !ary !n !i
236237
| i >= n = ()
237-
| otherwise = rnf (index ary i) `seq` go ary n (i+1)
238+
| (# x #) <- index# ary i
239+
= rnf x `seq` go ary n (i+1)
240+
-- We use index# just in case GHC can't see that the
241+
-- relevant rnf is strict, or in case it actually isn't.
238242
{-# INLINE rnfArray #-}
239243

240244
-- | Create a new mutable array of specified size, in the specified
@@ -388,7 +392,9 @@ updateM ary idx b =
388392
-- applying a function to it. Evaluates the element to WHNF before
389393
-- inserting it into the array.
390394
updateWith' :: Array e -> Int -> (e -> e) -> Array e
391-
updateWith' ary idx f = update ary idx $! f (index ary idx)
395+
updateWith' ary idx f
396+
| (# x #) <- index# ary idx
397+
= update ary idx $! f x
392398
{-# INLINE updateWith' #-}
393399

394400
-- | /O(1)/ Update the element at the given position in this array,

Data/HashMap/Base.hs

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,7 @@ instance (Show k, Show v) => Show (HashMap k v) where
267267

268268
instance Traversable (HashMap k) where
269269
traverse f = traverseWithKey (const f)
270+
{-# INLINABLE traverse #-}
270271

271272
#if MIN_VERSION_base(4,9,0)
272273
instance Eq2 HashMap where
@@ -1387,7 +1388,9 @@ unionArrayBy f b1 b2 ary1 ary2 = A.run $ do
13871388
| m > b' = return ()
13881389
| b' .&. m == 0 = go i i1 i2 (m `unsafeShiftL` 1)
13891390
| ba .&. m /= 0 = do
1390-
A.write mary i $! f (A.index ary1 i1) (A.index ary2 i2)
1391+
x1 <- A.indexM ary1 i1
1392+
x2 <- A.indexM ary2 i2
1393+
A.write mary i $! f x1 x2
13911394
go (i+1) (i1+1) (i2+1) (m `unsafeShiftL` 1)
13921395
| b1 .&. m /= 0 = do
13931396
A.write mary i =<< A.indexM ary1 i1
@@ -1657,7 +1660,7 @@ filterMapAux onLeaf onColl = go
16571660
return $! Collision h ary2
16581661
| otherwise -> do ary2 <- A.trim mary j
16591662
return $! Collision h ary2
1660-
| Just el <- onColl (A.index ary i)
1663+
| Just el <- onColl $! A.index ary i
16611664
= A.write mary j el >> step ary mary (i+1) (j+1) n
16621665
| otherwise = step ary mary (i+1) j n
16631666
{-# INLINE filterMapAux #-}
@@ -1830,7 +1833,9 @@ update16M ary idx b = do
18301833

18311834
-- | /O(n)/ Update the element at the given position in this array, by applying a function to it.
18321835
update16With' :: A.Array e -> Int -> (e -> e) -> A.Array e
1833-
update16With' ary idx f = update16 ary idx $! f (A.index ary idx)
1836+
update16With' ary idx f
1837+
| (# x #) <- A.index# ary idx
1838+
= update16 ary idx $! f x
18341839
{-# INLINE update16With' #-}
18351840

18361841
-- | Unsafely clone an array of 16 elements. The length of the input

0 commit comments

Comments
 (0)