@@ -119,7 +119,8 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type
119119
120120 -- * Memory Management
121121 , blockSize
122− , arrayChunkSize
122+ , arrayChunkBytes
123+ , allocBytesToElemCount
123124 , realloc
124125 , resize
125126 , resizeExp
@@ -187,7 +188,6 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type
187188 -- , appendSliceFrom
188189
189190 -- * Utilities
190− , bytesToElemCount
191191 , memcpy
192192 , memcmp
193193 , c_memchr
@@ -227,7 +227,7 @@ import Streamly.Internal.Data.SVar.Type (adaptState)
227227import Streamly.Internal.Data.Unfold.Type (Unfold(..))
228228
229229import Streamly.Internal.System.IO
230− (defaultChunkSize, mkChunkSize, unsafeInlineIO)
230+ (defaultChunkSize, arrayPayloadSize, unsafeInlineIO)
231231
232232#ifdef DEVBUILD
233233import qualified Data.Foldable as F
@@ -272,13 +272,13 @@ foreign import ccall unsafe "string.h memcmp" c_memcmp
272272 :: Ptr Word8 -> Ptr Word8 -> CSize -> IO CInt
273273
274274-- | Given a 'Storable' type (unused first arg) and a number of bytes, return
275−-- how many elements of that type will fit in after rounding the bytes up to
276−-- the element size.
275+-- how many elements of that type will completely fit in those bytes.
276+--
277277{-# INLINE bytesToElemCount #-}
278278bytesToElemCount :: Storable a => a -> Int -> Int
279279bytesToElemCount x n =
280280 let elemSize = sizeOf x
281− in n + elemSize - 1 `div` elemSize
281+ in n `div` elemSize
282282
283283-- XXX we are converting Int to CSize
284284memcpy :: Ptr Word8 -> Ptr Word8 -> Int -> IO ()
@@ -588,13 +588,30 @@ roundUpToPower2 :: Int -> Int
588588roundUpToPower2 = undefined
589589-}
590590
591−-- XXX remove mkChunkSize, it is not required.
591+-- | @allocBytesToBytes elem allocatedBytes@ returns the array size in bytes
592+-- such that the real allocation is less than or equal to @allocatedBytes@,
593+-- unless @allocatedBytes@ is less than the size of one array element in which
594+-- case it returns one element's size.
595+--
596+{-# INLINE allocBytesToBytes #-}
597+allocBytesToBytes :: forall a. Storable a => a -> Int -> Int
598+allocBytesToBytes _ n =
599+ max (arrayPayloadSize n) (sizeOf (undefined :: a))
600+
601+-- | Given a 'Storable' type (unused first arg) and real allocation size
602+-- (including overhead), return how many elements of that type will completely
603+-- fit in it, returns at least 1.
604+--
605+{-# INLINE allocBytesToElemCount #-}
606+allocBytesToElemCount :: Storable a => a -> Int -> Int
607+allocBytesToElemCount x bytes =
608+ let n = bytesToElemCount x (allocBytesToBytes x bytes)
609+ in assert (n >= 1) n
610+
592611-- | The default chunk size by which the array creation routines increase the
593612-- size of the array when the array is grown linearly.
594−--
595−{-# INLINE arrayChunkSize #-}
596−arrayChunkSize :: Int
597−arrayChunkSize = mkChunkSize 1024
613+arrayChunkBytes :: Int
614+arrayChunkBytes = 1024
598615
599616-------------------------------------------------------------------------------
600617-- Snoc
@@ -708,7 +725,7 @@ snocWith allocSize arr x = liftIO $ do
708725
709726-- | The array is mutated to append an additional element to it. If there
710727-- is no reserved space available in the array then it is reallocated to grow
711−-- it by 'arrayChunkSize' rounded up to 'blockSize' when the size becomes more
728+-- it by 'arrayChunkBytes' rounded up to 'blockSize' when the size becomes more
712729-- than 'largeObjectThreshold'.
713730--
714731-- Note that the returned array may be a mutated version of the original array.
@@ -718,7 +735,7 @@ snocWith allocSize arr x = liftIO $ do
718735-- /Pre-release/
719736{-# INLINE snocLinear #-}
720737snocLinear :: forall m a. (MonadIO m, Storable a) => Array a -> a -> m (Array a)
721−snocLinear = snocWith (+ arrayChunkSize)
738+snocLinear = snocWith (+ allocBytesToBytes (undefined :: a) arrayChunkBytes)
722739
723740-- XXX round it to next power of 2.
724741--
@@ -1162,9 +1179,11 @@ arraysOf n (D.Stream step state) =
11621179-- XXX buffer to a list instead?
11631180-- | Buffer the stream into arrays in memory.
11641181{-# INLINE arrayStreamKFromStreamD #-}
1165−arrayStreamKFromStreamD :: (MonadIO m, Storable a) =>
1182+arrayStreamKFromStreamD :: forall m a. (MonadIO m, Storable a) =>
11661183 D.Stream m a -> m (K.Stream m (Array a))
1167−arrayStreamKFromStreamD m = D.foldr K.cons K.nil $ arraysOf defaultChunkSize m
1184+arrayStreamKFromStreamD =
1185+ let n = allocBytesToElemCount (undefined :: a) defaultChunkSize
1186+ in D.foldr K.cons K.nil . arraysOf n
11681187
11691188-------------------------------------------------------------------------------
11701189-- Streams of arrays - Flattening
@@ -1725,7 +1744,7 @@ writeWith elemCount =
17251744
17261745-- | Fold the whole input to a single array.
17271746--
1728−-- Same as 'writeWith' using an initial array size of 'arrayChunkSize' bytes
1747+-- Same as 'writeWith' using an initial array size of 'arrayChunkBytes' bytes
17291748-- rounded up to the element size.
17301749--
17311750-- /Caution! Do not use this on infinite streams./
@@ -1733,7 +1752,7 @@ writeWith elemCount =
17331752-- @since 0.7.0
17341753{-# INLINE write #-}
17351754write :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)
1736−write = writeWith (bytesToElemCount (undefined :: a) arrayChunkSize)
1755+write = writeWith (allocBytesToElemCount (undefined :: a) arrayChunkBytes)
17371756
17381757-------------------------------------------------------------------------------
17391758-- construct from streams, known size