5283713 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Re-arrange functions and update doc in Foreign.Mut.Type

5283713Adithya Kumar5 years ago

src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+92−84

compare with master

unifiedsplit

@@ -236,19 +236,6 @@ mkChunkSizeKB n = mkChunkSize (n * k)

236defaultChunkSize :: Int

237defaultChunkSize = mkChunkSizeKB 32

238

239−-- | Remove the free space from an Array.

240−shrinkToFit :: forall a. Storable a => Array a -> IO (Array a)

241−shrinkToFit arr@Array{..} = do

242− assert (aEnd <= aBound) (return ())

243− let start = unsafeForeignPtrToPtr aStart

244− let used = aEnd `minusPtr` start

245− waste = aBound `minusPtr` aEnd

246− -- if used == waste == 0 then do not realloc

247− -- if the wastage is more than 25% of the array then realloc

248− if used < 3 * waste

249− then realloc used arr

250− else return arr

251−

252-------------------------------------------------------------------------------

239-------------------------------------------------------------------------------

253-- Construction

240-- Construction

254-------------------------------------------------------------------------------

241-------------------------------------------------------------------------------

@@ -300,7 +287,7 @@ unsafeWithNewArray count f = do

300 unsafeWithForeignPtr (aStart arr) $ \p -> f p >> return arr

287 unsafeWithForeignPtr (aStart arr) $ \p -> f p >> return arr

301

288

302-------------------------------------------------------------------------------

289-------------------------------------------------------------------------------

303−-- snoc

290+-- Mutation

304-------------------------------------------------------------------------------

291-------------------------------------------------------------------------------

305

292

306{-# INLINE unsafeWriteIndex #-}

293{-# INLINE unsafeWriteIndex #-}

@@ -349,7 +336,7 @@ snoc arr@Array {..} x =

349 return $ arr {aEnd = aEnd `plusPtr` sizeOf (undefined :: a)}

336 return $ arr {aEnd = aEnd `plusPtr` sizeOf (undefined :: a)}

350

337

351-------------------------------------------------------------------------------

338-------------------------------------------------------------------------------

352−-- re-allocate

339+-- Resizing

353-------------------------------------------------------------------------------

340-------------------------------------------------------------------------------

354

341

355-- | Reallocate the array to the specified size in bytes. If the size is less

342-- | Reallocate the array to the specified size in bytes. If the size is less

@@ -375,8 +362,21 @@ reallocAligned alignSize newSize Array{..} = do

375realloc :: forall a. Storable a => Int -> Array a -> IO (Array a)

362realloc :: forall a. Storable a => Int -> Array a -> IO (Array a)

376realloc = reallocAligned (alignment (undefined :: a))

363realloc = reallocAligned (alignment (undefined :: a))

377

364

365+-- | Remove the free space from an Array.

366+shrinkToFit :: forall a. Storable a => Array a -> IO (Array a)

367+shrinkToFit arr@Array{..} = do

368+ assert (aEnd <= aBound) (return ())

369+ let start = unsafeForeignPtrToPtr aStart

370+ let used = aEnd `minusPtr` start

371+ waste = aBound `minusPtr` aEnd

372+ -- if used == waste == 0 then do not realloc

373+ -- if the wastage is more than 25% of the array then realloc

374+ if used < 3 * waste

375+ then realloc used arr

376+ else return arr

377+

378-------------------------------------------------------------------------------

379−-- Elimination

379+--

380-------------------------------------------------------------------------------

381

382-- | Return element at the specified index without checking the bounds.

@@ -392,6 +392,10 @@ unsafeIndexIO Array {..} i =

392 (return ())

393 peek elemOff

394

395+-------------------------------------------------------------------------------

396+-- Size

397+-------------------------------------------------------------------------------

398+

395-- | /O(1)/ Get the byte length of the array.

399-- | /O(1)/ Get the byte length of the array.

396--

400--

397-- @since 0.7.0

401-- @since 0.7.0

@@ -484,75 +488,9 @@ bufferChunks :: (MonadIO m, Storable a) =>

484bufferChunks m = D.foldr K.cons K.nil $ arraysOf defaultChunkSize m

488bufferChunks m = D.foldr K.cons K.nil $ arraysOf defaultChunkSize m

485

489

486-------------------------------------------------------------------------------

490-------------------------------------------------------------------------------

487−-- Streams of arrays - flatten

491+-- Streams of arrays - Flattening

488-------------------------------------------------------------------------------

492-------------------------------------------------------------------------------

489

493

490−data ReadUState a = ReadUState

491− {-# UNPACK #-} !(ForeignPtr a) -- foreign ptr with end of array pointer

492− {-# UNPACK #-} !(Ptr a) -- current pointer

493−

494−-- | Resumable unfold of an array.

495−--

496−{-# INLINE_NORMAL producer #-}

497−producer :: forall m a. (Monad m, Storable a) => Producer m (Array a) a

498−producer = Producer step inject extract

499− where

500−

501− inject (Array (ForeignPtr start contents) (Ptr end) _) =

502− return $ ReadUState (ForeignPtr end contents) (Ptr start)

503−

504− {-# INLINE_LATE step #-}

505− step (ReadUState fp@(ForeignPtr end _) p) | p == Ptr end =

506− let x = unsafeInlineIO $ touchForeignPtr fp

507− in x `seq` return D.Stop

508− step (ReadUState fp p) = do

509− -- unsafeInlineIO allows us to run this in Identity monad for pure

510− -- toList/foldr case which makes them much faster due to not

511− -- accumulating the list and fusing better with the pure consumers.

512− --

513− -- This should be safe as the array contents are guaranteed to be

514− -- evaluated/written to before we peek at them.

515− let !x = unsafeInlineIO $ peek p

516− return $ D.Yield x

517− (ReadUState fp (p `plusPtr` sizeOf (undefined :: a)))

518−

519− extract (ReadUState (ForeignPtr end contents) (Ptr p)) =

520− return $ Array (ForeignPtr p contents) (Ptr end) (Ptr end)

521−

522−-- | Unfold an array into a stream.

523−--

524−-- @since 0.7.0

525−{-# INLINE_NORMAL read #-}

526−read :: forall m a. (Monad m, Storable a) => Unfold m (Array a) a

527−read = Producer.simplify producer

528−

529−-- | Unfold an array into a stream in reverse order.

530−--

531−-- /Pre-release/

532−{-# INLINE_NORMAL readRev #-}

533−readRev :: forall m a. (Monad m, Storable a) => Unfold m (Array a) a

534−readRev = Unfold step inject

535− where

536−

537− inject (Array fp end _) =

538− let p = end `plusPtr` negate (sizeOf (undefined :: a))

539− in return $ ReadUState fp p

540−

541− {-# INLINE_LATE step #-}

542− step (ReadUState fp@(ForeignPtr start _) p) | p < Ptr start =

543− let x = unsafeInlineIO $ touchForeignPtr fp

544− in x `seq` return D.Stop

545− step (ReadUState fp p) = do

546− -- unsafeInlineIO allows us to run this in Identity monad for pure

547− -- toList/foldr case which makes them much faster due to not

548− -- accumulating the list and fusing better with the pure consumers.

549− --

550− -- This should be safe as the array contents are guaranteed to be

551− -- evaluated/written to before we peek at them.

552− let !x = unsafeInlineIO $ peek p

553− return $ D.Yield x

554− (ReadUState fp (p `plusPtr` negate (sizeOf (undefined :: a))))

555−

556data FlattenState s a =

494data FlattenState s a =

557 OuterLoop s

495 OuterLoop s

558 | InnerLoop s !(ForeignPtr a) !(Ptr a) !(Ptr a)

496 | InnerLoop s !(ForeignPtr a) !(Ptr a) !(Ptr a)

@@ -626,6 +564,76 @@ flattenArraysRev (D.Stream step state) = D.Stream step' (OuterLoop state)

626 return $ D.Yield x (InnerLoop st startf

564 return $ D.Yield x (InnerLoop st startf

627 (p `plusPtr` negate (sizeOf (undefined :: a))) end)

565 (p `plusPtr` negate (sizeOf (undefined :: a))) end)

628

566

567+-------------------------------------------------------------------------------

568+-- Unfolds

569+-------------------------------------------------------------------------------

570+

571+data ReadUState a = ReadUState

572+ {-# UNPACK #-} !(ForeignPtr a) -- foreign ptr with end of array pointer

573+ {-# UNPACK #-} !(Ptr a) -- current pointer

574+

575+-- | Resumable unfold of an array.

576+--

577+{-# INLINE_NORMAL producer #-}

578+producer :: forall m a. (Monad m, Storable a) => Producer m (Array a) a

579+producer = Producer step inject extract

580+ where

581+

582+ inject (Array (ForeignPtr start contents) (Ptr end) _) =

583+ return $ ReadUState (ForeignPtr end contents) (Ptr start)

584+

585+ {-# INLINE_LATE step #-}

586+ step (ReadUState fp@(ForeignPtr end _) p) | p == Ptr end =

587+ let x = unsafeInlineIO $ touchForeignPtr fp

588+ in x `seq` return D.Stop

589+ step (ReadUState fp p) = do

590+ -- unsafeInlineIO allows us to run this in Identity monad for pure

591+ -- toList/foldr case which makes them much faster due to not

592+ -- accumulating the list and fusing better with the pure consumers.

593+ --

594+ -- This should be safe as the array contents are guaranteed to be

595+ -- evaluated/written to before we peek at them.

596+ let !x = unsafeInlineIO $ peek p

597+ return $ D.Yield x

598+ (ReadUState fp (p `plusPtr` sizeOf (undefined :: a)))

599+

600+ extract (ReadUState (ForeignPtr end contents) (Ptr p)) =

601+ return $ Array (ForeignPtr p contents) (Ptr end) (Ptr end)

602+

603+-- | Unfold an array into a stream.

604+--

605+-- @since 0.7.0

606+{-# INLINE_NORMAL read #-}

607+read :: forall m a. (Monad m, Storable a) => Unfold m (Array a) a

608+read = Producer.simplify producer

609+

610+-- | Unfold an array into a stream in reverse order.

611+--

612+-- /Pre-release/

613+{-# INLINE_NORMAL readRev #-}

614+readRev :: forall m a. (Monad m, Storable a) => Unfold m (Array a) a

615+readRev = Unfold step inject

616+ where

617+

618+ inject (Array fp end _) =

619+ let p = end `plusPtr` negate (sizeOf (undefined :: a))

620+ in return $ ReadUState fp p

621+

622+ {-# INLINE_LATE step #-}

623+ step (ReadUState fp@(ForeignPtr start _) p) | p < Ptr start =

624+ let x = unsafeInlineIO $ touchForeignPtr fp

625+ in x `seq` return D.Stop

626+ step (ReadUState fp p) = do

627+ -- unsafeInlineIO allows us to run this in Identity monad for pure

628+ -- toList/foldr case which makes them much faster due to not

629+ -- accumulating the list and fusing better with the pure consumers.

630+ --

631+ -- This should be safe as the array contents are guaranteed to be

632+ -- evaluated/written to before we peek at them.

633+ let !x = unsafeInlineIO $ peek p

634+ return $ D.Yield x

635+ (ReadUState fp (p `plusPtr` negate (sizeOf (undefined :: a))))

636+

629-------------------------------------------------------------------------------

637-------------------------------------------------------------------------------

630-- to Lists and streams

638-- to Lists and streams

631-------------------------------------------------------------------------------

639-------------------------------------------------------------------------------

@@ -757,7 +765,7 @@ foldr :: Storable a => (a -> b -> b) -> b -> Array a -> b

757foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr

765foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr

758

766

759-------------------------------------------------------------------------------

767-------------------------------------------------------------------------------

760−-- Write Folds to fold streams into arrays

768+-- Folds

761-------------------------------------------------------------------------------

769-------------------------------------------------------------------------------

762

770

763{-# INLINE_NORMAL writeNAllocWith #-}

771{-# INLINE_NORMAL writeNAllocWith #-}