27d7ea8 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

expose the experimental Fold APIs via internal module

27d7ea8Harendra Kumar7 years ago

benchmark/Linear.hs+2−1

compare with master

unifiedsplit

@@ -20,6 +20,7 @@ import qualified LinearOps as Ops

21import Streamly

22import qualified Streamly.Fold as FL

23+import qualified Streamly.Internal as Internal

23--import qualified Streamly.Pipe as Pipe

24--import qualified Streamly.Pipe as Pipe

24import qualified Streamly.Memory.Array as A

25import qualified Streamly.Memory.Array as A

25import qualified Streamly.Prelude as S

26import qualified Streamly.Prelude as S

@@ -241,7 +242,7 @@ main =

241 ]

242 ]

242 , bgroup "folds-transforms"

243 , bgroup "folds-transforms"

243 [ benchIOSink "drain" (S.runFold FL.drain)

244 [ benchIOSink "drain" (S.runFold FL.drain)

244− , benchIOSink "lmap" (S.runFold (FL.lmap (+1) FL.drain))

245+ , benchIOSink "lmap" (S.runFold (Internal.lmap (+1) FL.drain))

245 {-, benchIOSink "pipe-mapM"

246 {-, benchIOSink "pipe-mapM"

246 (S.runFold (FL.transform (Pipe.mapM (\x -> return $ x + 1)) FL.drain))-}

247 (S.runFold (FL.transform (Pipe.mapM (\x -> return $ x + 1)) FL.drain))-}

247 ]

248 ]

src/Streamly/Fold.hs+4−206

@@ -134,12 +134,12 @@ module Streamly.Fold

134

135 -- ** Mapping

136 --, transform

137− , lmap

137+ -- , lmap

138 --, lsequence

139− , lmapM

139+ -- , lmapM

140

141 -- -- ** Filtering

142− , lfilter

142+ -- , lfilter

143 -- , lfilterM

144 -- , ldeleteBy

145 -- , luniq

@@ -247,13 +247,7 @@ module Streamly.Fold

247 )

248where

249

250−import Control.Concurrent (threadDelay, forkIO, killThread)

251−import Control.Concurrent.MVar (MVar, newMVar, takeMVar, putMVar)

252−import Control.Exception (SomeException(..), catch, mask)

253import Control.Monad (void)

250import Control.Monad (void)

254−import Control.Monad.Catch (throwM)

255−import Control.Monad.IO.Class (MonadIO(..))

256−import Control.Monad.Trans.Control (control)

257import Data.Map.Strict (Map)

251import Data.Map.Strict (Map)

258

252

259import Prelude

253import Prelude

@@ -267,8 +261,7 @@ import Prelude

267import qualified Data.Map.Strict as Map

261import qualified Data.Map.Strict as Map

268import qualified Prelude

262import qualified Prelude

269

263

270−import Streamly.SVar (MonadAsync)

264+import Streamly.Fold.Types (Fold(..), lmap)

271−import Streamly.Fold.Types (Fold(..))

272import Streamly.Pipe.Types (Pipe (..), PipeState(..))

265import Streamly.Pipe.Types (Pipe (..), PipeState(..))

273import Streamly.Strict

266import Streamly.Strict

274

267

@@ -336,27 +329,6 @@ _transform (Pipe pstep1 pstep2 pinitial) (Fold fstep finitial fextract) =

336

329

337 extract (Tuple' _ fs) = fextract fs

330 extract (Tuple' _ fs) = fextract fs

338

331

339−-- | @(lmap f fold)@ maps the function @f@ on the input of the fold.

340−--

341−-- >>> S.runFold (FL.lmap (\x -> x * x) FL.sum) (S.enumerateFromTo 1 100)

342−-- 338350

343−--

344−-- @since 0.7.0

345−{-# INLINABLE lmap #-}

346−lmap :: (a -> b) -> Fold m b r -> Fold m a r

347−lmap f (Fold step begin done) = Fold step' begin done

348− where

349− step' x a = step x (f a)

350−

351−-- | @(lmapM f fold)@ maps the monadic function @f@ on the input of the fold.

352−--

353−-- @since 0.7.0

354−{-# INLINABLE lmapM #-}

355−lmapM :: Monad m => (a -> m b) -> Fold m b r -> Fold m a r

356−lmapM f (Fold step begin done) = Fold step' begin done

357− where

358− step' x a = f a >>= step x

359−

360------------------------------------------------------------------------------

332------------------------------------------------------------------------------

361-- Utilities

333-- Utilities

362------------------------------------------------------------------------------

334------------------------------------------------------------------------------

@@ -774,65 +746,6 @@ and = Fold (\x a -> return $ x && a) (return True) return

774or :: Monad m => Fold m Bool Bool

746or :: Monad m => Fold m Bool Bool

775or = Fold (\x a -> return $ x || a) (return False) return

747or = Fold (\x a -> return $ x || a) (return False) return

776

748

777−------------------------------------------------------------------------------

778−-- Filtering

779−------------------------------------------------------------------------------

780−

781−-- | Include only those elements that pass a predicate.

782−--

783−-- >>> S.runFold (lfilter (> 5) FL.sum) [1..10]

784−-- 40

785−--

786−-- @since 0.7.0

787−{-# INLINABLE lfilter #-}

788−lfilter :: Monad m => (a -> Bool) -> Fold m a r -> Fold m a r

789−lfilter f (Fold step begin done) = Fold step' begin done

790− where

791− step' x a = if f a then step x a else return x

792−

793−-- | Like 'lfilter' but with a monadic predicate.

794−--

795−-- @since 0.7.0

796−{-# INLINABLE _lfilterM #-}

797−_lfilterM :: Monad m => (a -> m Bool) -> Fold m a r -> Fold m a r

798−_lfilterM f (Fold step begin done) = Fold step' begin done

799− where

800− step' x a = do

801− use <- f a

802− if use then step x a else return x

803−

804−-- | Take first 'n' elements from the stream and discard the rest.

805−--

806−-- @since 0.7.0

807−{-# INLINABLE _ltake #-}

808−_ltake :: Monad m => Int -> Fold m a b -> Fold m a b

809−_ltake n (Fold step initial done) = Fold step' initial' done'

810− where

811− initial' = fmap (Tuple' 0) initial

812− step' (Tuple' i r) a = do

813− if i < n

814− then do

815− res <- step r a

816− return $ Tuple' (i + 1) res

817− else return $ Tuple' i r

818− done' (Tuple' _ r) = done r

819−

820−-- | Takes elements from the input as long as the predicate succeeds.

821−--

822−-- @since 0.7.0

823−{-# INLINABLE _ltakeWhile #-}

824−_ltakeWhile :: Monad m => (a -> Bool) -> Fold m a b -> Fold m a b

825−_ltakeWhile predicate (Fold step initial done) = Fold step' initial' done'

826− where

827− initial' = fmap Left' initial

828− step' (Left' r) a = do

829− if predicate a

830− then fmap Left' $ step r a

831− else return (Right' r)

832− step' r _ = return r

833− done' (Left' r) = done r

834− done' (Right' r) = done r

835−

836------------------------------------------------------------------------------

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

837-- Distributing

750-- Distributing

838------------------------------------------------------------------------------

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

@@ -1231,23 +1144,6 @@ unzip = unzipWith id

1231-- Nesting

1144-- Nesting

1232------------------------------------------------------------------------------

1145------------------------------------------------------------------------------

1233--

1146--

1234−-- | Modify the fold such that when the fold is done, instead of returning the

1235−-- accumulator, it returns a fold. The returned fold starts from where we left

1236−-- i.e. it uses the last accumulator value as the initial value of the

1237−-- accumulator. Thus we can resume the fold later and feed it more input.

1238−--

1239−-- >> do

1240−-- > more <- S.runFold (FL.duplicate FL.sum) (S.enumerateFromTo 1 10)

1241−-- > evenMore <- S.runFold (FL.duplicate more) (S.enumerateFromTo 11 20)

1242−-- > S.runFold evenMore (S.enumerateFromTo 21 30)

1243−-- > 465

1244−--

1245−-- @since 0.7.0

1246−{-# INLINABLE _duplicate #-}

1247−_duplicate :: Applicative m => Fold m a b -> Fold m a (Fold m a b)

1248−_duplicate (Fold step begin done) =

1249− Fold step begin (\x -> pure (Fold step (pure x) done))

1250−

1251{-

1147{-

1252-- All the stream flattening transformations can also be applied to a fold

1148-- All the stream flattening transformations can also be applied to a fold

1253-- input stream.

1149-- input stream.

@@ -1281,101 +1177,3 @@ lchunksInRange :: Monad m

1281lchunksInRange low high (Fold step1 initial1 extract1)

1177lchunksInRange low high (Fold step1 initial1 extract1)

1282 (Fold step2 initial2 extract2) = undefined

1178 (Fold step2 initial2 extract2) = undefined

1283-}

1179-}

1284−

1285−-- | Group the input stream into groups of @n@ elements each and then fold each

1286−-- group using the provided fold function.

1287−--

1288−-- @

1289−--

1290−-- -----Fold m a b----|-Fold n a c-|-Fold n a c-|-...-|----Fold m a c

1291−--

1292−-- @

1293−--

1294−{-# INLINE _lchunksOf #-}

1295−_lchunksOf :: Monad m => Int -> Fold m a b -> Fold m b c -> Fold m a c

1296−_lchunksOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =

1297− Fold step' initial' extract'

1298−

1299− where

1300−

1301− initial' = (Tuple3' 0) <$> initial1 <*> initial2

1302− step' (Tuple3' i r1 r2) a = do

1303− if i < n

1304− then do

1305− res <- step1 r1 a

1306− return $ Tuple3' (i + 1) res r2

1307− else do

1308− res <- extract1 r1

1309− acc2 <- step2 r2 res

1310−

1311− i1 <- initial1

1312− acc1 <- step1 i1 a

1313− return $ Tuple3' 1 acc1 acc2

1314− extract' (Tuple3' _ _ r) = extract2 r

1315−

1316−-- | Group the input stream into windows of n second each and then fold each

1317−-- group using the provided fold function.

1318−--

1319−-- For example, we can copy and distribute a stream to multiple folds where

1320−-- each fold can group the input differently e.g. by one second, one minute and

1321−-- one hour windows respectively and fold each resulting stream of folds.

1322−--

1323−-- @

1324−--

1325−-- -----Fold m a b----|-Fold n a c-|-Fold n a c-|-...-|----Fold m a c

1326−--

1327−-- @

1328−{-# INLINE _lsessionsOf #-}

1329−_lsessionsOf :: MonadAsync m => Double -> Fold m a b -> Fold m b c -> Fold m a c

1330−_lsessionsOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =

1331− Fold step' initial' extract'

1332−

1333− where

1334−

1335− -- XXX MVar may be expensive we need a cheaper synch mechanism here

1336− initial' = do

1337− i1 <- initial1

1338− i2 <- initial2

1339− mv1 <- liftIO $ newMVar i1

1340− mv2 <- liftIO $ newMVar (Right i2)

1341− t <- control $ \run ->

1342− mask $ \restore -> do

1343− tid <- forkIO $ catch (restore $ void $ run (timerThread mv1 mv2))

1344− (handleChildException mv2)

1345− run (return tid)

1346− return $ Tuple3' t mv1 mv2

1347− step' acc@(Tuple3' _ mv1 _) a = do

1348− r1 <- liftIO $ takeMVar mv1

1349− res <- step1 r1 a

1350− liftIO $ putMVar mv1 res

1351− return acc

1352− extract' (Tuple3' tid _ mv2) = do

1353− r2 <- liftIO $ takeMVar mv2

1354− liftIO $ killThread tid

1355− case r2 of

1356− Left e -> throwM e

1357− Right x -> extract2 x

1358−

1359− timerThread mv1 mv2 = do

1360− liftIO $ threadDelay (round $ n * 1000000)

1361−

1362− r1 <- liftIO $ takeMVar mv1

1363− i1 <- initial1

1364− liftIO $ putMVar mv1 i1

1365−

1366− res1 <- extract1 r1

1367− r2 <- liftIO $ takeMVar mv2

1368− res <- case r2 of

1369− Left _ -> return r2

1370− Right x -> fmap Right $ step2 x res1

1371− liftIO $ putMVar mv2 res

1372− timerThread mv1 mv2

1373−

1374− handleChildException ::

1375− MVar (Either SomeException a) -> SomeException -> IO ()

1376− handleChildException mv2 e = do

1377− r2 <- takeMVar mv2

1378− let r = case r2 of

1379− Left _ -> r2

1380− Right _ -> Left e

1381− putMVar mv2 r

src/Streamly/Fold/Types.hs+222−1

@@ -1,5 +1,6 @@

1{-# LANGUAGE CPP #-}

2{-# LANGUAGE ExistentialQuantification #-}

3+{-# LANGUAGE FlexibleContexts #-}

4-- |

5-- |

5-- Module : Streamly.Fold.Types

6-- Module : Streamly.Fold.Types

@@ -12,12 +13,29 @@

13module Streamly.Fold.Types

14module Streamly.Fold.Types

14 ( Fold (..)

15 ( Fold (..)

16+ , lmap

17+ , lmapM

18+ , lfilter

19+ , lfilterM

20+ , ltake

21+ , ltakeWhile

22+ , lsessionsOf

23+ , lchunksOf

24+ , duplicate

15 )

25 )

16where

26where

18import Control.Applicative (liftA2)

28import Control.Applicative (liftA2)

29+import Control.Concurrent (threadDelay, forkIO, killThread)

30+import Control.Concurrent.MVar (MVar, newMVar, takeMVar, putMVar)

31+import Control.Exception (SomeException(..), catch, mask)

32+import Control.Monad (void)

33+import Control.Monad.Catch (throwM)

34+import Control.Monad.IO.Class (MonadIO(..))

35+import Control.Monad.Trans.Control (control)

19import Data.Semigroup (Semigroup(..))

36import Data.Semigroup (Semigroup(..))

20−import Streamly.Strict (Tuple'(..))

37+import Streamly.Strict (Tuple'(..), Tuple3'(..), Either'(..))

38+import Streamly.SVar (MonadAsync)

22------------------------------------------------------------------------------

40------------------------------------------------------------------------------

23-- Monadic left folds

41-- Monadic left folds

@@ -170,3 +188,206 @@ instance (Monad m, Floating b) => Floating (Fold m a b) where

170

188

171 {-# INLINE logBase #-}

189 {-# INLINE logBase #-}

172 logBase = liftA2 logBase

190 logBase = liftA2 logBase

191+

192+------------------------------------------------------------------------------

193+-- Internal APIs

194+------------------------------------------------------------------------------

195+

196+-- | @(lmap f fold)@ maps the function @f@ on the input of the fold.

197+--

198+-- >>> S.runFold (FL.lmap (\x -> x * x) FL.sum) (S.enumerateFromTo 1 100)

199+-- 338350

200+--

201+-- @since 0.7.0

202+{-# INLINABLE lmap #-}

203+lmap :: (a -> b) -> Fold m b r -> Fold m a r

204+lmap f (Fold step begin done) = Fold step' begin done

205+ where

206+ step' x a = step x (f a)

207+

208+-- | @(lmapM f fold)@ maps the monadic function @f@ on the input of the fold.

209+--

210+-- @since 0.7.0

211+{-# INLINABLE lmapM #-}

212+lmapM :: Monad m => (a -> m b) -> Fold m b r -> Fold m a r

213+lmapM f (Fold step begin done) = Fold step' begin done

214+ where

215+ step' x a = f a >>= step x

216+

217+------------------------------------------------------------------------------

218+-- Filtering

219+------------------------------------------------------------------------------

220+

221+-- | Include only those elements that pass a predicate.

222+--

223+-- >>> S.runFold (lfilter (> 5) FL.sum) [1..10]

224+-- 40

225+--

226+-- @since 0.7.0

227+{-# INLINABLE lfilter #-}

228+lfilter :: Monad m => (a -> Bool) -> Fold m a r -> Fold m a r

229+lfilter f (Fold step begin done) = Fold step' begin done

230+ where

231+ step' x a = if f a then step x a else return x

232+

233+-- | Like 'lfilter' but with a monadic predicate.

234+--

235+-- @since 0.7.0

236+{-# INLINABLE lfilterM #-}

237+lfilterM :: Monad m => (a -> m Bool) -> Fold m a r -> Fold m a r

238+lfilterM f (Fold step begin done) = Fold step' begin done

239+ where

240+ step' x a = do

241+ use <- f a

242+ if use then step x a else return x

243+

244+-- | Take first 'n' elements from the stream and discard the rest.

245+--

246+-- @since 0.7.0

247+{-# INLINABLE ltake #-}

248+ltake :: Monad m => Int -> Fold m a b -> Fold m a b

249+ltake n (Fold step initial done) = Fold step' initial' done'

250+ where

251+ initial' = fmap (Tuple' 0) initial

252+ step' (Tuple' i r) a = do

253+ if i < n

254+ then do

255+ res <- step r a

256+ return $ Tuple' (i + 1) res

257+ else return $ Tuple' i r

258+ done' (Tuple' _ r) = done r

259+

260+-- | Takes elements from the input as long as the predicate succeeds.

261+--

262+-- @since 0.7.0

263+{-# INLINABLE ltakeWhile #-}

264+ltakeWhile :: Monad m => (a -> Bool) -> Fold m a b -> Fold m a b

265+ltakeWhile predicate (Fold step initial done) = Fold step' initial' done'

266+ where

267+ initial' = fmap Left' initial

268+ step' (Left' r) a = do

269+ if predicate a

270+ then fmap Left' $ step r a

271+ else return (Right' r)

272+ step' r _ = return r

273+ done' (Left' r) = done r

274+ done' (Right' r) = done r

275+

276+------------------------------------------------------------------------------

277+-- Nesting

278+------------------------------------------------------------------------------

279+--

280+-- | Modify the fold such that when the fold is done, instead of returning the

281+-- accumulator, it returns a fold. The returned fold starts from where we left

282+-- i.e. it uses the last accumulator value as the initial value of the

283+-- accumulator. Thus we can resume the fold later and feed it more input.

284+--

285+-- >> do

286+-- > more <- S.runFold (FL.duplicate FL.sum) (S.enumerateFromTo 1 10)

287+-- > evenMore <- S.runFold (FL.duplicate more) (S.enumerateFromTo 11 20)

288+-- > S.runFold evenMore (S.enumerateFromTo 21 30)

289+-- > 465

290+--

291+-- @since 0.7.0

292+{-# INLINABLE duplicate #-}

293+duplicate :: Applicative m => Fold m a b -> Fold m a (Fold m a b)

294+duplicate (Fold step begin done) =

295+ Fold step begin (\x -> pure (Fold step (pure x) done))

296+

297+-- | Group the input stream into groups of @n@ elements each and then fold each

298+-- group using the provided fold function.

299+--

300+-- @

301+--

302+-- -----Fold m a b----|-Fold n a c-|-Fold n a c-|-...-|----Fold m a c

303+--

304+-- @

305+--

306+{-# INLINE lchunksOf #-}

307+lchunksOf :: Monad m => Int -> Fold m a b -> Fold m b c -> Fold m a c

308+lchunksOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =

309+ Fold step' initial' extract'

310+

311+ where

312+

313+ initial' = (Tuple3' 0) <$> initial1 <*> initial2

314+ step' (Tuple3' i r1 r2) a = do

315+ if i < n

316+ then do

317+ res <- step1 r1 a

318+ return $ Tuple3' (i + 1) res r2

319+ else do

320+ res <- extract1 r1

321+ acc2 <- step2 r2 res

322+

323+ i1 <- initial1

324+ acc1 <- step1 i1 a

325+ return $ Tuple3' 1 acc1 acc2

326+ extract' (Tuple3' _ _ r) = extract2 r

327+

328+-- | Group the input stream into windows of n second each and then fold each

329+-- group using the provided fold function.

330+--

331+-- For example, we can copy and distribute a stream to multiple folds where

332+-- each fold can group the input differently e.g. by one second, one minute and

333+-- one hour windows respectively and fold each resulting stream of folds.

334+--

335+-- @

336+--

337+-- -----Fold m a b----|-Fold n a c-|-Fold n a c-|-...-|----Fold m a c

338+--

339+-- @

340+{-# INLINE lsessionsOf #-}

341+lsessionsOf :: MonadAsync m => Double -> Fold m a b -> Fold m b c -> Fold m a c

342+lsessionsOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =

343+ Fold step' initial' extract'

344+

345+ where

346+

347+ -- XXX MVar may be expensive we need a cheaper synch mechanism here

348+ initial' = do

349+ i1 <- initial1

350+ i2 <- initial2

351+ mv1 <- liftIO $ newMVar i1

352+ mv2 <- liftIO $ newMVar (Right i2)

353+ t <- control $ \run ->

354+ mask $ \restore -> do

355+ tid <- forkIO $ catch (restore $ void $ run (timerThread mv1 mv2))

356+ (handleChildException mv2)

357+ run (return tid)

358+ return $ Tuple3' t mv1 mv2

359+ step' acc@(Tuple3' _ mv1 _) a = do

360+ r1 <- liftIO $ takeMVar mv1

361+ res <- step1 r1 a

362+ liftIO $ putMVar mv1 res

363+ return acc

364+ extract' (Tuple3' tid _ mv2) = do

365+ r2 <- liftIO $ takeMVar mv2

366+ liftIO $ killThread tid

367+ case r2 of

368+ Left e -> throwM e

369+ Right x -> extract2 x

370+

371+ timerThread mv1 mv2 = do

372+ liftIO $ threadDelay (round $ n * 1000000)

373+

374+ r1 <- liftIO $ takeMVar mv1

375+ i1 <- initial1

376+ liftIO $ putMVar mv1 i1

377+

378+ res1 <- extract1 r1

379+ r2 <- liftIO $ takeMVar mv2

380+ res <- case r2 of

381+ Left _ -> return r2

382+ Right x -> fmap Right $ step2 x res1

383+ liftIO $ putMVar mv2 res

384+ timerThread mv1 mv2

385+

386+ handleChildException ::

387+ MVar (Either SomeException a) -> SomeException -> IO ()

388+ handleChildException mv2 e = do

389+ r2 <- takeMVar mv2

390+ let r = case r2 of

391+ Left _ -> r2

392+ Right _ -> Left e

393+ putMVar mv2 r

src/Streamly/Internal.hs+20−5

@@ -7,10 +7,14 @@

7-- Stability : experimental

8-- Portability : GHC

9--

10−-- This module is only for internal use. There is no warranty for the routines

10+-- This module is only for internal or experimental APIs. The constructors and

11−-- in this module to work correctly, please use at your own risk. The

12−-- constructors and routines exposed through this module are likely to change

11+-- routines exposed through this module are likely to change or to be removed

13−-- or to be removed in future without notice.

12+-- at any time in future without notice. There is no warranty for the routines

13+-- in this module to work correctly, please use at your own risk.

14+--

15+-- If you think some of the routines exported through this module should be

16+-- included in the officially supported modules please raise an issue at

17+-- https://github.com/composewell/streamly/issues .

14--

18--

15module Streamly.Internal

19module Streamly.Internal

16 ( Fold (..)

20 ( Fold (..)

@@ -18,9 +22,20 @@ module Streamly.Internal

19 -- * Streamly.Prelude

23 -- * Streamly.Prelude

20 , tapAsync

24 , tapAsync

25+

26+ -- * Streamly.Fold

27+ , lmap

28+ , lmapM

29+ , lfilter

30+ , lfilterM

31+ , ltake

32+ , ltakeWhile

33+ , lsessionsOf

34+ , lchunksOf

35+ , duplicate

21 )

36 )

22where

37where

24−import Streamly.Fold.Types (Fold(..))

39+import Streamly.Fold.Types

25import Streamly.Streams.Combinators (inspectMode)

40import Streamly.Streams.Combinators (inspectMode)

26import Streamly.Streams.Parallel (tapAsync)

41import Streamly.Streams.Parallel (tapAsync)

src/Streamly/Prelude.hs+1

@@ -657,6 +657,7 @@ import Streamly.Strict

657

658import qualified Streamly.Memory.Array as A

659import qualified Streamly.Fold as FL

660+import qualified Streamly.Fold.Types as FL

660import qualified Streamly.Streams.Prelude as P

661import qualified Streamly.Streams.Prelude as P

661import qualified Streamly.Streams.StreamK as K

662import qualified Streamly.Streams.StreamK as K

662import qualified Streamly.Streams.StreamD as D

663import qualified Streamly.Streams.StreamD as D