dc7a543 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add chunksOf to generic arrays

dc7a543Adithya Kumar3 years ago

core/src/Streamly/Internal/Data/Array/Generic.hs+11

compare with master

unifiedsplit

@@ -24,6 +24,8 @@ module Streamly.Internal.Data.Array.Generic

24 , fromListN

25 , fromList

27+ , chunksOf

28+

27 -- * Elimination

29 -- * Elimination

28 , length

30 , length

29 , reader

31 , reader

@@ -127,6 +129,15 @@ fromPureStream x =

127fromByteStr# :: Addr# -> Array Word8

129fromByteStr# :: Addr# -> Array Word8

128fromByteStr# addr = fromPureStream (D.fromByteStr# addr)

130fromByteStr# addr = fromPureStream (D.fromByteStr# addr)

129

131

132+-------------------------------------------------------------------------------

133+-- Stream Ops

134+-------------------------------------------------------------------------------

135+

136+{-# INLINE_NORMAL chunksOf #-}

137+chunksOf :: forall m a. MonadIO m

138+ => Int -> Stream m a -> Stream m (Array a)

139+chunksOf n strm = fmap unsafeFreeze $ MArray.chunksOf n strm

140+

130-------------------------------------------------------------------------------

141-------------------------------------------------------------------------------

131-- Construction - from streams

142-- Construction - from streams

132-------------------------------------------------------------------------------

143-------------------------------------------------------------------------------

core/src/Streamly/Internal/Data/MutArray/Generic.hs+87−8

@@ -132,7 +132,7 @@ module Streamly.Internal.Data.MutArray.Generic

132 -- multidimensional array representations.

133

134 -- ** Construct from streams

135− -- , chunksOf

135+ , chunksOf

136 -- , arrayStreamKFromStreamD

137 -- , writeChunks

138

@@ -182,6 +182,7 @@ import Streamly.Internal.Data.Fold.Type (Fold(..))

182import Streamly.Internal.Data.Producer.Type (Producer (..))

183import Streamly.Internal.Data.Unfold.Type (Unfold(..))

184import Streamly.Internal.Data.Stream.StreamD.Type (Stream)

185+import Streamly.Internal.Data.SVar.Type (adaptState)

185

186

186import qualified Streamly.Internal.Data.Fold.Type as FL

187import qualified Streamly.Internal.Data.Fold.Type as FL

187import qualified Streamly.Internal.Data.Producer as Producer

188import qualified Streamly.Internal.Data.Producer as Producer

@@ -263,6 +264,21 @@ nil = new 0

263-- Random writes

264-- Random writes

264-------------------------------------------------------------------------------

265-------------------------------------------------------------------------------

265

266

267+-- | Write the given element to the given index of the 'MutableArray#'. Does not

268+-- check if the index is out of bounds of the array.

269+--

270+-- /Pre-release/

271+{-# INLINE putIndexUnsafeWith #-}

272+putIndexUnsafeWith :: MonadIO m => Int -> MutableArray# RealWorld a -> a -> m ()

273+putIndexUnsafeWith n _arrContents# x =

274+ liftIO

275+ $ IO

276+ $ \s# ->

277+ case n of

278+ I# n# ->

279+ let s1# = writeArray# _arrContents# n# x s#

280+ in (# s1#, () #)

281+

266-- | Write the given element to the given index of the array. Does not check if

282-- | Write the given element to the given index of the array. Does not check if

267-- the index is out of bounds of the array.

283-- the index is out of bounds of the array.

268--

284--

@@ -271,13 +287,7 @@ nil = new 0

271putIndexUnsafe :: forall m a. MonadIO m => Int -> MutArray a -> a -> m ()

287putIndexUnsafe :: forall m a. MonadIO m => Int -> MutArray a -> a -> m ()

272putIndexUnsafe i MutArray {..} x =

288putIndexUnsafe i MutArray {..} x =

273 assert (i >= 0 && i < arrLen)

289 assert (i >= 0 && i < arrLen)

274− (liftIO

290+ putIndexUnsafeWith (i + arrStart) arrContents# x

275− $ IO

276− $ \s# ->

277− case i + arrStart of

278− I# n# ->

279− let s1# = writeArray# arrContents# n# x s#

280− in (# s1#, () #))

281

291

282invalidIndex :: String -> Int -> a

292invalidIndex :: String -> Int -> a

283invalidIndex label i =

293invalidIndex label i =

@@ -666,6 +676,75 @@ fromPureStream :: MonadIO m => Stream Identity a -> m (MutArray a)

666fromPureStream xs =

676fromPureStream xs =

667 liftIO $ D.fold write $ D.morphInner (return . runIdentity) xs

677 liftIO $ D.fold write $ D.morphInner (return . runIdentity) xs

668

678

679+-------------------------------------------------------------------------------

680+-- Chunking

681+-------------------------------------------------------------------------------

682+

683+data GroupState s a start end bound

684+ = GroupStart s

685+ | GroupBuffer s (MutableArray# RealWorld a) start end bound

686+ | GroupYield

687+ (MutableArray# RealWorld a)

688+ start

689+ end

690+ bound

691+ (GroupState s a start end bound)

692+ | GroupFinish

693+

694+-- | @chunksOf n stream@ groups the input stream into a stream of

695+-- arrays of size n.

696+--

697+-- @chunksOf n = foldMany (MutArray.writeN n)@

698+--

699+-- /Pre-release/

700+{-# INLINE_NORMAL chunksOf #-}

701+chunksOf :: forall m a. MonadIO m

702+ => Int -> D.Stream m a -> D.Stream m (MutArray a)

703+-- XXX the idiomatic implementation leads to large regression in the D.reverse'

704+-- benchmark. It seems it has difficulty producing optimized code when

705+-- converting to StreamK. Investigate GHC optimizations.

706+-- chunksOf n = D.foldMany (writeN n)

707+chunksOf n (D.Stream step state) =

708+ D.Stream step' (GroupStart state)

709+

710+ where

711+

712+ -- start is always 0

713+ -- end and len are always equal

714+

715+ {-# INLINE_LATE step' #-}

716+ step' _ (GroupStart st) = do

717+ when (n <= 0) $

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

719+ error $ "Streamly.Internal.Data.Array.Generic.Mut.Type.chunksOf: "

720+ ++ "the size of arrays [" ++ show n

721+ ++ "] must be a natural number"

722+ (MutArray contents start end bound :: MutArray a) <- new n

723+ return $ D.Skip (GroupBuffer st contents start end bound)

724+

725+ step' gst (GroupBuffer st contents start end bound) = do

726+ r <- step (adaptState gst) st

727+ case r of

728+ D.Yield x s -> do

729+ putIndexUnsafeWith end contents x

730+ let end1 = end + 1

731+ return $

732+ if end1 >= bound

733+ then D.Skip

734+ (GroupYield

735+ contents start end1 bound (GroupStart s))

736+ else D.Skip (GroupBuffer s contents start end1 bound)

737+ D.Skip s ->

738+ return $ D.Skip (GroupBuffer s contents start end bound)

739+ D.Stop ->

740+ return

741+ $ D.Skip (GroupYield contents start end bound GroupFinish)

742+

743+ step' _ (GroupYield contents start end bound next) =

744+ return $ D.Yield (MutArray contents start end bound) next

745+

746+ step' _ GroupFinish = return D.Stop

747+

669-------------------------------------------------------------------------------

748-------------------------------------------------------------------------------

670-- Unfolds

749-- Unfolds

671-------------------------------------------------------------------------------

750-------------------------------------------------------------------------------