From 170042346c20523fbe067c100cf916c8cb5fd7f3 Mon Sep 17 00:00:00 2001 From: Roman Leshchinskiy Date: Fri, 11 Jul 2008 17:19:22 +0000 Subject: [PATCH] New combinators --- Data/Vector/Base.hs | 22 ++++++++-- Data/Vector/Stream.hs | 116 ++++++++++++++++++++++++++++++++++++++++++++------ 2 files changed, 122 insertions(+), 16 deletions(-) diff --git a/Data/Vector/Base.hs b/Data/Vector/Base.hs index f404709..494c603 100644 --- a/Data/Vector/Base.hs +++ b/Data/Vector/Base.hs @@ -47,9 +47,13 @@ infixr ++ {-# INLINE (++) #-} v ++ w = unstream (stream v Stream.++ stream w) -filter :: Base v a => (a -> Bool) -> v a -> v a -{-# INLINE filter #-} -filter f = unstream . Stream.filter f . stream +take :: Base v a => Int -> v a -> v a +{-# INLINE take #-} +take n = unstream . Stream.take n . stream + +drop :: Base v a => Int -> v a -> v a +{-# INLINE drop #-} +drop n = unstream . Stream.drop n . stream map :: (Base v a, Base v b) => (a -> b) -> v a -> v b {-# INLINE map #-} @@ -59,6 +63,18 @@ zipWith :: (Base v a, Base v b, Base v c) => (a -> b -> c) -> v a -> v b -> v c {-# INLINE zipWith #-} zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys)) +filter :: Base v a => (a -> Bool) -> v a -> v a +{-# INLINE filter #-} +filter f = unstream . Stream.filter f . stream + +takeWhile :: Base v a => (a -> Bool) -> v a -> v a +{-# INLINE takeWhile #-} +takeWhile f = unstream . Stream.takeWhile f . stream + +dropWhile :: Base v a => (a -> Bool) -> v a -> v a +{-# INLINE dropWhile #-} +dropWhile f = unstream . Stream.dropWhile f . stream + foldl' :: Base v b => (a -> b -> a) -> a -> v b -> a {-# INLINE foldl' #-} foldl' f z = Stream.foldl' f z . stream diff --git a/Data/Vector/Stream.hs b/Data/Vector/Stream.hs index 782dfa2..ad1b2c7 100644 --- a/Data/Vector/Stream.hs +++ b/Data/Vector/Stream.hs @@ -7,14 +7,19 @@ module Data.Vector.Stream ( size, sized, unfold, toList, fromList, empty, singleton, replicate, (++), - map, filter, zipWith, + take, drop, + map, zipWith, + filter, takeWhile, dropWhile, foldr, foldl, foldl', mapM_, foldM ) where import Data.Vector.Stream.Size -import Prelude hiding ( replicate, (++), map, filter, zipWith, +import Prelude hiding ( replicate, (++), + take, drop, + map, zipWith, + filter, takeWhile, dropWhile, foldr, foldl, mapM_ ) @@ -52,6 +57,9 @@ fromList xs = Stream step xs Unknown step (x:xs) = Yield x xs step [] = Done +-- Construction +-- ------------ + empty :: Stream a {-# INLINE_STREAM empty #-} empty = Stream (const Done) () (Exact 0) @@ -86,6 +94,46 @@ Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb) Skip sb' -> Skip (Right sb') Done -> Done +-- Substreams +-- ---------- + +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 + +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 -> b) -> Stream a -> Stream b {-# INLINE_STREAM map #-} map f (Stream step s n) = Stream step' s n @@ -96,17 +144,6 @@ map f (Stream step s n) = Stream step' s n Skip s' -> Skip s' Done -> Done -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 - zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c {-# INLINE_STREAM zipWith #-} zipWith f (Stream stepa sa na) (Stream stepb sb nb) @@ -123,6 +160,59 @@ zipWith f (Stream stepa sa na) (Stream stepb sb nb) Skip sb' -> Skip (sa, sb', Just x) Done -> Done +-- Filtering +-- --------- + +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 + +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 + +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 + +-- Folding +-- ------- + foldl :: (a -> b -> a) -> a -> Stream b -> a {-# INLINE_STREAM foldl #-} foldl f z (Stream step s _) = foldl_go z s -- 1.9.1