1192c13 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add array transform/fold operations

1192c13Harendra Kumar7 years ago

src/Streamly/Fold.hs+4

compare with master

unifiedsplit

@@ -444,6 +444,10 @@ import qualified Streamly.Streams.Prelude as P

444-- Running a Fold

445------------------------------------------------------------------------------

446

447+-- XXX rename to foldStream to be consistent with foldArray for arrays? Should

448+-- we move it in the stream module (Prelude as of now) or should we move

449+-- foldArray as well here?

450+

447-- | Fold a stream using the supplied monadic fold.

451-- | Fold a stream using the supplied monadic fold.

448--

452--

449-- >>> FL.foldl' FL.sum (S.enumerateFromTo 1 100)

453-- >>> FL.foldl' FL.sum (S.enumerateFromTo 1 100)

src/Streamly/Mem/Array.hs+57−7

@@ -128,9 +128,16 @@ module Streamly.Mem.Array

128 , writeSlice

129 , writeSliceRev

130 -}

131+ -- * Immutable Transformations

132+ , transformWith

133+

134+ -- * Folds

135+ , foldWith

136+ , foldArray

131 )

137 )

132where

138where

133

139

140+import Control.Monad (when)

134import Control.Monad.IO.Class (MonadIO(..))

141import Control.Monad.IO.Class (MonadIO(..))

135import Data.Functor.Identity (Identity)

142import Data.Functor.Identity (Identity)

136import Foreign.ForeignPtr (withForeignPtr)

143import Foreign.ForeignPtr (withForeignPtr)

@@ -169,7 +176,9 @@ newArray len = undefined

169-- @since 0.7.0

176-- @since 0.7.0

170{-# INLINE writeN #-}

177{-# INLINE writeN #-}

171writeN :: (MonadIO m, Storable a) => Int -> SerialT m a -> m (Array a)

178writeN :: (MonadIO m, Storable a) => Int -> SerialT m a -> m (Array a)

179+writeN n m = do

180+ if n < 0 then error "writeN: negative write count specified" else return ()

172−writeN n m = fromStreamDN n $ D.toStreamD m

181+ fromStreamDN n $ D.toStreamD m

173

182

174-------------------------------------------------------------------------------

183-------------------------------------------------------------------------------

175-- Elimination

184-- Elimination

@@ -386,11 +395,17 @@ _toArraysOf n = FL.lchunksOf n (toArrayN n) FL.toStream

386-- XXX check if GHC's memory allocator is efficient enough. We can try the C

395-- XXX check if GHC's memory allocator is efficient enough. We can try the C

387-- malloc to compare against.

396-- malloc to compare against.

388

397

398+{-# INLINE bytesToCount #-}

399+bytesToCount :: Storable a => a -> Int -> Int

400+bytesToCount x n =

401+ let elemSize = sizeOf x

402+ in n + elemSize - 1 `div` elemSize

403+

389{-# INLINE toArrayMinChunk #-}

404{-# INLINE toArrayMinChunk #-}

390toArrayMinChunk :: forall m a. (MonadIO m, Storable a)

405toArrayMinChunk :: forall m a. (MonadIO m, Storable a)

391 => Int -> Fold m a (Array a)

406 => Int -> Fold m a (Array a)

392-- toArrayMinChunk n = FL.mapM spliceArrays $ toArraysOf n

407-- toArrayMinChunk n = FL.mapM spliceArrays $ toArraysOf n

393−toArrayMinChunk chunkSize = Fold step initial extract

408+toArrayMinChunk elemCount = Fold step initial extract

394

409

395 where

410 where

396

411

@@ -398,7 +413,9 @@ toArrayMinChunk chunkSize = Fold step initial extract

398 liftIO $ poke end x

413 liftIO $ poke end x

399 return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound

414 return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound

400

415

416+ initial = do

417+ when (elemCount < 0) $ error "toArrayMinChunk: elemCount is negative"

401− initial = liftIO $ A.newArray (mkChunkSize chunkSize)

418+ liftIO $ A.newArray elemCount

402 step arr@(Array _ end bound) x | end == bound = do

419 step arr@(Array _ end bound) x | end == bound = do

403 arr1 <- liftIO $ reallocDouble 1 arr

420 arr1 <- liftIO $ reallocDouble 1 arr

404 insertElem arr1 x

421 insertElem arr1 x

@@ -412,7 +429,7 @@ toArrayMinChunk chunkSize = Fold step initial extract

412-- @since 0.7.0

429-- @since 0.7.0

413{-# INLINE toArray #-}

430{-# INLINE toArray #-}

414toArray :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)

431toArray :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)

415−toArray = toArrayMinChunk 1024

432+toArray = toArrayMinChunk (bytesToCount (undefined :: a) (mkChunkSize 1024))

416

433

417-- | Convert a stream of arrays into a stream of their elements.

434-- | Convert a stream of arrays into a stream of their elements.

418--

435--

@@ -467,10 +484,11 @@ _spliceArraysRealloced s = do

467 return $ dst `plusPtr` len

484 return $ dst `plusPtr` len

468

485

469{-# INLINE spliceArraysBuffered #-}

486{-# INLINE spliceArraysBuffered #-}

470−spliceArraysBuffered :: (MonadIO m, Storable a)

487+spliceArraysBuffered :: forall m a. (MonadIO m, Storable a)

471 => SerialT m (Array a) -> m (Array a)

488 => SerialT m (Array a) -> m (Array a)

472spliceArraysBuffered s = do

489spliceArraysBuffered s = do

490+ idst <- liftIO $ A.newArray (bytesToCount (undefined :: a)

473− idst <- liftIO $ A.newArray (mkChunkSizeKB 4)

491+ (mkChunkSizeKB 4))

474 arr <- S.foldlM' appendArr idst s

492 arr <- S.foldlM' appendArr idst s

475 liftIO $ shrinkToFit arr

493 liftIO $ shrinkToFit arr

476

494

@@ -508,8 +526,40 @@ spliceArrays = spliceArraysBuffered

508--

526--

509{-# INLINE write #-}

527{-# INLINE write #-}

510write :: (MonadIO m, Storable a) => SerialT m a -> m (Array a)

528write :: (MonadIO m, Storable a) => SerialT m a -> m (Array a)

529+write = FL.foldl' toArray

511−write m = A.fromStreamD $ D.toStreamD m

530+-- write m = A.fromStreamD $ D.toStreamD m

512

531

513-- XXX efficiently compare two streams of arrays. Two streams can have chunks

532-- XXX efficiently compare two streams of arrays. Two streams can have chunks

514-- of different sizes, we can handle that in the stream comparison abstraction.

533-- of different sizes, we can handle that in the stream comparison abstraction.

515-- This could be useful e.g. to fast compare whether two files differ.

534-- This could be useful e.g. to fast compare whether two files differ.

535+

536+-------------------------------------------------------------------------------

537+-- Transform via stream operations

538+-------------------------------------------------------------------------------

539+

540+-- for non-length changing operations we can use the original length for

541+-- allocation. If we can predict the length then we can use the prediction for

542+-- new allocation. Otherwise we can use a hint and adjust dynamically.

543+

544+-- | Transform an array into another array using a stream transformation

545+-- operation.

546+--

547+-- @since 0.7.0

548+{-# INLINE transformWith #-}

549+transformWith :: (MonadIO m, Storable a, Storable b)

550+ => (SerialT m a -> SerialT m b) -> Array a -> m (Array b)

551+transformWith f arr = FL.foldl' (toArrayMinChunk (length arr)) $ f (read arr)

552+

553+-- | Fold an array using a 'Fold'.

554+--

555+-- @since 0.7.0

556+{-# INLINE foldArray #-}

557+foldArray :: (MonadIO m, Storable a) => Fold m a b -> Array a -> m b

558+foldArray f arr = FL.foldl' f (read arr)

559+

560+-- | Fold an array using a stream fold operation.

561+--

562+-- @since 0.7.0

563+{-# INLINE foldWith #-}

564+foldWith :: (MonadIO m, Storable a) => (SerialT m a -> m b) -> Array a -> m b

565+foldWith f arr = f (read arr)

src/Streamly/Mem/Array/Types.hs+5−4

@@ -34,7 +34,7 @@ module Streamly.Mem.Array.Types

34 , fromList

35 , fromListN

36 , fromStreamDN

37− , fromStreamD

37+ -- , fromStreamD

38 , fromStreamDArraysOf

39 , flattenArrays

40 , flattenArraysRev

@@ -384,7 +384,9 @@ toArrayN n = Fold step initial extract

384

385 where

386

387+ initial = do

388+ if n < 0 then error "toArrayN: negative count specified" else return ()

387− initial = liftIO $ newArray n

389+ liftIO $ newArray n

388 step arr@(Array _ end bound) _ | end == bound = return arr

390 step arr@(Array _ end bound) _ | end == bound = return arr

389 step (Array start end bound) x = do

391 step (Array start end bound) x = do

390 liftIO $ poke end x

392 liftIO $ poke end x

@@ -652,7 +654,7 @@ allocOverhead :: Int

652allocOverhead = 2 * sizeOf (undefined :: Int)

654allocOverhead = 2 * sizeOf (undefined :: Int)

653

655

654mkChunkSize :: Int -> Int

656mkChunkSize :: Int -> Int

655−mkChunkSize n = n - allocOverhead

657+mkChunkSize n = let size = n - allocOverhead in max size 0

656

658

657mkChunkSizeKB :: Int -> Int

659mkChunkSizeKB :: Int -> Int

658mkChunkSizeKB n = mkChunkSize (n * k)

660mkChunkSizeKB n = mkChunkSize (n * k)

@@ -663,4 +665,3 @@ mkChunkSizeKB n = mkChunkSize (n * k)

663-- bytes, so that the actual allocation is 32KB.

665-- bytes, so that the actual allocation is 32KB.

664defaultChunkSize :: Int

666defaultChunkSize :: Int

665defaultChunkSize = mkChunkSizeKB 32

667defaultChunkSize = mkChunkSizeKB 32

666−