Avoid boxing arrays

Previously, we used a lot of things that looked like

    runST (act >>= unsafeFreeze)

The trouble is that GHC can't unbox past the `runRW#` primitive
that `runST` is based on. So this actually allocates an `Array`
constructor which we will then throw away immediately.

The way to avoid this is to call `runRW#` manually to produce
a raw `SmallArray#`, then apply the `Array` constructor on the
outside.

Author: treeowl

Data/HashMap/Array.hs

@@ -36,7 +36,6 @@ module Data.HashMap.Array
     , unsafeThaw
     , unsafeSameArray
     , run
-    , run2
     , copy
     , copyM
 
@@ -59,6 +58,11 @@ import Control.Applicative (Applicative (..), (<$>))
 import Control.Applicative (liftA2)
 import Control.DeepSeq
 import GHC.Exts(Int(..), Int#, reallyUnsafePtrEquality#, tagToEnum#, unsafeCoerce#, State#)
+#if MIN_VERSION_base(4,10,0)
+import GHC.Exts (runRW#)
+#elif MIN_VERSION_base(4,9,0)
+import GHC.Base (runRW#)
+#endif
 import GHC.ST (ST(..))
 import Control.Monad.ST (stToIO)
 
@@ -256,13 +260,17 @@ new_ :: Int -> ST s (MArray s a)
 new_ n = new n undefinedElem
 
 singleton :: a -> Array a
-singleton x = runST (singletonM x)
+singleton x = run (singletonMut x)
 {-# INLINE singleton #-}
 
 singletonM :: a -> ST s (Array a)
 singletonM x = new 1 x >>= unsafeFreeze
 {-# INLINE singletonM #-}
 
+singletonMut :: a -> ST s (MArray s a)
+singletonMut x = new 1 x
+{-# INLINE singletonMut #-}
+
 pair :: a -> a -> Array a
 pair x y = run $ do
     ary <- new 2 x
@@ -314,15 +322,19 @@ unsafeThaw ary
 {-# INLINE unsafeThaw #-}
 
 run :: (forall s . ST s (MArray s e)) -> Array e
-run act = runST $ act >>= unsafeFreeze
+#if MIN_VERSION_base(4,9,0)
+-- GHC can't unbox across the runRW# boundary, so we apply the Array constructor
+-- on the outside.
+run (ST act) =
+  case runRW# $ \s ->
+    case act s of { (# s', MArray mary #) ->
+      unsafeFreezeArray# mary s' } of
+    (# _, ary #) -> Array ary
+#else
+run act = runST (act >>= unsafeFreeze)
+#endif
 {-# INLINE run #-}
 
-run2 :: (forall s. ST s (MArray s e, a)) -> (Array e, a)
-run2 k = runST (do
-                 (marr,b) <- k
-                 arr <- unsafeFreeze marr
-                 return (arr,b))
-
 -- | Unsafely copy the elements of an array. Array bounds are not checked.
 copy :: Array e -> Int -> MArray s e -> Int -> Int -> ST s ()
 copy !src !_sidx@(I# sidx#) !dst !_didx@(I# didx#) _n@(I# n#) =
@@ -357,34 +369,38 @@ trim mary n = cloneM mary 0 n >>= unsafeFreeze
 -- | /O(n)/ Insert an element at the given position in this array,
 -- increasing its size by one.
 insert :: Array e -> Int -> e -> Array e
-insert ary idx b = runST (insertM ary idx b)
+insert ary idx b = run (insertMut ary idx b)
 {-# INLINE insert #-}
 
 -- | /O(n)/ Insert an element at the given position in this array,
 -- increasing its size by one.
 insertM :: Array e -> Int -> e -> ST s (Array e)
-insertM ary idx b =
+insertM ary idx b = insertMut ary idx b >>= unsafeFreeze
+{-# INLINE insertM #-}
+
+insertMut :: Array e -> Int -> e -> ST s (MArray s e)
+insertMut ary idx b =
     CHECK_BOUNDS("insertM", count + 1, idx)
         do mary <- new_ (count+1)
            copy ary 0 mary 0 idx
            write mary idx b
            copy ary idx mary (idx+1) (count-idx)
-           unsafeFreeze mary
+           return mary
   where !count = length ary
-{-# INLINE insertM #-}
+{-# INLINE insertMut #-}
 
 -- | /O(n)/ Update the element at the given position in this array.
 update :: Array e -> Int -> e -> Array e
-update ary idx b = runST (updateM ary idx b)
+update ary idx b = run (updateM ary idx b)
 {-# INLINE update #-}
 
 -- | /O(n)/ Update the element at the given position in this array.
-updateM :: Array e -> Int -> e -> ST s (Array e)
+updateM :: Array e -> Int -> e -> ST s (MArray s e)
 updateM ary idx b =
     CHECK_BOUNDS("updateM", count, idx)
         do mary <- thaw ary 0 count
            write mary idx b
-           unsafeFreeze mary
+           return mary
   where !count = length ary
 {-# INLINE updateM #-}
 
@@ -442,18 +458,18 @@ thaw !ary !_o@(I# o#) !n@(I# n#) =
 -- | /O(n)/ Delete an element at the given position in this array,
 -- decreasing its size by one.
 delete :: Array e -> Int -> Array e
-delete ary idx = runST (deleteM ary idx)
+delete ary idx = run (deleteM ary idx)
 {-# INLINE delete #-}
 
 -- | /O(n)/ Delete an element at the given position in this array,
 -- decreasing its size by one.
-deleteM :: Array e -> Int -> ST s (Array e)
+deleteM :: Array e -> Int -> ST s (MArray s e)
 deleteM ary idx = do
     CHECK_BOUNDS("deleteM", count, idx)
         do mary <- new_ (count-1)
            copy ary 0 mary 0 idx
            copy ary (idx+1) mary idx (count-(idx+1))
-           unsafeFreeze mary
+           return mary
   where !count = length ary
 {-# INLINE deleteM #-}
 
@@ -463,9 +479,10 @@ map f = \ ary ->
     in run $ do
         mary <- new_ n
         go ary mary 0 n
+        return mary
   where
     go ary mary i n
-        | i >= n    = return mary
+        | i >= n    = return ()
         | otherwise = do
              x <- indexM ary i
              write mary i $ f x
@@ -479,9 +496,10 @@ map' f = \ ary ->
     in run $ do
         mary <- new_ n
         go ary mary 0 n
+        return mary
   where
     go ary mary i n
-        | i >= n    = return mary
+        | i >= n    = return ()
         | otherwise = do
              x <- indexM ary i
              write mary i $! f x
@@ -494,8 +512,9 @@ fromList n xs0 =
         run $ do
             mary <- new_ n
             go xs0 mary 0
+            return mary
   where
-    go [] !mary !_   = return mary
+    go [] !_ !_   = return ()
     go (x:xs) mary i = do write mary i x
                           go xs mary (i+1)
 

Data/HashMap/Base.hs

@@ -90,7 +90,6 @@ module Data.HashMap.Base
     , two
     , unionArrayBy
     , update16
-    , update16M
     , update16With'
     , updateOrConcatWith
     , updateOrConcatWithKey
@@ -653,7 +652,10 @@ collision h !e1 !e2 =
 
 -- | Create a 'BitmapIndexed' or 'Full' node.
 bitmapIndexedOrFull :: Bitmap -> A.Array (HashMap k v) -> HashMap k v
-bitmapIndexedOrFull b ary
+-- I don't know if it ever matters in context (once inlined),
+-- but the Core for this function looks a lot nicer if we force
+-- the array argument manually. I don't know why that is.
+bitmapIndexedOrFull b !ary
     | b == fullNodeMask = Full ary
     | otherwise         = BitmapIndexed b ary
 {-# INLINE bitmapIndexedOrFull #-}
@@ -1394,7 +1396,7 @@ unionWithKey f = go 0
 -- | Strict in the result of @f@.
 unionArrayBy :: (a -> a -> a) -> Bitmap -> Bitmap -> A.Array a -> A.Array a
              -> A.Array a
-unionArrayBy f b1 b2 ary1 ary2 = A.run $ do
+unionArrayBy f !b1 !b2 !ary1 !ary2 = A.run $ do
     let b' = b1 .|. b2
     mary <- A.new_ (popCount b')
     -- iterate over nonzero bits of b1 .|. b2
@@ -1836,16 +1838,16 @@ updateOrConcatWithKey f ary1 ary2 = A.run $ do
 
 -- | /O(n)/ Update the element at the given position in this array.
 update16 :: A.Array e -> Int -> e -> A.Array e
-update16 ary idx b = runST (update16M ary idx b)
+update16 ary idx b = A.run (update16Mut ary idx b)
 {-# INLINE update16 #-}
 
 -- | /O(n)/ Update the element at the given position in this array.
-update16M :: A.Array e -> Int -> e -> ST s (A.Array e)
-update16M ary idx b = do
+update16Mut :: A.Array e -> Int -> e -> ST s (A.MArray s e)
+update16Mut ary idx b = do
     mary <- clone16 ary
     A.write mary idx b
-    A.unsafeFreeze mary
-{-# INLINE update16M #-}
+    return mary
+{-# INLINE update16Mut #-}
 
 -- | /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

unordered-containers.cabal

@@ -200,7 +200,7 @@ benchmark benchmarks
     base >= 4.8.0,
     bytestring,
     containers,
-    criterion >= 1.0 && < 1.3,
+    criterion >= 1.0,
     deepseq >= 1.1,
     deepseq-generics,
     hashable >= 1.0.1.1,