@@ -67,6 +67,8 @@ module Streamly.Internal.Memory.Array.Types
6767 , toArrayMinChunk
6868 , writeN
6969 , writeNUnsafe
70+ , writeNAligned
71+ , writeNAlignedUnmanaged
7072 , write
7173 , writeAligned
7274
@@ -234,16 +236,13 @@ bytesToElemCount x n =
234236-- Construction
235237-------------------------------------------------------------------------------
236238
237−-- | Allocate an array that can hold 'count' items. The memory of the array is
238−-- uninitialized.
239−--
240−-- Note that this is internal routine, the reference to this array cannot be
241−-- given out until the array has been written to and frozen.
242−{-# INLINE newArray #-}
243−newArray :: forall a. Storable a => Int -> IO (Array a)
244−newArray count = do
239+-- | allocate a new array using the provided allocator function.
240+{-# INLINE newArrayAlignedAllocWith #-}
241+newArrayAlignedAllocWith :: forall a. Storable a
242+ => (Int -> Int -> IO (ForeignPtr a)) -> Int -> Int -> IO (Array a)
243+newArrayAlignedAllocWith alloc alignSize count = do
245244 let size = count * sizeOf (undefined :: a)
246− fptr <- Malloc.mallocForeignPtrAlignedBytes size (alignment (undefined :: a))
245+ fptr <- alloc size alignSize
247246 let p = unsafeForeignPtrToPtr fptr
248247 return $ Array
249248 { aStart = fptr
@@ -251,6 +250,30 @@ newArray count = do
251250 , aBound = p `plusPtr` size
252251 }
253252
253+-- | Allocate a new array aligned to the specified alignmend and using
254+-- unmanaged pinned memory. The memory will not be automatically freed by GHC.
255+-- This could be useful in allocate once global data structures. Use carefully
256+-- as incorrect use can lead to memory leak.
257+{-# INLINE newArrayAlignedUnmanaged #-}
258+newArrayAlignedUnmanaged :: forall a. Storable a => Int -> Int -> IO (Array a)
259+newArrayAlignedUnmanaged =
260+ newArrayAlignedAllocWith Malloc.mallocForeignPtrAlignedUnmanagedBytes
261+
262+{-# INLINE newArrayAligned #-}
263+newArrayAligned :: forall a. Storable a => Int -> Int -> IO (Array a)
264+newArrayAligned = newArrayAlignedAllocWith Malloc.mallocForeignPtrAlignedBytes
265+
266+-- XXX can unaligned allocation be more efficient when alignment is not needed?
267+--
268+-- | Allocate an array that can hold 'count' items. The memory of the array is
269+-- uninitialized.
270+--
271+-- Note that this is internal routine, the reference to this array cannot be
272+-- given out until the array has been written to and frozen.
273+{-# INLINE newArray #-}
274+newArray :: forall a. Storable a => Int -> IO (Array a)
275+newArray = newArrayAligned (alignment (undefined :: a))
276+
254277-- | Allocate an Array of the given size and run an IO action passing the array
255278-- start pointer.
256279{-# INLINE withNewArray #-}
@@ -315,6 +338,7 @@ reallocAligned alignSize newSize Array{..} = do
315338 , aBound = pNew `plusPtr` newSize
316339 }
317340
341+-- XXX can unaligned allocation be more efficient when alignment is not needed?
318342{-# INLINABLE realloc #-}
319343realloc :: forall a. Storable a => Int -> Array a -> IO (Array a)
320344realloc = reallocAligned (alignment (undefined :: a))
@@ -491,23 +515,53 @@ foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
491515-- Instances
492516-------------------------------------------------------------------------------
493517
494−-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
495−-- 'Array'.
496−--
497−-- @since 0.7.0
498−{-# INLINE_NORMAL writeN #-}
499−writeN :: forall m a. (MonadIO m, Storable a) => Int -> Fold m a (Array a)
500−writeN n = Fold step initial extract
518+{-# INLINE_NORMAL writeNAllocWith #-}
519+writeNAllocWith :: forall m a. (MonadIO m, Storable a)
520+ => (Int -> IO (Array a)) -> Int -> Fold m a (Array a)
521+writeNAllocWith alloc n = Fold step initial extract
501522
502523 where
503524
504− initial = liftIO $ newArray (max n 0)
525+ initial = liftIO $ alloc (max n 0)
505526 step arr@(Array _ end bound) _ | end == bound = return arr
506527 step (Array start end bound) x = do
507528 liftIO $ poke end x
508529 return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound
530+ -- XXX note that shirkToFit does not maintain alignment, in case we are
531+ -- using aligned allocation.
509532 extract = return -- liftIO . shrinkToFit
510533
534+-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
535+-- 'Array'.
536+--
537+-- @since 0.7.0
538+{-# INLINE_NORMAL writeN #-}
539+writeN :: forall m a. (MonadIO m, Storable a) => Int -> Fold m a (Array a)
540+writeN = writeNAllocWith newArray
541+
542+-- | @writeNAligned n@ folds a maximum of @n@ elements from the input
543+-- stream to an 'Array' aligned to the given size.
544+--
545+-- /Internal/
546+--
547+{-# INLINE_NORMAL writeNAligned #-}
548+writeNAligned :: forall m a. (MonadIO m, Storable a)
549+ => Int -> Int -> Fold m a (Array a)
550+writeNAligned alignSize = writeNAllocWith (newArrayAligned alignSize)
551+
552+-- | @writeNAlignedUnmanaged n@ folds a maximum of @n@ elements from the input
553+-- stream to an 'Array' aligned to the given size and using unmanaged memory.
554+-- This could be useful to allocate memory that we need to allocate only once
555+-- in the lifetime of the program.
556+--
557+-- /Internal/
558+--
559+{-# INLINE_NORMAL writeNAlignedUnmanaged #-}
560+writeNAlignedUnmanaged :: forall m a. (MonadIO m, Storable a)
561+ => Int -> Int -> Fold m a (Array a)
562+writeNAlignedUnmanaged alignSize =
563+ writeNAllocWith (newArrayAlignedUnmanaged alignSize)
564+
511565data ArrayUnsafe a = ArrayUnsafe
512566 {-# UNPACK #-} !(ForeignPtr a) -- first address
513567 {-# UNPACK #-} !(Ptr a) -- first unused address
@@ -560,7 +614,7 @@ toArrayMinChunk alignSize elemCount = Fold step initial extract
560614
561615 initial = do
562616 when (elemCount < 0) $ error "toArrayMinChunk: elemCount is negative"
563− liftIO $ newArray elemCount
617+ liftIO $ newArrayAligned alignSize elemCount
564618 step arr@(Array start end bound) x | end == bound = do
565619 let p = unsafeForeignPtrToPtr start
566620 oldSize = end `minusPtr` p