00e9793 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add support for aligned array creation.

00e9793Harendra Kumar7 years ago

src/Streamly/Memory/Array.hs+4−2

compare with master

unifiedsplit

@@ -408,9 +408,11 @@ runPipe f arr = P.runPipe (toArrayMinChunk (length arr)) $ f (A.read arr)

408--

409-- @since 0.7.0

410{-# INLINE _runTransform #-}

411−_runTransform :: (MonadIO m, Storable a, Storable b)

411+_runTransform :: forall m a b. (MonadIO m, Storable a, Storable b)

412 => (SerialT m a -> SerialT m b) -> Array a -> m (Array b)

413+_runTransform f arr =

413−_runTransform f arr = P.runFold (A.toArrayMinChunk (length arr)) $ f (A.read arr)

414+ P.runFold (A.toArrayMinChunk (alignment (undefined :: a)) (length arr))

415+ $ f (A.read arr)

414

416

415-- | Fold an array using a 'Fold'.

417-- | Fold an array using a 'Fold'.

416--

418--

src/Streamly/Memory/Array/Types.hs+29−9

@@ -67,6 +67,7 @@ module Streamly.Memory.Array.Types

67 , writeN

68 , writeNUnsafe

69 , write

70+ , writeAligned

70 , read

71 , read

71 , readU

72 , readU

@@ -300,14 +301,13 @@ snoc arr@Array {..} x = do

300

301

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

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

302-- than the original array the array gets truncated.

303-- than the original array the array gets truncated.

303−{-# NOINLINE realloc #-}

304+{-# NOINLINE reallocAligned #-}

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

305+reallocAligned :: Int -> Int -> Array a -> IO (Array a)

305−realloc newSize Array{..} = do

306+reallocAligned alignSize newSize Array{..} = do

306 assert (aEnd <= aBound) (return ())

307 assert (aEnd <= aBound) (return ())

307 let oldStart = unsafeForeignPtrToPtr aStart

308 let oldStart = unsafeForeignPtrToPtr aStart

308 let size = aEnd `minusPtr` oldStart

309 let size = aEnd `minusPtr` oldStart

309− newPtr <- Malloc.mallocForeignPtrAlignedBytes

310+ newPtr <- Malloc.mallocForeignPtrAlignedBytes newSize alignSize

310− newSize (alignment (undefined :: a))

311 withForeignPtr newPtr $ \pNew -> do

312 memcpy (castPtr pNew) (castPtr oldStart) size

313 touchForeignPtr aStart

@@ -317,6 +317,10 @@ realloc newSize Array{..} = do

317 , aBound = pNew `plusPtr` newSize

318 }

319

320+{-# INLINABLE realloc #-}

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

322+realloc = reallocAligned (alignment (undefined :: a))

323+

320-- | Remove the free space from an Array.

324-- | Remove the free space from an Array.

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

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

322shrinkToFit arr@Array{..} = do

326shrinkToFit arr@Array{..} = do

@@ -574,9 +578,9 @@ writeNUnsafe n = Fold step initial extract

574

578

575{-# INLINE_NORMAL toArrayMinChunk #-}

579{-# INLINE_NORMAL toArrayMinChunk #-}

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

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

577− => Int -> Fold m a (Array a)

581+ => Int -> Int -> Fold m a (Array a)

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

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

579−toArrayMinChunk elemCount = Fold step initial extract

583+toArrayMinChunk alignSize elemCount = Fold step initial extract

580

584

581 where

585 where

582

586

@@ -591,7 +595,7 @@ toArrayMinChunk elemCount = Fold step initial extract

591 let p = unsafeForeignPtrToPtr start

595 let p = unsafeForeignPtrToPtr start

592 oldSize = end `minusPtr` p

596 oldSize = end `minusPtr` p

593 newSize = max (oldSize * 2) 1

597 newSize = max (oldSize * 2) 1

594− arr1 <- liftIO $ realloc newSize arr

598+ arr1 <- liftIO $ reallocAligned alignSize newSize arr

595 insertElem arr1 x

599 insertElem arr1 x

596 step arr x = insertElem arr x

600 step arr x = insertElem arr x

597 extract = liftIO . shrinkToFit

601 extract = liftIO . shrinkToFit

@@ -603,9 +607,25 @@ toArrayMinChunk elemCount = Fold step initial extract

603-- @since 0.7.0

607-- @since 0.7.0

604{-# INLINE write #-}

608{-# INLINE write #-}

605write :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)

609write :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)

610+write = toArrayMinChunk (alignment (undefined :: a))

606−write = toArrayMinChunk (bytesToElemCount (undefined :: a)

611+ (bytesToElemCount (undefined :: a)

607 (mkChunkSize 1024))

612 (mkChunkSize 1024))

608

613

614+-- | Like 'write' but the array memory is aligned according to the specified

615+-- alignment size. This could be useful when we have specific alignment, for

616+-- example, cache aligned arrays for lookup table etc.

617+--

618+-- /Caution! Do not use this on infinite streams./

619+--

620+-- @since 0.7.0

621+{-# INLINE writeAligned #-}

622+writeAligned :: forall m a. (MonadIO m, Storable a)

623+ => Int -> Fold m a (Array a)

624+writeAligned alignSize =

625+ toArrayMinChunk alignSize

626+ (bytesToElemCount (undefined :: a)

627+ (mkChunkSize 1024))

628+

609{-# INLINE_NORMAL fromStreamDN #-}

629{-# INLINE_NORMAL fromStreamDN #-}

610fromStreamDN :: forall m a. (MonadIO m, Storable a)

630fromStreamDN :: forall m a. (MonadIO m, Storable a)

611 => Int -> D.Stream m a -> m (Array a)

631 => Int -> D.Stream m a -> m (Array a)