Adding manual worker/wrapper transforms to Data.Map

19 Aug 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

Adding manual worker/wrapper transforms to Data.Map

Johan Tibell