2{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE ExistentialQuantification #-}
4{-# LANGUAGE MagicHash #-}
5{-# LANGUAGE RecordWildCards #-}
6{-# LANGUAGE ScopedTypeVariables #-}
7{-# LANGUAGE TypeFamilies #-}
8{-# LANGUAGE UnboxedTuples #-}
13-- Module : Streamly.Mem.Array.Types
14-- Copyright : (c) 2019 Composewell Technologies
17-- Maintainer : harendra.kumar@gmail.com
18-- Stability : experimental
21module Streamly.Mem.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.
132-- XXX Can this cost be alleviated in GHC-8.10 specialization fix?
134-- XXX Another way to not require the Storable constraint in array operations
135-- is to store the elemSize in the array at construction and use that instead
136-- of using sizeOf. Need to charaterize perf cost of this.
138-- XXX rename the fields to "start, next, end".
145 { aStart :: {-# UNPACK #-} !(ForeignPtr a) -- first address
146 , aEnd :: {-# UNPACK #-} !(Ptr a) -- first unused address
147 , aBound :: {-# UNPACK #-} !(Ptr a) -- first address beyond allocated memory
150-------------------------------------------------------------------------------
152-------------------------------------------------------------------------------
154foreign import ccall unsafe "string.h memcpy" c_memcpy
155 :: Ptr Word8 -> Ptr Word8 -> CSize -> IO (Ptr Word8)
157foreign import ccall unsafe "string.h strlen" c_strlen
158 :: CString -> IO CSize
160-- XXX we are converting Int to CSize
161memcpy :: Ptr Word8 -> Ptr Word8 -> Int -> IO ()
162memcpy dst src len = c_memcpy dst src (fromIntegral len) >> return ()
164foreign import ccall unsafe "string.h memcmp" c_memcmp
165 :: Ptr Word8 -> Ptr Word8 -> CSize -> IO CInt
167-- XXX we are converting Int to CSize
168-- return True if the memory locations have identical contents
170memcmp :: Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
172 r <- c_memcmp p1 p2 (fromIntegral len)
175{-# INLINE unsafeInlineIO #-}
176unsafeInlineIO :: IO a -> a
177unsafeInlineIO (IO m) = case m realWorld# of (# _, r #) -> r
179-------------------------------------------------------------------------------
181-------------------------------------------------------------------------------
183-- | Allocate an array that can hold 'count' items. The memory of the array is
186-- Note that this is internal routine, the reference to this array cannot be
187-- given out until the array has been written to and frozen.
188{-# INLINE newArray #-}
189newArray :: forall a. Storable a => Int -> IO (Array a)
191 let size = count * sizeOf (undefined :: a)
192 fptr <- mallocPlainForeignPtrAlignedBytes size (alignment (undefined :: a))
193 let p = unsafeForeignPtrToPtr fptr
197 , aBound = p `plusPtr` size
200-- | Allocate an Array of the given size and run an IO action passing the array
202{-# INLINE withNewArray #-}
203withNewArray :: forall a. Storable a => Int -> (Ptr a -> IO ()) -> IO (Array a)
204withNewArray count f = do
205 arr <- newArray count
206 withForeignPtr (aStart arr) $ \p -> f p >> return arr
208-- XXX grow the array when we are beyond bound.
210-- Internal routine for when the array is being created. Appends one item at
211-- the end of the array. Useful when sequentially writing a stream to the
212-- array. DOES NOT CHECK THE ARRAY BOUNDS.
213{-# INLINE unsafeAppend #-}
214unsafeAppend :: forall a. Storable a => Array a -> a -> IO (Array a)
215unsafeAppend arr@Array{..} x = do
216 when (aEnd == aBound) $
217 error "BUG: unsafeAppend: writing beyond array bounds"
219 touchForeignPtr aStart
220 return $ arr {aEnd = aEnd `plusPtr` (sizeOf (undefined :: a))}
222-- | Remove the free space from an Array.
223shrinkToFit :: forall a. Storable a => Array a -> IO (Array a)
224shrinkToFit arr@Array{..} = do
225 assert (aEnd <= aBound) (return ())
228 let oldStart = unsafeForeignPtrToPtr aStart
229 let size = aEnd `minusPtr` oldStart
230 newPtr <- mallocPlainForeignPtrAlignedBytes
231 size (alignment (undefined :: a))
232 withForeignPtr newPtr $ \pNew -> do
233 memcpy (castPtr pNew) (castPtr oldStart) size
234 touchForeignPtr aStart
235 let end = pNew `plusPtr` size
243-- XXX when converting an array of Word8 from a literal string we can simply
244-- refer to the literal string. Is it possible to write rules such that
245-- fromList Word8 can be rewritten so that GHC does not first convert the
246-- literal to [Char] and then we convert it back to an Array Word8?
248-- Note that the address must be a read-only address (meant to be used for
249-- read-only string literals) because we are sharing it, any modification to
250-- the original address would change our array. That's why this function is
252{-# INLINE _fromCStringAddrUnsafe #-}
253_fromCStringAddrUnsafe :: Addr# -> IO (Array Word8)
254_fromCStringAddrUnsafe addr# = do
255 ptr <- newForeignPtr_ (castPtr cstr)
257 let n = fromIntegral len
258 let p = unsafeForeignPtrToPtr ptr
259 let end = p `plusPtr` n
269-------------------------------------------------------------------------------
271-------------------------------------------------------------------------------
273-- | Return element at the specified index without checking the bounds.
275-- Unsafe because it does not check the bounds of the array.
276{-# INLINE_NORMAL unsafeIndexIO #-}
277unsafeIndexIO :: forall a. Storable a => Array a -> Int -> IO a
278unsafeIndexIO Array {..} i =
279 withForeignPtr aStart $ \p -> do
280 let elemSize = sizeOf (undefined :: a)
281 elemOff = p `plusPtr` (elemSize * i)
282 assert (i >= 0 && elemOff `plusPtr` elemSize <= aEnd)
286-- | Return element at the specified index without checking the bounds.
287{-# INLINE_NORMAL unsafeIndex #-}
288unsafeIndex :: forall a. Storable a => Array a -> Int -> a
289unsafeIndex arr i = let !r = unsafeInlineIO $ unsafeIndexIO arr i in r
291-- | /O(1)/ Get the length of the array.
295length :: forall a. Storable a => Array a -> Int
297 let p = unsafeForeignPtrToPtr aStart
298 aLen = aEnd `minusPtr` p
299 in assert (aLen >= 0) (aLen `div` sizeOf (undefined :: a))
301{-# INLINE_NORMAL toStreamD #-}
302toStreamD :: forall m a. (Monad m, Storable a) => Array a -> D.Stream m a
304 let p = unsafeForeignPtrToPtr aStart
309 {-# INLINE_LATE step #-}
310 step _ p | p == aEnd = return D.Stop
312 -- unsafeInlineIO allows us to run this in Identity monad for pure
313 -- toList/foldr case which makes them much faster due to not
314 -- accumulating the list and fusing better with the pure consumers.
316 -- This should be safe as the array contents are guaranteed to be
317 -- evaluated/written to before we peek at them.
318 let !x = unsafeInlineIO $ do
320 touchForeignPtr aStart
322 return $ D.Yield x (p `plusPtr` (sizeOf (undefined :: a)))
324{-# INLINE_NORMAL toStreamDRev #-}
325toStreamDRev :: forall m a. (Monad m, Storable a) => Array a -> D.Stream m a
326toStreamDRev Array{..} =
327 let p = aEnd `plusPtr` negate (sizeOf (undefined :: a))
332 {-# INLINE_LATE step #-}
333 step _ p | p < unsafeForeignPtrToPtr aStart = return D.Stop
335 -- See comments in toStreamD for why we use unsafeInlineIO
336 let !x = unsafeInlineIO $ do
338 touchForeignPtr aStart
340 return $ D.Yield x (p `plusPtr` negate (sizeOf (undefined :: a)))
342{-# INLINE_NORMAL foldl' #-}
343foldl' :: forall a b. Storable a => (b -> a -> b) -> b -> Array a -> b
344foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
346{-# INLINE_NORMAL foldr #-}
347foldr :: Storable a => (a -> b -> b) -> b -> Array a -> b
348foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
350-------------------------------------------------------------------------------
352-------------------------------------------------------------------------------
354-- | @toArrayN n@ folds a maximum of @n@ elements from the input stream to an
358{-# INLINE_NORMAL toArrayN #-}
359toArrayN :: forall m a. (MonadIO m, Storable a) => Int -> Fold m a (Array a)
360toArrayN n = Fold step initial extract
364 initial = liftIO $ newArray n
365 step arr@(Array _ end bound) _ | end == bound = return arr
366 step (Array start end bound) x = do
368 return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound
372-- | Fold the input to a pure buffered stream (List) of arrays.
373{-# INLINE toArrays #-}
374toArrays :: (Monad m, Storable a) => Int -> Fold m a (List (Array a))
375toArrays n = Fold step initial extract
377-- This can be implemented by combining the List of arrays from toArrays into a
379-- | Fold the whole input to a single array.
380{-# INLINE toArray #-}
381toArray :: forall m a. (Monad m, Storable a) => Fold m a (Array a)
382toArray = Fold step begin done
385{-# INLINE_NORMAL fromStreamDN #-}
386fromStreamDN :: forall m a. (MonadIO m, Storable a)
387 => Int -> D.Stream m a -> m (Array a)
388fromStreamDN limit str = do
389 arr <- liftIO $ newArray limit
390 end <- D.foldlM' write (aEnd arr) $ D.take limit str
391 return $ arr {aEnd = end}
397 return $ ptr `plusPtr` sizeOf (undefined :: a)
399-- | @fromStreamArraysOf n stream@ groups the input stream into a stream of
401{-# INLINE fromStreamDArraysOf #-}
402fromStreamDArraysOf :: (MonadIO m, Storable a)
403 => Int -> D.Stream m a -> D.Stream m (Array a)
404fromStreamDArraysOf n str = D.groupsOf n (toArrayN n) str
406-- XXX concatMap does not seem to have the best possible performance so we have
407-- a custom way to concat arrays.
408data FlattenState s a =
410 | InnerLoop s (ForeignPtr a) (Ptr a) (Ptr a)
412{-# INLINE_NORMAL flattenArrays #-}
413flattenArrays :: forall m a. (MonadIO m, Storable a)
414 => D.Stream m (Array a) -> D.Stream m a
415flattenArrays (D.Stream step state) = D.Stream step' (OuterLoop state)
419 {-# INLINE_LATE step' #-}
420 step' gst (OuterLoop st) = do
421 r <- step (adaptState gst) st
423 D.Yield Array{..} s ->
424 let p = unsafeForeignPtrToPtr aStart
425 in D.Skip (InnerLoop s aStart p aEnd)
426 D.Skip s -> D.Skip (OuterLoop s)
429 step' _ (InnerLoop st _ p end) | p == end =
430 return $ D.Skip $ OuterLoop st
432 step' _ (InnerLoop st startf p end) = do
435 touchForeignPtr startf
437 return $ D.Yield x (InnerLoop st startf
438 (p `plusPtr` (sizeOf (undefined :: a))) end)
440{-# INLINE_NORMAL flattenArraysRev #-}
441flattenArraysRev :: forall m a. (MonadIO m, Storable a)
442 => D.Stream m (Array a) -> D.Stream m a
443flattenArraysRev (D.Stream step state) = D.Stream step' (OuterLoop state)
447 {-# INLINE_LATE step' #-}
448 step' gst (OuterLoop st) = do
449 r <- step (adaptState gst) st
451 D.Yield Array{..} s ->
452 let p = aEnd `plusPtr` negate (sizeOf (undefined :: a))
453 -- XXX we do not need aEnd
454 in D.Skip (InnerLoop s aStart p aEnd)
455 D.Skip s -> D.Skip (OuterLoop s)
458 step' _ (InnerLoop st start p _) | p < unsafeForeignPtrToPtr start =
459 return $ D.Skip $ OuterLoop st
461 step' _ (InnerLoop st startf p end) = do
464 touchForeignPtr startf
466 return $ D.Yield x (InnerLoop st startf
467 (p `plusPtr` negate (sizeOf (undefined :: a))) end)
469-- CAUTION: a very large number (millions) of arrays can degrade performance
470-- due to GC overhead because we need to buffer the arrays before we flatten
473-- We could take the approach of doubling the memory allocation on each
474-- overflow. This would result in more or less the same amount of copying as in
475-- the chunking approach. However, if we have to shrink in the end then it may
476-- result in an extra copy of the entire data.
478{-# INLINE fromStreamD #-}
479fromStreamD :: (MonadIO m, Storable a) => D.Stream m a -> m (Array a)
481 let s = fromStreamDArraysOf defaultChunkSize m
482 buffered <- D.foldr K.cons K.nil s
483 len <- K.foldl' (+) 0 (K.map length buffered)
484 fromStreamDN len $ flattenArrays $ D.fromStreamK buffered
486-- | Convert an 'Array' into a list.
489{-# INLINABLE toList #-}
490toList :: Storable a => Array a -> [a]
493instance (Show a, Storable a) => Show (Array a) where
494 {-# INLINE showsPrec #-}
495 showsPrec _ = shows . toList
497-- | Create an 'Array' from the first N elements of a list. The array is
498-- allocated to size N, if the list terminates before N elements then the
499-- array may hold less than N elements.
502{-# INLINABLE fromListN #-}
503fromListN :: Storable a => Int -> [a] -> Array a
504fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
506-- | Create an 'Array' from a list. The list must be of finite size.
509{-# INLINABLE fromList #-}
510fromList :: Storable a => [a] -> Array a
511fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
513instance (Storable a, Read a, Show a) => Read (Array a) where
514 {-# INLINE readPrec #-}
518 readListPrec = readListPrecDefault
520instance (a ~ Char) => IsString (Array a) where
521 {-# INLINE fromString #-}
522 fromString = fromList
524-- GHC versions 8.0 and below cannot derive IsList
525instance Storable a => IsList (Array a) where
526 type (Item (Array a)) = a
527 {-# INLINE fromList #-}
529 {-# INLINE fromListN #-}
530 fromListN = fromListN
531 {-# INLINE toList #-}
535arrcmp :: Array a -> Array a -> Bool
537 let !res = unsafeInlineIO $ do
538 let ptr1 = unsafeForeignPtrToPtr $ aStart arr1
539 let ptr2 = unsafeForeignPtrToPtr $ aStart arr2
540 let len1 = aEnd arr1 `minusPtr` ptr1
541 let len2 = aEnd arr2 `minusPtr` ptr2
545 r <- memcmp (castPtr ptr1) (castPtr ptr2) len1
546 touchForeignPtr $ aStart arr1
547 touchForeignPtr $ aStart arr2
552-- XXX we are assuming that Storable equality means element equality. This may
553-- or may not be correct? arrcmp is 40% faster compared to stream equality.
554instance (Storable a, Eq a) => Eq (Array a) where
557 -- arr1 == arr2 = runIdentity $ D.eqBy (==) (toStreamD arr1) (toStreamD arr2)
559instance (Storable a, NFData a) => NFData (Array a) where
561 rnf = foldl' (\_ x -> rnf x) ()
563instance (Storable a, Ord a) => Ord (Array a) where
564 {-# INLINE compare #-}
565 compare arr1 arr2 = unsafePerformIO $
566 D.cmpBy compare (toStreamD arr1) (toStreamD arr2)
568 -- Default definitions defined in base do not have an INLINE on them, so we
569 -- replicate them here with an INLINE.
571 x < y = case compare x y of { LT -> True; _ -> False }
574 x <= y = case compare x y of { GT -> False; _ -> True }
577 x > y = case compare x y of { GT -> True; _ -> False }
580 x >= y = case compare x y of { LT -> False; _ -> True }
582 -- These two default methods use '<=' rather than 'compare'
583 -- because the latter is often more expensive
585 max x y = if x <= y then y else x
588 min x y = if x <= y then x else y
591-- Definitions using the Storable constraint from the Array type. These are to
592-- make the Foldable instance possible though it is much slower (7x slower).
594{-# INLINE_NORMAL toStreamD_ #-}
595toStreamD_ :: forall m a. MonadIO m => Int -> Array a -> D.Stream m a
596toStreamD_ size Array{..} =
597 let p = unsafeForeignPtrToPtr aStart
602 {-# INLINE_LATE step #-}
603 step _ p | p == aEnd = return D.Stop
607 touchForeignPtr aStart
609 return $ D.Yield x (p `plusPtr` size)
611{-# INLINE_NORMAL _foldr #-}
612_foldr :: forall a b. (a -> b -> b) -> b -> Array a -> b
613_foldr f z arr@Array {..} =
614 let !n = sizeOf (undefined :: a)
615 in unsafePerformIO $ D.foldr f z $ toStreamD_ n arr
617-- | Note that the 'Foldable' instance is 7x slower than the direct
619instance Foldable Array where
623-------------------------------------------------------------------------------
624-- Semigroup and Monoid
625-------------------------------------------------------------------------------
627-------------------------------------------------------------------------------
629-------------------------------------------------------------------------------
631-- | GHC memory management allocation header overhead
633allocOverhead = 2 * sizeOf (undefined :: Int)
635-- | Default maximum buffer size in bytes, for reading from and writing to IO
636-- devices, the value is 32KB minus GHC allocation overhead, which is a few
637-- bytes, so that the actual allocation is 32KB.
638defaultChunkSize :: Int
639defaultChunkSize = 32 * k - allocOverhead
642mkChunkSizeKB :: Int -> Int
643mkChunkSizeKB 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