Monadic streams
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Thu, 7 Aug 2008 04:37:05 +0000 (04:37 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Thu, 7 Aug 2008 04:37:05 +0000 (04:37 +0000)
Data/Vector/MStream.hs [new file with mode: 0644]
Data/Vector/MVector.hs
Data/Vector/MVector/Mut.hs
vector.cabal

diff --git a/Data/Vector/MStream.hs b/Data/Vector/MStream.hs
new file mode 100644 (file)
index 0000000..3d36eb9
--- /dev/null
@@ -0,0 +1,83 @@
+{-# LANGUAGE ExistentialQuantification #-}
+
+#include "phases.h"
+
+module Data.Vector.MStream (
+  MStream,
+
+  sized,
+
+  unfoldM, foldM,
+
+  map, mapM, filter, filterM
+) where
+
+import Data.Vector.Stream ( Step(..) )
+import Data.Vector.Stream.Size
+
+import Control.Monad  ( liftM )
+import Prelude hiding ( map, mapM, filter )
+
+data MStream m a = forall s. MStream (s -> m (Step s a)) s Size
+
+sized :: MStream m a -> Size -> MStream m a
+{-# INLINE_STREAM sized #-}
+sized (MStream step s _) sz = MStream step s sz
+
+unfoldM :: Monad m => (s -> m (Maybe (a, s))) -> s -> MStream m a
+{-# INLINE_STREAM unfoldM #-}
+unfoldM f s = MStream step s Unknown
+  where
+    {-# INLINE step #-}
+    step s = do
+               r <- f s
+               case r of
+                 Just (x, s') -> return $ Yield x s'
+                 Nothing      -> return $ Done
+
+map :: Monad m => (a -> b) -> MStream m a -> MStream m b
+{-# INLINE map #-}
+map f = mapM (return . f)
+
+mapM :: Monad m => (a -> m b) -> MStream m a -> MStream m b
+{-# INLINE_STREAM mapM #-}
+mapM f (MStream step s n) = MStream step' s n
+  where
+    {-# INLINE step' #-}
+    step' s = do
+                r <- step s
+                case r of
+                  Yield x s' -> liftM  (`Yield` s') (f x)
+                  Skip    s' -> return (Skip    s')
+                  Done       -> return Done
+
+filter :: Monad m => (a -> Bool) -> MStream m a -> MStream m a
+{-# INLINE filter #-}
+filter f = filterM (return . f)
+
+filterM :: Monad m => (a -> m Bool) -> MStream m a -> MStream m a
+{-# INLINE_STREAM filterM #-}
+filterM f (MStream step s n) = MStream step' s (toMax n)
+  where
+    {-# INLINE step' #-}
+    step' s = do
+                r <- step s
+                case r of
+                  Yield x s' -> do
+                                  b <- f x
+                                  if b then return $ Yield x s'
+                                       else return $ Skip    s'
+                  Skip    s' -> return $ Skip s'
+                  Done       -> return $ Done
+
+foldM :: Monad m => (a -> b -> m a) -> a -> MStream m b -> m a
+{-# INLINE_STREAM foldM #-}
+foldM m z (MStream step s _) = foldM_go z s
+  where
+    foldM_go z s = do
+                     r <- step s
+                     case r of
+                       Yield x s' -> do { z' <- m z x; foldM_go z' s' }
+                       Skip    s' -> foldM_go z s'
+                       Done       -> return z
+
index ac52920..1107f73 100644 (file)
@@ -23,6 +23,8 @@ module Data.Vector.MVector (
 
 import qualified Data.Vector.Stream      as Stream
 import           Data.Vector.Stream      ( Stream )
+import qualified Data.Vector.MStream     as MStream
+import           Data.Vector.MStream     ( MStream )
 import           Data.Vector.Stream.Size
 
 import Control.Monad.ST ( ST )
@@ -166,23 +168,24 @@ grow :: MVector v m a => v a -> Int -> m (v a)
 grow v by = assert (by >= 0)
           $ unsafeGrow v by
 
-mstream :: MVector v m a => v a -> Stream (m a)
+mstream :: MVector v m a => v a -> MStream m a
 {-# INLINE mstream #-}
-mstream v = v `seq` (Stream.unfold get 0 `Stream.sized` Exact n)
+mstream v = v `seq` (MStream.unfoldM get 0 `MStream.sized` Exact n)
   where
     n = length v
 
     {-# INLINE get #-}
-    get i | i < n     = Just (unsafeRead v i, i+1)
-          | otherwise = Nothing
+    get i | i < n     = do x <- unsafeRead v i
+                           return $ Just (x, i+1)
+          | otherwise = return $ Nothing
 
-munstream :: MVector v m a => v a -> Stream (m a) -> m (v a)
+munstream :: MVector v m a => v a -> MStream m a -> m (v a)
 {-# INLINE munstream #-}
 munstream v s = v `seq` do
-                          n' <- Stream.foldM put 0 s
+                          n' <- MStream.foldM put 0 s
                           return $ slice v 0 n'
   where
-    put i m = do { write v i =<< m; return (i+1) }
+    put i x = do { write v i x; return (i+1) }
 
 -- | Create a new mutable vector and fill it with elements from the 'Stream'.
 -- The vector will grow logarithmically if the 'Size' hint of the 'Stream' is
index 7fa060e..0b0e76e 100644 (file)
@@ -12,6 +12,9 @@ import           Data.Vector.MVector ( MVector )
 import           Data.Vector.Stream ( Stream )
 import qualified Data.Vector.Stream as Stream
 
+import           Data.Vector.MStream ( MStream )
+import qualified Data.Vector.MStream as MStream
+
 import Control.Monad  ( liftM )
 import Prelude hiding ( reverse, map )
 
@@ -29,8 +32,7 @@ unstream :: Stream a -> Mut a
 {-# INLINE_STREAM unstream #-}
 unstream s = Mut (MVector.unstream s)
 
-restream :: (forall m. Monad m => Stream (m a) -> Stream (m a))
-          -> Mut a -> Mut a
+restream :: (forall m. Monad m => MStream m a -> MStream m a) -> Mut a -> Mut a
 {-# INLINE_STREAM restream #-}
 restream f (Mut p) = Mut (
   do
@@ -40,8 +42,8 @@ restream f (Mut p) = Mut (
 {-# RULES
 
 "restream/restream [Mut]"
-  forall (f :: forall m. Monad m => Stream (m a) -> Stream (m a))
-         (g :: forall m. Monad m => Stream (m a) -> Stream (m a)) p .
+  forall (f :: forall m. Monad m => MStream m a -> MStream m a)
+         (g :: forall m. Monad m => MStream m a -> MStream m a) p .
   restream f (restream g p) = restream (f . g) p
 
  #-}
@@ -56,5 +58,9 @@ reverse m = trans m (MVector.reverse)
 
 map :: (a -> a) -> Mut a -> Mut a
 {-# INLINE map #-}
-map f = restream (Stream.map (liftM f))
+map f = restream (MStream.map f)
+
+filter :: (a -> Bool) -> Mut a -> Mut a
+{-# INLINE filter #-}
+filter f = restream (MStream.filter f)
 
index 49f760f..fd81eb4 100644 (file)
@@ -24,6 +24,7 @@ Library
   Exposed-Modules:
         Data.Vector.Stream.Size
         Data.Vector.Stream
+        Data.Vector.MStream
 
         Data.Vector.MVector
         Data.Vector.MVector.Mut