@@ -67,6 +67,7 @@ module Streamly.Memory.Array.Types
6767 , writeN
6868 , writeNUnsafe
6969 , write
70+ , writeAligned
7071 , read
7172 , readU
7273
@@ -300,14 +301,13 @@ snoc arr@Array {..} x = do
300301
301302-- | Reallocate the array to the specified size in bytes. If the size is less
302303-- than the original array the array gets truncated.
303−{-# NOINLINE realloc #-}
304−realloc :: forall a. Storable a => Int -> Array a -> IO (Array a)
305−realloc newSize Array{..} = do
304+{-# NOINLINE reallocAligned #-}
305+reallocAligned :: Int -> Int -> Array a -> IO (Array a)
306+reallocAligned alignSize newSize Array{..} = do
306307 assert (aEnd <= aBound) (return ())
307308 let oldStart = unsafeForeignPtrToPtr aStart
308309 let size = aEnd `minusPtr` oldStart
309− newPtr <- Malloc.mallocForeignPtrAlignedBytes
310− newSize (alignment (undefined :: a))
310+ newPtr <- Malloc.mallocForeignPtrAlignedBytes newSize alignSize
311311 withForeignPtr newPtr $ \pNew -> do
312312 memcpy (castPtr pNew) (castPtr oldStart) size
313313 touchForeignPtr aStart
@@ -317,6 +317,10 @@ realloc newSize Array{..} = do
317317 , aBound = pNew `plusPtr` newSize
318318 }
319319
320+{-# INLINABLE realloc #-}
321+realloc :: forall a. Storable a => Int -> Array a -> IO (Array a)
322+realloc = reallocAligned (alignment (undefined :: a))
323+
320324-- | Remove the free space from an Array.
321325shrinkToFit :: forall a. Storable a => Array a -> IO (Array a)
322326shrinkToFit arr@Array{..} = do
@@ -574,9 +578,9 @@ writeNUnsafe n = Fold step initial extract
574578
575579{-# INLINE_NORMAL toArrayMinChunk #-}
576580toArrayMinChunk :: forall m a. (MonadIO m, Storable a)
577− => Int -> Fold m a (Array a)
581+ => Int -> Int -> Fold m a (Array a)
578582-- toArrayMinChunk n = FL.mapM spliceArrays $ toArraysOf n
579−toArrayMinChunk elemCount = Fold step initial extract
583+toArrayMinChunk alignSize elemCount = Fold step initial extract
580584
581585 where
582586
@@ -591,7 +595,7 @@ toArrayMinChunk elemCount = Fold step initial extract
591595 let p = unsafeForeignPtrToPtr start
592596 oldSize = end `minusPtr` p
593597 newSize = max (oldSize * 2) 1
594− arr1 <- liftIO $ realloc newSize arr
598+ arr1 <- liftIO $ reallocAligned alignSize newSize arr
595599 insertElem arr1 x
596600 step arr x = insertElem arr x
597601 extract = liftIO . shrinkToFit
@@ -603,9 +607,25 @@ toArrayMinChunk elemCount = Fold step initial extract
603607-- @since 0.7.0
604608{-# INLINE write #-}
605609write :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)
606−write = toArrayMinChunk (bytesToElemCount (undefined :: a)
610+write = toArrayMinChunk (alignment (undefined :: a))
611+ (bytesToElemCount (undefined :: a)
607612 (mkChunkSize 1024))
608613
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+
609629{-# INLINE_NORMAL fromStreamDN #-}
610630fromStreamDN :: forall m a. (MonadIO m, Storable a)
611631 => Int -> D.Stream m a -> m (Array a)