Surprising strictness properties of pre-order fold over a Data.Map

18 Aug 2010

      Hi,

I was adding a strict pre-order fold to the Data.Map module and I ran into
this slightly surprising behavior. Modeled on foldl' for lists I defined

    foldlWithKey' :: (b -> k -> a -> b) -> b -> Map k a -> b
    foldlWithKey' f z0 m = go z0 m
      where
        go z Tip              = z
        go z (Bin _ kx x l r) = let x' = f (go z l) kx x in x' `seq` go x' r
    {-# INLINE foldlWithKey' #-}

and, in a separate module, I defined this test

    module Test (test) where

    import qualified Data.Map as M

    test :: M.Map Int Int -> Int
    test m = M.foldlWithKey' (\n k v -> n + k + v) 0 m

which generates this core:

    test1 :: Int
    test1 = I# 0

    test_go2 :: Int -> Data.Map.Map Int Int -> Int
    test_go2 =
      \ (z_ani :: Int)
        (ds_anj :: Data.Map.Map Int Int) ->
        case ds_anj of _ {
          Data.Map.Tip -> z_ani;
          Data.Map.Bin _ kx_anp x_anq l_anr r_ans ->
            case test_go2 z_ani l_anr of _ { I# x1_ao5 ->
            case kx_anp of _ { I# y_ao9 ->
            case x_anq of _ { I# y1_Xot ->
            test_go2
              (I# (+# (+# x1_ao5 y_ao9) y1_Xot))
              r_ans
            }
            }
            }
        }

    test :: Data.Map.Map Int Int -> Int
    test =
      \ (m_ajo :: Data.Map.Map Int Int) ->
        test_go2 test1 m_ajo

Note how the accumulator 'z' is not unboxed in the loop. I don't quite
understand why that is. I do know I can get the core that I want by defining

    foldlWithKey2' :: (b -> k -> a -> b) -> b -> Map k a -> b
    foldlWithKey2' f z0 m = go z0 m
      where
        go z _ | z `seq` False = undefined
        go z Tip              = z
        go z (Bin _ kx x l r) = go (f (go z l) kx x) r
    {-# INLINE foldlWithKey2' #-}

and

    module Test (test2) where

    import qualified Data.Map as M

    test2 :: M.Map Int Int -> Int
    test2 m = M.foldlWithKey2' (\n k v -> n + k + v) 0 m

you get this core:

    test2_$s$wgo2 :: Data.Map.Map Int Int -> Int# -> Int#
    test2_$s$wgo2 =
      \ (sc_soS :: Data.Map.Map Int Int)
        (sc1_soT :: Int#) ->
        case sc_soS of _ {
          Data.Map.Tip -> sc1_soT;
          Data.Map.Bin _ kx_anK x_anL l_anM r_anN ->
            case test2_$s$wgo2 l_anM sc1_soT of ww_sou { __DEFAULT ->
            case kx_anK of _ { I# y_ao9 ->
            case x_anL of _ { I# y1_Xow ->
            test2_$s$wgo2
              r_anN (+# (+# ww_sou y_ao9) y1_Xow)
            }
            }
            }
        }

    $wtest2 :: Data.Map.Map Int Int -> Int#
    $wtest2 =
      \ (w_sox :: Data.Map.Map Int Int) ->
        test2_$s$wgo2 w_sox 0

    test2 :: Data.Map.Map Int Int -> Int
    test2 =
      __inline_me (\ (w_sox :: Data.Map.Map
                                 Int Int) ->
                     case $wtest2 w_sox of ww_soA { __DEFAULT ->
                     I# ww_soA
                     })

Note that the accumulator is now unboxed.

Could someone please explain the difference. I would like to be able to to
understand when I would get the former or the latter by looking at the
Haskell source.

Cheers,
Johan

Surprising strictness properties of pre-order fold over a Data.Map

Johan Tibell