2{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE ExistentialQuantification #-}
4{-# LANGUAGE MagicHash #-}
5{-# LANGUAGE RecordWildCards #-}
6{-# LANGUAGE ScopedTypeVariables #-}
7{-# LANGUAGE TypeFamilies #-}
8{-# LANGUAGE UnboxedTuples #-}
13-- Module : Streamly.Array
14-- Copyright : (c) 2019 Composewell Technologies
17-- Maintainer : harendra.kumar@gmail.com
18-- Stability : experimental
21module Streamly.Array.Types
60import Control.Exception (assert)
61import Control.DeepSeq (NFData(..))
62import Control.Monad (when)
63import Control.Monad.IO.Class (MonadIO(..))
64import Data.Functor.Identity (runIdentity)
65import Data.Word (Word8)
66import Foreign.C.String (CString)
67import Foreign.C.Types (CSize(..), CInt(..))
68import Foreign.ForeignPtr
69 (ForeignPtr, withForeignPtr, touchForeignPtr)
70import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
71import Foreign.Ptr (plusPtr, minusPtr, castPtr)
72import Foreign.Storable (Storable(..))
73import Prelude hiding (length, foldr)
74import Text.Read (readPrec, readListPrec, readListPrecDefault)
76import GHC.Base (Addr#, realWorld#)
77import GHC.Exts (IsList, IsString(..))
78import GHC.ForeignPtr (mallocPlainForeignPtrAlignedBytes, newForeignPtr_)
79import GHC.IO (IO(IO), unsafePerformIO)
80import GHC.Ptr (Ptr(..))
82import Streamly.Fold.Types (Fold(..))
83import Streamly.SVar (adaptState)
85import qualified Streamly.Streams.StreamD.Type as D
86import qualified Streamly.Streams.StreamK as K
87import qualified GHC.Exts as Exts
90import qualified Data.Foldable as F
93-------------------------------------------------------------------------------
95-------------------------------------------------------------------------------
97-- There are two goals that we want to fulfill with arrays. One, holding large
98-- amounts of data in non-GC memory and the ability to pass this data to and
99-- from the operating system without an extra copy overhead. Two, allow random
100-- access to elements based on their indices. The first one falls in the
101-- category of storage buffers while the second one falls in the category of
102-- maps/multisets/hashmaps.
104-- For the first requirement we use an array of Storables and store it in a
105-- ForeignPtr. We can have both immutable and mutable variants of this array
106-- using wrappers over the same underlying type.
108-- For the second requirement, we need a separate type for arrays of
109-- polymorphic values, for example vectors of handler functions, lookup tables.
110-- We can call this type a "vector" in contrast to arrays. It should not
111-- require Storable instance for the type. In that case we need to use an
112-- Array# instead of a ForeignPtr. This type of array would not reduce the GC
113-- overhead much because each element of the array still needs to be scanned by
114-- the GC. However, it can store polymorphic elements and allow random access
115-- to those. But in most cases random access means storage, and it means we
116-- need to avoid GC scanning except in cases of trivially small storage. One
117-- way to achieve that would be to put the array in a Compact region. However,
118-- if and when we mutate this, we will have to use a manual GC copying out to
119-- another CR and freeing the old one.
121-------------------------------------------------------------------------------
123-------------------------------------------------------------------------------
125-- We require that an array stores only Storable. Arrays are used for buffering
126-- data while streams are used for processing. If you want something to be
127-- buffered it better be Storable so that we can store it in non-GC memory.
129-- We can use a "Storable" constraint in the Array type and the Constraint can
130-- be automatically provided to a function that pattern matches on the Array
131-- type. However, it has huge performance cost, so we do not use it.
133-- XXX Another way to not require the Storable constraint in array operations
134-- is to store the elemSize in the array at construction and use that instead
135-- of using sizeOf. Need to charaterize perf cost of this.
137-- XXX rename the fields to "start, next, end".
144 { aStart :: {-# UNPACK #-} !(ForeignPtr a) -- first address
145 , aEnd :: {-# UNPACK #-} !(Ptr a) -- first unused address
146 , aBound :: {-# UNPACK #-} !(Ptr a) -- first address beyond allocated memory
149-------------------------------------------------------------------------------
151-------------------------------------------------------------------------------
153foreign import ccall unsafe "string.h memcpy" c_memcpy
154 :: Ptr Word8 -> Ptr Word8 -> CSize -> IO (Ptr Word8)
156foreign import ccall unsafe "string.h strlen" c_strlen
157 :: CString -> IO CSize
159-- XXX we are converting Int to CSize
160memcpy :: Ptr Word8 -> Ptr Word8 -> Int -> IO ()
161memcpy dst src len = c_memcpy dst src (fromIntegral len) >> return ()
163foreign import ccall unsafe "string.h memcmp" c_memcmp
164 :: Ptr Word8 -> Ptr Word8 -> CSize -> IO CInt
166-- XXX we are converting Int to CSize
167-- return True if the memory locations have identical contents
169memcmp :: Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
171 r <- c_memcmp p1 p2 (fromIntegral len)
174{-# INLINE unsafeInlineIO #-}
175unsafeInlineIO :: IO a -> a
176unsafeInlineIO (IO m) = case m realWorld# of (# _, r #) -> r
178-------------------------------------------------------------------------------
180-------------------------------------------------------------------------------
182-- | Allocate an array that can hold 'count' items. The memory of the array is
185-- Note that this is internal routine, the reference to this array cannot be
186-- given out until the array has been written to and frozen.
187{-# INLINE newArray #-}
188newArray :: forall a. Storable a => Int -> IO (Array a)
190 let size = count * sizeOf (undefined :: a)
191 fptr <- mallocPlainForeignPtrAlignedBytes size (alignment (undefined :: a))
192 let p = unsafeForeignPtrToPtr fptr
196 , aBound = p `plusPtr` size
199-- | Allocate an Array of the given size and run an IO action passing the array
201{-# INLINE withNewArray #-}
202withNewArray :: forall a. Storable a => Int -> (Ptr a -> IO ()) -> IO (Array a)
203withNewArray count f = do
204 arr <- newArray count
205 withForeignPtr (aStart arr) $ \p -> f p >> return arr
207-- XXX grow the array when we are beyond bound.
209-- Internal routine for when the array is being created. Appends one item at
210-- the end of the array. Useful when sequentially writing a stream to the
211-- array. DOES NOT CHECK THE ARRAY BOUNDS.
212{-# INLINE unsafeAppend #-}
213unsafeAppend :: forall a. Storable a => Array a -> a -> IO (Array a)
214unsafeAppend arr@Array{..} x = do
215 when (aEnd == aBound) $
216 error "BUG: unsafeAppend: writing beyond array bounds"
218 touchForeignPtr aStart
219 return $ arr {aEnd = aEnd `plusPtr` (sizeOf (undefined :: a))}
221-- | Remove the free space from an Array.
222shrinkToFit :: forall a. Storable a => Array a -> IO (Array a)
223shrinkToFit arr@Array{..} = do
224 assert (aEnd <= aBound) (return ())
227 let oldStart = unsafeForeignPtrToPtr aStart
228 let size = aEnd `minusPtr` oldStart
229 newPtr <- mallocPlainForeignPtrAlignedBytes
230 size (alignment (undefined :: a))
231 withForeignPtr newPtr $ \pNew -> do
232 memcpy (castPtr pNew) (castPtr oldStart) size
233 touchForeignPtr aStart
234 let end = pNew `plusPtr` size
242-- XXX when converting an array of Word8 from a literal string we can simply
243-- refer to the literal string. Is it possible to write rules such that
244-- fromList Word8 can be rewritten so that GHC does not first convert the
245-- literal to [Char] and then we convert it back to an Array Word8?
247-- Note that the address must be a read-only address (meant to be used for
248-- read-only string literals) because we are sharing it, any modification to
249-- the original address would change our array. That's why this function is
251{-# INLINE _fromCStringAddrUnsafe #-}
252_fromCStringAddrUnsafe :: Addr# -> IO (Array Word8)
253_fromCStringAddrUnsafe addr# = do
254 ptr <- newForeignPtr_ (castPtr cstr)
256 let n = fromIntegral len
257 let p = unsafeForeignPtrToPtr ptr
258 let end = p `plusPtr` n
268-------------------------------------------------------------------------------
270-------------------------------------------------------------------------------
272-- | Return element at the specified index without checking the bounds.
274-- Unsafe because it does not check the bounds of the array.
275{-# INLINE_NORMAL unsafeIndexIO #-}
276unsafeIndexIO :: forall a. Storable a => Array a -> Int -> IO a
277unsafeIndexIO Array {..} i =
278 withForeignPtr aStart $ \p -> do
279 let elemSize = sizeOf (undefined :: a)
280 elemOff = p `plusPtr` (elemSize * i)
281 assert (i >= 0 && elemOff `plusPtr` elemSize <= aEnd)
285-- | Return element at the specified index without checking the bounds.
286{-# INLINE_NORMAL unsafeIndex #-}
287unsafeIndex :: forall a. Storable a => Array a -> Int -> a
288unsafeIndex arr i = let !r = unsafeInlineIO $ unsafeIndexIO arr i in r
290-- | /O(1)/ Get the length of the array.
294length :: forall a. Storable a => Array a -> Int
296 let p = unsafeForeignPtrToPtr aStart
297 aLen = aEnd `minusPtr` p
298 in assert (aLen >= 0) (aLen `div` sizeOf (undefined :: a))
300{-# INLINE_NORMAL toStreamD #-}
301toStreamD :: forall m a. (Monad m, Storable a) => Array a -> D.Stream m a
303 let p = unsafeForeignPtrToPtr aStart
308 {-# INLINE_LATE step #-}
309 step _ p | p == aEnd = return D.Stop
311 -- unsafeInlineIO allows us to run this in Identity monad for pure
312 -- toList/foldr case which makes them much faster due to not
313 -- accumulating the list and fusing better with the pure consumers.
315 -- This should be safe as the array contents are guaranteed to be
316 -- evaluated/written to before we peek at them.
317 let !x = unsafeInlineIO $ do
319 touchForeignPtr aStart
321 return $ D.Yield x (p `plusPtr` (sizeOf (undefined :: a)))
323{-# INLINE_NORMAL toStreamDRev #-}
324toStreamDRev :: forall m a. (Monad m, Storable a) => Array a -> D.Stream m a
325toStreamDRev Array{..} =
326 let p = aEnd `plusPtr` negate (sizeOf (undefined :: a))
331 {-# INLINE_LATE step #-}
332 step _ p | p < unsafeForeignPtrToPtr aStart = return D.Stop
334 -- See comments in toStreamD for why we use unsafeInlineIO
335 let !x = unsafeInlineIO $ do
337 touchForeignPtr aStart
339 return $ D.Yield x (p `plusPtr` negate (sizeOf (undefined :: a)))
341{-# INLINE_NORMAL foldl' #-}
342foldl' :: forall a b. Storable a => (b -> a -> b) -> b -> Array a -> b
343foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
345{-# INLINE_NORMAL foldr #-}
346foldr :: Storable a => (a -> b -> b) -> b -> Array a -> b
347foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
349-------------------------------------------------------------------------------
351-------------------------------------------------------------------------------
353-- | @toArrayN n@ folds a maximum of @n@ elements from the input stream to an
357{-# INLINE_NORMAL toArrayN #-}
358toArrayN :: forall m a. (MonadIO m, Storable a) => Int -> Fold m a (Array a)
359toArrayN n = Fold step initial extract
363 initial = liftIO $ newArray n
364 step arr@(Array _ end bound) _ | end == bound = return arr
365 step (Array start end bound) x = do
367 return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound
371-- | Fold the input to a pure buffered stream (List) of arrays.
372{-# INLINE toArrays #-}
373toArrays :: (Monad m, Storable a) => Int -> Fold m a (List (Array a))
374toArrays n = Fold step initial extract
376-- This can be implemented by combining the List of arrays from toArrays into a
378-- | Fold the whole input to a single array.
379{-# INLINE toArray #-}
380toArray :: forall m a. (Monad m, Storable a) => Fold m a (Array a)
381toArray = Fold step begin done
384{-# INLINE_NORMAL fromStreamDN #-}
385fromStreamDN :: forall m a. (MonadIO m, Storable a)
386 => Int -> D.Stream m a -> m (Array a)
387fromStreamDN limit str = do
388 arr <- liftIO $ newArray limit
389 end <- D.foldlM' write (aEnd arr) $ D.take limit str
390 return $ arr {aEnd = end}
396 return $ ptr `plusPtr` sizeOf (undefined :: a)
398-- | @fromStreamArraysOf n stream@ groups the input stream into a stream of
400{-# INLINE fromStreamDArraysOf #-}
401fromStreamDArraysOf :: (MonadIO m, Storable a)
402 => Int -> D.Stream m a -> D.Stream m (Array a)
403fromStreamDArraysOf n str = D.groupsOf n (toArrayN n) str
405-- XXX concatMap does not seem to have the best possible performance so we have
406-- a custom way to concat arrays.
407data FlattenState s a =
409 | InnerLoop s (ForeignPtr a) (Ptr a) (Ptr a)
411{-# INLINE_NORMAL flattenArrays #-}
412flattenArrays :: forall m a. (MonadIO m, Storable a)
413 => D.Stream m (Array a) -> D.Stream m a
414flattenArrays (D.Stream step state) = D.Stream step' (OuterLoop state)
418 {-# INLINE_LATE step' #-}
419 step' gst (OuterLoop st) = do
420 r <- step (adaptState gst) st
422 D.Yield Array{..} s ->
423 let p = unsafeForeignPtrToPtr aStart
424 in D.Skip (InnerLoop s aStart p aEnd)
425 D.Skip s -> D.Skip (OuterLoop s)
428 step' _ (InnerLoop st _ p end) | p == end =
429 return $ D.Skip $ OuterLoop st
431 step' _ (InnerLoop st startf p end) = do
434 touchForeignPtr startf
436 return $ D.Yield x (InnerLoop st startf
437 (p `plusPtr` (sizeOf (undefined :: a))) end)
439{-# INLINE_NORMAL flattenArraysRev #-}
440flattenArraysRev :: forall m a. (MonadIO m, Storable a)
441 => D.Stream m (Array a) -> D.Stream m a
442flattenArraysRev (D.Stream step state) = D.Stream step' (OuterLoop state)
446 {-# INLINE_LATE step' #-}
447 step' gst (OuterLoop st) = do
448 r <- step (adaptState gst) st
450 D.Yield Array{..} s ->
451 let p = aEnd `plusPtr` negate (sizeOf (undefined :: a))
452 -- XXX we do not need aEnd
453 in D.Skip (InnerLoop s aStart p aEnd)
454 D.Skip s -> D.Skip (OuterLoop s)
457 step' _ (InnerLoop st start p _) | p < unsafeForeignPtrToPtr start =
458 return $ D.Skip $ OuterLoop st
460 step' _ (InnerLoop st startf p end) = do
463 touchForeignPtr startf
465 return $ D.Yield x (InnerLoop st startf
466 (p `plusPtr` negate (sizeOf (undefined :: a))) end)
468-- CAUTION: a very large number (millions) of arrays can degrade performance
469-- due to GC overhead because we need to buffer the arrays before we flatten
472{-# INLINE fromStreamD #-}
473fromStreamD :: (MonadIO m, Storable a) => D.Stream m a -> m (Array a)
475 let s = fromStreamDArraysOf defaultChunkSize m
476 buffered <- D.foldr K.cons K.nil s
477 len <- K.foldl' (+) 0 (K.map length buffered)
478 fromStreamDN len $ flattenArrays $ D.fromStreamK buffered
480-- | Convert an 'Array' into a list.
483{-# INLINABLE toList #-}
484toList :: Storable a => Array a -> [a]
487instance (Show a, Storable a) => Show (Array a) where
488 {-# INLINE showsPrec #-}
489 showsPrec _ = shows . toList
491-- | Create an 'Array' from the first N elements of a list. The array is
492-- allocated to size N, if the list terminates before N elements then the
493-- array may hold less than N elements.
496{-# INLINABLE fromListN #-}
497fromListN :: Storable a => Int -> [a] -> Array a
498fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
500-- | Create an 'Array' from a list. The list must be of finite size.
503{-# INLINABLE fromList #-}
504fromList :: Storable a => [a] -> Array a
505fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
507instance (Storable a, Read a, Show a) => Read (Array a) where
508 {-# INLINE readPrec #-}
512 readListPrec = readListPrecDefault
514instance (a ~ Char) => IsString (Array a) where
515 {-# INLINE fromString #-}
516 fromString = fromList
518-- GHC versions 8.0 and below cannot derive IsList
519instance Storable a => IsList (Array a) where
520 type (Item (Array a)) = a
521 {-# INLINE fromList #-}
523 {-# INLINE fromListN #-}
524 fromListN = fromListN
525 {-# INLINE toList #-}
529arrcmp :: Array a -> Array a -> Bool
531 let !res = unsafeInlineIO $ do
532 let ptr1 = unsafeForeignPtrToPtr $ aStart arr1
533 let ptr2 = unsafeForeignPtrToPtr $ aStart arr2
534 let len1 = aEnd arr1 `minusPtr` ptr1
535 let len2 = aEnd arr2 `minusPtr` ptr2
539 r <- memcmp (castPtr ptr1) (castPtr ptr2) len1
540 touchForeignPtr $ aStart arr1
541 touchForeignPtr $ aStart arr2
546-- XXX we are assuming that Storable equality means element equality. This may
547-- or may not be correct? arrcmp is 40% faster compared to stream equality.
548instance (Storable a, Eq a) => Eq (Array a) where
551 -- arr1 == arr2 = runIdentity $ D.eqBy (==) (toStreamD arr1) (toStreamD arr2)
553instance (Storable a, NFData a) => NFData (Array a) where
555 rnf = foldl' (\_ x -> rnf x) ()
557instance (Storable a, Ord a) => Ord (Array a) where
558 {-# INLINE compare #-}
559 compare arr1 arr2 = unsafePerformIO $
560 D.cmpBy compare (toStreamD arr1) (toStreamD arr2)
562 -- Default definitions defined in base do not have an INLINE on them, so we
563 -- replicate them here with an INLINE.
565 x < y = case compare x y of { LT -> True; _ -> False }
568 x <= y = case compare x y of { GT -> False; _ -> True }
571 x > y = case compare x y of { GT -> True; _ -> False }
574 x >= y = case compare x y of { LT -> False; _ -> True }
576 -- These two default methods use '<=' rather than 'compare'
577 -- because the latter is often more expensive
579 max x y = if x <= y then y else x
582 min x y = if x <= y then x else y
585-- Definitions using the Storable constraint from the Array type. These are to
586-- make the Foldable instance possible though it is much slower (7x slower).
588{-# INLINE_NORMAL toStreamD_ #-}
589toStreamD_ :: forall m a. MonadIO m => Int -> Array a -> D.Stream m a
590toStreamD_ size Array{..} =
591 let p = unsafeForeignPtrToPtr aStart
596 {-# INLINE_LATE step #-}
597 step _ p | p == aEnd = return D.Stop
601 touchForeignPtr aStart
603 return $ D.Yield x (p `plusPtr` size)
605{-# INLINE_NORMAL _foldr #-}
606_foldr :: forall a b. (a -> b -> b) -> b -> Array a -> b
607_foldr f z arr@Array {..} =
608 let !n = sizeOf (undefined :: a)
609 in unsafePerformIO $ D.foldr f z $ toStreamD_ n arr
611-- | Note that the 'Foldable' instance is 7x slower than the direct
613instance Foldable Array where
617-------------------------------------------------------------------------------
618-- Semigroup and Monoid
619-------------------------------------------------------------------------------
621-------------------------------------------------------------------------------
623-------------------------------------------------------------------------------
625-- | GHC memory management allocation header overhead
627allocOverhead = 2 * sizeOf (undefined :: Int)
629-- | Default maximum buffer size in bytes, for reading from and writing to IO
630-- devices, the value is 32KB minus GHC allocation overhead, which is a few
631-- bytes, so that the actual allocation is 32KB.
632defaultChunkSize :: Int
633defaultChunkSize = 32 * k - allocOverhead
636mkChunkSizeKB :: Int -> Int
637mkChunkSizeKB n = n * k - allocOverhead
keysesc to close
- j k
- move down, up
- enter
- open
- esc
- back, or up a level
- 1 2 3 4
- readme, files, commits, refs
- /
- filter
- y
- copy the clone command
- c
- compare the chosen commit with another (or shift-click both)
- ?
- show or hide this list