on ea342a4 help sign in
627 lines 23 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE MagicHash #-}
4{-# LANGUAGE RecordWildCards #-}
5{-# LANGUAGE UnboxedTuples #-}
6{-# LANGUAGE ScopedTypeVariables #-}
8#include "inline.hs"
10-- |
11-- Module : Streamly.Mem.Array
12-- Copyright : (c) 2019 Harendra Kumar
13--
14-- License : BSD3
15-- Maintainer : harendra.kumar@gmail.com
16-- Stability : experimental
17-- Portability : GHC
18--
19-- Arrays are the computing duals of streams. Streams are good at sequential
20-- access and immutable transformations of in-transit data whereas arrays are
21-- good at random access and in-place transformations of buffered data. Unlike
22-- streams which are potentially infinite, arrays are necessarily /finite/.
23-- Arrays can be used as an efficient interface between streams and external
24-- storage systems like memory, files and network. Streams and arrays complete
25-- each other to provide a general purpose computing system. The design of
26-- streamly as a general purpose computing framework is centered around these
27-- two fundamental aspects of computing and storage.
28--
29-- Even though the implementation allows for in-place mutation, the current
30-- implementation of arrays is immutable. Arrays are of fixed size. The size of
31-- an array is fixed at creation and cannot be changed later without copying.
32-- Arrays are chunks of memory that hold a sequence of 'Storable' values of a
33-- given type. They are designed to store serializable data, they cannot store
34-- non-serializable data like functions, this is reflected by the 'Storable'
35-- constraint. For efficient buffering of data, arrays use pinned memory and
36-- therefore can hold arbitrary number of elements with a single movable
37-- pointer visible to GC, therefore, adding no pressure to GC. Moreover,
38-- pinned memory allows communication with foreign consumers and producers
39-- (e.g. file or network IO) without copying the data.
40--
41-- By design, there are no transformation operations provided in this module,
42-- only stream conversion and IO routines are provided. An array is purely a
43-- data store, therefore, it is not a 'Functor'. Arrays can be operated upon
44-- /efficiently/ by converting them into a stream, applying the desired
45-- transformations on the stream and then converting it back to an array.
46-- 'Foldable' instance is not provided because the implementation would be much
47-- less efficient compared to folding via streams. 'Semigroup' and 'Monoid'
48-- instances are deliberately not provided to avoid misuse; concatenating
49-- arrays using binary operations can be highly inefficient. Instead, use
50-- 'spliceArrays' to concatenate N arrays at once.
51--
52-- To summarize:
53--
54-- * Arrays are finite and fixed in size
55-- * provide /O(1)/ access to elements
56-- * store only data and not functions
57-- * provide efficient IO interfacing
58--
59-- 'ByteString' data type from the 'bytestring' package and the 'Text' data
60-- type from the 'text' package are just special cases of arrays. 'ByteString'
61-- is equivalent to @Array Word8@ and 'Text' is equivalent to a @utf16@ encoded
62-- @Array Word8@. All the 'bytestring' and 'text' operations can be performed
63-- on arrays with equivalent or better performance by converting them to and
64-- from streams.
65--
66-- This module is designed to be imported qualified:
67--
68-- > import qualified Streamly.Array as A
70-- Each array is one pointer visible to the GC. Too many small arrays (e.g.
71-- single byte) are only as good as holding those elements in a Haskell list.
72-- However, small arrays can be compacted into large ones to reduce the
73-- overhead. To hold 32GB memory in 32k sized buffers we need 1 million arrays
74-- if we use one array for each chunk. This is still significant to add
75-- pressure to GC. However, we can create arrays of arrays (trees) to scale to
76-- arbitrarily large amounts of memory but still using small chunks of
77-- contiguous memory.
79module Streamly.Mem.Array
80 (
81 Array
83 -- , defaultChunkSize
85 -- * Construction
86 -- | When performance matters, the fastest way to generate an array is
87 -- 'writeN'. For regular use, 'IsList' and 'IsString' instances can be
88 -- used to conveniently construct arrays from literal values.
89 -- 'OverloadedLists' extension or 'fromList' can be used to construct an
90 -- array from a list literal. Similarly, 'OverloadedStrings' extension or
91 -- 'fromList' can be used to construct an array from a string literal.
93 -- Pure List APIs
94 , A.fromListN
95 , A.fromList
97 -- Monadic APIs
98 -- , newArray
99 , writeN
100 , write
102 -- Stream Folds
103 , A.toArrayN
104 -- , toArrays
105 , toArray
107 -- * Elimination
108 -- 'GHC.Exts.toList' from "GHC.Exts" can be used to convert an array to a
109 -- list.
111 , A.toList
112 , A.read
113 , readRev
115 -- * Random Access
116 , length
117 -- , (!!)
119 , readIndex
120 {-
121 , readSlice
122 , readSliceRev
123 -}
125 , writeIndex
126 {-
127 , writeSlice
128 , writeSliceRev
129 -}
131 -- * Streams of arrays
132 , arraysOf
134 -- Streams of arrays
135 , flattenArrays
136 -- , flattenArraysRev
137 , packArrays
138 , packArraysChunksOf
139 , unlinesArraysBy
140 , splitArraysOn
142 -- * Immutable Transformations
143 , transformWith
145 -- * Folding Arrays
146 , foldWith
147 , foldArray
148 )
149where
151import Control.Monad (when)
152import Control.Monad.IO.Class (MonadIO(..))
153-- import Data.Functor.Identity (Identity)
154import Data.Word (Word8)
155import Foreign.ForeignPtr (withForeignPtr)
156import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
157import Foreign.Ptr (minusPtr, plusPtr, castPtr)
158import Foreign.Storable (Storable(..))
159import Prelude hiding (length, null, last, map, (!!), read)
161import Streamly.Mem.Array.Types (Array(..), length)
162import Streamly.Streams.Serial (SerialT)
163import Streamly.Streams.StreamK.Type (IsStream)
164import Streamly.Fold.Types (Fold(..))
166import qualified Streamly.Fold as FL
167import qualified Streamly.Mem.Array.Types as A
168import qualified Streamly.Prelude as S
169import qualified Streamly.Streams.StreamD as D
170import qualified Streamly.Streams.Prelude as P
172-------------------------------------------------------------------------------
173-- Construction
174-------------------------------------------------------------------------------
176{-
177-- | Create a new uninitialized array of given length.
178--
179-- @since 0.7.0
180newArray :: (MonadIO m, Storable a) => Int -> m (Array a)
181newArray len = undefined
182-}
184-- | Create an 'Array' from the first N elements of a stream. The array is
185-- allocated to size N, if the stream terminates before N elements then the
186-- array may hold less than N elements.
187--
188-- @since 0.7.0
189{-# INLINE writeN #-}
190writeN :: (MonadIO m, Storable a) => Int -> SerialT m a -> m (Array a)
191writeN n m = do
192 if n < 0 then error "writeN: negative write count specified" else return ()
193 A.fromStreamDN n $ D.toStreamD m
195-------------------------------------------------------------------------------
196-- Elimination
197-------------------------------------------------------------------------------
199-- | Convert an 'Array' into a stream in reverse order.
200--
201-- @since 0.7.0
202{-# INLINE_EARLY readRev #-}
203readRev :: (Monad m, IsStream t, Storable a) => Array a -> t m a
204readRev = D.fromStreamD . A.toStreamDRev
205-- XXX add fallback to StreamK rule
206-- {-# RULES "Streamly.Array.readRev fallback to StreamK" [1]
207-- forall a. S.toStreamK (readRev a) = K.revFromArray a #-}
209{-# INLINE _null #-}
210_null :: Storable a => Array a -> Bool
211_null arr = length arr == 0
213{-# INLINE _last #-}
214_last :: (MonadIO m, Storable a) => Array a -> m (Maybe a)
215_last arr = readIndex arr (length arr - 1)
217-------------------------------------------------------------------------------
218-- Random Access
219-------------------------------------------------------------------------------
221-------------------------------------------------------------------------------
222-- Searching
223-------------------------------------------------------------------------------
225{-
226-- | Perform a binary search in the array to find an element.
227bsearch :: a -> Array a -> Maybe Bool
228bsearch = undefined
230-- | Perform a binary search in the array to find an element index.
231{-# INLINE elemIndex #-}
232bsearchIndex :: a -> Array a -> Maybe Int
233bsearchIndex elem arr = undefined
235-- find/findIndex etc can potentially be implemented more efficiently on arrays
236-- compared to streams by using SIMD instructions.
238find :: (a -> Bool) -> Array a -> Bool
239find = undefined
241findIndex :: (a -> Bool) -> Array a -> Maybe Int
242findIndex = undefined
244findIndices :: (a -> Bool) -> Array a -> Array Int
245findIndices = undefined
246-}
248-------------------------------------------------------------------------------
249-- Folds
250-------------------------------------------------------------------------------
252-- XXX We can potentially use SIMD instructions on arrays to fold faster.
254-------------------------------------------------------------------------------
255-- Slice and splice
256-------------------------------------------------------------------------------
258{-
259slice :: Int -> Int -> Array a
260slice begin end arr = undefined
262splitAt :: Int -> Array a -> (Array a, Array a)
263splitAt i arr = undefined
265-- XXX This operation can be performed efficiently via streams.
266-- | Append two arrays together to create a single array.
267splice :: Array a -> Array a -> Array a
268splice arr1 arr2 = undefined
270-------------------------------------------------------------------------------
271-- In-place mutation APIs
272-------------------------------------------------------------------------------
274-- | Partition an array into two halves using a partitioning predicate. The
275-- first half retains values where the predicate is 'False' and the second half
276-- retains values where the predicate is 'True'.
277{-# INLINE partitionBy #-}
278partitionBy :: (a -> Bool) -> Array a -> (Array a, Array a)
279partitionBy f arr = undefined
281-- | Shuffle corresponding elements from two arrays using a shuffle function.
282-- If the shuffle function returns 'False' then do nothing otherwise swap the
283-- elements. This can be used in a bottom up fold to shuffle or reorder the
284-- elements.
285shuffleBy :: (a -> a -> m Bool) -> Array a -> Array a -> m (Array a)
286shuffleBy f arr1 arr2 = undefined
288-- XXX we can also make the folds partial by stopping at a certain level.
289--
290-- | Perform a top down hierarchical recursive partitioning fold of items in
291-- the container using the given function as the partition function.
292--
293-- This will perform a quick sort if the partition function is
294-- 'partitionBy (< pivot)'.
295--
296-- @since 0.7.0
297{-# INLINABLE foldtWith #-}
298foldtWith :: Int -> (Array a -> Array a -> m (Array a)) -> Array a -> m (Array a)
299foldtWith level f = undefined
301-- | Perform a pairwise bottom up fold recursively merging the pairs. Level
302-- indicates the level in the tree where the fold would stop.
303--
304-- This will perform a random shuffle if the shuffle function is random.
305-- If we stop at level 0 and repeatedly apply the function then we can do a
306-- bubble sort.
307foldbWith :: Int -> (Array a -> Array a -> m (Array a)) -> Array a -> m (Array a)
308foldbWith level f = undefined
309-}
311-- XXX consider the bulk update/accumulation/permutation APIs from vector.
313-------------------------------------------------------------------------------
314-- Random reads and writes
315-------------------------------------------------------------------------------
317-- | /O(1)/ Lookup the element at the given index.
318--
319-- @since 0.7.0
320{-# INLINE readIndex #-}
321readIndex :: (MonadIO m, Storable a) => Array a -> Int -> m (Maybe a)
322readIndex arr i =
323 if i < 0 || i > length arr - 1
324 then return Nothing
325 else liftIO $ withForeignPtr (aStart arr) $ \p ->
326 fmap Just $ peekElemOff p i
328{-
329-- | @readSlice arr i count@ streams a slice of the array @arr@ starting
330-- at index @i@ and reading up to @count@ elements in the forward direction
331-- ending at the index @i + count - 1@.
332--
333-- @since 0.7.0
334{-# INLINE readSlice #-}
335readSlice :: (IsStream t, Monad m, Storable a)
336 => Array a -> Int -> Int -> t m a
337readSlice arr i len = undefined
339-- | @readSliceRev arr i count@ streams a slice of the array @arr@ starting at
340-- index @i@ and reading up to @count@ elements in the reverse direction ending
341-- at the index @i - count + 1@.
342--
343-- @since 0.7.0
344{-# INLINE readSliceRev #-}
345readSliceRev :: (IsStream t, Monad m, Storable a)
346 => Array a -> Int -> Int -> t m a
347readSliceRev arr i len = undefined
348-}
350-- | /O(1)/ Write the given element at the given index in the array.
351-- Performs in-place mutation of the array.
352--
353-- @since 0.7.0
354{-# INLINE writeIndex #-}
355writeIndex :: (MonadIO m, Storable a) => Array a -> Int -> a -> m ()
356writeIndex arr i a = do
357 let maxIndex = length arr - 1
358 if i < 0
359 then error "writeIndex: negative array index"
360 else if i > maxIndex
361 then error $ "writeIndex: specified array index " ++ show i
362 ++ " is beyond the maximum index " ++ show maxIndex
363 else
364 liftIO $ withForeignPtr (aStart arr) $ \p ->
365 pokeElemOff p i a
367{-
368-- | @writeSlice arr i count stream@ writes a stream to the array @arr@
369-- starting at index @i@ and writing up to @count@ elements in the forward
370-- direction ending at the index @i + count - 1@.
371--
372-- @since 0.7.0
373{-# INLINE writeSlice #-}
374writeSlice :: (IsStream t, Monad m, Storable a)
375 => Array a -> Int -> Int -> t m a -> m ()
376writeSlice arr i len s = undefined
378-- | @writeSliceRev arr i count stream@ writes a stream to the array @arr@
379-- starting at index @i@ and writing up to @count@ elements in the reverse
380-- direction ending at the index @i - count + 1@.
381--
382-- @since 0.7.0
383{-# INLINE writeSliceRev #-}
384writeSliceRev :: (IsStream t, Monad m, Storable a)
385 => Array a -> Int -> Int -> t m a -> m ()
386writeSliceRev arr i len s = undefined
387-}
389-------------------------------------------------------------------------------
390-- Streams of Arrays
391-------------------------------------------------------------------------------
393-- exponentially increasing sizes of the chunks upto the max limit.
394-- XXX this will be easier to implement with parsers/terminating folds
395-- With this we should be able to reduce the number of chunks/allocations.
396-- The reallocation/copy based toArray can also be implemented using this.
397--
398{-
399{-# INLINE toArraysInRange #-}
400toArraysInRange :: (IsStream t, MonadIO m, Storable a)
401 => Int -> Int -> Fold m (Array a) b -> Fold m a b
402toArraysInRange low high (Fold step initial extract) =
403-}
405{-
406-- | Fold the input to a pure buffered stream (List) of arrays.
407{-# INLINE _toArraysOf #-}
408_toArraysOf :: (MonadIO m, Storable a)
409 => Int -> Fold m a (SerialT Identity (Array a))
410_toArraysOf n = FL.lchunksOf n (A.toArrayN n) FL.toStream
411-}
413-- XXX The realloc based implementation needs to make one extra copy if we use
414-- shrinkToFit. On the other hand, the stream of arrays implementation may
415-- buffer the array chunk pointers in memory but it does not have to shrink as
416-- we know the exact size in the end. However, memory copying does not seems to
417-- be as expensive as the allocations. Therefore, we need to reduce the number
418-- of allocations instead. Also, the size of allocations matters, right sizing
419-- an allocation even at the cost of copying sems to help. Should be measured
420-- on a big stream with heavy calls to toArray to see the effect.
421--
422-- XXX check if GHC's memory allocator is efficient enough. We can try the C
423-- malloc to compare against.
425{-# INLINE bytesToCount #-}
426bytesToCount :: Storable a => a -> Int -> Int
427bytesToCount x n =
428 let elemSize = sizeOf x
429 in n + elemSize - 1 `div` elemSize
431{-# INLINE_NORMAL toArrayMinChunk #-}
432toArrayMinChunk :: forall m a. (MonadIO m, Storable a)
433 => Int -> Fold m a (Array a)
434-- toArrayMinChunk n = FL.mapM spliceArrays $ toArraysOf n
435toArrayMinChunk elemCount = Fold step initial extract
437 where
439 insertElem (Array start end bound) x = do
440 liftIO $ poke end x
441 return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound
443 initial = do
444 when (elemCount < 0) $ error "toArrayMinChunk: elemCount is negative"
445 liftIO $ A.newArray elemCount
446 step arr@(Array start end bound) x | end == bound = do
447 let p = unsafeForeignPtrToPtr start
448 oldSize = end `minusPtr` p
449 newSize = max (oldSize * 2) 1
450 arr1 <- liftIO $ A.realloc newSize arr
451 insertElem arr1 x
452 step arr x = insertElem arr x
453 extract = liftIO . A.shrinkToFit
455-- | Fold the whole input to a single array.
456--
457-- /Caution! Do not use this on infinite streams./
458--
459-- @since 0.7.0
460{-# INLINE toArray #-}
461toArray :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)
462toArray = toArrayMinChunk (bytesToCount (undefined :: a) (A.mkChunkSize 1024))
464-- | Convert a stream of arrays into a stream of their elements.
465--
466-- Same as the following but more efficient:
467--
468-- > flattenArrays = S.concatMap A.read
469--
470-- @since 0.7.0
471{-# INLINE flattenArrays #-}
472flattenArrays :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
473flattenArrays m = D.fromStreamD $ A.flattenArrays (D.toStreamD m)
475-- XXX should we have a reverseArrays API to reverse the stream of arrays
476-- instead?
477--
478-- | Convert a stream of arrays into a stream of their elements reversing the
479-- contents of each array before flattening.
480--
481-- @since 0.7.0
482{-# INLINE _flattenArraysRev #-}
483_flattenArraysRev :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
484_flattenArraysRev m = D.fromStreamD $ A.flattenArraysRev (D.toStreamD m)
486-- XXX use an Array instead as separator? Or use a separate unlinesArraysBySeq
487-- API for that?
488--
489-- | Flatten a stream of arrays appending the given element after each
490-- array.
491--
492-- @since 0.7.0
493{-# INLINE unlinesArraysBy #-}
494unlinesArraysBy :: (MonadIO m, IsStream t, Storable a)
495 => a -> t m (Array a) -> t m a
496unlinesArraysBy x = D.fromStreamD . A.unlines x . D.toStreamD
498-- | Split a stream of arrays on a given separator byte, dropping the separator
499-- and coalescing all the arrays between two separators into a single array.
500--
501-- @since 0.7.0
502{-# INLINE splitArraysOn #-}
503splitArraysOn
504 :: (IsStream t, MonadIO m)
505 => Word8
506 -> t m (Array Word8)
507 -> t m (Array Word8)
508splitArraysOn byte s = D.fromStreamD $ A.splitOn byte $ D.toStreamD s
510-- | Coalesce adajcent arrays in incoming stream to form bigger arrays of a
511-- maximum specified size.
512--
513-- @since 0.7.0
514{-# INLINE packArraysChunksOf #-}
515packArraysChunksOf :: (MonadIO m, Storable a)
516 => Int -> SerialT m (Array a) -> SerialT m (Array a)
517packArraysChunksOf n xs =
518 D.fromStreamD $ A.packArraysChunksOf n (D.toStreamD xs)
520-- | Groups the elements in an input stream into arrays of given size.
521--
522-- Same as the following but more efficient:
523--
524-- > arraysOf n = FL.groupsOf n (FL.toArrayN n)
525--
526-- @since 0.7.0
527{-# INLINE arraysOf #-}
528arraysOf :: (IsStream t, MonadIO m, Storable a)
529 => Int -> t m a -> t m (Array a)
530arraysOf n str =
531 D.fromStreamD $ A.fromStreamDArraysOf n (D.toStreamD str)
533-- XXX Both of these implementations of splicing seem to perform equally well.
534-- We need to perform benchmarks over a range of sizes though.
536-- CAUTION! length must more than equal to lengths of all the arrays in the
537-- stream.
538{-# INLINE spliceArraysLenUnsafe #-}
539spliceArraysLenUnsafe :: (MonadIO m, Storable a)
540 => Int -> SerialT m (Array a) -> m (Array a)
541spliceArraysLenUnsafe len buffered = do
542 arr <- liftIO $ A.newArray len
543 end <- S.foldlM' writeArr (aEnd arr) buffered
544 return $ arr {aEnd = end}
546 where
548 writeArr dst Array{..} =
549 liftIO $ withForeignPtr aStart $ \src -> do
550 let count = aEnd `minusPtr` src
551 A.memcpy (castPtr dst) (castPtr src) count
552 return $ dst `plusPtr` count
554{-# INLINE _spliceArraysBuffered #-}
555_spliceArraysBuffered :: (MonadIO m, Storable a)
556 => SerialT m (Array a) -> m (Array a)
557_spliceArraysBuffered s = do
558 buffered <- P.foldr S.cons S.nil s
559 len <- S.sum (S.map length buffered)
560 spliceArraysLenUnsafe len s
562{-# INLINE spliceArraysRealloced #-}
563spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
564 => SerialT m (Array a) -> m (Array a)
565spliceArraysRealloced s = do
566 idst <- liftIO $ A.newArray (bytesToCount (undefined :: a)
567 (A.mkChunkSizeKB 4))
569 arr <- S.foldlM' A.spliceWithDoubling idst s
570 liftIO $ A.shrinkToFit arr
572-- | Given a stream of arrays, splice them all together to generate a single
573-- array. The stream must be /finite/.
574--
575-- @since 0.7.0
576{-# INLINE packArrays #-}
577packArrays :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
578packArrays = spliceArraysRealloced
579-- spliceArrays = _spliceArraysBuffered
581-- | Create an 'Array' from a stream. This is useful when we want to create a
582-- single array from a stream of unknown size. 'writeN' is at least twice
583-- as efficient when the size is already known.
584--
585-- Note that if the input stream is too large memory allocation for the array
586-- may fail. When the stream size is not known, `arraysOf` followed by
587-- processing of indvidual arrays in the resulting stream should be preferred.
588--
589{-# INLINE write #-}
590write :: (MonadIO m, Storable a) => SerialT m a -> m (Array a)
591write = FL.foldl' toArray
592-- write m = A.fromStreamD $ D.toStreamD m
594-- XXX efficiently compare two streams of arrays. Two streams can have chunks
595-- of different sizes, we can handle that in the stream comparison abstraction.
596-- This could be useful e.g. to fast compare whether two files differ.
598-------------------------------------------------------------------------------
599-- Transform via stream operations
600-------------------------------------------------------------------------------
602-- for non-length changing operations we can use the original length for
603-- allocation. If we can predict the length then we can use the prediction for
604-- new allocation. Otherwise we can use a hint and adjust dynamically.
606-- | Transform an array into another array using a stream transformation
607-- operation.
608--
609-- @since 0.7.0
610{-# INLINE transformWith #-}
611transformWith :: (MonadIO m, Storable a, Storable b)
612 => (SerialT m a -> SerialT m b) -> Array a -> m (Array b)
613transformWith f arr = FL.foldl' (toArrayMinChunk (length arr)) $ f (A.read arr)
615-- | Fold an array using a 'Fold'.
616--
617-- @since 0.7.0
618{-# INLINE foldArray #-}
619foldArray :: (MonadIO m, Storable a) => Fold m a b -> Array a -> m b
620foldArray f arr = FL.foldl' f (A.read arr)
622-- | Fold an array using a stream fold operation.
623--
624-- @since 0.7.0
625{-# INLINE foldWith #-}
626foldWith :: (MonadIO m, Storable a) => (SerialT m a -> m b) -> Array a -> m b
627foldWith f arr = f (A.read arr)