Reduce code size with unboxed unary tuples

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 haskell-unordered-containers#64

Author: treeowl

Data/HashMap/Base.hs

@@ -107,6 +107,8 @@ module Data.HashMap.Base
     , insertModifying
     , ptrEq
     , adjust#
+    , unionWithKey#
+    , unsafeInsertWith
     ) where
 
 #if __GLASGOW_HASKELL__ < 710
@@ -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' #-}
 
@@ -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 #-}
 
@@ -882,7 +884,7 @@ insertModifyingArr x f k0 ary0 = go k0 ary0 0 (A.length ary0)
 
 -- | In-place update version of insertWith
 unsafeInsertWith :: forall k v. (Eq k, Hashable k)
-                 => (v -> v -> v) -> k -> v -> HashMap k v
+                 => (v -> v -> (# v #)) -> k -> v -> HashMap k v
                  -> HashMap k v
 unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
   where
@@ -891,7 +893,7 @@ unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
     go !h !k x !_ Empty = return $! Leaf h (L k x)
     go h k x s (Leaf hy l@(L ky y))
         | hy == h = if ky == k
-                    then return $! Leaf h (L k (f x y))
+                    then case f x y of (# v #) -> return $! Leaf h (L k v)
                     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)
@@ -1256,6 +1258,9 @@ 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 m = unionWithKey# (\k v1 v2 -> (# f k v1 v2 #)) m
+{-# INLINE unionWithKey #-}
+{-
 unionWithKey f = go 0
   where
     -- empty vs. anything
@@ -1337,6 +1342,93 @@ unionWithKey f = go 0
         m1 = mask h1 s
         m2 = mask h2 s
 {-# 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 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 (\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)
+        | otherwise = goDifferentHash s h1 h2 t1 t2
+    -- branch vs. branch
+    go s (BitmapIndexed b1 ary1) (BitmapIndexed b2 ary2) =
+        let b'   = b1 .|. b2
+            ary' = unionArrayBy (go (s+bitsPerSubkey)) b1 b2 ary1 ary2
+        in bitmapIndexedOrFull b' ary'
+    go s (BitmapIndexed b1 ary1) (Full ary2) =
+        let ary' = unionArrayBy (go (s+bitsPerSubkey)) b1 fullNodeMask ary1 ary2
+        in Full ary'
+    go s (Full ary1) (BitmapIndexed b2 ary2) =
+        let ary' = unionArrayBy (go (s+bitsPerSubkey)) fullNodeMask b2 ary1 ary2
+        in Full ary'
+    go s (Full ary1) (Full ary2) =
+        let ary' = unionArrayBy (go (s+bitsPerSubkey)) fullNodeMask fullNodeMask
+                   ary1 ary2
+        in Full ary'
+    -- leaf vs. branch
+    go s (BitmapIndexed b1 ary1) t2
+        | b1 .&. m2 == 0 = let ary' = A.insert ary1 i t2
+                               b'   = b1 .|. m2
+                           in bitmapIndexedOrFull b' ary'
+        | otherwise      = let ary' = A.updateWith' ary1 i $ \st1 ->
+                                   go (s+bitsPerSubkey) st1 t2
+                           in BitmapIndexed b1 ary'
+        where
+          h2 = leafHashCode t2
+          m2 = mask h2 s
+          i = sparseIndex b1 m2
+    go s t1 (BitmapIndexed b2 ary2)
+        | b2 .&. m1 == 0 = let ary' = A.insert ary2 i $! t1
+                               b'   = b2 .|. m1
+                           in bitmapIndexedOrFull b' ary'
+        | otherwise      = let ary' = A.updateWith' ary2 i $ \st2 ->
+                                   go (s+bitsPerSubkey) t1 st2
+                           in BitmapIndexed b2 ary'
+      where
+        h1 = leafHashCode t1
+        m1 = mask h1 s
+        i = sparseIndex b2 m1
+    go s (Full ary1) t2 =
+        let h2   = leafHashCode t2
+            i    = index h2 s
+            ary' = update16With' ary1 i $ \st1 -> go (s+bitsPerSubkey) st1 t2
+        in Full ary'
+    go s t1 (Full ary2) =
+        let h1   = leafHashCode t1
+            i    = index h1 s
+            ary' = update16With' ary2 i $ \st2 -> go (s+bitsPerSubkey) t1 st2
+        in Full ary'
+
+    leafHashCode (Leaf h _) = h
+    leafHashCode (Collision h _) = h
+    leafHashCode _ = error "leafHashCode"
+
+    goDifferentHash s h1 h2 t1 t2
+        | m1 == m2  = BitmapIndexed m1 (A.singleton $! go (s+bitsPerSubkey) t1 t2)
+        | m1 <  m2  = BitmapIndexed (m1 .|. m2) (A.pair t1 t2)
+        | otherwise = BitmapIndexed (m1 .|. m2) (A.pair t2 t1)
+      where
+        m1 = mask h1 s
+        m2 = mask h2 s
+{-# INLINE unionWithKey# #-}
+
 
 -- | Strict in the result of @f@.
 unionArrayBy :: (a -> a -> a) -> Bitmap -> Bitmap -> A.Array a -> A.Array a
@@ -1667,7 +1759,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) -> unsafeInsertWith (\x y -> (# f x y #)) k v m) empty
 {-# INLINE fromListWith #-}
 
 ------------------------------------------------------------------------
@@ -1719,12 +1811,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
@@ -1736,15 +1828,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
@@ -1763,7 +1855,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