[GHC] #16357: Add `oneShot` to the implementation of foldlM

[GHC]

#16357: Add `oneShot` to the implementation of foldlM
-------------------------------------+-------------------------------------
           Reporter:  autotaker      |             Owner:  (none)
               Type:  bug            |            Status:  new
           Priority:  normal         |         Milestone:
          Component:                 |           Version:  8.9
  libraries/base                     |
           Keywords:                 |  Operating System:  Unknown/Multiple
       Architecture:                 |   Type of failure:  None/Unknown
  Unknown/Multiple                   |
          Test Case:                 |        Blocked By:
           Blocking:                 |   Related Tickets:
Differential Rev(s):                 |         Wiki Page:
-------------------------------------+-------------------------------------
 The current (473632d7671619ee08a2a0025aa22bd4f79eca2d) implementation of
 `Data.Foldable.foldlM` is the like this
 {{{#!hs
 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.
 {{{#!hs
 f :: Int -> IO Int
 f = foldlM (\a b -> pure $! (a + b)) 0 (filter even [1..n])
 }}}

 Generated core:
 {{{#!hs
 -- 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,
 {{{#!hs
 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.
 {{{#!hs
 -- 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 #)
       }
 }}}

-- 
Ticket URL: <http://ghc.haskell.org/trac/ghc/ticket/16357>
GHC <http://www.haskell.org/ghc/>
The Glasgow Haskell Compiler