git.iko.soy/streamly
on 95ebb7b help sign in
streamly/src/Streamly/Internal/Memory/ArrayStream.hs
388 lines 14 kB history raw
1-- |
2-- Module : Streamly.Internal.Memory.ArrayStream
3-- Copyright : (c) 2019 Composewell Technologies
4-- License : BSD3-3-Clause
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
8--
9-- Combinators to efficiently manipulate streams of immutable arrays.
10--
11module Streamly.Internal.Memory.ArrayStream
12 (
13 -- * Creation
14 arraysOf
16 -- * Flattening to elements
17 , concat
18 , concatRev
19 , interpose
20 , interposeSuffix
21 , intercalateSuffix
22 , unlines
24 -- * Elimination
25 , toArray
27 -- * Compaction
28 , lpackArraysChunksOf
29#if !defined(mingw32_HOST_OS)
30 , groupIOVecsOf
31#endif
32 , compact
34 -- * Splitting
35 , splitOn
36 , splitOnSuffix
37 )
38where
40#include "inline.hs"
42import Control.Monad.IO.Class (MonadIO(..))
43import Data.Word (Word8)
44import Foreign.ForeignPtr (withForeignPtr, touchForeignPtr)
45import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
46import Foreign.Ptr (minusPtr, plusPtr, castPtr)
47import Foreign.Storable (Storable(..))
48import GHC.ForeignPtr (ForeignPtr(..))
49import GHC.Ptr (Ptr(..))
50import Prelude hiding (length, null, last, map, (!!), read, concat, unlines)
52#if !defined(mingw32_HOST_OS)
53import Streamly.FileSystem.FDIO (IOVec(..))
54#endif
55import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..), length)
56import Streamly.Internal.Data.Fold.Types (Fold(..))
57import Streamly.Internal.Data.Stream.Serial (SerialT)
58import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
59import Streamly.Internal.Data.SVar (adaptState)
61import qualified Streamly.Internal.Data.Array.Storable.Foreign as A
62import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
63import qualified Streamly.Internal.Data.Array.Storable.Foreign.Mut.Types as MA
64import qualified Streamly.Internal.Data.Array.Stream.Mut.Foreign as AS
65import qualified Streamly.Internal.Data.Fold.Types as FL
66import qualified Streamly.Internal.Data.Stream.IsStream as S
67import qualified Streamly.Internal.Data.Stream.StreamD as D
69-------------------------------------------------------------------------------
70-- Generation
71-------------------------------------------------------------------------------
73-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
74-- @n@ elements each.
75--
76-- > arraysOf n = Stream.chunksOf n (Array.writeN n)
77--
78-- /Internal/
79{-# INLINE arraysOf #-}
80arraysOf :: (IsStream t, MonadIO m, Storable a)
81 => Int -> t m a -> t m (Array a)
82arraysOf n str = D.fromStreamD $ A.arraysOf n (D.toStreamD str)
84-------------------------------------------------------------------------------
85-- Append
86-------------------------------------------------------------------------------
88-- XXX efficiently compare two streams of arrays. Two streams can have chunks
89-- of different sizes, we can handle that in the stream comparison abstraction.
90-- This could be useful e.g. to fast compare whether two files differ.
92-- | Convert a stream of arrays into a stream of their elements.
93--
94-- Same as the following but more efficient:
95--
96-- > concat = S.concatMap A.read
97--
98-- @since 0.7.0
99{-# INLINE concat #-}
100concat :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
101-- concat m = D.fromStreamD $ A.flattenArrays (D.toStreamD m)
102-- concat m = D.fromStreamD $ D.concatMap A.toStreamD (D.toStreamD m)
103concat m = D.fromStreamD $ D.concatUnfold A.read (D.toStreamD m)
105-- XXX should we have a reverseArrays API to reverse the stream of arrays
106-- instead?
107--
108-- | Convert a stream of arrays into a stream of their elements reversing the
109-- contents of each array before flattening.
110--
111-- @since 0.7.0
112{-# INLINE concatRev #-}
113concatRev :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
114concatRev m = D.fromStreamD $ A.flattenArraysRev (D.toStreamD m)
116-------------------------------------------------------------------------------
117-- Intersperse and append
118-------------------------------------------------------------------------------
120-- | Flatten a stream of arrays after inserting the given element between
121-- arrays.
122--
123-- /Internal/
124{-# INLINE interpose #-}
125interpose :: (MonadIO m, IsStream t, Storable a) => a -> t m (Array a) -> t m a
126interpose x = S.interpose x A.read
128{-# INLINE intercalateSuffix #-}
129intercalateSuffix :: (MonadIO m, IsStream t, Storable a)
130 => Array a -> t m (Array a) -> t m a
131intercalateSuffix arr = S.intercalateSuffix arr A.read
133-- | Flatten a stream of arrays appending the given element after each
134-- array.
135--
136-- @since 0.7.0
137{-# INLINE interposeSuffix #-}
138interposeSuffix :: (MonadIO m, IsStream t, Storable a)
139 => a -> t m (Array a) -> t m a
140-- interposeSuffix x = D.fromStreamD . A.unlines x . D.toStreamD
141interposeSuffix x = S.interposeSuffix x A.read
143data FlattenState s a =
144 OuterLoop s
145 | InnerLoop s !(ForeignPtr a) !(Ptr a) !(Ptr a)
147{-# INLINE_NORMAL unlines #-}
148unlines :: forall m a. (MonadIO m, Storable a)
149 => a -> D.Stream m (Array a) -> D.Stream m a
150unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
151 where
152 {-# INLINE_LATE step' #-}
153 step' gst (OuterLoop st) = do
154 r <- step (adaptState gst) st
155 return $ case r of
156 D.Yield Array{..} s ->
157 let p = unsafeForeignPtrToPtr aStart
158 in D.Skip (InnerLoop s aStart p aEnd)
159 D.Skip s -> D.Skip (OuterLoop s)
160 D.Stop -> D.Stop
162 step' _ (InnerLoop st _ p end) | p == end =
163 return $ D.Yield sep $ OuterLoop st
165 step' _ (InnerLoop st startf p end) = do
166 x <- liftIO $ do
167 r <- peek p
168 touchForeignPtr startf
169 return r
170 return $ D.Yield x (InnerLoop st startf
171 (p `plusPtr` sizeOf (undefined :: a)) end)
173-------------------------------------------------------------------------------
174-- Compact
175-------------------------------------------------------------------------------
177{-# INLINE_NORMAL packArraysChunksOf #-}
178packArraysChunksOf :: (MonadIO m, Storable a)
179 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
180packArraysChunksOf n str =
181 D.map A.unsafeFreeze $ AS.packArraysChunksOf n $ D.map A.unsafeThaw str
183-- XXX instead of writing two different versions of this operation, we should
184-- write it as a pipe.
185{-# INLINE_NORMAL lpackArraysChunksOf #-}
186lpackArraysChunksOf :: (MonadIO m, Storable a)
187 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
188lpackArraysChunksOf n fld =
189 FL.map A.unsafeThaw $ AS.lpackArraysChunksOf n (FL.map A.unsafeFreeze fld)
191#if !defined(mingw32_HOST_OS)
193-- | @groupIOVecsOf maxBytes maxEntries@ groups arrays in the incoming stream
194-- to create a stream of 'IOVec' arrays with a maximum of @maxBytes@ bytes in
195-- each array and a maximum of @maxEntries@ entries in each array.
196--
197-- @since 0.7.0
198{-# INLINE_NORMAL groupIOVecsOf #-}
199groupIOVecsOf :: MonadIO m
200 => Int -> Int -> D.Stream m (Array a) -> D.Stream m (Array IOVec)
201groupIOVecsOf n maxIOVLen str =
202 D.map A.unsafeFreeze
203 $ AS.groupIOVecsOf n maxIOVLen
204 $ D.map A.unsafeThaw str
205#endif
207-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
208-- maximum specified size in bytes.
209--
210-- @since 0.7.0
211{-# INLINE compact #-}
212compact :: (MonadIO m, Storable a)
213 => Int -> SerialT m (Array a) -> SerialT m (Array a)
214compact n xs = D.fromStreamD $ packArraysChunksOf n (D.toStreamD xs)
216-------------------------------------------------------------------------------
217-- Split
218-------------------------------------------------------------------------------
220data SplitState s arr
221 = Initial s
222 | Buffering s arr
223 | Splitting s arr
224 | Yielding arr (SplitState s arr)
225 | Finishing
227-- | Split a stream of arrays on a given separator byte, dropping the separator
228-- and coalescing all the arrays between two separators into a single array.
229--
230-- @since 0.7.0
231{-# INLINE_NORMAL _splitOn #-}
232_splitOn
233 :: MonadIO m
234 => Word8
235 -> D.Stream m (Array Word8)
236 -> D.Stream m (Array Word8)
237_splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
239 where
241 {-# INLINE_LATE step' #-}
242 step' gst (Initial st) = do
243 r <- step gst st
244 case r of
245 D.Yield arr s -> do
246 (arr1, marr2) <- A.breakOn byte arr
247 return $ case marr2 of
248 Nothing -> D.Skip (Buffering s arr1)
249 Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
250 D.Skip s -> return $ D.Skip (Initial s)
251 D.Stop -> return D.Stop
253 step' gst (Buffering st buf) = do
254 r <- step gst st
255 case r of
256 D.Yield arr s -> do
257 (arr1, marr2) <- A.breakOn byte arr
258 buf' <- A.spliceTwo buf arr1
259 return $ case marr2 of
260 Nothing -> D.Skip (Buffering s buf')
261 Just x -> D.Skip (Yielding buf' (Splitting s x))
262 D.Skip s -> return $ D.Skip (Buffering s buf)
263 D.Stop -> return $
264 if A.byteLength buf == 0
265 then D.Stop
266 else D.Skip (Yielding buf Finishing)
268 step' _ (Splitting st buf) = do
269 (arr1, marr2) <- A.breakOn byte buf
270 return $ case marr2 of
271 Nothing -> D.Skip $ Buffering st arr1
272 Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
274 step' _ (Yielding arr next) = return $ D.Yield arr next
275 step' _ Finishing = return D.Stop
277-- | Split a stream of arrays on a given separator byte, dropping the separator
278-- and coalescing all the arrays between two separators into a single array.
279--
280-- @since 0.7.0
281{-# INLINE splitOn #-}
282splitOn
283 :: (IsStream t, MonadIO m)
284 => Word8
285 -> t m (Array Word8)
286 -> t m (Array Word8)
287splitOn byte s =
288 D.fromStreamD $ D.splitInnerBy (A.breakOn byte) A.spliceTwo $ D.toStreamD s
290{-# INLINE splitOnSuffix #-}
291splitOnSuffix
292 :: (IsStream t, MonadIO m)
293 => Word8
294 -> t m (Array Word8)
295 -> t m (Array Word8)
296-- splitOn byte s = D.fromStreamD $ A.splitOn byte $ D.toStreamD s
297splitOnSuffix byte s =
298 D.fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.spliceTwo $ D.toStreamD s
300-------------------------------------------------------------------------------
301-- Elimination
302-------------------------------------------------------------------------------
304-- XXX Both of these implementations of splicing seem to perform equally well.
305-- We need to perform benchmarks over a range of sizes though.
307-- CAUTION! length must more than equal to lengths of all the arrays in the
308-- stream.
309{-# INLINE spliceArraysLenUnsafe #-}
310spliceArraysLenUnsafe :: (MonadIO m, Storable a)
311 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
312spliceArraysLenUnsafe len buffered = do
313 arr <- liftIO $ MA.newArray len
314 end <- S.foldlM' writeArr (return $ MA.aEnd arr) buffered
315 return $ arr {MA.aEnd = end}
317 where
319 writeArr dst (MA.Array as ae _) =
320 liftIO $ withForeignPtr as $ \src -> do
321 let count = ae `minusPtr` src
322 A.memcpy (castPtr dst) (castPtr src) count
323 return $ dst `plusPtr` count
325{-# INLINE _spliceArrays #-}
326_spliceArrays :: (MonadIO m, Storable a)
327 => SerialT m (Array a) -> m (Array a)
328_spliceArrays s = do
329 buffered <- S.foldr S.cons S.nil s
330 len <- S.sum (S.map length buffered)
331 arr <- liftIO $ MA.newArray len
332 end <- S.foldlM' writeArr (return $ MA.aEnd arr) s
333 return $ A.unsafeFreeze $ arr {MA.aEnd = end}
335 where
337 writeArr dst (Array as ae) =
338 liftIO $ withForeignPtr as $ \src -> do
339 let count = ae `minusPtr` src
340 A.memcpy (castPtr dst) (castPtr src) count
341 return $ dst `plusPtr` count
343{-# INLINE _spliceArraysBuffered #-}
344_spliceArraysBuffered :: (MonadIO m, Storable a)
345 => SerialT m (Array a) -> m (Array a)
346_spliceArraysBuffered s = do
347 buffered <- S.foldr S.cons S.nil s
348 len <- S.sum (S.map length buffered)
349 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
351{-# INLINE spliceArraysRealloced #-}
352spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
353 => SerialT m (Array a) -> m (Array a)
354spliceArraysRealloced s = do
355 let idst = liftIO $ MA.newArray (A.bytesToElemCount (undefined :: a)
356 (A.mkChunkSizeKB 4))
358 arr <- S.foldlM' MA.spliceWithDoubling idst (S.map A.unsafeThaw s)
359 liftIO $ A.unsafeFreeze <$> MA.shrinkToFit arr
361-- | Given a stream of arrays, splice them all together to generate a single
362-- array. The stream must be /finite/.
363--
364-- @since 0.7.0
365{-# INLINE toArray #-}
366toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
367toArray = spliceArraysRealloced
368-- spliceArrays = _spliceArraysBuffered
370-- exponentially increasing sizes of the chunks upto the max limit.
371-- XXX this will be easier to implement with parsers/terminating folds
372-- With this we should be able to reduce the number of chunks/allocations.
373-- The reallocation/copy based toArray can also be implemented using this.
374--
375{-
376{-# INLINE toArraysInRange #-}
377toArraysInRange :: (IsStream t, MonadIO m, Storable a)
378 => Int -> Int -> Fold m (Array a) b -> Fold m a b
379toArraysInRange low high (Fold step initial extract) =
380-}
382{-
383-- | Fold the input to a pure buffered stream (List) of arrays.
384{-# INLINE _toArraysOf #-}
385_toArraysOf :: (MonadIO m, Storable a)
386 => Int -> Fold m a (SerialT Identity (Array a))
387_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
388-}