on bebdf08 help sign in
1105 lines 39 kB history raw
1{-# LANGUAGE CPP #-}
2{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE ExistentialQuantification #-}
4{-# LANGUAGE MagicHash #-}
5{-# LANGUAGE RecordWildCards #-}
6{-# LANGUAGE ScopedTypeVariables #-}
7{-# LANGUAGE TypeFamilies #-}
8{-# LANGUAGE UnboxedTuples #-}
9{-# LANGUAGE FlexibleContexts #-}
10#include "inline.hs"
12-- |
13-- Module : Streamly.Mem.Array.Types
14-- Copyright : (c) 2019 Composewell Technologies
15--
16-- License : BSD3
17-- Maintainer : harendra.kumar@gmail.com
18-- Stability : experimental
19-- Portability : GHC
20--
21module Streamly.Mem.Array.Types
22 (
23 Array (..)
25 -- * Construction
26 , withNewArray
27 , newArray
28 , unsafeSnoc
29 , snoc
30 , spliceWithDoubling
32 , fromList
33 , fromListN
34 , fromStreamDN
35 -- , fromStreamD
37 -- * Streams of arrays
38 , fromStreamDArraysOf
39 , flattenArrays
40 , flattenArraysRev
41 , packArraysChunksOf
42#if !defined(mingw32_HOST_OS)
43 , groupIOVecsOf
44#endif
45 , splitOn
47 -- * Elimination
48 , unsafeIndexIO
49 , unsafeIndex
50 , length
51 , byteLength
52 , byteCapacity
53 , foldl'
54 , foldr
55 , splitAt
57 , toStreamD
58 , toStreamDRev
59 , toStreamK
60 , toStreamKRev
61 , toList
62 , toArrayN
63 , read
65 -- * Utilities
66 , defaultChunkSize
67 , mkChunkSize
68 , mkChunkSizeKB
69 , unsafeInlineIO
70 , realloc
71 , shrinkToFit
72 , memcpy
73 , memcmp
75 , unlines
76 )
77where
79import Control.Exception (assert)
80import Control.DeepSeq (NFData(..))
81import Control.Monad (when)
82import Control.Monad.IO.Class (MonadIO(..))
83import Data.Functor.Identity (runIdentity)
84import Data.Word (Word8)
85import Foreign.C.String (CString)
86import Foreign.C.Types (CSize(..), CInt(..))
87import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, touchForeignPtr)
88import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
89import Foreign.Ptr (plusPtr, minusPtr, castPtr, nullPtr)
90import Foreign.Storable (Storable(..))
91import Prelude hiding (length, foldr, read, unlines, splitAt)
92import Text.Read (readPrec, readListPrec, readListPrecDefault)
94import GHC.Base (Addr#, realWorld#)
95import GHC.Exts (IsList, IsString(..))
96import GHC.ForeignPtr (mallocPlainForeignPtrAlignedBytes, newForeignPtr_)
97import GHC.IO (IO(IO), unsafePerformIO)
98import GHC.Ptr (Ptr(..))
100import Streamly.Fold.Types (Fold(..))
101import Streamly.SVar (adaptState)
103#if !defined(mingw32_HOST_OS)
104import Streamly.FileSystem.FDIO (IOVec(..))
105#endif
107import qualified Streamly.Streams.StreamD.Type as D
108import qualified Streamly.Streams.StreamK as K
109import qualified GHC.Exts as Exts
111#ifdef DEVBUILD
112import qualified Data.Foldable as F
113#endif
115#if MIN_VERSION_base(4,10,0)
116import Foreign.ForeignPtr (plusForeignPtr)
117#else
118import GHC.Base (Int(..), plusAddr#)
119import GHC.ForeignPtr (ForeignPtr(..))
120plusForeignPtr :: ForeignPtr a -> Int -> ForeignPtr b
121plusForeignPtr (ForeignPtr addr c) (I# d) = ForeignPtr (plusAddr# addr d) c
122#endif
124-------------------------------------------------------------------------------
125-- Design Notes
126-------------------------------------------------------------------------------
128-- There are two goals that we want to fulfill with arrays. One, holding large
129-- amounts of data in non-GC memory and the ability to pass this data to and
130-- from the operating system without an extra copy overhead. Two, allow random
131-- access to elements based on their indices. The first one falls in the
132-- category of storage buffers while the second one falls in the category of
133-- maps/multisets/hashmaps.
134--
135-- For the first requirement we use an array of Storables and store it in a
136-- ForeignPtr. We can have both immutable and mutable variants of this array
137-- using wrappers over the same underlying type.
138--
139-- For the second requirement, we need a separate type for arrays of
140-- polymorphic values, for example vectors of handler functions, lookup tables.
141-- We can call this type a "vector" in contrast to arrays. It should not
142-- require Storable instance for the type. In that case we need to use an
143-- Array# instead of a ForeignPtr. This type of array would not reduce the GC
144-- overhead much because each element of the array still needs to be scanned by
145-- the GC. However, it can store polymorphic elements and allow random access
146-- to those. But in most cases random access means storage, and it means we
147-- need to avoid GC scanning except in cases of trivially small storage. One
148-- way to achieve that would be to put the array in a Compact region. However,
149-- if and when we mutate this, we will have to use a manual GC copying out to
150-- another CR and freeing the old one.
152-------------------------------------------------------------------------------
153-- Array Data Type
154-------------------------------------------------------------------------------
156-- We require that an array stores only Storable. Arrays are used for buffering
157-- data while streams are used for processing. If you want something to be
158-- buffered it better be Storable so that we can store it in non-GC memory.
159--
160-- We can use a "Storable" constraint in the Array type and the Constraint can
161-- be automatically provided to a function that pattern matches on the Array
162-- type. However, it has huge performance cost, so we do not use it.
163-- XXX Can this cost be alleviated in GHC-8.10 specialization fix?
164--
165-- XXX Another way to not require the Storable constraint in array operations
166-- is to store the elemSize in the array at construction and use that instead
167-- of using sizeOf. Need to charaterize perf cost of this.
168--
169-- XXX rename the fields to "start, next, end".
170--
171data Array a =
172#ifdef DEVBUILD
173 Storable a =>
174#endif
175 Array
176 { aStart :: {-# UNPACK #-} !(ForeignPtr a) -- first address
177 , aEnd :: {-# UNPACK #-} !(Ptr a) -- first unused address
178 , aBound :: {-# UNPACK #-} !(Ptr a) -- first address beyond allocated memory
179 }
181-------------------------------------------------------------------------------
182-- Utility functions
183-------------------------------------------------------------------------------
185foreign import ccall unsafe "string.h memcpy" c_memcpy
186 :: Ptr Word8 -> Ptr Word8 -> CSize -> IO (Ptr Word8)
188foreign import ccall unsafe "string.h strlen" c_strlen
189 :: CString -> IO CSize
191foreign import ccall unsafe "string.h memchr" c_memchr
192 :: Ptr Word8 -> Word8 -> CSize -> IO (Ptr Word8)
194-- XXX we are converting Int to CSize
195memcpy :: Ptr Word8 -> Ptr Word8 -> Int -> IO ()
196memcpy dst src len = c_memcpy dst src (fromIntegral len) >> return ()
198foreign import ccall unsafe "string.h memcmp" c_memcmp
199 :: Ptr Word8 -> Ptr Word8 -> CSize -> IO CInt
201-- XXX we are converting Int to CSize
202-- return True if the memory locations have identical contents
203{-# INLINE memcmp #-}
204memcmp :: Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
205memcmp p1 p2 len = do
206 r <- c_memcmp p1 p2 (fromIntegral len)
207 return $ r == 0
209{-# INLINE unsafeInlineIO #-}
210unsafeInlineIO :: IO a -> a
211unsafeInlineIO (IO m) = case m realWorld# of (# _, r #) -> r
213-------------------------------------------------------------------------------
214-- Construction
215-------------------------------------------------------------------------------
217-- | Allocate an array that can hold 'count' items. The memory of the array is
218-- uninitialized.
219--
220-- Note that this is internal routine, the reference to this array cannot be
221-- given out until the array has been written to and frozen.
222{-# INLINE newArray #-}
223newArray :: forall a. Storable a => Int -> IO (Array a)
224newArray count = do
225 let size = count * sizeOf (undefined :: a)
226 fptr <- mallocPlainForeignPtrAlignedBytes size (alignment (undefined :: a))
227 let p = unsafeForeignPtrToPtr fptr
228 return $ Array
229 { aStart = fptr
230 , aEnd = p
231 , aBound = p `plusPtr` size
232 }
234-- | Allocate an Array of the given size and run an IO action passing the array
235-- start pointer.
236{-# INLINE withNewArray #-}
237withNewArray :: forall a. Storable a => Int -> (Ptr a -> IO ()) -> IO (Array a)
238withNewArray count f = do
239 arr <- newArray count
240 withForeignPtr (aStart arr) $ \p -> f p >> return arr
242-- XXX grow the array when we are beyond bound.
243--
244-- Internal routine for when the array is being created. Appends one item at
245-- the end of the array. Useful when sequentially writing a stream to the
246-- array. DOES NOT CHECK THE ARRAY BOUNDS.
247{-# INLINE unsafeSnoc #-}
248unsafeSnoc :: forall a. Storable a => Array a -> a -> IO (Array a)
249unsafeSnoc arr@Array{..} x = do
250 when (aEnd == aBound) $
251 error "BUG: unsafeSnoc: writing beyond array bounds"
252 poke aEnd x
253 touchForeignPtr aStart
254 return $ arr {aEnd = aEnd `plusPtr` (sizeOf (undefined :: a))}
256{-# INLINE snoc #-}
257snoc :: forall a. Storable a => Array a -> a -> IO (Array a)
258snoc arr@Array {..} x = do
259 if (aEnd == aBound)
260 then do
261 let oldStart = unsafeForeignPtrToPtr aStart
262 size = aEnd `minusPtr` oldStart
263 newSize = (size + (sizeOf (undefined :: a)))
264 newPtr <- mallocPlainForeignPtrAlignedBytes newSize (alignment (undefined :: a))
265 withForeignPtr newPtr $ \pNew -> do
266 memcpy (castPtr pNew) (castPtr oldStart) size
267 poke (pNew `plusPtr` size) x
268 touchForeignPtr aStart
269 return $ Array
270 { aStart = newPtr
271 , aEnd = pNew `plusPtr` (size + sizeOf (undefined :: a))
272 , aBound = pNew `plusPtr` newSize
273 }
274 else do
275 poke aEnd x
276 touchForeignPtr aStart
277 return $ arr {aEnd = aEnd `plusPtr` (sizeOf (undefined :: a))}
279-- | Reallocate the array to the specified size in bytes. If the size is less
280-- than the original array the array gets truncated.
281{-# NOINLINE realloc #-}
282realloc :: forall a. Storable a => Int -> Array a -> IO (Array a)
283realloc newSize Array{..} = do
284 assert (aEnd <= aBound) (return ())
285 let oldStart = unsafeForeignPtrToPtr aStart
286 let size = aEnd `minusPtr` oldStart
287 newPtr <- mallocPlainForeignPtrAlignedBytes
288 newSize (alignment (undefined :: a))
289 withForeignPtr newPtr $ \pNew -> do
290 memcpy (castPtr pNew) (castPtr oldStart) size
291 touchForeignPtr aStart
292 return $ Array
293 { aStart = newPtr
294 , aEnd = pNew `plusPtr` size
295 , aBound = pNew `plusPtr` newSize
296 }
298-- | Remove the free space from an Array.
299shrinkToFit :: forall a. Storable a => Array a -> IO (Array a)
300shrinkToFit arr@Array{..} = do
301 assert (aEnd <= aBound) (return ())
302 if aEnd /= aBound
303 then do
304 let oldStart = unsafeForeignPtrToPtr aStart
305 let size = aEnd `minusPtr` oldStart
306 realloc size arr
307 else return arr
309-- XXX when converting an array of Word8 from a literal string we can simply
310-- refer to the literal string. Is it possible to write rules such that
311-- fromList Word8 can be rewritten so that GHC does not first convert the
312-- literal to [Char] and then we convert it back to an Array Word8?
313--
314-- Note that the address must be a read-only address (meant to be used for
315-- read-only string literals) because we are sharing it, any modification to
316-- the original address would change our array. That's why this function is
317-- unsafe.
318{-# INLINE _fromCStringAddrUnsafe #-}
319_fromCStringAddrUnsafe :: Addr# -> IO (Array Word8)
320_fromCStringAddrUnsafe addr# = do
321 ptr <- newForeignPtr_ (castPtr cstr)
322 len <- c_strlen cstr
323 let n = fromIntegral len
324 let p = unsafeForeignPtrToPtr ptr
325 let end = p `plusPtr` n
326 return $ Array
327 { aStart = ptr
328 , aEnd = end
329 , aBound = end
330 }
331 where
332 cstr :: CString
333 cstr = Ptr addr#
335-------------------------------------------------------------------------------
336-- Elimination
337-------------------------------------------------------------------------------
339-- | Return element at the specified index without checking the bounds.
340--
341-- Unsafe because it does not check the bounds of the array.
342{-# INLINE_NORMAL unsafeIndexIO #-}
343unsafeIndexIO :: forall a. Storable a => Array a -> Int -> IO a
344unsafeIndexIO Array {..} i =
345 withForeignPtr aStart $ \p -> do
346 let elemSize = sizeOf (undefined :: a)
347 elemOff = p `plusPtr` (elemSize * i)
348 assert (i >= 0 && elemOff `plusPtr` elemSize <= aEnd)
349 (return ())
350 peek elemOff
352-- | Return element at the specified index without checking the bounds.
353{-# INLINE_NORMAL unsafeIndex #-}
354unsafeIndex :: forall a. Storable a => Array a -> Int -> a
355unsafeIndex arr i = let !r = unsafeInlineIO $ unsafeIndexIO arr i in r
357-- | /O(1)/ Get the byte length of the array.
358--
359-- @since 0.7.0
360{-# INLINE byteLength #-}
361byteLength :: Array a -> Int
362byteLength Array{..} =
363 let p = unsafeForeignPtrToPtr aStart
364 len = aEnd `minusPtr` p
365 in assert (len >= 0) len
367-- | /O(1)/ Get the length of the array i.e. the number of elements in the
368-- array.
369--
370-- @since 0.7.0
371{-# INLINE length #-}
372length :: forall a. Storable a => Array a -> Int
373length arr = byteLength arr `div` sizeOf (undefined :: a)
375{-# INLINE byteCapacity #-}
376byteCapacity :: Array a -> Int
377byteCapacity Array{..} =
378 let p = unsafeForeignPtrToPtr aStart
379 len = aBound `minusPtr` p
380 in assert (len >= 0) len
382{-# INLINE_NORMAL toStreamD #-}
383toStreamD :: forall m a. (Monad m, Storable a) => Array a -> D.Stream m a
384toStreamD Array{..} =
385 let p = unsafeForeignPtrToPtr aStart
386 in D.Stream step p
388 where
390 {-# INLINE_LATE step #-}
391 step _ p | p == aEnd = return D.Stop
392 step _ p = do
393 -- unsafeInlineIO allows us to run this in Identity monad for pure
394 -- toList/foldr case which makes them much faster due to not
395 -- accumulating the list and fusing better with the pure consumers.
396 --
397 -- This should be safe as the array contents are guaranteed to be
398 -- evaluated/written to before we peek at them.
399 let !x = unsafeInlineIO $ do
400 r <- peek p
401 touchForeignPtr aStart
402 return r
403 return $ D.Yield x (p `plusPtr` (sizeOf (undefined :: a)))
405{-# INLINE toStreamK #-}
406toStreamK :: forall t m a. (K.IsStream t, Storable a) => Array a -> t m a
407toStreamK Array{..} =
408 let p = unsafeForeignPtrToPtr aStart
409 in go p
411 where
413 go p | p == aEnd = K.nil
414 | otherwise =
415 -- See Note in toStreamD.
416 let !x = unsafeInlineIO $ do
417 r <- peek p
418 touchForeignPtr aStart
419 return r
420 in x `K.cons` go (p `plusPtr` (sizeOf (undefined :: a)))
422{-# INLINE_NORMAL toStreamDRev #-}
423toStreamDRev :: forall m a. (Monad m, Storable a) => Array a -> D.Stream m a
424toStreamDRev Array{..} =
425 let p = aEnd `plusPtr` negate (sizeOf (undefined :: a))
426 in D.Stream step p
428 where
430 {-# INLINE_LATE step #-}
431 step _ p | p < unsafeForeignPtrToPtr aStart = return D.Stop
432 step _ p = do
433 -- See comments in toStreamD for why we use unsafeInlineIO
434 let !x = unsafeInlineIO $ do
435 r <- peek p
436 touchForeignPtr aStart
437 return r
438 return $ D.Yield x (p `plusPtr` negate (sizeOf (undefined :: a)))
440{-# INLINE toStreamKRev #-}
441toStreamKRev :: forall t m a. (K.IsStream t, Storable a) => Array a -> t m a
442toStreamKRev Array {..} =
443 let p = aEnd `plusPtr` negate (sizeOf (undefined :: a))
444 in go p
446 where
448 go p | p < unsafeForeignPtrToPtr aStart = K.nil
449 | otherwise =
450 let !x = unsafeInlineIO $ do
451 r <- peek p
452 touchForeignPtr aStart
453 return r
454 in x `K.cons` go (p `plusPtr` negate (sizeOf (undefined :: a)))
456{-# INLINE_NORMAL foldl' #-}
457foldl' :: forall a b. Storable a => (b -> a -> b) -> b -> Array a -> b
458foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
460{-# INLINE_NORMAL foldr #-}
461foldr :: Storable a => (a -> b -> b) -> b -> Array a -> b
462foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
464-------------------------------------------------------------------------------
465-- Instances
466-------------------------------------------------------------------------------
468-- | @toArrayN n@ folds a maximum of @n@ elements from the input stream to an
469-- 'Array'.
470--
471-- @since 0.7.0
472{-# INLINE_NORMAL toArrayN #-}
473toArrayN :: forall m a. (MonadIO m, Storable a) => Int -> Fold m a (Array a)
474toArrayN n = Fold step initial extract
476 where
478 initial = do
479 if n < 0 then error "toArrayN: negative count specified" else return ()
480 liftIO $ newArray n
481 step arr@(Array _ end bound) _ | end == bound = return arr
482 step (Array start end bound) x = do
483 liftIO $ poke end x
484 return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound
485 extract = liftIO . shrinkToFit
487{-# INLINE_NORMAL fromStreamDN #-}
488fromStreamDN :: forall m a. (MonadIO m, Storable a)
489 => Int -> D.Stream m a -> m (Array a)
490fromStreamDN limit str = do
491 arr <- liftIO $ newArray limit
492 end <- D.foldlM' write (aEnd arr) $ D.take limit str
493 return $ arr {aEnd = end}
495 where
497 write ptr x = do
498 liftIO $ poke ptr x
499 return $ ptr `plusPtr` sizeOf (undefined :: a)
501data GroupState s start end bound
502 = GroupStart s
503 | GroupBuffer s start end bound
504 | GroupYield start end bound (GroupState s start end bound)
505 | GroupFinish
507-- | @fromStreamArraysOf n stream@ groups the input stream into a stream of
508-- arrays of size n.
509{-# INLINE_NORMAL fromStreamDArraysOf #-}
510fromStreamDArraysOf :: forall m a. (MonadIO m, Storable a)
511 => Int -> D.Stream m a -> D.Stream m (Array a)
512-- fromStreamDArraysOf n str = D.groupsOf n (toArrayN n) str
513fromStreamDArraysOf n (D.Stream step state) =
514 D.Stream step' (GroupStart state)
516 where
518 {-# INLINE_LATE step' #-}
519 step' _ (GroupStart st) = do
520 when (n <= 0) $
521 -- XXX we can pass the module string from the higher level API
522 error $ "Streamly.Mem.Array.Types.fromStreamDArraysOf: the size of "
523 ++ "arrays [" ++ show n ++ "] must be a natural number"
524 Array start end bound <- liftIO $ newArray n
525 return $ D.Skip (GroupBuffer st start end bound)
527 step' gst (GroupBuffer st start end bound) = do
528 r <- step (adaptState gst) st
529 case r of
530 D.Yield x s -> do
531 liftIO $ poke end x
532 let end' = end `plusPtr` sizeOf (undefined :: a)
533 return $
534 if end' >= bound
535 then D.Skip (GroupYield start end' bound (GroupStart s))
536 else D.Skip (GroupBuffer s start end' bound)
537 D.Skip s -> return $ D.Skip (GroupBuffer s start end bound)
538 D.Stop -> return $ D.Skip (GroupYield start end bound GroupFinish)
540 step' _ (GroupYield start end bound next) =
541 return $ D.Yield (Array start end bound) next
543 step' _ GroupFinish = return D.Stop
545-- XXX concatMap does not seem to have the best possible performance so we have
546-- a custom way to concat arrays.
547data FlattenState s a =
548 OuterLoop s
549 | InnerLoop s (ForeignPtr a) (Ptr a) (Ptr a)
551{-# INLINE_NORMAL flattenArrays #-}
552flattenArrays :: forall m a. (MonadIO m, Storable a)
553 => D.Stream m (Array a) -> D.Stream m a
554flattenArrays (D.Stream step state) = D.Stream step' (OuterLoop state)
556 where
558 {-# INLINE_LATE step' #-}
559 step' gst (OuterLoop st) = do
560 r <- step (adaptState gst) st
561 return $ case r of
562 D.Yield Array{..} s ->
563 let p = unsafeForeignPtrToPtr aStart
564 in D.Skip (InnerLoop s aStart p aEnd)
565 D.Skip s -> D.Skip (OuterLoop s)
566 D.Stop -> D.Stop
568 step' _ (InnerLoop st _ p end) | p == end =
569 return $ D.Skip $ OuterLoop st
571 step' _ (InnerLoop st startf p end) = do
572 x <- liftIO $ do
573 r <- peek p
574 touchForeignPtr startf
575 return r
576 return $ D.Yield x (InnerLoop st startf
577 (p `plusPtr` (sizeOf (undefined :: a))) end)
579{-# INLINE_NORMAL flattenArraysRev #-}
580flattenArraysRev :: forall m a. (MonadIO m, Storable a)
581 => D.Stream m (Array a) -> D.Stream m a
582flattenArraysRev (D.Stream step state) = D.Stream step' (OuterLoop state)
584 where
586 {-# INLINE_LATE step' #-}
587 step' gst (OuterLoop st) = do
588 r <- step (adaptState gst) st
589 return $ case r of
590 D.Yield Array{..} s ->
591 let p = aEnd `plusPtr` negate (sizeOf (undefined :: a))
592 -- XXX we do not need aEnd
593 in D.Skip (InnerLoop s aStart p aEnd)
594 D.Skip s -> D.Skip (OuterLoop s)
595 D.Stop -> D.Stop
597 step' _ (InnerLoop st start p _) | p < unsafeForeignPtrToPtr start =
598 return $ D.Skip $ OuterLoop st
600 step' _ (InnerLoop st startf p end) = do
601 x <- liftIO $ do
602 r <- peek p
603 touchForeignPtr startf
604 return r
605 return $ D.Yield x (InnerLoop st startf
606 (p `plusPtr` negate (sizeOf (undefined :: a))) end)
608-- CAUTION: a very large number (millions) of arrays can degrade performance
609-- due to GC overhead because we need to buffer the arrays before we flatten
610-- all the arrays.
611--
612-- We could take the approach of doubling the memory allocation on each
613-- overflow. This would result in more or less the same amount of copying as in
614-- the chunking approach. However, if we have to shrink in the end then it may
615-- result in an extra copy of the entire data.
616--
617{-# INLINE fromStreamD #-}
618fromStreamD :: (MonadIO m, Storable a) => D.Stream m a -> m (Array a)
619fromStreamD m = do
620 let s = fromStreamDArraysOf defaultChunkSize m
621 buffered <- D.foldr K.cons K.nil s
622 len <- K.foldl' (+) 0 (K.map length buffered)
623 fromStreamDN len $ flattenArrays $ D.fromStreamK buffered
625-- | Convert an 'Array' into a list.
626--
627-- @since 0.7.0
628{-# INLINABLE toList #-}
629toList :: Storable a => Array a -> [a]
630toList = foldr (:) []
632instance (Show a, Storable a) => Show (Array a) where
633 {-# INLINE showsPrec #-}
634 showsPrec _ = shows . toList
636-- | Create an 'Array' from the first N elements of a list. The array is
637-- allocated to size N, if the list terminates before N elements then the
638-- array may hold less than N elements.
639--
640-- @since 0.7.0
641{-# INLINABLE fromListN #-}
642fromListN :: Storable a => Int -> [a] -> Array a
643fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
645-- | Create an 'Array' from a list. The list must be of finite size.
646--
647-- @since 0.7.0
648{-# INLINABLE fromList #-}
649fromList :: Storable a => [a] -> Array a
650fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
652instance (Storable a, Read a, Show a) => Read (Array a) where
653 {-# INLINE readPrec #-}
654 readPrec = do
655 xs <- readPrec
656 return (fromList xs)
657 readListPrec = readListPrecDefault
659instance (a ~ Char) => IsString (Array a) where
660 {-# INLINE fromString #-}
661 fromString = fromList
663-- GHC versions 8.0 and below cannot derive IsList
664instance Storable a => IsList (Array a) where
665 type (Item (Array a)) = a
666 {-# INLINE fromList #-}
667 fromList = fromList
668 {-# INLINE fromListN #-}
669 fromListN = fromListN
670 {-# INLINE toList #-}
671 toList = toList
673{-# INLINE arrcmp #-}
674arrcmp :: Array a -> Array a -> Bool
675arrcmp arr1 arr2 =
676 let !res = unsafeInlineIO $ do
677 let ptr1 = unsafeForeignPtrToPtr $ aStart arr1
678 let ptr2 = unsafeForeignPtrToPtr $ aStart arr2
679 let len1 = aEnd arr1 `minusPtr` ptr1
680 let len2 = aEnd arr2 `minusPtr` ptr2
682 if len1 == len2
683 then do
684 r <- memcmp (castPtr ptr1) (castPtr ptr2) len1
685 touchForeignPtr $ aStart arr1
686 touchForeignPtr $ aStart arr2
687 return r
688 else return False
689 in res
691-- XXX we are assuming that Storable equality means element equality. This may
692-- or may not be correct? arrcmp is 40% faster compared to stream equality.
693instance (Storable a, Eq a) => Eq (Array a) where
694 {-# INLINE (==) #-}
695 (==) = arrcmp
696 -- arr1 == arr2 = runIdentity $ D.eqBy (==) (toStreamD arr1) (toStreamD arr2)
698instance (Storable a, NFData a) => NFData (Array a) where
699 {-# INLINE rnf #-}
700 rnf = foldl' (\_ x -> rnf x) ()
702instance (Storable a, Ord a) => Ord (Array a) where
703 {-# INLINE compare #-}
704 compare arr1 arr2 = unsafePerformIO $
705 D.cmpBy compare (toStreamD arr1) (toStreamD arr2)
707 -- Default definitions defined in base do not have an INLINE on them, so we
708 -- replicate them here with an INLINE.
709 {-# INLINE (<) #-}
710 x < y = case compare x y of { LT -> True; _ -> False }
712 {-# INLINE (<=) #-}
713 x <= y = case compare x y of { GT -> False; _ -> True }
715 {-# INLINE (>) #-}
716 x > y = case compare x y of { GT -> True; _ -> False }
718 {-# INLINE (>=) #-}
719 x >= y = case compare x y of { LT -> False; _ -> True }
721 -- These two default methods use '<=' rather than 'compare'
722 -- because the latter is often more expensive
723 {-# INLINE max #-}
724 max x y = if x <= y then y else x
726 {-# INLINE min #-}
727 min x y = if x <= y then x else y
729#ifdef DEVBUILD
730-- Definitions using the Storable constraint from the Array type. These are to
731-- make the Foldable instance possible though it is much slower (7x slower).
732--
733{-# INLINE_NORMAL toStreamD_ #-}
734toStreamD_ :: forall m a. MonadIO m => Int -> Array a -> D.Stream m a
735toStreamD_ size Array{..} =
736 let p = unsafeForeignPtrToPtr aStart
737 in D.Stream step p
739 where
741 {-# INLINE_LATE step #-}
742 step _ p | p == aEnd = return D.Stop
743 step _ p = do
744 x <- liftIO $ do
745 r <- peek p
746 touchForeignPtr aStart
747 return r
748 return $ D.Yield x (p `plusPtr` size)
750{-# INLINE_NORMAL _foldr #-}
751_foldr :: forall a b. (a -> b -> b) -> b -> Array a -> b
752_foldr f z arr@Array {..} =
753 let !n = sizeOf (undefined :: a)
754 in unsafePerformIO $ D.foldr f z $ toStreamD_ n arr
756-- | Note that the 'Foldable' instance is 7x slower than the direct
757-- operations.
758instance Foldable Array where
759 foldr = _foldr
760#endif
762-- | Convert an 'Array' into a stream.
763--
764-- @since 0.7.0
765{-# INLINE_EARLY read #-}
766read :: (Monad m, K.IsStream t, Storable a) => Array a -> t m a
767read = D.fromStreamD . toStreamD
768-- XXX add fallback to StreamK rule
769-- {-# RULES "Streamly.Array.read fallback to StreamK" [1]
770-- forall a. S.readK (read a) = K.fromArray a #-}
772-------------------------------------------------------------------------------
773-- Semigroup and Monoid
774-------------------------------------------------------------------------------
776-------------------------------------------------------------------------------
777-- IO
778-------------------------------------------------------------------------------
780-- | GHC memory management allocation header overhead
781allocOverhead :: Int
782allocOverhead = 2 * sizeOf (undefined :: Int)
784mkChunkSize :: Int -> Int
785mkChunkSize n = let size = n - allocOverhead in max size 0
787mkChunkSizeKB :: Int -> Int
788mkChunkSizeKB n = mkChunkSize (n * k)
789 where k = 1024
791-- | Default maximum buffer size in bytes, for reading from and writing to IO
792-- devices, the value is 32KB minus GHC allocation overhead, which is a few
793-- bytes, so that the actual allocation is 32KB.
794defaultChunkSize :: Int
795defaultChunkSize = mkChunkSizeKB 32
797{-# INLINE_NORMAL unlines #-}
798unlines :: forall m a. (MonadIO m, Storable a)
799 => a -> D.Stream m (Array a) -> D.Stream m a
800unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
801 where
802 {-# INLINE_LATE step' #-}
803 step' gst (OuterLoop st) = do
804 r <- step (adaptState gst) st
805 return $ case r of
806 D.Yield Array{..} s ->
807 let p = unsafeForeignPtrToPtr aStart
808 in D.Skip (InnerLoop s aStart p aEnd)
809 D.Skip s -> D.Skip (OuterLoop s)
810 D.Stop -> D.Stop
812 step' _ (InnerLoop st _ p end) | p == end =
813 return $ D.Yield sep $ OuterLoop st
815 step' _ (InnerLoop st startf p end) = do
816 x <- liftIO $ do
817 r <- peek p
818 touchForeignPtr startf
819 return r
820 return $ D.Yield x (InnerLoop st startf
821 (p `plusPtr` (sizeOf (undefined :: a))) end)
823{-# INLINE spliceTwo #-}
824spliceTwo :: (MonadIO m, Storable a) => Array a -> Array a -> m (Array a)
825spliceTwo arr1 arr2 = do
826 let src1 = unsafeForeignPtrToPtr (aStart arr1)
827 src2 = unsafeForeignPtrToPtr (aStart arr2)
828 len1 = aEnd arr1 `minusPtr` src1
829 len2 = aEnd arr2 `minusPtr` src2
831 arr <- liftIO $ newArray (len1 + len2)
832 let dst = unsafeForeignPtrToPtr (aStart arr)
834 liftIO $ do
835 memcpy (castPtr dst) (castPtr src1) len1
836 touchForeignPtr (aStart arr1)
837 memcpy (castPtr (dst `plusPtr` len1)) (castPtr src2) len2
838 touchForeignPtr (aStart arr2)
839 return arr { aEnd = dst `plusPtr` (len1 + len2) }
841-- Splice a new array into a pre-reserved array. The user must ensure that
842-- there is enough space in the array.
843{-# INLINE spliceWith #-}
844spliceWith :: (MonadIO m) => Array a -> Array a -> m (Array a)
845spliceWith dst@(Array _ end bound) src = liftIO $ do
846 let srcLen = byteLength src
847 if end `plusPtr` srcLen > bound
848 then error "Bug: spliceIntoUnsafe: Not enough space in the target array"
849 else
850 withForeignPtr (aStart dst) $ \_ -> do
851 withForeignPtr (aStart src) $ \psrc -> do
852 let pdst = aEnd dst
853 memcpy (castPtr pdst) (castPtr psrc) srcLen
854 return $ dst { aEnd = pdst `plusPtr` srcLen }
856-- Splice a new array into a preallocated array, doubling the space if there is
857-- no space in the target array.
858{-# INLINE spliceWithDoubling #-}
859spliceWithDoubling :: (MonadIO m, Storable a)
860 => Array a -> Array a -> m (Array a)
861spliceWithDoubling dst@(Array start end bound) src = do
862 assert (end <= bound) (return ())
863 let srcLen = aEnd src `minusPtr` unsafeForeignPtrToPtr (aStart src)
865 dst1 <-
866 if end `plusPtr` srcLen >= bound
867 then do
868 let oldStart = unsafeForeignPtrToPtr start
869 oldSize = end `minusPtr` oldStart
870 newSize = max (oldSize * 2) (oldSize + srcLen)
871 liftIO $ realloc newSize dst
872 else return dst
873 spliceWith dst1 src
875data SpliceState s arr
876 = SpliceInitial s
877 | SpliceBuffering s arr
878 | SpliceYielding arr (SpliceState s arr)
879 | SpliceFinish
881-- XXX can use general grouping combinators to achieve this?
882-- | Coalesce adajcent arrays in incoming stream to form bigger arrays of a
883-- maximum specified size.
884--
885-- @since 0.7.0
886{-# INLINE_NORMAL packArraysChunksOf #-}
887packArraysChunksOf :: (MonadIO m, Storable a)
888 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
889packArraysChunksOf n (D.Stream step state) =
890 D.Stream step' (SpliceInitial state)
892 where
894 {-# INLINE_LATE step' #-}
895 step' gst (SpliceInitial st) = do
896 when (n <= 0) $
897 -- XXX we can pass the module string from the higher level API
898 error $ "Streamly.Mem.Array.Types.packArraysChunksOf: the size of "
899 ++ "arrays [" ++ show n ++ "] must be a natural number"
900 r <- step gst st
901 case r of
902 D.Yield arr s -> return $
903 let len = byteLength arr
904 in if len >= n
905 then D.Skip (SpliceYielding arr (SpliceInitial s))
906 else D.Skip (SpliceBuffering s arr)
907 D.Skip s -> return $ D.Skip (SpliceInitial s)
908 D.Stop -> return $ D.Stop
910 step' gst (SpliceBuffering st buf) = do
911 r <- step gst st
912 case r of
913 D.Yield arr s -> do
914 let len = byteLength buf + byteLength arr
915 if len > n
916 then return $
917 D.Skip (SpliceYielding buf (SpliceBuffering s arr))
918 else do
919 buf' <- if byteCapacity buf < n
920 then liftIO $ realloc n buf
921 else return buf
922 buf'' <- spliceWith buf' arr
923 return $ D.Skip (SpliceBuffering s buf'')
924 D.Skip s -> return $ D.Skip (SpliceBuffering s buf)
925 D.Stop -> return $ D.Skip (SpliceYielding buf SpliceFinish)
927 step' _ SpliceFinish = return D.Stop
929 step' _ (SpliceYielding arr next) = return $ D.Yield arr next
931#if !defined(mingw32_HOST_OS)
932data GatherState s arr
933 = GatherInitial s
934 | GatherBuffering s arr Int
935 | GatherYielding arr (GatherState s arr)
936 | GatherFinish
938-- | @groupIOVecsOf maxBytes maxEntries@ groups arrays in the incoming stream
939-- to create a stream of 'IOVec' arrays with a maximum of @maxBytes@ bytes in
940-- each array and a maximum of @maxEntries@ entries in each array.
941--
942-- @since 0.7.0
943{-# INLINE_NORMAL groupIOVecsOf #-}
944groupIOVecsOf :: MonadIO m
945 => Int -> Int -> D.Stream m (Array a) -> D.Stream m (Array IOVec)
946groupIOVecsOf n maxIOVLen (D.Stream step state) =
947 D.Stream step' (GatherInitial state)
949 where
951 {-# INLINE_LATE step' #-}
952 step' gst (GatherInitial st) = do
953 when (n <= 0) $
954 -- XXX we can pass the module string from the higher level API
955 error $ "Streamly.Mem.Array.Types.groupIOVecsOf: the size of "
956 ++ "groups [" ++ show n ++ "] must be a natural number"
957 when (maxIOVLen <= 0) $
958 -- XXX we can pass the module string from the higher level API
959 error $ "Streamly.Mem.Array.Types.groupIOVecsOf: the number of "
960 ++ "IOVec entries [" ++ show n ++ "] must be a natural number"
961 r <- step (adaptState gst) st
962 case r of
963 D.Yield arr s -> do
964 let p = unsafeForeignPtrToPtr (aStart arr)
965 len = byteLength arr
966 iov <- liftIO $ newArray maxIOVLen
967 iov' <- liftIO $ unsafeSnoc iov (IOVec (castPtr p)
968 (fromIntegral len))
969 if len >= n
970 then return $ D.Skip (GatherYielding iov' (GatherInitial s))
971 else return $ D.Skip (GatherBuffering s iov' len)
972 D.Skip s -> return $ D.Skip (GatherInitial s)
973 D.Stop -> return $ D.Stop
975 step' gst (GatherBuffering st iov len) = do
976 r <- step (adaptState gst) st
977 case r of
978 D.Yield arr s -> do
979 let p = unsafeForeignPtrToPtr (aStart arr)
980 alen = byteLength arr
981 len' = len + alen
982 if len' > n || length iov >= maxIOVLen
983 then do
984 iov' <- liftIO $ newArray maxIOVLen
985 iov'' <- liftIO $ unsafeSnoc iov' (IOVec (castPtr p)
986 (fromIntegral alen))
987 return $ D.Skip (GatherYielding iov
988 (GatherBuffering s iov'' alen))
989 else do
990 iov' <- liftIO $ unsafeSnoc iov (IOVec (castPtr p)
991 (fromIntegral alen))
992 return $ D.Skip (GatherBuffering s iov' len')
993 D.Skip s -> return $ D.Skip (GatherBuffering s iov len)
994 D.Stop -> return $ D.Skip (GatherYielding iov GatherFinish)
996 step' _ GatherFinish = return D.Stop
998 step' _ (GatherYielding iov next) = return $ D.Yield iov next
999#endif
1001-- | Create two slices of an array without copying the original array. The
1002-- specified index @i@ is the first index of the second slice.
1004-- @since 0.7.0
1005splitAt :: forall a. Storable a => Int -> Array a -> (Array a, Array a)
1006splitAt i arr@Array{..} =
1007 let maxIndex = length arr - 1
1008 in if i < 0
1009 then error "sliceAt: negative array index"
1010 else if i > maxIndex
1011 then error $ "sliceAt: specified array index " ++ show i
1012 ++ " is beyond the maximum index " ++ show maxIndex
1013 else let off = i * sizeOf (undefined :: a)
1014 p = unsafeForeignPtrToPtr aStart `plusPtr` off
1015 in ( Array
1016 { aStart = aStart
1017 , aEnd = p
1018 , aBound = p
1020 , Array
1021 { aStart = aStart `plusForeignPtr` off
1022 , aEnd = aEnd
1023 , aBound = aBound
1027-- Drops the separator byte
1028{-# INLINE breakOn #-}
1029breakOn :: MonadIO m
1030 => Word8 -> Array Word8 -> m (Array Word8, Maybe (Array Word8))
1031breakOn sep arr@Array{..} = liftIO $ do
1032 let p = unsafeForeignPtrToPtr aStart
1033 loc <- c_memchr p sep (fromIntegral $ aEnd `minusPtr` p)
1034 return $
1035 if loc == nullPtr
1036 then (arr, Nothing)
1037 else
1038 ( Array
1039 { aStart = aStart
1040 , aEnd = loc
1041 , aBound = loc
1043 , Just $ Array
1044 { aStart = aStart `plusForeignPtr` (loc `minusPtr` p + 1)
1045 , aEnd = aEnd
1046 , aBound = aBound
1050data SplitState s arr
1051 = Initial s
1052 | Buffering s arr
1053 | Splitting s arr
1054 | Yielding arr (SplitState s arr)
1055 | Finishing
1057-- | Split a stream of arrays on a given separator byte, dropping the separator
1058-- and coalescing all the arrays between two separators into a single array.
1060-- @since 0.7.0
1061{-# INLINE_NORMAL splitOn #-}
1062splitOn
1063 :: MonadIO m
1064 => Word8
1065 -> D.Stream m (Array Word8)
1066 -> D.Stream m (Array Word8)
1067splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
1069 where
1071 {-# INLINE_LATE step' #-}
1072 step' gst (Initial st) = do
1073 r <- step gst st
1074 case r of
1075 D.Yield arr s -> do
1076 (arr1, marr2) <- breakOn byte arr
1077 return $ case marr2 of
1078 Nothing -> D.Skip (Buffering s arr1)
1079 Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
1080 D.Skip s -> return $ D.Skip (Initial s)
1081 D.Stop -> return $ D.Stop
1083 step' gst (Buffering st buf) = do
1084 r <- step gst st
1085 case r of
1086 D.Yield arr s -> do
1087 (arr1, marr2) <- breakOn byte arr
1088 buf' <- spliceTwo buf arr1
1089 return $ case marr2 of
1090 Nothing -> D.Skip (Buffering s buf')
1091 Just x -> D.Skip (Yielding buf' (Splitting s x))
1092 D.Skip s -> return $ D.Skip (Buffering s buf)
1093 D.Stop -> return $
1094 if byteLength buf == 0
1095 then D.Stop
1096 else D.Skip (Yielding buf Finishing)
1098 step' _ (Splitting st buf) = do
1099 (arr1, marr2) <- breakOn byte buf
1100 return $ case marr2 of
1101 Nothing -> D.Skip $ Buffering st arr1
1102 Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
1104 step' _ (Yielding arr next) = return $ D.Yield arr next
1105 step' _ Finishing = return $ D.Stop