@@ -67,6 +67,8 @@ module Streamly.Internal.Memory.Array.Types
67 , toArrayMinChunk
67 , toArrayMinChunk
69 , writeNUnsafe
69 , writeNUnsafe
71+ , writeNAlignedUnmanaged
71 , writeAligned
73 , writeAligned
@@ -234,16 +236,13 @@ bytesToElemCount x n =
234-- Construction
236-- Construction
235-------------------------------------------------------------------------------
237-------------------------------------------------------------------------------
237−-- | Allocate an array that can hold 'count' items. The memory of the array is
239+-- | allocate a new array using the provided allocator function.
238−-- uninitialized.
240+{-# INLINE newArrayAlignedAllocWith #-}
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)
241+newArrayAlignedAllocWith :: forall a. Storable a
244−newArray count = do
242+ => (Int -> Int -> IO (ForeignPtr a)) -> Int -> Int -> IO (Array a)
243+newArrayAlignedAllocWith alloc alignSize count = do
245 let size = count * sizeOf (undefined :: a)
244 let size = count * sizeOf (undefined :: a)
246− fptr <- Malloc.mallocForeignPtrAlignedBytes size (alignment (undefined :: a))
245+ fptr <- alloc size alignSize
247 let p = unsafeForeignPtrToPtr fptr
246 let p = unsafeForeignPtrToPtr fptr
248 return $ Array
247 return $ Array
249 { aStart = fptr
248 { aStart = fptr
@@ -251,6 +250,30 @@ newArray count = do
251 , aBound = p `plusPtr` size
250 , aBound = p `plusPtr` size
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
262+{-# INLINE newArrayAligned #-}
263+newArrayAligned :: forall a. Storable a => Int -> Int -> IO (Array a)
264+newArrayAligned = newArrayAlignedAllocWith Malloc.mallocForeignPtrAlignedBytes
266+-- XXX can unaligned allocation be more efficient when alignment is not needed?
268+-- | Allocate an array that can hold 'count' items. The memory of the array is
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))
254-- | Allocate an Array of the given size and run an IO action passing the array
277-- | Allocate an Array of the given size and run an IO action passing the array
255-- start pointer.
278-- start pointer.
256{-# INLINE withNewArray #-}
279{-# INLINE withNewArray #-}
@@ -315,6 +338,7 @@ reallocAligned alignSize newSize Array{..} = do
315 , aBound = pNew `plusPtr` newSize
338 , aBound = pNew `plusPtr` newSize
341+-- XXX can unaligned allocation be more efficient when alignment is not needed?
318{-# INLINABLE realloc #-}
342{-# INLINABLE realloc #-}
319realloc :: forall a. Storable a => Int -> Array a -> IO (Array a)
343realloc :: forall a. Storable a => Int -> Array a -> IO (Array a)
320realloc = reallocAligned (alignment (undefined :: a))
344realloc = reallocAligned (alignment (undefined :: a))
@@ -491,23 +515,53 @@ foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
491-- Instances
515-- Instances
492-------------------------------------------------------------------------------
516-------------------------------------------------------------------------------
494−-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
498−{-# INLINE_NORMAL writeN #-}
518+{-# INLINE_NORMAL writeNAllocWith #-}
499−writeN :: forall m a. (MonadIO m, Storable a) => Int -> Fold m a (Array a)
519+writeNAllocWith :: forall m a. (MonadIO m, Storable a)
520+ => (Int -> IO (Array a)) -> Int -> Fold m a (Array a)
500−writeN n = Fold step initial extract
521+writeNAllocWith alloc n = Fold step initial extract
504− initial = liftIO $ newArray (max n 0)
525+ initial = liftIO $ alloc (max n 0)
505 step arr@(Array _ end bound) _ | end == bound = return arr
526 step arr@(Array _ end bound) _ | end == bound = return arr
506 step (Array start end bound) x = do
527 step (Array start end bound) x = do
507 liftIO $ poke end x
528 liftIO $ poke end x
508 return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound
529 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.
509 extract = return -- liftIO . shrinkToFit
532 extract = return -- liftIO . shrinkToFit
534+-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
538+{-# INLINE_NORMAL writeN #-}
539+writeN :: forall m a. (MonadIO m, Storable a) => Int -> Fold m a (Array a)
540+writeN = writeNAllocWith newArray
542+-- | @writeNAligned n@ folds a maximum of @n@ elements from the input
543+-- stream to an 'Array' aligned to the given size.
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)
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.
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)
511data ArrayUnsafe a = ArrayUnsafe
565data ArrayUnsafe a = ArrayUnsafe
512 {-# UNPACK #-} !(ForeignPtr a) -- first address
566 {-# UNPACK #-} !(ForeignPtr a) -- first address
513 {-# UNPACK #-} !(Ptr a) -- first unused address
567 {-# UNPACK #-} !(Ptr a) -- first unused address
@@ -560,7 +614,7 @@ toArrayMinChunk alignSize elemCount = Fold step initial extract
561 initial = do
615 initial = do
562 when (elemCount < 0) $ error "toArrayMinChunk: elemCount is negative"
616 when (elemCount < 0) $ error "toArrayMinChunk: elemCount is negative"
563− liftIO $ newArray elemCount
617+ liftIO $ newArrayAligned alignSize elemCount
564 step arr@(Array start end bound) x | end == bound = do
618 step arr@(Array start end bound) x | end == bound = do
565 let p = unsafeForeignPtrToPtr start
619 let p = unsafeForeignPtrToPtr start
566 oldSize = end `minusPtr` p
620 oldSize = end `minusPtr` p