/StreamFusionReassoc.hs
Haskell | 89 lines | 44 code | 17 blank | 28 comment | 0 complexity | d6ff5d7c9e0f9a6aace1ae009ad3027e MD5 | raw file
Possible License(s): BSD-3-Clause
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))