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