/StreamFusionReassoc.hs

 1{-# LANGUAGE ExistentialQuantification, BangPatterns, TypeOperators #-}
 2import Prelude hiding (enumFromTo, concatMap, replicate)
 3
 4data Stream a = forall s. Stream !(s -> Step a s)  -- a stepper function
 5                                 !s                -- an initial state
 6
 7-- | A stream step.
 8--
 9--   A step either ends a stream, skips a value, or yields a value
10--
11data Step a s = Yield a !s
12              | Skip    !s
13              | Done
14
15
16-- | Construct an abstract stream from a list.
17stream :: [a] -> Stream a
18stream xs0 = Stream next xs0
19  where
20    {-# INLINE next #-}
21    next []     = Done
22    next (x:xs) = Yield x xs
23{-# INLINE [0] stream #-}
24
25-- | Flatten a stream back into a list.
26unstream :: Stream a -> [a]
27unstream (Stream next s0) = unfold_unstream s0
28  where
29    unfold_unstream !s = case next s of
30      Done       -> []
31      Skip    s' -> unfold_unstream s'
32      Yield x s' -> x : unfold_unstream s'
33{-# INLINE [0] unstream #-}
34
35--
36-- /The/ stream fusion rule
37--
38
39{-# RULES
40"STREAM stream/unstream fusion" forall s.
41    stream (unstream s) = s
42  #-}
43
44
45{-# INLINE replicate #-}
46replicate n x = unstream (replicateS n x)
47
48{-# INLINE [0] replicateS #-}
49replicateS :: Int -> a -> Stream a
50replicateS n x = Stream next n
51  where
52    {-# INLINE next #-}
53    next !i | i <= 0    = Done
54            | otherwise = Yield x (i-1)
55
56{-# INLINE enumFromTo #-}
57enumFromTo x y = unstream (enumFromToS x y)
58
59{-# INLINE [0] enumFromToS #-}
60enumFromToS x y = Stream step x
61  where
62    {-# INLINE step #-}
63    step x | x <= y    = Yield x (x + 1)
64           | otherwise = Done
65
66data a :!: b = !a :!: !b
67
68{-# INLINE concatMap #-}
69concatMap f xs = unstream (concatMapS (stream . f) (stream xs))
70
71{-# INLINE [0] concatMapS #-}
72concatMapS :: (a -> Stream b) -> Stream a -> Stream b
73concatMapS f (Stream next0 s0) = Stream next (s0 :!: Nothing)
74  where
75    {-# INLINE next #-}
76    next (s :!: Nothing) = case next0 s of
77      Done       -> Done
78      Skip    s' -> Skip (s' :!: Nothing)
79      Yield x s' -> Skip (s' :!: Just (f x))
80
81    next (s :!: Just (Stream g t)) = case g t of
82      Done       -> Skip    (s :!: Nothing)
83      Skip    t' -> Skip    (s :!: Just (Stream g t'))
84      Yield x t' -> Yield x (s :!: Just (Stream g t'))
85
86-- [1,1,2,2,3,3,4,4,5,5,2,2,3,3,4,4,5,5,3,3,4,4,5,5,4,4,5,5,5,5]
87main = do
88    print $ concatMap (\y -> replicate 2 y) (concatMap (\x -> enumFromTo x 5) (enumFromTo 1 (5 :: Int)))
89    --print $ concatMap (\x -> concatMap (\y -> replicate 2 y) (enumFromTo x 5)) (enumFromTo 1 (5 :: Int))