minimal stream fusion

fusion.hs

{-# LANGUAGE ExistentialQuantification #-}
module Main (main) where

import Prelude hiding (map, filter, foldl, sum)
import System.Environment (getArgs)

data Stream a = forall s. Stream (s -> Step a s) s
data Step a s = Done
              | Yield a s
              | Skip s

-- | map for Stream
mapS :: (a -> b) -> Stream a -> Stream b
mapS f (Stream next0 s0) = Stream next s0
  where
    next s = case next0 s of
      Done       -> Done
      Skip s'    -> Skip s'
      Yield x s' -> Yield (f x) s'

-- | filter for Stream
filterS :: (a -> Bool) -> Stream a -> Stream a
filterS p (Stream next0 s0) = Stream next s0
  where
    next s = case next0 s of
      Done                   -> Done
      Skip s'                -> Skip s'
      Yield x s' | p x       -> Yield x s'
                 | otherwise -> Skip s'

-- | foldl for Stream
foldlS :: (b -> a -> b) -> b -> Stream a -> b
foldlS f z0 (Stream next0 s0) = go z0 s0
  where
      go z s = case next0 s of
        Done       -> z
        Skip s'    -> go z s'
        Yield x s' -> go (z `f` x) s'

stream :: [a] -> Stream a
{-# INLINE [1] stream #-}
stream = Stream next
  where
    next []     = Done
    next (x:xs) = Yield x xs

unstream :: Stream a -> [a]
{-# INLINE [1] unstream #-}
unstream (Stream next0 s0) = unfold s0
  where
    unfold s = case next0 s of
      Done       -> []
      Skip s'    -> unfold s'
      Yield x s' -> x : unfold s'

{-# RULES
"stream/unstream" [2] forall s. stream (unstream s) = s
  #-}

map :: (a -> b) -> [a] -> [b]
{-# INLINE map #-}
map f = unstream . mapS f . stream

filter :: (a -> Bool) -> [a] -> [a]
{-# INLINE filter #-}
filter p = unstream . filterS p . stream

foldl :: (b -> a -> b) -> b -> [a] -> b
{-# INLINE foldl #-}
foldl f z = foldlS f z . stream

sum :: Num a => [a] -> a
sum = foldl (+) 0

main :: IO ()
main = do
  n:d:_ <- fmap (map read) getArgs
  let xs = [k * m | k <- [1..n], m <- [1..k]] :: [Int]
  -- boring computation
  print $ sum . map (\x -> x - d) . filter (> 0) . map (+ d) $ xs