1aa909b · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add concurrent zipping API

1aa909bHarendra Kumar9 years ago

src/Asyncly.hs+5

compare with master

unifiedsplit

@@ -146,9 +146,14 @@ module Asyncly

146 -- * Special folds

147 , take

148 , drop

149+

150+ -- * Zipping

149 , zipWith

151 , zipWith

150 , zipWithM

152 , zipWithM

151 , ZipSerial (..)

153 , ZipSerial (..)

154+ , zipAsyncWithM

155+ , zipAsyncWith

156+ , ZipAsync (..)

152

157

153 -- * Re-exports

158 -- * Re-exports

154 , Monoid (..)

159 , Monoid (..)

src/Asyncly/AsyncT.hs+57−2

@@ -27,6 +27,9 @@ module Asyncly.AsyncT

27 , zipWithM

28 , zipWith

29 , ZipSerial (..)

30+ , zipAsyncWithM

31+ , zipAsyncWith

32+ , ZipAsync (..)

31 , interleave

34 , interleave

32 , (<=>)

35 , (<=>)

@@ -711,7 +714,7 @@ drop n m = AsyncT $ \ctx stp yld -> do

711-- Zipping Streams

714-- Zipping Streams

712------------------------------------------------------------------------------

715------------------------------------------------------------------------------

713

716

714−-- | Zip two AsyncT streams serially using a monadic function.

717+-- | Zip two AsyncT streams serially using a monadic zipping function.

715zipWithM :: MonadAsync m =>

718zipWithM :: MonadAsync m =>

716 (a -> b -> AsyncT m c) -> AsyncT m a -> AsyncT m b -> AsyncT m c

719 (a -> b -> AsyncT m c) -> AsyncT m a -> AsyncT m b -> AsyncT m c

717zipWithM f m1 m2 = AsyncT $ \_ stp yld -> do

720zipWithM f m1 m2 = AsyncT $ \_ stp yld -> do

@@ -724,7 +727,7 @@ zipWithM f m1 m2 = AsyncT $ \_ stp yld -> do

724 yield1 a (Just ra) = merge a ra

727 yield1 a (Just ra) = merge a ra

725 (runAsyncT m1) Nothing stp yield1

728 (runAsyncT m1) Nothing stp yield1

726

729

727−-- | Zip two AsyncT streams serially using a pure function.

730+-- | Zip two AsyncT streams serially using a pure zipping function.

728zipWith :: MonadAsync m

731zipWith :: MonadAsync m

729 => (a -> b -> c) -> AsyncT m a -> AsyncT m b -> AsyncT m c

732 => (a -> b -> c) -> AsyncT m a -> AsyncT m b -> AsyncT m c

730zipWith f m1 m2 = AsyncT $ \_ stp yld -> do

733zipWith f m1 m2 = AsyncT $ \_ stp yld -> do

@@ -746,6 +749,58 @@ instance MonadAsync m => Applicative (ZipSerial m) where

746 pure a = ZipSerial $ cycle1 (pure a)

749 pure a = ZipSerial $ cycle1 (pure a)

747 (ZipSerial xs) <*> (ZipSerial ys) = ZipSerial (zipWith id xs ys)

750 (ZipSerial xs) <*> (ZipSerial ys) = ZipSerial (zipWith id xs ys)

748

751

752+-- | The computation specified in the argument is pushed to a new thread and

753+-- the context is computation pulls output from the thread.

754+pushOneToCtx :: MonadAsync m => CtxType -> AsyncT m a -> m (Context m a)

755+pushOneToCtx ctype m = do

756+ ctx <- liftIO $

757+ case ctype of

758+ CtxType _ FIFO -> do

759+ c <- getFifoCtx ctype

760+ return c

761+ CtxType _ LIFO -> do

762+ c <- getLifoCtx ctype

763+ return c

764+ -- Note: We must have all the work on the queue before sending the

765+ -- pushworker, otherwise the pushworker may exit before we even get a

766+ -- chance to push.

767+ liftIO $ (enqueue ctx) m

768+ pushWorker ctx

769+ return ctx

770+

771+async :: MonadAsync m => AsyncT m a -> m (AsyncT m a)

772+async m = do

773+ ctx <- pushOneToCtx (CtxType Disjunction LIFO) m

774+ return $ pullFromCtx ctx

775+

776+-- | Zip two AsyncT streams asyncly (i.e. both the streams are generated

777+-- concurrently) using a monadic zipping function.

778+zipAsyncWithM :: MonadAsync m

779+ => (a -> b -> AsyncT m c) -> AsyncT m a -> AsyncT m b -> AsyncT m c

780+zipAsyncWithM f m1 m2 = AsyncT $ \_ stp yld -> do

781+ ma <- async m1

782+ mb <- async m2

783+ (runAsyncT (zipWithM f ma mb)) Nothing stp yld

784+

785+-- | Zip two AsyncT streams asyncly (i.e. both the streams are generated

786+-- concurrently) using a pure zipping function.

787+zipAsyncWith :: MonadAsync m

788+ => (a -> b -> c) -> AsyncT m a -> AsyncT m b -> AsyncT m c

789+zipAsyncWith f m1 m2 = AsyncT $ \_ stp yld -> do

790+ ma <- async m1

791+ mb <- async m2

792+ (runAsyncT (zipWith f ma mb)) Nothing stp yld

793+

794+-- | Wrapper around AsyncT type with a parallel zipping Applicative instance

795+--

796+-- > f <$> ZipAsync xs1 <*> ... <*> ZipAsync xsN

797+newtype ZipAsync m a = ZipAsync {getZipAsync :: AsyncT m a}

798+ deriving (Functor)

799+

800+instance MonadAsync m => Applicative (ZipAsync m) where

801+ pure a = ZipAsync $ cycle1 (pure a)

802+ (ZipAsync xs) <*> (ZipAsync ys) = ZipAsync (zipAsyncWith id xs ys)

803+

749------------------------------------------------------------------------------

804------------------------------------------------------------------------------

750-- Utilities

805-- Utilities

751------------------------------------------------------------------------------

806------------------------------------------------------------------------------

test/Main.hs+21−8

@@ -1,3 +1,6 @@

1+{-# LANGUAGE FlexibleContexts #-}

2+{-# LANGUAGE RankNTypes #-}

1module Main (main) where

4module Main (main) where

3import Control.Concurrent (threadDelay)

6import Control.Concurrent (threadDelay)

@@ -198,26 +201,36 @@ main = hspec $ do

198 describe "Miscellaneous combined examples" mixedOps

201 describe "Miscellaneous combined examples" mixedOps

199

202

200 describe "Transformation" $ transformOps (<>)

203 describe "Transformation" $ transformOps (<>)

201− describe "Serial zipping" $ serialZip

204+ describe "Serial zipping" $

202−

205+ zipOps zipWith zipWithM ZipSerial getZipSerial

206+ describe "Async zipping" $

207+ zipOps zipAsyncWith zipAsyncWithM ZipAsync getZipAsync

208+

209+zipOps :: Applicative (f IO)

210+ => (forall a b c. (a -> b -> c)

211+ -> AsyncT IO a -> AsyncT IO b -> AsyncT IO c)

212+ -> (forall a b c. (a -> b -> AsyncT IO c)

213+ -> AsyncT IO a -> AsyncT IO b -> AsyncT IO c)

214+ -> (forall a. AsyncT IO a -> f IO a)

215+ -> (forall a. f IO a -> AsyncT IO a)

203−serialZip :: Spec

216+ -> Spec

204−serialZip = do

217+zipOps z zM app getz = do

205 it "zipWith" $

218 it "zipWith" $

206 let s1 = foldMapWith (<|>) return [1..10]

219 let s1 = foldMapWith (<|>) return [1..10]

207 s2 = foldMapWith (<>) return [1..]

220 s2 = foldMapWith (<>) return [1..]

208− in toList (zipWith (+) s1 s2)

221+ in toList (z (+) s1 s2)

209 `shouldReturn` ([2,4..20] :: [Int])

222 `shouldReturn` ([2,4..20] :: [Int])

210

223

211 it "zipWithM" $

224 it "zipWithM" $

212 let s1 = foldMapWith (<|>) return [1..10]

225 let s1 = foldMapWith (<|>) return [1..10]

213 s2 = foldMapWith (<>) return [1..]

226 s2 = foldMapWith (<>) return [1..]

214− in toList (zipWithM (\a b -> return (a + b)) s1 s2)

227+ in toList (zM (\a b -> return (a + b)) s1 s2)

215 `shouldReturn` ([2,4..20] :: [Int])

228 `shouldReturn` ([2,4..20] :: [Int])

216

229

217− it "Applicative ZipSerial" $

230+ it "Applicative zip" $

218 let s1 = foldMapWith (<|>) return [1..10]

231 let s1 = foldMapWith (<|>) return [1..10]

219 s2 = foldMapWith (<>) return [1..]

232 s2 = foldMapWith (<>) return [1..]

220− in toList (getZipSerial ((+) <$> ZipSerial s1 <*> ZipSerial s2))

233+ in toList (getz ((+) <$> app s1 <*> app s2))

221 `shouldReturn` ([2,4..20] :: [Int])

234 `shouldReturn` ([2,4..20] :: [Int])

222

235

223timed :: Int -> AsyncT IO Int

236timed :: Int -> AsyncT IO Int