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/Internal.hs

@@ -120,7 +120,7 @@ module Data.HashMap.Internal
     , two
     , unionArrayBy
     , update32
-    , update32M
+    , update32Mut
     , update32With'
     , updateOrConcatWith
     , updateOrConcatWithKey
@@ -740,7 +740,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 #-}
@@ -1615,7 +1618,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
@@ -2218,16 +2221,16 @@ subsetArray cmpV ary1 ary2 = A.length ary1 <= A.length ary2 && A.all inAry2 ary1
 
 -- | /O(n)/ Update the element at the given position in this array.
 update32 :: A.Array e -> Int -> e -> A.Array e
-update32 ary idx b = runST (update32M ary idx b)
+update32 ary idx b = A.run (update32Mut ary idx b)
 {-# INLINE update32 #-}
 
 -- | /O(n)/ Update the element at the given position in this array.
-update32M :: A.Array e -> Int -> e -> ST s (A.Array e)
-update32M ary idx b = do
+update32Mut :: A.Array e -> Int -> e -> ST s (A.MArray s e)
+update32Mut ary idx b = do
     mary <- clone ary
     A.write mary idx b
-    A.unsafeFreeze mary
-{-# INLINE update32M #-}
+    return mary
+{-# INLINE update32Mut #-}
 
 -- | /O(n)/ Update the element at the given position in this array, by applying a function to it.
 update32With' :: A.Array e -> Int -> (e -> e) -> A.Array e

Data/HashMap/Internal/Array.hs

@@ -86,7 +86,7 @@ import GHC.Exts            (Int (..), SmallArray#, SmallMutableArray#,
                             cloneSmallMutableArray#, copySmallArray#,
                             copySmallMutableArray#, indexSmallArray#,
                             newSmallArray#, readSmallArray#,
-                            reallyUnsafePtrEquality#, sizeofSmallArray#,
+                            reallyUnsafePtrEquality#, runRW#, sizeofSmallArray#,
                             sizeofSmallMutableArray#, tagToEnum#,
                             thawSmallArray#, unsafeCoerce#,
                             unsafeFreezeSmallArray#, unsafeThawSmallArray#,
@@ -205,13 +205,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
@@ -263,7 +267,17 @@ 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 #) ->
+      unsafeFreezeSmallArray# mary s' } of
+    (# _, ary #) -> Array ary
+#else
+run act = runST (act >>= unsafeFreeze)
+#endif
 {-# INLINE run #-}
 
 -- | Unsafely copy the elements of an array. Array bounds are not checked.
@@ -300,33 +314,37 @@ 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) b
            copy ary 0 mary 0 idx
            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 #-}
 
@@ -421,18 +439,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 #-}
 
@@ -442,9 +460,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
@@ -458,9 +477,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
@@ -473,8 +493,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)