Potential missed unboxing

Summary

Consider a tree and an inorder foldr on it.

data Tree a = Tip | Bin a !(Tree a) !(Tree a)

foldrTree :: (a -> b -> b) -> b -> Tree a -> b
foldrTree f z0 t = go t z0
  where
    go Tip z = z
    go (Bin x l r) z = go l (f x (go r z))

For a CPS-style fold,

import GHC.Exts

len :: Tree a -> Int
len t = foldrTree (\_ k -> oneShot (\ !acc -> k (acc + 1))) id t 0

GHC 9.8.2 -O2 generates Core like this

Rec {
-- RHS size: {terms: 24, types: 18, coercions: 0, joins: 0/0}
M.len2 [Occ=LoopBreaker]
  :: forall {a}. Tree a -> (Int -> Int) -> Int -> Int
[GblId, Arity=3, Str=<1L><1C(1,L)><L>, Unf=OtherCon []]
M.len2
  = \ (@a_aD2)
      (ds_dEF :: Tree a_aD2)
      (z_ayd :: Int -> Int)
      (eta_B0 [OS=OneShot] :: Int) ->
      case ds_dEF of {
        Tip -> z_ayd eta_B0;
        Bin x_aye l_ayf r_ayg ->
          M.len2
            @a_aD2
            l_ayf
            (\ (v_B2 [OS=OneShot] :: Int) ->
               case v_B2 of { I# ipv_sEK ->
               M.len2 @a_aD2 r_ayg z_ayd (GHC.Types.I# (+# ipv_sEK 1#))
               })
            eta_B0
      }
end Rec }

-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
M.len1 :: Int
[GblId,
 Unf=Unf{Src=<vanilla>, TopLvl=True,
         Value=True, ConLike=True, WorkFree=True, Expandable=True,
         Guidance=IF_ARGS [] 10 10}]
M.len1 = GHC.Types.I# 0#

-- RHS size: {terms: 6, types: 5, coercions: 0, joins: 0/0}
len :: forall a. Tree a -> Int
[GblId,
 Arity=1,
 Str=<1L>,
 Unf=Unf{Src=<vanilla>, TopLvl=True,
         Value=True, ConLike=True, WorkFree=True, Expandable=True,
         Guidance=IF_ARGS [0] 40 0}]
len
  = \ (@a_aD2) (t_ayi :: Tree a_aD2) ->
      M.len2 @a_aD2 t_ayi (id @Int) M.len1

Notice that the accumulator eta_B0 remains boxed in M.len2. Would it be possible to unbox this?

Note that it is only a small change to do so manually.

lenAlt :: Tree a -> Int
lenAlt t = foldrTree (\_ k -> oneShot (\acc# -> k (acc# +# 1#))) I# t 0#
-- len t = foldrTree (\_ k -> oneShot (\ !acc -> k (acc + 1))) id t 0

The Core for this is now

Rec {
-- RHS size: {terms: 18, types: 14, coercions: 0, joins: 0/0}
M.lenAlt_$spoly_go [Occ=LoopBreaker]
  :: forall {a}. Tree a -> Int# -> Int
[GblId[StrictWorker([!])], Arity=2, Str=<1L><L>, Unf=OtherCon []]
M.lenAlt_$spoly_go
  = \ (@a_aAn)
      (sc_sBJ :: Tree a_aAn)
      (eta_B0 [OS=OneShot] :: Int#) ->
      case sc_sBJ of {
        Tip -> GHC.Types.I# eta_B0;
        Bin x_aye l_ayf r_ayg ->
          poly_go_rCg
            @a_aAn
            l_ayf
            (\ (v_B2 [OS=OneShot] :: Int#) ->
               M.lenAlt_$spoly_go @a_aAn r_ayg (+# v_B2 1#))
            eta_B0
      }

-- RHS size: {terms: 20, types: 16, coercions: 0, joins: 0/0}
poly_go_rCg :: forall {a}. Tree a -> (Int# -> Int) -> Int# -> Int
[GblId, Arity=3, Str=<1L><1C(1,L)><L>, Unf=OtherCon []]
poly_go_rCg
  = \ (@a_aAn)
      (ds_dB8 :: Tree a_aAn)
      (z_ayd :: Int# -> Int)
      (eta_B0 [OS=OneShot] :: Int#) ->
      case ds_dB8 of {
        Tip -> z_ayd eta_B0;
        Bin x_aye l_ayf r_ayg ->
          poly_go_rCg
            @a_aAn
            l_ayf
            (\ (v_B2 [OS=OneShot] :: Int#) ->
               poly_go_rCg @a_aAn r_ayg z_ayd (+# v_B2 1#))
            eta_B0
      }
end Rec }

-- RHS size: {terms: 5, types: 4, coercions: 0, joins: 0/0}
lenAlt :: forall a. Tree a -> Int
[GblId,
 Arity=1,
 Str=<1L>,
 Unf=Unf{Src=<vanilla>, TopLvl=True,
         Value=True, ConLike=True, WorkFree=True, Expandable=True,
         Guidance=IF_ARGS [0] 30 0}]
lenAlt
  = \ (@a_aAn) (t_ayi :: Tree a_aAn) ->
      M.lenAlt_$spoly_go @a_aAn t_ayi 0#

Haskell Playground link: https://play.haskell.org/saved/VdH9F8Xc

Steps to reproduce

Compile and observe the generated Core.

Expected behavior

The accumulator would be unboxed.

Environment

• GHC version used: 9.8.2