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
59import Control.Exception (assert)
60import Control.DeepSeq (NFData(..))
61import Control.Monad (when)
62import Data.Functor.Identity (runIdentity)
63import Data.Word (Word8)
64import Foreign.C.String (CString)
65import Foreign.C.Types (CSize(..), CInt(..))
66import Foreign.ForeignPtr
67 (ForeignPtr, withForeignPtr, touchForeignPtr)
68import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
69import Foreign.Ptr (plusPtr, minusPtr, castPtr)
70import Foreign.Storable (Storable(..))
71import Prelude hiding (length, foldr)
72import Text.Read (readPrec, readListPrec, readListPrecDefault)
74import GHC.Base (Addr#, realWorld#)
75import GHC.Exts (IsList, IsString(..))
76import GHC.ForeignPtr (mallocPlainForeignPtrAlignedBytes, newForeignPtr_)
77import GHC.IO (IO(IO), unsafeDupablePerformIO)
78import GHC.Ptr (Ptr(..))
80import Streamly.Fold.Types (Fold(..))
81import Streamly.SVar (adaptState)
83import qualified Streamly.Streams.StreamD.Type as D
84import qualified Streamly.Streams.StreamK as K
85import qualified GHC.Exts as Exts
87-------------------------------------------------------------------------------
89-------------------------------------------------------------------------------
91-- There are two goals that we want to fulfill with arrays. One, holding large
92-- amounts of data in non-GC memory and the ability to pass this data to and
93-- from the operating system without an extra copy overhead. Two, allow random
94-- access to elements based on their indices. The first one falls in the
95-- category of storage buffers while the second one falls in the category of
96-- maps/multisets/hashmaps.
98-- For the first requirement we use an array of Storables and store it in a
99-- ForeignPtr. We can have both immutable and mutable variants of this array
100-- using wrappers over the same underlying type.
102-- For the second requirement, we need a separate type for arrays of
103-- polymorphic values, for example vectors of handler functions, lookup tables.
104-- We can call this type a "vector" in contrast to arrays. It should not
105-- require Storable instance for the type. In that case we need to use an
106-- Array# instead of a ForeignPtr. This type of array would not reduce the GC
107-- overhead much because each element of the array still needs to be scanned by
108-- the GC. However, it can store polymorphic elements and allow random access
109-- to those. But in most cases random access means storage, and it means we
110-- need to avoid GC scanning except in cases of trivially small storage. One
111-- way to achieve that would be to put the array in a Compact region. However,
112-- if and when we mutate this, we will have to use a manual GC copying out to
113-- another CR and freeing the old one.
115-------------------------------------------------------------------------------
117-------------------------------------------------------------------------------
119-- We require that an array stores only Storable. Arrays are used for buffering
120-- data while streams are used for processing. If you want something to be
121-- buffered it better be Storable so that we can store it in non-GC memory.
123-- We can use a "Storable" constraint in the Array type and the Constraint can
124-- be automatically provided to a function that pattern matches on the Array
125-- type. However, it has huge performance cost, so we do not use it.
127-- XXX Another way to not require the Storable constraint in array operations
128-- is to store the elemSize in the array at construction and use that instead
129-- of using sizeOf. Need to charaterize perf cost of this.
131-- XXX rename the fields to "start, next, end".
138 { aStart :: {-# UNPACK #-} !(ForeignPtr a) -- first address
139 , aEnd :: {-# UNPACK #-} !(Ptr a) -- first unused address
140 , aBound :: {-# UNPACK #-} !(Ptr a) -- first address beyond allocated memory
143-------------------------------------------------------------------------------
145-------------------------------------------------------------------------------
147foreign import ccall unsafe "string.h memcpy" c_memcpy
148 :: Ptr Word8 -> Ptr Word8 -> CSize -> IO (Ptr Word8)
150foreign import ccall unsafe "string.h strlen" c_strlen
151 :: CString -> IO CSize
153-- XXX we are converting Int to CSize
154memcpy :: Ptr Word8 -> Ptr Word8 -> Int -> IO ()
155memcpy dst src len = c_memcpy dst src (fromIntegral len) >> return ()
157foreign import ccall unsafe "string.h memcmp" c_memcmp
158 :: Ptr Word8 -> Ptr Word8 -> CSize -> IO CInt
160-- XXX we are converting Int to CSize
161-- return True if the memory locations have identical contents
163memcmp :: Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
165 r <- c_memcmp p1 p2 (fromIntegral len)
168{-# INLINE unsafeInlineIO #-}
169unsafeInlineIO :: IO a -> a
170unsafeInlineIO (IO m) = case m realWorld# of (# _, r #) -> r
172-------------------------------------------------------------------------------
174-------------------------------------------------------------------------------
176-- | Allocate an array that can hold 'count' items. The memory of the array is
179-- Note that this is internal routine, the reference to this array cannot be
180-- given out until the array has been written to and frozen.
181{-# INLINE unsafeNew #-}
182unsafeNew :: forall a. Storable a => Int -> IO (Array a)
184 let size = count * sizeOf (undefined :: a)
185 fptr <- mallocPlainForeignPtrAlignedBytes size (alignment (undefined :: a))
186 let p = unsafeForeignPtrToPtr fptr
190 , aBound = p `plusPtr` size
193-- | Allocate an Array of the given size and run an IO action passing the array
195{-# INLINE withNewArray #-}
196withNewArray :: forall a. Storable a => Int -> (Ptr a -> IO ()) -> IO (Array a)
197withNewArray count f = do
198 arr <- unsafeNew count
199 withForeignPtr (aStart arr) $ \p -> f p >> return arr
201-- XXX grow the array when we are beyond bound.
203-- Internal routine for when the array is being created. Appends one item at
204-- the end of the array. Useful when sequentially writing a stream to the
205-- array. DOES NOT CHECK THE ARRAY BOUNDS.
206{-# INLINE unsafeAppend #-}
207unsafeAppend :: forall a. Storable a => Array a -> a -> IO (Array a)
208unsafeAppend arr@Array{..} x = do
209 when (aEnd == aBound) $
210 error "BUG: unsafeAppend: writing beyond array bounds"
212 touchForeignPtr aStart
213 return $ arr {aEnd = aEnd `plusPtr` (sizeOf (undefined :: a))}
215-- | Remove the free space from an Array.
216shrinkToFit :: forall a. Storable a => Array a -> IO (Array a)
217shrinkToFit arr@Array{..} = do
218 assert (aEnd <= aBound) (return ())
221 let oldStart = unsafeForeignPtrToPtr aStart
222 let size = aEnd `minusPtr` oldStart
223 newPtr <- mallocPlainForeignPtrAlignedBytes
224 size (alignment (undefined :: a))
225 withForeignPtr newPtr $ \pNew -> do
226 memcpy (castPtr pNew) (castPtr oldStart) size
227 touchForeignPtr aStart
228 let end = pNew `plusPtr` size
236-- XXX when converting an array of Word8 from a literal string we can simply
237-- refer to the literal string. Is it possible to write rules such that
238-- fromList Word8 can be rewritten so that GHC does not first convert the
239-- literal to [Char] and then we convert it back to an Array Word8?
241-- Note that the address must be a read-only address (meant to be used for
242-- read-only string literals) because we are sharing it, any modification to
243-- the original address would change our array. That's why this function is
245{-# INLINE _fromCStringAddrUnsafe #-}
246_fromCStringAddrUnsafe :: Addr# -> IO (Array Word8)
247_fromCStringAddrUnsafe addr# = do
248 ptr <- newForeignPtr_ (castPtr cstr)
250 let n = fromIntegral len
251 let p = unsafeForeignPtrToPtr ptr
252 let end = p `plusPtr` n
262-------------------------------------------------------------------------------
264-------------------------------------------------------------------------------
266-- | Return element at the specified index without checking the bounds.
267{-# INLINE_NORMAL unsafeIndex #-}
268unsafeIndex :: forall a. Storable a => Array a -> Int -> a
269unsafeIndex Array {..} i =
270 let !r = unsafeInlineIO $
271 withForeignPtr aStart $ \p -> do
272 let elemSize = sizeOf (undefined :: a)
273 elemOff = p `plusPtr` (elemSize * i)
274 assert (i >= 0 && elemOff `plusPtr` elemSize <= aEnd)
279-- | /O(1)/ Get the length of the array.
283length :: forall a. Storable a => Array a -> Int
285 let p = unsafeForeignPtrToPtr aStart
286 aLen = aEnd `minusPtr` p
287 in assert (aLen >= 0) (aLen `div` sizeOf (undefined :: a))
289{-# INLINE_NORMAL toStreamD #-}
290toStreamD :: forall m a. (Monad m, Storable a) => Array a -> D.Stream m a
292 let p = unsafeForeignPtrToPtr aStart
297 {-# INLINE_LATE step #-}
298 step _ p | p == aEnd = return D.Stop
300 let !x = unsafeInlineIO $ do
302 touchForeignPtr aStart
304 return $ D.Yield x (p `plusPtr` (sizeOf (undefined :: a)))
306{-# INLINE_NORMAL toStreamDRev #-}
307toStreamDRev :: forall m a. (Monad m, Storable a) => Array a -> D.Stream m a
308toStreamDRev Array{..} =
309 let p = aEnd `plusPtr` negate (sizeOf (undefined :: a))
314 {-# INLINE_LATE step #-}
315 step _ p | p < unsafeForeignPtrToPtr aStart = return D.Stop
317 let !x = unsafeInlineIO $ do
319 touchForeignPtr aStart
321 return $ D.Yield x (p `plusPtr` negate (sizeOf (undefined :: a)))
323{-# INLINE_NORMAL foldl' #-}
324foldl' :: forall a b. Storable a => (b -> a -> b) -> b -> Array a -> b
325foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
327{-# INLINE_NORMAL foldr #-}
328foldr :: Storable a => (a -> b -> b) -> b -> Array a -> b
329foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
331-------------------------------------------------------------------------------
333-------------------------------------------------------------------------------
335-- | @toArrayN n@ folds a maximum of @n@ elements from the input stream to an
339{-# INLINE_NORMAL toArrayN #-}
340toArrayN :: forall m a. (Monad m, Storable a) => Int -> Fold m a (Array a)
341toArrayN n = Fold step initial extract
346 let !arr = unsafeDupablePerformIO $ unsafeNew n
348 step arr@(Array _ end bound) _ | end == bound = return arr
349 step (Array start end bound) x =
350 let !_ = unsafeInlineIO $ poke end x
351 in return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound
355-- | Fold the input to a pure buffered stream (List) of arrays.
356{-# INLINE toArrays #-}
357toArrays :: (Monad m, Storable a) => Int -> Fold m a (List (Array a))
358toArrays n = Fold step initial extract
360-- This can be implemented by combining the List of arrays from toArrays into a
362-- | Fold the whole input to a single array.
363{-# INLINE toArray #-}
364toArray :: forall m a. (Monad m, Storable a) => Fold m a (Array a)
365toArray = Fold step begin done
368{-# INLINE_NORMAL fromStreamDN #-}
369fromStreamDN :: forall m a. (Monad m, Storable a)
370 => Int -> D.Stream m a -> m (Array a)
371fromStreamDN limit str = do
372 let !arr = unsafeDupablePerformIO $ unsafeNew limit
373 end <- D.foldl' write (aEnd arr) $ D.take limit str
374 return $ arr {aEnd = end}
379 let !_ = unsafeInlineIO $ poke ptr x
380 in ptr `plusPtr` sizeOf (undefined :: a)
382-- | @fromStreamArraysOf n stream@ groups the input stream into a stream of
384{-# INLINE fromStreamDArraysOf #-}
385fromStreamDArraysOf :: (Monad m, Storable a)
386 => Int -> D.Stream m a -> D.Stream m (Array a)
387fromStreamDArraysOf n str = D.groupsOf n (toArrayN n) str
389-- XXX concatMap does not seem to have the best possible performance so we have
390-- a custom way to concat arrays.
391data CAState s a = CAParent s | CANested s (ForeignPtr a) (Ptr a) (Ptr a)
393{-# INLINE_NORMAL flattenArrays #-}
394flattenArrays :: forall m a. (Monad m, Storable a)
395 => D.Stream m (Array a) -> D.Stream m a
396flattenArrays (D.Stream step state) = D.Stream step' (CAParent state)
400 {-# INLINE_LATE step' #-}
401 step' gst (CAParent st) = do
402 r <- step (adaptState gst) st
404 D.Yield Array{..} s ->
405 let p = unsafeForeignPtrToPtr aStart
406 in D.Skip (CANested s aStart p aEnd)
407 D.Skip s -> D.Skip (CAParent s)
410 step' _ (CANested st _ p end) | p == end =
411 return $ D.Skip $ CAParent st
413 step' _ (CANested st startf p end) = do
414 let !x = unsafeInlineIO $ do
416 touchForeignPtr startf
418 return $ D.Yield x (CANested st startf
419 (p `plusPtr` (sizeOf (undefined :: a))) end)
421{-# INLINE_NORMAL flattenArraysRev #-}
422flattenArraysRev :: forall m a. (Monad m, Storable a)
423 => D.Stream m (Array a) -> D.Stream m a
424flattenArraysRev (D.Stream step state) = D.Stream step' (CAParent state)
428 {-# INLINE_LATE step' #-}
429 step' gst (CAParent st) = do
430 r <- step (adaptState gst) st
432 D.Yield Array{..} s ->
433 let p = aEnd `plusPtr` negate (sizeOf (undefined :: a))
434 -- XXX we do not need aEnd
435 in D.Skip (CANested s aStart p aEnd)
436 D.Skip s -> D.Skip (CAParent s)
439 step' _ (CANested st start p _) | p < unsafeForeignPtrToPtr start =
440 return $ D.Skip $ CAParent st
442 step' _ (CANested st startf p end) = do
443 let !x = unsafeInlineIO $ do
445 touchForeignPtr startf
447 return $ D.Yield x (CANested st startf
448 (p `plusPtr` negate (sizeOf (undefined :: a))) end)
450-- CAUTION: a very large number (millions) of arrays can degrade performance
451-- due to GC overhead because we need to buffer the arrays before we flatten
454{-# INLINE fromStreamD #-}
455fromStreamD :: (Monad m, Storable a) => D.Stream m a -> m (Array a)
457 let s = fromStreamDArraysOf defaultChunkSize m
458 buffered <- D.foldr K.cons K.nil s
459 len <- K.foldl' (+) 0 (K.map length buffered)
460 fromStreamDN len $ flattenArrays $ D.fromStreamK buffered
462-- | Convert an 'Array' into a list.
465{-# INLINABLE toList #-}
466toList :: Storable a => Array a -> [a]
469instance (Show a, Storable a) => Show (Array a) where
470 {-# INLINE showsPrec #-}
471 showsPrec _ = shows . toList
473-- | Create an 'Array' from the first N elements of a list. The array is
474-- allocated to size N, if the list terminates before N elements then the
475-- array may hold less than N elements.
478{-# INLINABLE fromListN #-}
479fromListN :: Storable a => Int -> [a] -> Array a
480fromListN n xs = runIdentity $ fromStreamDN n $ D.fromList xs
482-- | Create an 'Array' from a list. The list must be of finite size.
485{-# INLINABLE fromList #-}
486fromList :: Storable a => [a] -> Array a
487fromList xs = runIdentity $ fromStreamD $ D.fromList xs
489instance (Storable a, Read a, Show a) => Read (Array a) where
490 {-# INLINE readPrec #-}
494 readListPrec = readListPrecDefault
496instance (a ~ Char) => IsString (Array a) where
497 {-# INLINE fromString #-}
498 fromString = fromList
500-- GHC versions 8.0 and below cannot derive IsList
501instance Storable a => IsList (Array a) where
502 type (Item (Array a)) = a
503 {-# INLINE fromList #-}
505 {-# INLINE fromListN #-}
506 fromListN = fromListN
507 {-# INLINE toList #-}
511arrcmp :: Array a -> Array a -> Bool
513 let !res = unsafeInlineIO $ do
514 let ptr1 = unsafeForeignPtrToPtr $ aStart arr1
515 let ptr2 = unsafeForeignPtrToPtr $ aStart arr2
516 let len1 = aEnd arr1 `minusPtr` ptr1
517 let len2 = aEnd arr2 `minusPtr` ptr2
521 r <- memcmp (castPtr ptr1) (castPtr ptr2) len1
522 touchForeignPtr $ aStart arr1
523 touchForeignPtr $ aStart arr2
528-- XXX we are assuming that Storable equality means element equality. This may
529-- or may not be correct? arrcmp is 40% faster compared to stream equality.
530instance (Storable a, Eq a) => Eq (Array a) where
533 -- arr1 == arr2 = runIdentity $ D.eqBy (==) (toStreamD arr1) (toStreamD arr2)
535instance (Storable a, NFData a) => NFData (Array a) where
537 rnf = foldl' (\_ x -> rnf x) ()
539instance (Storable a, Ord a) => Ord (Array a) where
540 {-# INLINE compare #-}
541 compare arr1 arr2 = runIdentity $
542 D.cmpBy compare (toStreamD arr1) (toStreamD arr2)
544 -- Default definitions defined in base do not have an INLINE on them, so we
545 -- replicate them here with an INLINE.
547 x < y = case compare x y of { LT -> True; _ -> False }
550 x <= y = case compare x y of { GT -> False; _ -> True }
553 x > y = case compare x y of { GT -> True; _ -> False }
556 x >= y = case compare x y of { LT -> False; _ -> True }
558 -- These two default methods use '<=' rather than 'compare'
559 -- because the latter is often more expensive
561 max x y = if x <= y then y else x
564 min x y = if x <= y then x else y
567-- Definitions using the Storable constraint from the Array type. These are to
568-- make the Foldable instance possible though it is much slower (7x slower).
570{-# INLINE_NORMAL toStreamD_ #-}
571toStreamD_ :: forall m a. Monad m => Int -> Array a -> D.Stream m a
572toStreamD_ size Array{..} =
573 let p = unsafeForeignPtrToPtr aStart
578 {-# INLINE_LATE step #-}
579 step _ p | p == aEnd = return D.Stop
581 let !x = unsafeInlineIO $ do
583 touchForeignPtr aStart
585 return $ D.Yield x (p `plusPtr` size)
587{-# INLINE_NORMAL _foldr #-}
588_foldr :: forall a b. (a -> b -> b) -> b -> Array a -> b
589_foldr f z arr@Array {..} =
590 let !n = sizeOf (undefined :: a)
591 in runIdentity $ D.foldr f z $ toStreamD_ n arr
593-- | Note that the 'Foldable' instance is 7x slower than the direct
595instance Foldable Array where
599-------------------------------------------------------------------------------
600-- Semigroup and Monoid
601-------------------------------------------------------------------------------
603-------------------------------------------------------------------------------
605-------------------------------------------------------------------------------
607-- | GHC memory management allocation header overhead
609allocOverhead = 2 * sizeOf (undefined :: Int)
611-- | Default maximum buffer size in bytes, for reading from and writing to IO
612-- devices, the value is 32KB minus GHC allocation overhead, which is a few
613-- bytes, so that the actual allocation is 32KB.
614defaultChunkSize :: Int
615defaultChunkSize = 32 * k - allocOverhead
618mkChunkSizeKB :: Int -> Int
619mkChunkSizeKB 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