Adding manual worker/wrapper transforms to Data.Map

Johan Tibell johan.tibell at gmail.com
Thu Aug 19 05:38:10 EDT 2010

```Hi all,

I tried doing the "standard" worker/wrapper transform to some functions in
Data.Map. For example, by transforming

insertWith' :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWith' f k x m
= insertWithKey' (\_ x' y' -> f x' y') k x m

-- | Same as 'insertWithKey', but the combining function is applied
strictly.
insertWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a ->
Map k a
insertWithKey' f kx x t0
= case t of
Tip -> singleton kx \$! x
Bin sy ky y l r
-> case compare kx ky of
LT -> balance ky y (insertWithKey' f kx x l) r
GT -> balance ky y l (insertWithKey' f kx x r)
EQ -> let x' = f kx x y in seq x' (Bin sy kx x' l r)

to

insertWith' :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWith' f k x m
= insertWithKey' (\_ x' y' -> f x' y') k x m
{-# INLINE insertWith' #-}

insertWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a ->
Map k a
insertWithKey' f kx x t0 = kx `seq` go t0
where
go t = case t of
Tip -> singleton kx \$! x
Bin sy ky y l r
-> case compare kx ky of
LT -> balance ky y (go l) r
GT -> balance ky y l (go r)
EQ -> let x' = f kx x y in seq x' (Bin sy kx x' l r)
{-# INLINE insertWithKey' #-}

I got a 16% speedup on this benchmark:

{-# LANGUAGE BangPatterns #-}
module Main where

import Criterion.Main
import qualified Data.Map as M

main = defaultMain
[ bench "insertWith20k/size" \$ whnf (M.size . insertWith) n
]
where
-- Number of elements
n = 20000

insertWith :: Int -> M.Map Int Int
insertWith max = go 0 M.empty
where
go :: Int -> M.Map Int Int -> M.Map Int Int
go n !m
| n >= max  = m
| otherwise = go (n + 1) \$ M.insertWith' (+) (n `mod` 20) n m

There are lots of other functions in Data.Map that could benefit from the
same transform, in particular some of the folds.

Does anyone see a reason for me to not go ahead and try to create a patch
that performs this transformation on all functions that could benefit from
it? I would include a Criterion benchmark that shows the gains.

Cheers,
Johan
-------------- next part --------------
An HTML attachment was scrubbed...