Commit d45693a5 authored by David Feuer's avatar David Feuer Committed by Joachim Breitner

Make scanl fuse; add scanl'

Summary:
Make scanl a good producer and a good consumer for fold/build
fusion. Add strictly-accumulating scanl', which is required for
Data.List.inits.

Reviewers: nomeata, austin

Reviewed By: austin

Subscribers: spacekitteh, thomie, carter, ezyang, simonmar

Differential Revision: https://phabricator.haskell.org/D314

GHC Trac Issues: #9356
parent 205b1032
......@@ -22,7 +22,7 @@ module GHC.List (
map, (++), filter, concat,
head, last, tail, init, uncons, null, length, (!!),
foldl, scanl, scanl1, foldr, foldr1, scanr, scanr1,
foldl, scanl, scanl1, scanl', foldr, foldr1, scanr, scanr1,
iterate, repeat, replicate, cycle,
take, drop, splitAt, takeWhile, dropWhile, span, break,
reverse, and, or,
......@@ -200,10 +200,33 @@ foldl k z0 xs = foldr (\(v::a) (fn::b->b) (z::b) -> fn (k z v)) (id :: b -> b) x
--
-- > last (scanl f z xs) == foldl f z xs.
-- This peculiar arrangement is necessary to prevent scanl being rewritten in
-- its own right-hand side.
{-# NOINLINE [1] scanl #-}
scanl :: (b -> a -> b) -> b -> [a] -> [b]
scanl f q ls = q : (case ls of
[] -> []
x:xs -> scanl f (f q x) xs)
scanl = scanlGo
where
scanlGo :: (b -> a -> b) -> b -> [a] -> [b]
scanlGo f q ls = q : (case ls of
[] -> []
x:xs -> scanlGo f (f q x) xs)
-- Note [scanl rewrite rules]
{-# RULES
"scanl" [~1] forall f a bs . scanl f a bs =
build (\c n -> a `c` foldr (scanlFB f c) (constScanl n) bs a)
"scanlList" [1] forall f (a::a) bs .
foldr (scanlFB f (:)) (constScanl []) bs a = tail (scanl f a bs)
#-}
{-# INLINE [0] scanlFB #-}
scanlFB :: (b -> a -> b) -> (b -> c -> c) -> a -> (b -> c) -> b -> c
scanlFB f c = \b g x -> let b' = f x b in b' `c` g b'
{-# INLINE [0] constScanl #-}
constScanl :: a -> b -> a
constScanl = const
-- | 'scanl1' is a variant of 'scanl' that has no starting value argument:
--
......@@ -213,6 +236,64 @@ scanl1 :: (a -> a -> a) -> [a] -> [a]
scanl1 f (x:xs) = scanl f x xs
scanl1 _ [] = []
-- | A strictly accumulating version of 'scanl'
{-# NOINLINE [1] scanl' #-}
scanl' :: (b -> a -> b) -> b -> [a] -> [b]
-- This peculiar form is needed to prevent scanl' from being rewritten
-- in its own right hand side.
scanl' = scanlGo'
where
scanlGo' :: (b -> a -> b) -> b -> [a] -> [b]
scanlGo' f q ls = q `seq` q : (case ls of
[] -> []
x:xs -> scanlGo' f (f q x) xs)
-- Note [scanl rewrite rules]
{-# RULES
"scanl'" [~1] forall f a bs . scanl' f a bs =
build (\c n -> a `c` foldr (scanlFB' f c) (flipSeqScanl' n) bs a)
"scanlList'" [1] forall f a bs .
foldr (scanlFB' f (:)) (flipSeqScanl' []) bs a = tail (scanl' f a bs)
#-}
{-# INLINE [0] scanlFB' #-}
scanlFB' :: (b -> a -> b) -> (b -> c -> c) -> a -> (b -> c) -> b -> c
scanlFB' f c = \b g x -> let b' = f x b in b' `seq` b' `c` g b'
{-# INLINE [0] flipSeqScanl' #-}
flipSeqScanl' :: a -> b -> a
flipSeqScanl' = flip seq
{-
Note [scanl rewrite rules]
~~~~~~~~~~~~~~~~~~~~~~~~~~
In most cases, when we rewrite a form to one that can fuse, we try to rewrite it
back to the original form if it does not fuse. For scanl, we do something a
little different. In particular, we rewrite
scanl f a bs
to
build (\c n -> a `c` foldr (scanlFB f c) (constScanl n) bs a)
When build is inlined, this becomes
a : foldr (scanlFB f (:)) (constScanl []) bs a
To rewrite this form back to scanl, we would need a rule that looked like
forall f a bs. a : foldr (scanlFB f (:)) (constScanl []) bs a = scanl f a bs
The problem with this rule is that it has (:) at its head. This would have the
effect of changing the way the inliner looks at (:), not only here but
everywhere. In most cases, this makes no difference, but in some cases it
causes it to come to a different decision about whether to inline something.
Based on nofib benchmarks, this is bad for performance. Therefore, we instead
match on everything past the :, which is just the tail of scanl.
-}
-- foldr, foldr1, scanr, and scanr1 are the right-to-left duals of the
-- above functions.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment