vector-0.1: Data/Vector/Stream.hs
{-# LANGUAGE ExistentialQuantification #-}
-- |
-- Module : Data.Vector.Stream.Size
-- Copyright : (c) Roman Leshchinskiy 2008
-- License : BSD-style
--
-- Maintainer : rl@cse.unsw.edu.au
-- Stability : experimental
-- Portability : non-portable
--
-- Fusible streams
--
#include "phases.h"
module Data.Vector.Stream (
-- * Types
Step(..), Stream(..),
-- * Size hints
size, sized,
-- * Length information
length, null,
-- * Construction
empty, singleton, cons, snoc, replicate, (++),
-- * Accessing individual elements
head, last, (!!),
-- * Substreams
extract, init, tail, take, drop,
-- * Mapping and zipping
map, zipWith,
-- * Filtering
filter, takeWhile, dropWhile,
-- * Searching
elem, notElem, find, findIndex,
-- * Folding
foldl, foldl1, foldl', foldl1', foldr, foldr1,
-- * Unfolding
unfold,
-- * Conversion to/from lists
toList, fromList,
-- * Monadic combinators
mapM_, foldM
) where
import Data.Vector.Stream.Size
import Prelude hiding ( length, null,
replicate, (++),
head, last, (!!),
init, tail, take, drop,
map, zipWith,
filter, takeWhile, dropWhile,
elem, notElem,
foldl, foldl1, foldr, foldr1,
mapM_ )
data Step s a = Yield a s
| Skip s
| Done
-- | The type of fusible streams
data Stream a = forall s. Stream (s -> Step s a) s Size
-- | 'Size' hint of a 'Stream'
size :: Stream a -> Size
{-# INLINE size #-}
size (Stream _ _ sz) = sz
-- | Attach a 'Size' hint to a 'Stream'
sized :: Stream a -> Size -> Stream a
{-# INLINE_STREAM sized #-}
sized (Stream step s _) sz = Stream step s sz
-- | Unfold
unfold :: (s -> Maybe (a, s)) -> s -> Stream a
{-# INLINE_STREAM unfold #-}
unfold f s = Stream step s Unknown
where
{-# INLINE step #-}
step s = case f s of
Just (x, s') -> Yield x s'
Nothing -> Done
-- | Convert a 'Stream' to a list
toList :: Stream a -> [a]
{-# INLINE toList #-}
toList s = foldr (:) [] s
-- | Create a 'Stream' from a list
fromList :: [a] -> Stream a
{-# INLINE_STREAM fromList #-}
fromList xs = Stream step xs Unknown
where
step (x:xs) = Yield x xs
step [] = Done
-- Length
-- ------
-- | Length of a 'Stream'
length :: Stream a -> Int
{-# INLINE_STREAM length #-}
length s = foldl' (\n _ -> n+1) 0 s
-- | Check if a 'Stream' is empty
null :: Stream a -> Bool
{-# INLINE_STREAM null #-}
null s = foldr (\_ _ -> False) True s
-- Construction
-- ------------
-- | Empty 'Stream'
empty :: Stream a
{-# INLINE_STREAM empty #-}
empty = Stream (const Done) () (Exact 0)
-- | Singleton 'Stream'
singleton :: a -> Stream a
{-# INLINE_STREAM singleton #-}
singleton x = Stream step True (Exact 1)
where
{-# INLINE step #-}
step True = Yield x False
step False = Done
-- | Replicate a value to a given length
replicate :: Int -> a -> Stream a
{-# INLINE_STREAM replicate #-}
replicate n x = Stream step n (Exact (max n 0))
where
{-# INLINE step #-}
step i | i > 0 = Yield x (i-1)
| otherwise = Done
-- | Prepend an element
cons :: a -> Stream a -> Stream a
{-# INLINE cons #-}
cons x s = singleton x ++ s
-- | Append an element
snoc :: Stream a -> a -> Stream a
{-# INLINE snoc #-}
snoc s x = s ++ singleton x
infixr 5 ++
-- | Concatenate two 'Stream's
(++) :: Stream a -> Stream a -> Stream a
{-# INLINE_STREAM (++) #-}
Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
where
step (Left sa) = case stepa sa of
Yield x sa' -> Yield x (Left sa')
Skip sa' -> Skip (Left sa')
Done -> Skip (Right sb)
step (Right sb) = case stepb sb of
Yield x sb' -> Yield x (Right sb')
Skip sb' -> Skip (Right sb')
Done -> Done
-- Accessing elements
-- ------------------
-- | First element of the 'Stream' or error if empty
head :: Stream a -> a
{-# INLINE_STREAM head #-}
head (Stream step s _) = head_loop s
where
head_loop s = case step s of
Yield x _ -> x
Skip s' -> head_loop s'
Done -> error "Data.Vector.Stream.head: empty stream"
-- | Last element of the 'Stream' or error if empty
last :: Stream a -> a
{-# INLINE_STREAM last #-}
last (Stream step s _) = last_loop0 s
where
last_loop0 s = case step s of
Yield x s' -> last_loop1 x s'
Skip s' -> last_loop0 s'
Done -> error "Data.Vector.Stream.last: empty stream"
last_loop1 x s = case step s of
Yield y s' -> last_loop1 y s'
Skip s' -> last_loop1 x s'
Done -> x
-- | Element at the given position
(!!) :: Stream a -> Int -> a
{-# INLINE (!!) #-}
s !! i = head (drop i s)
-- Substreams
-- ----------
-- | Extract a substream of the given length starting at the given position.
extract :: Stream a -> Int -- ^ starting index
-> Int -- ^ length
-> Stream a
{-# INLINE extract #-}
extract s i n = take n (drop i s)
-- | All but the last element
init :: Stream a -> Stream a
{-# INLINE_STREAM init #-}
init (Stream step s sz) = Stream step' (Nothing, s) (sz - 1)
where
{-# INLINE step' #-}
step' (Nothing, s) = case step s of
Yield x s' -> Skip (Just x, s')
Skip s' -> Skip (Nothing, s')
Done -> Done -- FIXME: should be an error
step' (Just x, s) = case step s of
Yield y s' -> Yield x (Just y, s')
Skip s' -> Skip (Just x, s')
Done -> Done
-- | All but the first element
tail :: Stream a -> Stream a
{-# INLINE_STREAM tail #-}
tail (Stream step s sz) = Stream step' (Left s) (sz - 1)
where
{-# INLINE step' #-}
step' (Left s) = case step s of
Yield x s' -> Skip (Right s')
Skip s' -> Skip (Left s')
Done -> Done -- FIXME: should be error?
step' (Right s) = case step s of
Yield x s' -> Yield x (Right s')
Skip s' -> Skip (Right s')
Done -> Done
-- | The first @n@ elements
take :: Int -> Stream a -> Stream a
{-# INLINE_STREAM take #-}
take n (Stream step s sz) = Stream step' (s, 0) (smaller (Exact n) sz)
where
{-# INLINE step' #-}
step' (s, i) | i < n = case step s of
Yield x s' -> Yield x (s', i+1)
Skip s' -> Skip (s', i)
Done -> Done
step' (s, i) = Done
data Drop s = Drop_Drop s Int | Drop_Keep s
-- | All but the first @n@ elements
drop :: Int -> Stream a -> Stream a
{-# INLINE_STREAM drop #-}
drop n (Stream step s sz) = Stream step' (Drop_Drop s 0) (sz - Exact n)
where
{-# INLINE step' #-}
step' (Drop_Drop s i) | i < n = case step s of
Yield x s' -> Skip (Drop_Drop s' (i+1))
Skip s' -> Skip (Drop_Drop s' i)
Done -> Done
| otherwise = Skip (Drop_Keep s)
step' (Drop_Keep s) = case step s of
Yield x s' -> Yield x (Drop_Keep s')
Skip s' -> Skip (Drop_Keep s')
Done -> Done
-- Mapping/zipping
-- ---------------
instance Functor Stream where
{-# INLINE_STREAM fmap #-}
fmap = map
-- | Map a function over a 'Stream'
map :: (a -> b) -> Stream a -> Stream b
{-# INLINE_STREAM map #-}
map f (Stream step s n) = Stream step' s n
where
{-# INLINE step' #-}
step' s = case step s of
Yield x s' -> Yield (f x) s'
Skip s' -> Skip s'
Done -> Done
-- | Zip two 'Stream's with the given function
zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
{-# INLINE_STREAM zipWith #-}
zipWith f (Stream stepa sa na) (Stream stepb sb nb)
= Stream step (sa, sb, Nothing) (smaller na nb)
where
{-# INLINE step #-}
step (sa, sb, Nothing) = case stepa sa of
Yield x sa' -> Skip (sa', sb, Just x)
Skip sa' -> Skip (sa', sb, Nothing)
Done -> Done
step (sa, sb, Just x) = case stepb sb of
Yield y sb' -> Yield (f x y) (sa, sb', Nothing)
Skip sb' -> Skip (sa, sb', Just x)
Done -> Done
-- Filtering
-- ---------
-- | Drop elements which do not satisfy the predicate
filter :: (a -> Bool) -> Stream a -> Stream a
{-# INLINE_STREAM filter #-}
filter f (Stream step s n) = Stream step' s (toMax n)
where
{-# INLINE step' #-}
step' s = case step s of
Yield x s' | f x -> Yield x s'
| otherwise -> Skip s'
Skip s' -> Skip s'
Done -> Done
-- | Longest prefix of elements that satisfy the predicate
takeWhile :: (a -> Bool) -> Stream a -> Stream a
{-# INLINE_STREAM takeWhile #-}
takeWhile f (Stream step s n) = Stream step' s (toMax n)
where
{-# INLINE step' #-}
step' s = case step s of
Yield x s' | f x -> Yield x s'
| otherwise -> Done
Skip s' -> Skip s'
Done -> Done
data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
-- | Drop the longest prefix of elements that satisfy the predicate
dropWhile :: (a -> Bool) -> Stream a -> Stream a
{-# INLINE_STREAM dropWhile #-}
dropWhile f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
where
-- NOTE: we jump through hoops here to have only one Yield; local data
-- declarations would be nice!
{-# INLINE step' #-}
step' (DropWhile_Drop s)
= case step s of
Yield x s' | f x -> Skip (DropWhile_Drop s')
| otherwise -> Skip (DropWhile_Yield x s')
Skip s' -> Skip (DropWhile_Drop s')
Done -> Done
step' (DropWhile_Yield x s) = Yield x (DropWhile_Next s)
step' (DropWhile_Next s) = case step s of
Yield x s' -> Skip (DropWhile_Yield x s')
Skip s' -> Skip (DropWhile_Next s')
Done -> Done
-- Searching
-- ---------
infix 4 `elem`
-- | Check whether the 'Stream' contains an element
elem :: Eq a => a -> Stream a -> Bool
{-# INLINE_STREAM elem #-}
elem x (Stream step s _) = elem_loop s
where
elem_loop s = case step s of
Yield y s' | x == y -> True
| otherwise -> elem_loop s'
Skip s' -> elem_loop s'
Done -> False
infix 4 `notElem`
-- | Inverse of `elem`
notElem :: Eq a => a -> Stream a -> Bool
{-# INLINE notElem #-}
notElem x = not . elem x
-- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
-- such element exists.
find :: (a -> Bool) -> Stream a -> Maybe a
{-# INLINE_STREAM find #-}
find f (Stream step s _) = find_loop s
where
find_loop s = case step s of
Yield x s' | f x -> Just x
| otherwise -> find_loop s'
Skip s' -> find_loop s'
Done -> Nothing
-- | Yield 'Just' the index of the first element matching the predicate or
-- 'Nothing' if no such element exists.
findIndex :: (a -> Bool) -> Stream a -> Maybe Int
{-# INLINE_STREAM findIndex #-}
findIndex f (Stream step s _) = findIndex_loop s 0
where
findIndex_loop s i = case step s of
Yield x s' | f x -> Just i
| otherwise -> findIndex_loop s' (i+1)
Skip s' -> findIndex_loop s' i
Done -> Nothing
-- Folding
-- -------
-- | Left fold
foldl :: (a -> b -> a) -> a -> Stream b -> a
{-# INLINE_STREAM foldl #-}
foldl f z (Stream step s _) = foldl_go z s
where
foldl_go z s = case step s of
Yield x s' -> foldl_go (f z x) s'
Skip s' -> foldl_go z s'
Done -> z
-- | Left fold on non-empty 'Stream's
foldl1 :: (a -> a -> a) -> Stream a -> a
{-# INLINE_STREAM foldl1 #-}
foldl1 f (Stream step s sz) = foldl1_loop s
where
foldl1_loop s = case step s of
Yield x s' -> foldl f x (Stream step s' (sz - 1))
Skip s' -> foldl1_loop s'
Done -> error "Data.Vector.Stream.foldl1: empty stream"
-- | Left fold with strict accumulator
foldl' :: (a -> b -> a) -> a -> Stream b -> a
{-# INLINE_STREAM foldl' #-}
foldl' f z (Stream step s _) = foldl_go z s
where
foldl_go z s = z `seq`
case step s of
Yield x s' -> foldl_go (f z x) s'
Skip s' -> foldl_go z s'
Done -> z
-- | Left fold on non-empty 'Stream's with strict accumulator
foldl1' :: (a -> a -> a) -> Stream a -> a
{-# INLINE_STREAM foldl1' #-}
foldl1' f (Stream step s sz) = foldl1'_loop s
where
foldl1'_loop s = case step s of
Yield x s' -> foldl' f x (Stream step s' (sz - 1))
Skip s' -> foldl1'_loop s'
Done -> error "Data.Vector.Stream.foldl1': empty stream"
-- | Right fold
foldr :: (a -> b -> b) -> b -> Stream a -> b
{-# INLINE_STREAM foldr #-}
foldr f z (Stream step s _) = foldr_go s
where
foldr_go s = case step s of
Yield x s' -> f x (foldr_go s')
Skip s' -> foldr_go s'
Done -> z
-- | Right fold on non-empty 'Stream's
foldr1 :: (a -> a -> a) -> Stream a -> a
{-# INLINE_STREAM foldr1 #-}
foldr1 f (Stream step s sz) = foldr1_loop s
where
foldr1_loop s = case step s of
Yield x s' -> foldr f x (Stream step s' (sz - 1))
Skip s' -> foldr1_loop s'
Done -> error "Data.Vector.Stream.foldr1: empty stream"
-- Comparisons
-- -----------
eq :: Eq a => Stream a -> Stream a -> Bool
{-# INLINE_STREAM eq #-}
eq (Stream step1 s1 _) (Stream step2 s2 _) = eq_loop0 s1 s2
where
eq_loop0 s1 s2 = case step1 s1 of
Yield x s1' -> eq_loop1 x s1' s2
Skip s1' -> eq_loop0 s1' s2
Done -> null (Stream step2 s2 Unknown)
eq_loop1 x s1 s2 = case step2 s2 of
Yield y s2' -> x == y && eq_loop0 s1 s2'
Skip s2' -> eq_loop1 x s1 s2'
Done -> False
cmp :: Ord a => Stream a -> Stream a -> Ordering
{-# INLINE_STREAM cmp #-}
cmp (Stream step1 s1 _) (Stream step2 s2 _) = cmp_loop0 s1 s2
where
cmp_loop0 s1 s2 = case step1 s1 of
Yield x s1' -> cmp_loop1 x s1' s2
Skip s1' -> cmp_loop0 s1' s2
Done -> if null (Stream step2 s2 Unknown)
then EQ else LT
cmp_loop1 x s1 s2 = case step2 s2 of
Yield y s2' -> case x `compare` y of
EQ -> cmp_loop0 s1 s2'
c -> c
Skip s2' -> cmp_loop1 x s1 s2'
Done -> GT
instance Eq a => Eq (Stream a) where
{-# INLINE (==) #-}
(==) = eq
instance Ord a => Ord (Stream a) where
{-# INLINE compare #-}
compare = cmp
-- Monadic combinators
-- -------------------
-- | Apply a monadic action to each element of the stream
mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
{-# INLINE_STREAM mapM_ #-}
mapM_ m (Stream step s _) = mapM_go s
where
mapM_go s = case step s of
Yield x s' -> do { m x; mapM_go s' }
Skip s' -> mapM_go s'
Done -> return ()
-- | Monadic fold
foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
{-# INLINE_STREAM foldM #-}
foldM m z (Stream step s _) = foldM_go z s
where
foldM_go z s = case step s of
Yield x s' -> do { z' <- m z x; foldM_go z' s' }
Skip s' -> foldM_go z s'
Done -> return z