Let's make takeWhile fuse!

David Feuer david.feuer at gmail.com
Tue Jul 22 02:29:32 UTC 2014


`takeWhile` doesn't do the stream fusion thing. This definition seems
to fix that, at least to a great extent. I don't know how to write the
necessary rules to make this definition be used in the right places.

takeWhile :: forall a . (a -> Bool) -> [a] -> [a]
takeWhile p xs = build tw'
  where
    tw' :: forall b . (a -> b -> b) -> b -> b
    tw' kons knil = foldr go knil xs
      where
        go x rest | p x       = x `kons` rest
                  | otherwise = knil


Tests (performed on 7.8.3, but the definition hasn't changed in HEAD):

In the trivial

main = print $ length $ takeWhile (< 10000000) [(1::Int) .. 20000000]

it fused completely, allocating virtually nothing, whereas using
Prelude.takeWhile (on 7.8.3) required a whopping 1,360,049,352

In the more complex

main = print $ length $ filter (\x -> x `rem` 3 == 0) $ takeWhile (<
10000000) $ filter even  [(1::Int) .. 20000000]

it fused partially, allocating 266,716,048 bytes.

The current GHC definition allocates 746,716,000 bytes in this case
(strangely, less than it used for the simple test)!

David Feuer


More information about the Libraries mailing list