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

1 files changed, 85 insertions, 4 deletions

diff --git a/libraries/base/GHC/List.lhs b/libraries/base/GHC/List.lhs
index 51f68abe51..6137249bcd 100644
--- a/libraries/base/GHC/List.lhs
+++ b/libraries/base/GHC/List.lhs
@@ -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.