9ef94b0 · streamly · git.iko.soy

git.iko.soy/streamly

Folds with injection of initial state

commit9ef94b0e395f66c8eceda1f24daff96f62830633 authorHarendra Kumar <harendra.kumar@gmail.com> dateMon 11 Nov 2019, 19:06 (7 years ago) parenta4c2ae4

7 files changed+155−2

compare with master

unifiedsplit

examples/Split.hs+40 src/Streamly/Internal/Data/Fold.hs+5 src/Streamly/Internal/Data/Fold/Types.hs+41 src/Streamly/Internal/Data/Stream/StreamD/Type.hs+46−1 src/Streamly/Internal/FileSystem/Handle.hs+14 src/Streamly/Internal/Prelude.hs+8−1 src/Streamly/Streams/StreamD.hs+1

examples/Split.hsnew file+40

@@ -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

4141 -- ** Full Folds

4242 , drain

4343 , drainBy

44+ , drainBy2

4445 , last

4546 , length

4647 , sum

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

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

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

379380

381+{-# INLINABLE drainBy2 #-}

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

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

384+

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

381386--

382387-- @since 0.7.0

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

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

1414

1515module Streamly.Internal.Data.Fold.Types

1616 ( Fold (..)

17+ , Fold2 (..)

18+ , simplify

1719 , toListRevF -- experimental

1820 , lmap

1921 , lmapM

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

2426 , ltakeWhile

2527 , lsessionsOf

2628 , lchunksOf

29+ , lchunksOf2

2730

2831 , duplicate

2932 , initialize

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

6467 -- | @Fold @ @ step @ @ initial @ @ extract@

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

6669

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+

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

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

6985 {-# INLINE fmap #-}

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

370386 acc2 <- step2 r2 res

371387 extract2 acc2

372388

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+

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

374415-- group using the provided fold function.

375416--

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

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

6363 , take

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

6565 , groupsOf

66+ , groupsOf2

6667 )

6768where

6869

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

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

7778

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

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

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

8081

8182import qualified Streamly.Streams.StreamK as K

8283

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

551552 return $ Yield r next

552553

553554 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

4343 -- ** Write to Handle

4444 -- Byte stream write (Folds)

4545 , write

46+ , write2

4647 -- , writeUtf8

4748 -- , writeUtf8ByLines

4849 -- , writeByFrames

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

117118

118119import Streamly (MonadAsync)

119120import Streamly.Data.Fold (Fold)

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

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

121123import Streamly.Internal.Memory.Array.Types

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

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

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

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

412414

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+

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

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

415421-- 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)

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

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

441447

448+{-# INLINE writeWithBufferOf2 #-}

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

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

451+

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

443453--

444454-- | 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)

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

451461write = writeWithBufferOf defaultChunkSize

452462

463+{-# INLINE write2 #-}

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

465+write2 = writeWithBufferOf2 defaultChunkSize

466+

453467{-

454468{-# INLINE write #-}

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

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

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

287287

288288 -- -- *** Chunks

289289 , chunksOf

290+ , chunksOf2

290291 , arraysOf

291292 , intervalsOf

292293

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

429430import qualified System.IO as IO

430431

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

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

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

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

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

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

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

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

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

25722573

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+

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

25742581-- @n@ elements each.

25752582--

src/Streamly/Streams/StreamD.hs+1

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

159159

160160 -- ** Grouping

161161 , groupsOf

162+ , groupsOf2

162163 , groupsBy

163164 , groupsRollingBy

164165