on 7cb856b help sign in
589 lines 19 kB history raw
1{-# LANGUAGE CPP #-}
2-- |
3-- Module : Streamly.Internal.Data.Array
4-- Copyright : (c) 2019 Composewell Technologies
5--
6-- License : BSD3
7-- Maintainer : streamly@composewell.com
8-- Stability : experimental
9-- Portability : GHC
10--
11module Streamly.Internal.Data.Array
12 (
13 -- * Setup
14 -- $setup
16 -- * Design Notes
17 -- $design
19 -- * The Array Type
20 module Streamly.Internal.Data.Array.Type
22 -- * Construction
23 -- Stream Folds
24 , fromStreamN
25 , fromStream
27 -- Monadic Folds
28 , writeLastN
30 -- * Elimination
31 -- ** Conversion
33 -- ** Unfolds
34 , reader
35 , readerUnsafe
36 , producer -- experimental
38 -- * Random Access
39 -- , (!!)
40 , getIndex
41 , getIndexRev
42 , last -- XXX getIndexLast?
43 , getIndices
44 , getIndicesFromThenTo
45 -- , getIndicesFrom -- read from a given position to the end of file
46 -- , getIndicesUpto -- read from beginning up to the given position
47 -- , getIndicesFromTo
48 -- , getIndicesFromRev -- read from a given position to the beginning of file
49 -- , getIndicesUptoRev -- read from end to the given position in file
51 -- * Size
52 , null
54 -- * Search
55 , binarySearch
56 , findIndicesOf
57 -- , findIndexOf
58 -- , find
60 -- * Casting
61 , cast
62 , asBytes
63 , castUnsafe
64 , asCStringUnsafe
66 -- * Subarrays
67 , getSliceUnsafe
68 -- , getSlice
69 , genSlicesFromLen
70 , getSlicesFromLen
71 , splitOn
73 -- * Streaming Operations
74 , streamTransform
76 -- ** Folding
77 , streamFold
78 , fold
80 -- ** Serialization
81 , encodeAs
82 , serialize
83 , pinnedSerialize
84 , deserialize
85 )
86where
88#include "assert.hs"
89#include "inline.hs"
90#include "ArrayMacros.h"
92import Control.Monad (when)
93import Control.Monad.IO.Class (MonadIO(..))
94import Data.Functor.Identity (Identity)
95import Data.Proxy (Proxy(..))
96import Data.Word (Word8)
97import Foreign.C.String (CString)
98import Foreign.Ptr (castPtr)
99import Foreign.Storable (Storable)
100import Streamly.Internal.Data.Unbox (Unbox(..))
101import Prelude hiding (length, null, last, map, (!!), read, concat)
103import Streamly.Internal.Data.MutByteArray.Type (PinnedState(..))
104import Streamly.Internal.Data.Serialize.Type (Serialize)
105import Streamly.Internal.Data.Fold.Type (Fold(..))
106import Streamly.Internal.Data.Producer.Type (Producer(..))
107import Streamly.Internal.Data.Stream (Stream)
108import Streamly.Internal.Data.Tuple.Strict (Tuple3Fused'(..))
109import Streamly.Internal.Data.Unfold.Type (Unfold(..))
110import Streamly.Internal.System.IO (unsafeInlineIO)
112import qualified Streamly.Internal.Data.Serialize.Type as Serialize
113import qualified Streamly.Internal.Data.MutByteArray.Type as MBA
114import qualified Streamly.Internal.Data.MutArray as MA
115import qualified Streamly.Internal.Data.Array.Type as A
116import qualified Streamly.Internal.Data.Fold as FL
117import qualified Streamly.Internal.Data.Producer.Type as Producer
118import qualified Streamly.Internal.Data.Producer as Producer
119import qualified Streamly.Internal.Data.Ring as RB
120import qualified Streamly.Internal.Data.Stream as D
121import qualified Streamly.Internal.Data.Stream as Stream
122import qualified Streamly.Internal.Data.Unfold as Unfold
124import Streamly.Internal.Data.Array.Type
126#include "DocTestDataArray.hs"
128-- $design
129--
130-- To summarize:
131--
132-- * Arrays are finite and fixed in size
133-- * provide /O(1)/ access to elements
134-- * store only data and not functions
135-- * provide efficient IO interfacing
136--
137-- 'Foldable' instance is not provided because the implementation would be much
138-- less efficient compared to folding via streams. 'Semigroup' and 'Monoid'
139-- instances should be used with care; concatenating arrays using binary
140-- operations can be highly inefficient. Instead, use
141-- 'Streamly.Internal.Data.Stream.Chunked.toArray' to concatenate N
142-- arrays at once.
143--
144-- Each array is one pointer visible to the GC. Too many small arrays (e.g.
145-- single byte) are only as good as holding those elements in a Haskell list.
146-- However, small arrays can be compacted into large ones to reduce the
147-- overhead. To hold 32GB memory in 32k sized buffers we need 1 million arrays
148-- if we use one array for each chunk. This is still significant to add
149-- pressure to GC.
151-------------------------------------------------------------------------------
152-- Construction
153-------------------------------------------------------------------------------
155-- | Create an 'Array' from the first N elements of a stream. The array is
156-- allocated to size N, if the stream terminates before N elements then the
157-- array may hold less than N elements.
158--
159-- /Pre-release/
160{-# INLINE fromStreamN #-}
161fromStreamN :: (MonadIO m, Unbox a) => Int -> Stream m a -> m (Array a)
162fromStreamN n m = do
163 when (n < 0) $ error "writeN: negative write count specified"
164 A.fromStreamDN n m
166-- | Create an 'Array' from a stream. This is useful when we want to create a
167-- single array from a stream of unknown size. 'writeN' is at least twice
168-- as efficient when the size is already known.
169--
170-- Note that if the input stream is too large memory allocation for the array
171-- may fail. When the stream size is not known, `chunksOf` followed by
172-- processing of indvidual arrays in the resulting stream should be preferred.
173--
174-- /Pre-release/
175{-# INLINE fromStream #-}
176fromStream :: (MonadIO m, Unbox a) => Stream m a -> m (Array a)
177fromStream = Stream.fold A.write
178-- write m = A.fromStreamD $ D.fromStreamK m
180-------------------------------------------------------------------------------
181-- Elimination
182-------------------------------------------------------------------------------
184{-# INLINE_NORMAL producer #-}
185producer :: forall m a. (Monad m, Unbox a) => Producer m (Array a) a
186producer =
187 Producer.translate A.unsafeThaw A.unsafeFreeze
188 $ MA.producerWith (return . unsafeInlineIO)
190-- | Unfold an array into a stream.
191--
192{-# INLINE_NORMAL reader #-}
193reader :: forall m a. (Monad m, Unbox a) => Unfold m (Array a) a
194reader = Producer.simplify producer
196-- | Unfold an array into a stream, does not check the end of the array, the
197-- user is responsible for terminating the stream within the array bounds. For
198-- high performance application where the end condition can be determined by
199-- a terminating fold.
200--
201-- Written in the hope that it may be faster than "read", however, in the case
202-- for which this was written, "read" proves to be faster even though the core
203-- generated with unsafeRead looks simpler.
204--
205-- /Pre-release/
206--
207{-# INLINE_NORMAL readerUnsafe #-}
208readerUnsafe :: forall m a. (Monad m, Unbox a) => Unfold m (Array a) a
209readerUnsafe = Unfold step inject
210 where
212 inject (Array contents start end) =
213 return (ArrayUnsafe contents end start)
215 {-# INLINE_LATE step #-}
216 step (ArrayUnsafe contents end p) = do
217 -- unsafeInlineIO allows us to run this in Identity monad for pure
218 -- toList/foldr case which makes them much faster due to not
219 -- accumulating the list and fusing better with the pure consumers.
220 --
221 -- This should be safe as the array contents are guaranteed to be
222 -- evaluated/written to before we peek at them.
223 let !x = unsafeInlineIO $ peekByteIndex p contents
224 let !p1 = INDEX_NEXT(p,a)
225 return $ D.Yield x (ArrayUnsafe contents end p1)
227-- |
228--
229-- >>> null arr = Array.byteLength arr == 0
230--
231-- /Pre-release/
232{-# INLINE null #-}
233null :: Array a -> Bool
234null arr = A.byteLength arr == 0
236-- | Like 'getIndex' but indexes the array in reverse from the end.
237--
238-- /Pre-release/
239{-# INLINE getIndexRev #-}
240getIndexRev :: forall a. Unbox a => Int -> Array a -> Maybe a
241getIndexRev i arr =
242 unsafeInlineIO
243 $ do
244 let elemPtr = RINDEX_OF(arrEnd arr, i, a)
245 if i >= 0 && elemPtr >= arrStart arr
246 then Just <$> peekByteIndex elemPtr (arrContents arr)
247 else return Nothing
249-- |
250--
251-- >>> last arr = Array.getIndexRev arr 0
252--
253-- /Pre-release/
254{-# INLINE last #-}
255last :: Unbox a => Array a -> Maybe a
256last = getIndexRev 0
258-------------------------------------------------------------------------------
259-- Folds with Array as the container
260-------------------------------------------------------------------------------
262-- | @writeLastN n@ folds a maximum of @n@ elements from the end of the input
263-- stream to an 'Array'.
264--
265{-# INLINE writeLastN #-}
266writeLastN ::
267 (Storable a, Unbox a, MonadIO m) => Int -> Fold m a (Array a)
268writeLastN n
269 | n <= 0 = fmap (const mempty) FL.drain
270 | otherwise = A.unsafeFreeze <$> Fold step initial done done
272 where
274 step (Tuple3Fused' rb rh i) a = do
275 rh1 <- liftIO $ RB.unsafeInsert rb rh a
276 return $ FL.Partial $ Tuple3Fused' rb rh1 (i + 1)
278 initial =
279 let f (a, b) = FL.Partial $ Tuple3Fused' a b (0 :: Int)
280 in fmap f $ liftIO $ RB.new n
282 done (Tuple3Fused' rb rh i) = do
283 arr <- liftIO $ MA.new n
284 foldFunc i rh snoc' arr rb
286 -- XXX We should write a read unfold for ring.
287 snoc' b a = liftIO $ MA.snocUnsafe b a
289 foldFunc i
290 | i < n = RB.unsafeFoldRingM
291 | otherwise = RB.unsafeFoldRingFullM
293-------------------------------------------------------------------------------
294-- Random Access
295-------------------------------------------------------------------------------
297-------------------------------------------------------------------------------
298-- Searching
299-------------------------------------------------------------------------------
301-- | Given a sorted array, perform a binary search to find the given element.
302-- Returns the index of the element if found.
303--
304-- /Unimplemented/
305{-# INLINE binarySearch #-}
306binarySearch :: a -> Array a -> Maybe Int
307binarySearch = undefined
309-- find/findIndex etc can potentially be implemented more efficiently on arrays
310-- compared to streams by using SIMD instructions.
311-- We can also return a bit array instead.
313-- | Perform a linear search to find all the indices where a given element is
314-- present in an array.
315--
316-- /Unimplemented/
317findIndicesOf :: (a -> Bool) -> Unfold Identity (Array a) Int
318findIndicesOf = undefined
320{-
321findIndexOf :: (a -> Bool) -> Array a -> Maybe Int
322findIndexOf p = Unfold.fold Fold.one . Stream.unfold (findIndicesOf p)
324find :: (a -> Bool) -> Array a -> Bool
325find = Unfold.fold Fold.null . Stream.unfold (findIndicesOf p)
326-}
328-------------------------------------------------------------------------------
329-- Folds
330-------------------------------------------------------------------------------
332-- XXX We can potentially use SIMD instructions on arrays to fold faster.
334-------------------------------------------------------------------------------
335-- Slice
336-------------------------------------------------------------------------------
338-- | /O(1)/ Slice an array in constant time.
339--
340-- Caution: The bounds of the slice are not checked.
341--
342-- /Unsafe/
343--
344-- /Pre-release/
345{-# INLINE getSliceUnsafe #-}
346getSliceUnsafe ::
347 forall a. Unbox a
348 => Int -- ^ starting index
349 -> Int -- ^ length of the slice
350 -> Array a
351 -> Array a
352getSliceUnsafe index len (Array contents start e) =
353 let size = SIZE_OF(a)
354 start1 = start + (index * size)
355 end1 = start1 + (len * size)
356 in assert (end1 <= e) (Array contents start1 end1)
358-- | Split the array into a stream of slices using a predicate. The element
359-- matching the predicate is dropped.
360--
361-- /Pre-release/
362{-# INLINE splitOn #-}
363splitOn :: (Monad m, Unbox a) =>
364 (a -> Bool) -> Array a -> Stream m (Array a)
365splitOn predicate arr =
366 fmap (\(i, len) -> getSliceUnsafe i len arr)
367 $ D.sliceOnSuffix predicate (A.toStreamD arr)
369{-# INLINE genSlicesFromLen #-}
370genSlicesFromLen :: forall m a. (Monad m, Unbox a)
371 => Int -- ^ from index
372 -> Int -- ^ length of the slice
373 -> Unfold m (Array a) (Int, Int)
374genSlicesFromLen from len =
375 Unfold.lmap A.unsafeThaw (MA.genSlicesFromLen from len)
377-- | Generate a stream of slices of specified length from an array, starting
378-- from the supplied array index. The last slice may be shorter than the
379-- requested length.
380--
381-- /Pre-release//
382{-# INLINE getSlicesFromLen #-}
383getSlicesFromLen :: forall m a. (Monad m, Unbox a)
384 => Int -- ^ from index
385 -> Int -- ^ length of the slice
386 -> Unfold m (Array a) (Array a)
387getSlicesFromLen from len =
388 fmap A.unsafeFreeze
389 $ Unfold.lmap A.unsafeThaw (MA.getSlicesFromLen from len)
391-------------------------------------------------------------------------------
392-- Random reads and writes
393-------------------------------------------------------------------------------
395-- XXX Change this to a partial function instead of a Maybe type? And use
396-- MA.getIndex instead.
397--
398-- | /O(1)/ Lookup the element at the given index. Index starts from 0.
399--
400{-# INLINE getIndex #-}
401getIndex :: forall a. Unbox a => Int -> Array a -> Maybe a
402getIndex i arr =
403 unsafeInlineIO
404 $ do
405 let elemPtr = INDEX_OF(arrStart arr, i, a)
406 if i >= 0 && INDEX_VALID(elemPtr, arrEnd arr, a)
407 then Just <$> peekByteIndex elemPtr (arrContents arr)
408 else return Nothing
410-- | Given a stream of array indices, read the elements on those indices from
411-- the supplied Array. An exception is thrown if an index is out of bounds.
412--
413-- This is the most general operation. We can implement other operations in
414-- terms of this:
415--
416-- @
417-- read =
418-- let u = lmap (\arr -> (0, length arr - 1)) Unfold.enumerateFromTo
419-- in Unfold.lmap f (getIndices arr)
420--
421-- readRev =
422-- let i = length arr - 1
423-- in Unfold.lmap f (getIndicesFromThenTo i (i - 1) 0)
424-- @
425--
426-- /Pre-release/
427{-# INLINE getIndices #-}
428getIndices :: (Monad m, Unbox a) => Stream m Int -> Unfold m (Array a) a
429getIndices m =
430 let unf = MA.getIndicesD (return . unsafeInlineIO) m
431 in Unfold.lmap A.unsafeThaw unf
433-- | Unfolds @(from, then, to, array)@ generating a finite stream whose first
434-- element is the array value from the index @from@ and the successive elements
435-- are from the indices in increments of @then@ up to @to@. Index enumeration
436-- can occur downwards or upwards depending on whether @then@ comes before or
437-- after @from@.
438--
439-- @
440-- getIndicesFromThenTo =
441-- let f (from, next, to, arr) =
442-- (Stream.enumerateFromThenTo from next to, arr)
443-- in Unfold.lmap f getIndices
444-- @
445--
446-- /Unimplemented/
447{-# INLINE getIndicesFromThenTo #-}
448getIndicesFromThenTo :: Unfold m (Int, Int, Int, Array a) a
449getIndicesFromThenTo = undefined
451-------------------------------------------------------------------------------
452-- Transform via stream operations
453-------------------------------------------------------------------------------
455-- for non-length changing operations we can use the original length for
456-- allocation. If we can predict the length then we can use the prediction for
457-- new allocation. Otherwise we can use a hint and adjust dynamically.
459{-
460-- | Transform an array into another array using a pipe transformation
461-- operation.
462--
463{-# INLINE runPipe #-}
464runPipe :: (MonadIO m, Unbox a, Unbox b)
465 => Pipe m a b -> Array a -> m (Array b)
466runPipe f arr = P.runPipe (toArrayMinChunk (length arr)) $ f (A.read arr)
467-}
469-- XXX For transformations that cannot change the number of elements e.g. "map"
470-- we can use a predetermined array length.
471--
472-- | Transform an array into another array using a stream transformation
473-- operation.
474--
475-- /Pre-release/
476{-# INLINE streamTransform #-}
477streamTransform :: forall m a b. (MonadIO m, Unbox a, Unbox b)
478 => (Stream m a -> Stream m b) -> Array a -> m (Array b)
479streamTransform f arr =
480 Stream.fold (A.writeWith (length arr)) $ f (A.read arr)
482-------------------------------------------------------------------------------
483-- Casts
484-------------------------------------------------------------------------------
486-- | Cast an array having elements of type @a@ into an array having elements of
487-- type @b@. The array size must be a multiple of the size of type @b@
488-- otherwise accessing the last element of the array may result into a crash or
489-- a random value.
490--
491-- /Pre-release/
492--
493castUnsafe ::
494#ifdef DEVBUILD
495 Unbox b =>
496#endif
497 Array a -> Array b
498castUnsafe (Array contents start end) =
499 Array contents start end
501-- | Cast an @Array a@ into an @Array Word8@.
502--
503--
504asBytes :: Array a -> Array Word8
505asBytes = castUnsafe
507-- | Cast an array having elements of type @a@ into an array having elements of
508-- type @b@. The length of the array should be a multiple of the size of the
509-- target element otherwise 'Nothing' is returned.
510--
511--
512cast :: forall a b. (Unbox b) => Array a -> Maybe (Array b)
513cast arr =
514 let len = A.byteLength arr
515 r = len `mod` SIZE_OF(b)
516 in if r /= 0
517 then Nothing
518 else Just $ castUnsafe arr
520-- | Convert an array of any type into a null terminated CString Ptr. If the
521-- array is unpinned it is first converted to a pinned array which requires a
522-- copy.
523--
524-- /Unsafe/
525--
526-- /O(n) Time: (creates a copy of the array)/
527--
528-- /Pre-release/
529--
530asCStringUnsafe :: Array a -> (CString -> IO b) -> IO b
531asCStringUnsafe arr act = do
532 let arr1 = asBytes arr <> A.fromList [0]
533 -- asPtrUnsafe makes sure the array is pinned
534 asPtrUnsafe arr1 $ \ptr -> act (castPtr ptr)
536-------------------------------------------------------------------------------
537-- Folds
538-------------------------------------------------------------------------------
540-- XXX We can directly use toStreamD and D.fold here.
542-- | Fold an array using a 'Fold'.
543--
544-- /Pre-release/
545{-# INLINE fold #-}
546fold :: forall m a b. (Monad m, Unbox a) => Fold m a b -> Array a -> m b
547fold f arr = Stream.fold f (A.read arr)
549-- | Fold an array using a stream fold operation.
550--
551-- /Pre-release/
552{-# INLINE streamFold #-}
553streamFold :: (Monad m, Unbox a) => (Stream m a -> m b) -> Array a -> m b
554streamFold f arr = f (A.read arr)
556--------------------------------------------------------------------------------
557-- Serialization
558--------------------------------------------------------------------------------
560{-# INLINE encodeAs #-}
561encodeAs :: forall a. Serialize a => PinnedState -> a -> Array Word8
562encodeAs ps a =
563 unsafeInlineIO $ do
564 let len = Serialize.size 0 a
565 mbarr <- MBA.newBytesAs ps len
566 off <- Serialize.serialize 0 mbarr a
567 assertM(len == off)
568 pure $ Array mbarr 0 off
570{-# INLINE serialize #-}
571serialize :: Serialize a => a -> Array Word8
572serialize = encodeAs Unpinned
574-- | Serialize a Haskell type to a pinned byte array. The array is allocated
575-- using pinned memory so that it can be used directly in OS APIs for writing
576-- to file or sending over the network.
577{-# INLINE pinnedSerialize #-}
578pinnedSerialize :: Serialize a => a -> Array Word8
579pinnedSerialize = encodeAs Pinned
581-- | Decode a Haskell type from a byte array containing its serialized
582-- representation.
583{-# INLINE deserialize #-}
584deserialize :: Serialize a => Array Word8 -> a
585deserialize arr@(Array {..}) = unsafeInlineIO $ do
586 let lenArr = length arr
587 (off, val) <- Serialize.deserialize arrStart arrContents (arrStart + lenArr)
588 assertM(off == arrStart + lenArr)
589 pure val