Skip to content
This repository
Browse code

Only call runST once in insertWith'

Improves performance by 10% on insert/Int.  Before this change we used
to (implicitly) call runST once per level of the tree, causing
closures to be allocated.
  • Loading branch information...
commit 58b7815a057b3575c58b746d5971d59d806b1333 1 parent 7c19a6d
Johan Tibell authored February 29, 2012
50  Data/HashMap/Array.hs
@@ -11,6 +11,7 @@ module Data.HashMap.Array
11 11
     , new
12 12
     , new_
13 13
     , singleton
  14
+    , singleton'
14 15
     , pair
15 16
 
16 17
       -- * Basic interface
@@ -22,8 +23,10 @@ module Data.HashMap.Array
22 23
     , index_
23 24
     , indexM_
24 25
     , update
  26
+    , update'
25 27
     , updateWith
26 28
     , insert
  29
+    , insert'
27 30
     , delete
28 31
 
29 32
     , unsafeFreeze
@@ -145,9 +148,13 @@ new_ :: Int -> ST s (MArray s a)
145 148
 new_ n = new n undefinedElem
146 149
 
147 150
 singleton :: a -> Array a
148  
-singleton x = run (new 1 x)
  151
+singleton x = runST (singleton' x)
149 152
 {-# INLINE singleton #-}
150 153
 
  154
+singleton' :: a -> ST s (Array a)
  155
+singleton' x = new 1 x >>= unsafeFreeze
  156
+{-# INLINE singleton' #-}
  157
+
151 158
 pair :: a -> a -> Array a
152 159
 pair x y = run $ do
153 160
     ary <- new 2 x
@@ -249,28 +256,37 @@ copyM !src !sidx !dst !didx n =
249 256
 -- | /O(n)/ Insert an element at the given position in this array,
250 257
 -- increasing its size by one.
251 258
 insert :: Array e -> Int -> e -> Array e
252  
-insert ary idx b =
253  
-    CHECK_BOUNDS("insert", count + 1, idx)
254  
-        run $ do
255  
-            mary <- new_ (count+1)
256  
-            copy ary 0 mary 0 idx
257  
-            write mary idx b
258  
-            copy ary idx mary (idx+1) (count-idx)
259  
-            return mary
260  
-  where !count = length ary
  259
+insert ary idx b = runST (insert' ary idx b)
261 260
 {-# INLINE insert #-}
262 261
 
  262
+-- | /O(n)/ Insert an element at the given position in this array,
  263
+-- increasing its size by one.
  264
+insert' :: Array e -> Int -> e -> ST s (Array e)
  265
+insert' ary idx b =
  266
+    CHECK_BOUNDS("insert'", count + 1, idx)
  267
+        do mary <- new_ (count+1)
  268
+           copy ary 0 mary 0 idx
  269
+           write mary idx b
  270
+           copy ary idx mary (idx+1) (count-idx)
  271
+           unsafeFreeze mary
  272
+  where !count = length ary
  273
+{-# INLINE insert' #-}
  274
+
263 275
 -- | /O(n)/ Update the element at the given position in this array.
264 276
 update :: Array e -> Int -> e -> Array e
265  
-update ary idx b =
266  
-    CHECK_BOUNDS("update", count, idx)
267  
-        run $ do
268  
-            mary <- thaw ary 0 count
269  
-            write mary idx b
270  
-            return mary
271  
-  where !count = length ary
  277
+update ary idx b = runST (update' ary idx b)
272 278
 {-# INLINE update #-}
273 279
 
  280
+-- | /O(n)/ Update the element at the given position in this array.
  281
+update' :: Array e -> Int -> e -> ST s (Array e)
  282
+update' ary idx b =
  283
+    CHECK_BOUNDS("update'", count, idx)
  284
+        do mary <- thaw ary 0 count
  285
+           write mary idx b
  286
+           unsafeFreeze mary
  287
+  where !count = length ary
  288
+{-# INLINE update' #-}
  289
+
274 290
 -- | /O(n)/ Update the element at the given positio in this array, by
275 291
 -- applying a function to it.  Evaluates the element to WHNF before
276 292
 -- inserting it into the array.
56  Data/HashMap/Base.hs
@@ -250,34 +250,36 @@ insertWith = insertWith'
250 250
 -- pay the cost of using the higher-order argument.
251 251
 insertWith' :: (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v
252 252
             -> HashMap k v
253  
-insertWith' f k0 v0 = go h0 k0 v0 0
  253
+insertWith' f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
254 254
   where
255 255
     h0 = hash k0
256  
-    go !h !k x !_ Empty = Leaf h (L k x)
  256
+    go !h !k x !_ Empty = return $! Leaf h (L k x)
257 257
     go h k x s t@(Leaf hy l@(L ky y))
258 258
         | hy == h = if ky == k
259  
-                    then Leaf h (L k (f x y))
260  
-                    else collision h l (L k x)
261  
-        | otherwise = go h k x s $ BitmapIndexed (bitpos hy s) (A.singleton t)
  259
+                    then return $! Leaf h (L k (f x y))
  260
+                    else return $! collision h l (L k x)
  261
+        | otherwise = do
  262
+            leaf <- A.singleton' t
  263
+            go h k x s $ BitmapIndexed (bitpos hy s) leaf
262 264
     go h k x s (BitmapIndexed b ary)
263  
-        | b .&. m == 0 = let l    = Leaf h (L k x)
264  
-                             ary' = A.insert ary i $! l
265  
-                             b'   = b .|. m
266  
-                         in bitmapIndexedOrFull b' ary'
267  
-        | otherwise = let  st   = A.index ary i
268  
-                           st'  = go h k x (s+bitsPerSubkey) st
269  
-                           ary' = A.update ary i $! st'
270  
-                      in BitmapIndexed b ary'
  265
+        | b .&. m == 0 = do
  266
+            ary' <- A.insert' ary i $! Leaf h (L k x)
  267
+            return $! bitmapIndexedOrFull (b .|. m) ary'
  268
+        | otherwise = do
  269
+            st <- A.index_ ary i
  270
+            st' <- go h k x (s+bitsPerSubkey) st
  271
+            ary' <- A.update' ary i st'
  272
+            return $! BitmapIndexed b ary'
271 273
       where m = bitpos h s
272 274
             i = index b m
273  
-    go h k x s (Full ary) =
274  
-        let i    = mask h s
275  
-            st   = A.index ary i
276  
-            st'  = go h k x (s+bitsPerSubkey) st
277  
-            ary' = update16 ary i $! st'
278  
-        in Full ary'
  275
+    go h k x s (Full ary) = do
  276
+        st <- A.index_ ary i
  277
+        st' <- go h k x (s+bitsPerSubkey) st
  278
+        ary' <- update16' ary i st'
  279
+        return $! Full ary'
  280
+      where i = mask h s
279 281
     go h k x s t@(Collision hy v)
280  
-        | h == hy   = Collision h (updateOrSnocWith f k x v)
  282
+        | h == hy   = return $! Collision h (updateOrSnocWith f k x v)
281 283
         | otherwise = go h k x s $ BitmapIndexed (bitpos hy s) (A.singleton t)
282 284
 {-# INLINE insertWith' #-}
283 285
 
@@ -792,13 +794,17 @@ updateOrConcatWith f ary1 ary2 = A.run $ do
792 794
 
793 795
 -- | /O(n)/ Update the element at the given position in this array.
794 796
 update16 :: A.Array e -> Int -> e -> A.Array e
795  
-update16 ary idx b =
796  
-    A.run $ do
797  
-        mary <- clone16 ary
798  
-        A.write mary idx b
799  
-        return mary
  797
+update16 ary idx b = runST (update16' ary idx b)
800 798
 {-# INLINE update16 #-}
801 799
 
  800
+-- | /O(n)/ Update the element at the given position in this array.
  801
+update16' :: A.Array e -> Int -> e -> ST s (A.Array e)
  802
+update16' ary idx b = do
  803
+    mary <- clone16 ary
  804
+    A.write mary idx b
  805
+    A.unsafeFreeze mary
  806
+{-# INLINE update16' #-}
  807
+
802 808
 -- | /O(n)/ Update the element at the given position in this array, by applying a function to it.
803 809
 update16With :: A.Array e -> Int -> (e -> e) -> A.Array e
804 810
 update16With ary idx f = update16 ary idx $! f (A.index ary idx)

0 notes on commit 58b7815

Please sign in to comment.
Something went wrong with that request. Please try again.