Add `oneShot` to the implementation of foldlM

The current (473632d7) implementation of Data.Foldable.foldlM is the like this

foldlM :: (Foldable t, Monad m) => (b -> a -> m b) -> b -> t a -> m b
foldlM f z0 xs = foldr c return xs z0
  -- See Note [List fusion and continuations in 'c']
  where c x k z = f z x >>= k
        {-# INLINE c #-}

It generates an inefficient core for the following example.

f :: Int -> IO Int
f n = foldlM (\a b -> pure $! (a + b)) 0 (filter even [1..n])

Generated core:

-- RHS size: {terms: 48, types: 22, coercions: 12, joins: 0/1}
Main.$wf [InlPrag=NOUSERINLINE[2]]
  :: GHC.Prim.Int#
     -> GHC.Prim.State# GHC.Prim.RealWorld
     -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #)
[GblId,
 Arity=2,
 Caf=NoCafRefs,
 Str=<L,U><L,U>,
 Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
         WorkFree=True, Expandable=True, Guidance=IF_ARGS [0 0] 216 30}]
Main.$wf
  = \ (ww_s6TZ :: GHC.Prim.Int#)
      (w_s6TW :: GHC.Prim.State# GHC.Prim.RealWorld) ->
      case GHC.Prim.># 1# ww_s6TZ of {
        __DEFAULT ->
          letrec {
            go_a5un [Occ=LoopBreaker] :: GHC.Prim.Int# -> Int -> IO Int
            [LclId, Arity=1, Str=<L,U>, Unf=OtherCon []]
            go_a5un
              = \ (x_a5uo :: GHC.Prim.Int#) ->
                  case GHC.Prim.remInt# x_a5uo 2# of {
                    __DEFAULT ->
                      case GHC.Prim.==# x_a5uo ww_s6TZ of {
                        __DEFAULT -> go_a5un (GHC.Prim.+# x_a5uo 1#);
                        1# ->
                          (GHC.Base.$fApplicativeIO4 @ Int)
                          `cast` (<Int>_R ->_R Sym (GHC.Types.N:IO[0] <Int>_R)
                                  :: (Int
                                      -> GHC.Prim.State# GHC.Prim.RealWorld
                                      -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #))
                                     ~R# (Int -> IO Int))
                      };
                    0# ->
                      Main.main_c
                        @ Int
                        (GHC.Types.I# x_a5uo)
                        (case GHC.Prim.==# x_a5uo ww_s6TZ of {
                           __DEFAULT -> go_a5un (GHC.Prim.+# x_a5uo 1#);
                           1# ->
                             (GHC.Base.$fApplicativeIO4 @ Int)
                             `cast` (<Int>_R ->_R Sym (GHC.Types.N:IO[0] <Int>_R)
                                     :: (Int
                                         -> GHC.Prim.State# GHC.Prim.RealWorld
                                         -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #))
                                        ~R# (Int -> IO Int))
                         })
                  }; } in
          ((go_a5un 1# Main.main4)
           `cast` (GHC.Types.N:IO[0] <Int>_R
                   :: IO Int
                      ~R# (GHC.Prim.State# GHC.Prim.RealWorld
                           -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #))))
            w_s6TW;
        1# -> (# w_s6TW, Main.main4 #)
      }

It seems that the main loop go_a5un is not eta-expanded.

I think problem is that oneShot is missing in the definition of foldlM.

When I changed the definition of foldlM as follows,

import GHC.Exts(oneShot)
foldlM :: (Foldable t, Monad m) => (b -> a -> m b) -> b -> t a -> m b
foldlM f z0 xs = foldr c return xs z0
  -- See Note [List fusion and continuations in 'c']
  where c x k = oneShot (\z -> f z x >>= k)
        {-# INLINE c #-}

Then, the main loop of the wf is eta-expaned as expected.

-- RHS size: {terms: 64, types: 46, coercions: 0, joins: 1/1}
Main.$wf [InlPrag=NOUSERINLINE[2]]
  :: GHC.Prim.Int#
     -> GHC.Prim.State# GHC.Prim.RealWorld
     -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #)
[GblId,
 Arity=2,
 Caf=NoCafRefs,
 Str=<L,U><L,U>,
 Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
         WorkFree=True, Expandable=True, Guidance=IF_ARGS [0 0] 136 30}]
Main.$wf
  = \ (ww_s6Xc :: GHC.Prim.Int#)
      (w_s6X9 :: GHC.Prim.State# GHC.Prim.RealWorld) ->
      case GHC.Prim.># 1# ww_s6Xc of {
        __DEFAULT ->
          joinrec {
            go_s6WG [Occ=LoopBreaker]
              :: GHC.Prim.Int#
                 -> Int
                 -> GHC.Prim.State# GHC.Prim.RealWorld
                 -> (# GHC.Prim.State# GHC.Prim.RealWorld, Int #)
            [LclId[JoinId(3)],
             Arity=3,
             Str=<L,U><L,U(U)><L,U>,
             Unf=OtherCon []]
            go_s6WG (x_a5xy :: GHC.Prim.Int#)
                    (eta_B2 :: Int)
                    (eta1_Xz :: GHC.Prim.State# GHC.Prim.RealWorld)
              = case GHC.Prim.remInt# x_a5xy 2# of {
                  __DEFAULT ->
                    case GHC.Prim.==# x_a5xy ww_s6Xc of {
                      __DEFAULT -> jump go_s6WG (GHC.Prim.+# x_a5xy 1#) eta_B2 eta1_Xz;
                      1# -> (# eta1_Xz, eta_B2 #)
                    };
                  0# ->
                    case eta_B2 of { GHC.Types.I# x1_a5t8 ->
                    case GHC.Prim.==# x_a5xy ww_s6Xc of {
                      __DEFAULT ->
                        jump go_s6WG
                          (GHC.Prim.+# x_a5xy 1#)
                          (GHC.Types.I# (GHC.Prim.+# x1_a5t8 x_a5xy))
                          eta1_Xz;
                      1# -> (# eta1_Xz, GHC.Types.I# (GHC.Prim.+# x1_a5t8 x_a5xy) #)
                    }
                    }
                }; } in
          jump go_s6WG 1# Main.main4 w_s6X9;
        1# -> (# w_s6X9, Main.main4 #)
      }

changed weight to 5

added Tbug Trac import core libraries labels

changed the description

Attached file T16357.hs ($1409).

This bug is caused by INLINE annotation for c which is introduced while fixing #8763 (closed).

After list fusion, the example program would be:

f (I# n) =
  letrec go i =
        case i `remInt#` 2 of
           1# -> case i ==# n of
             0# -> go (i +# 1)
             1# -> pure
           0# -> c (I# i) (case i ==# n of
             0# -> go (i +# 1)
             1# -> pure)
  in go 0# (I# 0)
c x k z = k $! (z + x)
{-# INLINE c #-}

We can see function c is lifted to the **top-level definition**. Besides, the continuation case i ==# n of { 0# -> go (i +# 1); 1# -> pure } is passed as a **higher-order** argument of c. Therefore, CallArity analyzer cannot find the correct call-arity of function go.

Without INLINE pragma for c, we have:

f (I# n) =
  letrec go i =
        case i `remInt#` 2 of
           1# -> case i ==# n of
             0# -> go (i +# 1)
             1# -> pure
           0# -> 
             let k = case i ==# n of
                     0# -> go (i +# 1)
                     1# -> pure in
             \z -> k $! I# (z + (I# i))
  in go 0# (I# 0)

We can see the continuation k is defined inside go. In this case, CallArity works good and finds the call-arity of go as 3 (not 2 because it includes State# Realworld).

If we define c x k = \z -> f z x >>= k (instead of c x k z = f z x >>= k) and add an INLINE pragma for c, CallArity works good because the continuation is defined inside go.

added runtime perf label

It sounds like you have a good handle on what is going on here, @autotaker. Would you care to pick this up and propose a patch?

Ok! I’ll do it in a few days.

added Pnormal label

@autotaker did you/do you have time to produce a patch? This seems great.

mentioned in issue #18993

mentioned in merge request !4827