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