9ef94b0 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Folds with injection of initial state

9ef94b0Harendra Kumar7 years ago

examples/Split.hsnew file+40

compare with master

unifiedsplit

@@ -0,0 +1,40 @@

1+import qualified Streamly.Internal.Prelude as S

2+import qualified Streamly.Data.Fold as FL

3+-- import qualified Streamly.Memory.Array as A

4+import qualified Streamly.Internal.Memory.ArrayStream as AS

5+import qualified Streamly.Internal.FileSystem.Handle as IFH

6+import qualified Streamly.FileSystem.Handle as FH

7+import qualified System.IO as FH

8+-- import qualified Streamly.FileSystem.FD as FH

9+-- import qualified Streamly.Data.Unicode.Stream as US

10+

11+import Control.Monad.IO.Class (liftIO)

12+import Control.Monad.State.Strict (StateT(..), get, put)

13+import Data.Char (ord)

14+import System.Environment (getArgs)

15+import System.IO (IOMode(..), hSeek, SeekMode(..))

16+import Data.Function ((&))

17+

18+newHandle :: StateT (Maybe (FH.Handle, Int)) IO FH.Handle

19+newHandle = do

20+ old <- get

21+ idx <- case old of

22+ Nothing -> return 0

23+ Just (h, i) -> liftIO (FH.hClose h) >> return (i + 1)

24+ h <- liftIO $ FH.openFile ("dst-xyz-" ++ show idx ++ ".txt") WriteMode

25+ put (Just (h, idx))

26+ return h

27+

28+splitFile :: FH.Handle -> IO ()

29+splitFile inHandle =

30+ S.unfold FH.read inHandle

31+ & S.liftInner

32+ & S.chunksOf2 (180 * 1024 * 1024) newHandle IFH.write2

33+ & S.evalStateT Nothing -- generate new handle for each iteration

34+ & S.drain

35+

36+main :: IO ()

37+main = do

38+ name <- fmap head getArgs

39+ src <- FH.openFile name ReadMode

40+ splitFile src

src/Streamly/Internal/Data/Fold.hs+5

@@ -41,6 +41,7 @@ module Streamly.Internal.Data.Fold

41 -- ** Full Folds

42 , drain

43 , drainBy

44+ , drainBy2

44 , last

45 , last

45 , length

46 , length

46 , sum

47 , sum

@@ -377,6 +378,10 @@ drain = Fold step begin done

377drainBy :: Monad m => (a -> m b) -> Fold m a ()

378drainBy :: Monad m => (a -> m b) -> Fold m a ()

378drainBy f = Fold (const (void . f)) (return ()) return

379drainBy f = Fold (const (void . f)) (return ()) return

379

380

381+{-# INLINABLE drainBy2 #-}

382+drainBy2 :: Monad m => (a -> m b) -> Fold2 m c a ()

383+drainBy2 f = Fold2 (const (void . f)) (\_ -> return ()) return

384+

380-- | Extract the last element of the input stream, if any.

385-- | Extract the last element of the input stream, if any.

381--

386--

382-- @since 0.7.0

387-- @since 0.7.0

src/Streamly/Internal/Data/Fold/Types.hs+41

@@ -14,6 +14,8 @@

15module Streamly.Internal.Data.Fold.Types

16 ( Fold (..)

17+ , Fold2 (..)

18+ , simplify

17 , toListRevF -- experimental

19 , toListRevF -- experimental

18 , lmap

20 , lmap

19 , lmapM

21 , lmapM

@@ -24,6 +26,7 @@ module Streamly.Internal.Data.Fold.Types

24 , ltakeWhile

26 , ltakeWhile

25 , lsessionsOf

27 , lsessionsOf

26 , lchunksOf

28 , lchunksOf

29+ , lchunksOf2

28 , duplicate

31 , duplicate

29 , initialize

32 , initialize

@@ -64,6 +67,19 @@ data Fold m a b =

64 -- | @Fold @ @ step @ @ initial @ @ extract@

67 -- | @Fold @ @ step @ @ initial @ @ extract@

65 forall s. Fold (s -> a -> m s) (m s) (s -> m b)

68 forall s. Fold (s -> a -> m s) (m s) (s -> m b)

70+-- Experimental type to provide a side input to the fold for generating the

71+-- initial state. For example, if we have to fold chunks of a stream and write

72+-- each chunk to a different file, then we can generate the file name using a

73+-- monadic action. This is a generalized version of 'Fold'.

74+--

75+data Fold2 m c a b =

76+ -- | @Fold @ @ step @ @ inject @ @ extract@

77+ forall s. Fold2 (s -> a -> m s) (c -> m s) (s -> m b)

78+

79+-- | Convert more general type 'Fold2' into a simpler type 'Fold'

80+simplify :: Fold2 m c a b -> c -> Fold m a b

81+simplify (Fold2 step inject extract) c = Fold step (inject c) extract

82+

67-- | Maps a function on the output of the fold (the type @b@).

83-- | Maps a function on the output of the fold (the type @b@).

68instance Applicative m => Functor (Fold m a) where

84instance Applicative m => Functor (Fold m a) where

69 {-# INLINE fmap #-}

85 {-# INLINE fmap #-}

@@ -370,6 +386,31 @@ lchunksOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =

370 acc2 <- step2 r2 res

386 acc2 <- step2 r2 res

371 extract2 acc2

387 extract2 acc2

372

388

389+{-# INLINE lchunksOf2 #-}

390+lchunksOf2 :: Monad m => Int -> Fold m a b -> Fold2 m x b c -> Fold2 m x a c

391+lchunksOf2 n (Fold step1 initial1 extract1) (Fold2 step2 inject2 extract2) =

392+ Fold2 step' inject' extract'

393+

394+ where

395+

396+ inject' x = (Tuple3' 0) <$> initial1 <*> inject2 x

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

398+ if i < n

399+ then do

400+ res <- step1 r1 a

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

402+ else do

403+ res <- extract1 r1

404+ acc2 <- step2 r2 res

405+

406+ i1 <- initial1

407+ acc1 <- step1 i1 a

408+ return $ Tuple3' 1 acc1 acc2

409+ extract' (Tuple3' _ r1 r2) = do

410+ res <- extract1 r1

411+ acc2 <- step2 r2 res

412+ extract2 acc2

413+

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

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

374-- group using the provided fold function.

415-- group using the provided fold function.

375--

416--

src/Streamly/Internal/Data/Stream/StreamD/Type.hs+46−1

@@ -63,6 +63,7 @@ module Streamly.Internal.Data.Stream.StreamD.Type

63 , take

64 , GroupState (..) -- for inspection testing

65 , groupsOf

66+ , groupsOf2

66 )

67 )

67where

68where

@@ -76,7 +77,7 @@ import GHC.Types (SPEC(..))

76import Prelude hiding (map, mapM, foldr, take, concatMap)

77import Prelude hiding (map, mapM, foldr, take, concatMap)

78import Streamly.Internal.Data.SVar (State(..), adaptState, defState)

79import Streamly.Internal.Data.SVar (State(..), adaptState, defState)

79−import Streamly.Internal.Data.Fold.Types (Fold(..))

80+import Streamly.Internal.Data.Fold.Types (Fold(..), Fold2(..))

81import qualified Streamly.Streams.StreamK as K

82import qualified Streamly.Streams.StreamK as K

@@ -551,3 +552,47 @@ groupsOf n (Fold fstep initial extract) (Stream step state) =

551 return $ Yield r next

552 return $ Yield r next

552

553

553 step' _ GroupFinish = return Stop

554 step' _ GroupFinish = return Stop

555+

556+{-# INLINE_NORMAL groupsOf2 #-}

557+groupsOf2

558+ :: Monad m

559+ => Int

560+ -> m c

561+ -> Fold2 m c a b

562+ -> Stream m a

563+ -> Stream m b

564+groupsOf2 n input (Fold2 fstep inject extract) (Stream step state) =

565+ n `seq` Stream step' (GroupStart state)

566+

567+ where

568+

569+ {-# INLINE_LATE step' #-}

570+ step' _ (GroupStart st) = do

571+ -- XXX shall we use the Natural type instead? Need to check performance

572+ -- implications.

573+ when (n <= 0) $

574+ -- XXX we can pass the module string from the higher level API

575+ error $ "Streamly.Internal.Data.Stream.StreamD.Type.groupsOf: the size of "

576+ ++ "groups [" ++ show n ++ "] must be a natural number"

577+ -- fs = fold state

578+ fs <- input >>= inject

579+ return $ Skip (GroupBuffer st fs 0)

580+

581+ step' gst (GroupBuffer st fs i) = do

582+ r <- step (adaptState gst) st

583+ case r of

584+ Yield x s -> do

585+ fs' <- fstep fs x

586+ let i' = i + 1

587+ return $

588+ if i' >= n

589+ then Skip (GroupYield fs' (GroupStart s))

590+ else Skip (GroupBuffer s fs' i')

591+ Skip s -> return $ Skip (GroupBuffer s fs i)

592+ Stop -> return $ Skip (GroupYield fs GroupFinish)

593+

594+ step' _ (GroupYield fs next) = do

595+ r <- extract fs

596+ return $ Yield r next

597+

598+ step' _ GroupFinish = return Stop

src/Streamly/Internal/FileSystem/Handle.hs+14

@@ -43,6 +43,7 @@ module Streamly.Internal.FileSystem.Handle

43 -- ** Write to Handle

44 -- Byte stream write (Folds)

45 , write

46+ , write2

46 -- , writeUtf8

47 -- , writeUtf8

47 -- , writeUtf8ByLines

48 -- , writeUtf8ByLines

48 -- , writeByFrames

49 -- , writeByFrames

@@ -117,6 +118,7 @@ import Prelude hiding (read)

117

118

118import Streamly (MonadAsync)

119import Streamly (MonadAsync)

119import Streamly.Data.Fold (Fold)

120import Streamly.Data.Fold (Fold)

121+import Streamly.Internal.Data.Fold.Types (Fold2(..))

120import Streamly.Internal.Data.Unfold.Types (Unfold(..))

122import Streamly.Internal.Data.Unfold.Types (Unfold(..))

121import Streamly.Internal.Memory.Array.Types

123import Streamly.Internal.Memory.Array.Types

122 (Array(..), writeNUnsafe, defaultChunkSize, shrinkToFit,

124 (Array(..), writeNUnsafe, defaultChunkSize, shrinkToFit,

@@ -410,6 +412,10 @@ fromBytes = fromBytesWithBufferOf defaultChunkSize

410writeChunks :: (MonadIO m, Storable a) => Handle -> Fold m (Array a) ()

412writeChunks :: (MonadIO m, Storable a) => Handle -> Fold m (Array a) ()

411writeChunks h = FL.drainBy (liftIO . writeArray h)

413writeChunks h = FL.drainBy (liftIO . writeArray h)

412

414

415+{-# INLINE writeChunks2 #-}

416+writeChunks2 :: (MonadIO m, Storable a) => Fold2 m Handle (Array a) ()

417+writeChunks2 = Fold2 (\h arr -> liftIO $ writeArray h arr >> return h) return (\_ -> return ())

418+

413-- | @writeChunksWithBufferOf bufsize handle@ writes a stream of arrays

419-- | @writeChunksWithBufferOf bufsize handle@ writes a stream of arrays

414-- to @handle@ after coalescing the adjacent arrays in chunks of @bufsize@.

420-- to @handle@ after coalescing the adjacent arrays in chunks of @bufsize@.

415-- We never split an array, if a single array is bigger than the specified size

421-- We never split an array, if a single array is bigger than the specified size

@@ -439,6 +445,10 @@ writeChunksWithBufferOf n h = lpackArraysChunksOf n (writeChunks h)

439writeWithBufferOf :: MonadIO m => Int -> Handle -> Fold m Word8 ()

445writeWithBufferOf :: MonadIO m => Int -> Handle -> Fold m Word8 ()

440writeWithBufferOf n h = FL.lchunksOf n (writeNUnsafe n) (writeChunks h)

446writeWithBufferOf n h = FL.lchunksOf n (writeNUnsafe n) (writeChunks h)

441

447

448+{-# INLINE writeWithBufferOf2 #-}

449+writeWithBufferOf2 :: MonadIO m => Int -> Fold2 m Handle Word8 ()

450+writeWithBufferOf2 n = FL.lchunksOf2 n (writeNUnsafe n) writeChunks2

451+

442-- > write = 'writeWithBufferOf' A.defaultChunkSize

452-- > write = 'writeWithBufferOf' A.defaultChunkSize

443--

453--

444-- | Write a byte stream to a file handle. Accumulates the input in chunks of

454-- | Write a byte stream to a file handle. Accumulates the input in chunks of

@@ -450,6 +460,10 @@ writeWithBufferOf n h = FL.lchunksOf n (writeNUnsafe n) (writeChunks h)

450write :: MonadIO m => Handle -> Fold m Word8 ()

460write :: MonadIO m => Handle -> Fold m Word8 ()

451write = writeWithBufferOf defaultChunkSize

461write = writeWithBufferOf defaultChunkSize

452

462

463+{-# INLINE write2 #-}

464+write2 :: MonadIO m => Fold2 m Handle Word8 ()

465+write2 = writeWithBufferOf2 defaultChunkSize

466+

453{-

467{-

454{-# INLINE write #-}

468{-# INLINE write #-}

455write :: (MonadIO m, Storable a) => Handle -> SerialT m a -> m ()

469write :: (MonadIO m, Storable a) => Handle -> SerialT m a -> m ()

src/Streamly/Internal/Prelude.hs+8−1

@@ -287,6 +287,7 @@ module Streamly.Internal.Prelude

287

288 -- -- *** Chunks

289 , chunksOf

290+ , chunksOf2

290 , arraysOf

291 , arraysOf

291 , intervalsOf

292 , intervalsOf

292

293

@@ -429,7 +430,7 @@ import qualified Prelude

429import qualified System.IO as IO

430import qualified System.IO as IO

430

431

431import Streamly.Streams.Enumeration (Enumerable(..), enumerate, enumerateTo)

432import Streamly.Streams.Enumeration (Enumerable(..), enumerate, enumerateTo)

432−import Streamly.Internal.Data.Fold.Types (Fold (..))

433+import Streamly.Internal.Data.Fold.Types (Fold (..), Fold2 (..))

433import Streamly.Internal.Data.Unfold.Types (Unfold)

434import Streamly.Internal.Data.Unfold.Types (Unfold)

434import Streamly.Internal.Memory.Array.Types (Array, writeNUnsafe)

435import Streamly.Internal.Memory.Array.Types (Array, writeNUnsafe)

435-- import Streamly.Memory.Ring (Ring)

436-- import Streamly.Memory.Ring (Ring)

@@ -2570,6 +2571,12 @@ chunksOf

2570 => Int -> Fold m a b -> t m a -> t m b

2571 => Int -> Fold m a b -> t m a -> t m b

2571chunksOf n f m = D.fromStreamD $ D.groupsOf n f (D.toStreamD m)

2572chunksOf n f m = D.fromStreamD $ D.groupsOf n f (D.toStreamD m)

2572

2573

2574+{-# INLINE chunksOf2 #-}

2575+chunksOf2

2576+ :: (IsStream t, Monad m)

2577+ => Int -> m c -> Fold2 m c a b -> t m a -> t m b

2578+chunksOf2 n action f m = D.fromStreamD $ D.groupsOf2 n action f (D.toStreamD m)

2579+

2573-- | @arraysOf n stream@ groups the elements in the input stream into arrays of

2580-- | @arraysOf n stream@ groups the elements in the input stream into arrays of

2574-- @n@ elements each.

2581-- @n@ elements each.

2575--

2582--

src/Streamly/Streams/StreamD.hs+1

@@ -159,6 +159,7 @@ module Streamly.Streams.StreamD

159

160 -- ** Grouping

161 , groupsOf

162+ , groupsOf2

162 , groupsBy

163 , groupsBy

163 , groupsRollingBy

164 , groupsRollingBy

164

165