git.iko.soy/streamly
on 4300a9f help sign in
streamly/core/src/Streamly/Internal/Data/Array/Mut/Stream.hs
323 lines 12 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Array.Mut.Stream
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 mutable arrays.
10--
11module Streamly.Internal.Data.Array.Mut.Stream
12 (
13 -- * Generation
14 chunksOf
16 -- * Compaction
17 , packArraysChunksOf
18 , SpliceState (..)
19 , lpackArraysChunksOf
20 , compact
21 , compactLE
22 , compactEQ
23 , compactGE
24 )
25where
27#include "inline.hs"
28#include "ArrayMacros.h"
30import Control.Monad.IO.Class (MonadIO(..))
31import Control.Monad (when)
32import Data.Bifunctor (first)
33import Data.Proxy (Proxy(..))
34import Streamly.Internal.Data.Unboxed (Unbox, sizeOf)
35import Streamly.Internal.Data.Array.Mut.Type (MutArray(..))
36import Streamly.Internal.Data.Fold.Type (Fold(..))
37import Streamly.Internal.Data.Parser (ParseError)
38import Streamly.Internal.Data.Stream.StreamD.Type (Stream)
39import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
41import qualified Streamly.Internal.Data.Array.Mut.Type as MArray
42import qualified Streamly.Internal.Data.Fold.Type as FL
43import qualified Streamly.Internal.Data.Stream.StreamD as D
44import qualified Streamly.Internal.Data.Parser.ParserD as ParserD
46-- | @chunksOf n stream@ groups the elements in the input stream into arrays of
47-- @n@ elements each.
48--
49-- Same as the following but may be more efficient:
50--
51-- > chunksOf n = Stream.foldMany (MArray.writeN n)
52--
53-- /Pre-release/
54{-# INLINE chunksOf #-}
55chunksOf :: (MonadIO m, Unbox a)
56 => Int -> Stream m a -> Stream m (MutArray a)
57chunksOf = MArray.chunksOf
59-------------------------------------------------------------------------------
60-- Compact
61-------------------------------------------------------------------------------
63data SpliceState s arr
64 = SpliceInitial s
65 | SpliceBuffering s arr
66 | SpliceYielding arr (SpliceState s arr)
67 | SpliceFinish
69-- XXX This can be removed once compactLEFold/compactLE are implemented.
70--
71-- | This mutates the first array (if it has space) to append values from the
72-- second one. This would work for immutable arrays as well because an
73-- immutable array never has space so a new array is allocated instead of
74-- mutating it.
75--
76-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
77-- maximum specified size. Note that if a single array is bigger than the
78-- specified size we do not split it to fit. When we coalesce multiple arrays
79-- if the size would exceed the specified size we do not coalesce therefore the
80-- actual array size may be less than the specified chunk size.
81--
82-- @since 0.7.0
83{-# INLINE_NORMAL packArraysChunksOf #-}
84packArraysChunksOf :: (MonadIO m, Unbox a)
85 => Int -> D.Stream m (MutArray a) -> D.Stream m (MutArray a)
86packArraysChunksOf n (D.Stream step state) =
87 D.Stream step' (SpliceInitial state)
89 where
91 {-# INLINE_LATE step' #-}
92 step' gst (SpliceInitial st) = do
93 when (n <= 0) $
94 -- XXX we can pass the module string from the higher level API
95 error $ "Streamly.Internal.Data.Array.Mut.Type.packArraysChunksOf: the size of "
96 ++ "arrays [" ++ show n ++ "] must be a natural number"
97 r <- step gst st
98 case r of
99 D.Yield arr s -> return $
100 let len = MArray.byteLength arr
101 in if len >= n
102 then D.Skip (SpliceYielding arr (SpliceInitial s))
103 else D.Skip (SpliceBuffering s arr)
104 D.Skip s -> return $ D.Skip (SpliceInitial s)
105 D.Stop -> return D.Stop
107 step' gst (SpliceBuffering st buf) = do
108 r <- step gst st
109 case r of
110 D.Yield arr s -> do
111 let len = MArray.byteLength buf + MArray.byteLength arr
112 if len > n
113 then return $
114 D.Skip (SpliceYielding buf (SpliceBuffering s arr))
115 else do
116 buf' <- if MArray.byteCapacity buf < n
117 then liftIO $ MArray.realloc n buf
118 else return buf
119 buf'' <- MArray.splice buf' arr
120 return $ D.Skip (SpliceBuffering s buf'')
121 D.Skip s -> return $ D.Skip (SpliceBuffering s buf)
122 D.Stop -> return $ D.Skip (SpliceYielding buf SpliceFinish)
124 step' _ SpliceFinish = return D.Stop
126 step' _ (SpliceYielding arr next) = return $ D.Yield arr next
128-- XXX Remove this once compactLEFold is implemented
129-- lpackArraysChunksOf = Fold.many compactLEFold
130--
131{-# INLINE_NORMAL lpackArraysChunksOf #-}
132lpackArraysChunksOf :: (MonadIO m, Unbox a)
133 => Int -> Fold m (MutArray a) () -> Fold m (MutArray a) ()
134lpackArraysChunksOf n (Fold step1 initial1 extract1) =
135 Fold step initial extract
137 where
139 initial = do
140 when (n <= 0) $
141 -- XXX we can pass the module string from the higher level API
142 error $ "Streamly.Internal.Data.Array.Mut.Type.packArraysChunksOf: the size of "
143 ++ "arrays [" ++ show n ++ "] must be a natural number"
145 r <- initial1
146 return $ first (Tuple' Nothing) r
148 extract (Tuple' Nothing r1) = extract1 r1
149 extract (Tuple' (Just buf) r1) = do
150 r <- step1 r1 buf
151 case r of
152 FL.Partial rr -> extract1 rr
153 FL.Done _ -> return ()
155 step (Tuple' Nothing r1) arr =
156 let len = MArray.byteLength arr
157 in if len >= n
158 then do
159 r <- step1 r1 arr
160 case r of
161 FL.Done _ -> return $ FL.Done ()
162 FL.Partial s -> do
163 extract1 s
164 res <- initial1
165 return $ first (Tuple' Nothing) res
166 else return $ FL.Partial $ Tuple' (Just arr) r1
168 step (Tuple' (Just buf) r1) arr = do
169 let len = MArray.byteLength buf + MArray.byteLength arr
170 buf' <- if MArray.byteCapacity buf < len
171 then liftIO $ MArray.realloc (max n len) buf
172 else return buf
173 buf'' <- MArray.splice buf' arr
175 -- XXX this is common in both the equations of step
176 if len >= n
177 then do
178 r <- step1 r1 buf''
179 case r of
180 FL.Done _ -> return $ FL.Done ()
181 FL.Partial s -> do
182 extract1 s
183 res <- initial1
184 return $ first (Tuple' Nothing) res
185 else return $ FL.Partial $ Tuple' (Just buf'') r1
187-- XXX Same as compactLE, to be removed once that is implemented.
188--
189-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
190-- maximum specified size in bytes.
191--
192-- /Internal/
193{-# INLINE compact #-}
194compact :: (MonadIO m, Unbox a)
195 => Int -> Stream m (MutArray a) -> Stream m (MutArray a)
196compact = packArraysChunksOf
198-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
199-- maximum specified size. Note that if a single array is bigger than the
200-- specified size we do not split it to fit. When we coalesce multiple arrays
201-- if the size would exceed the specified size we do not coalesce therefore the
202-- actual array size may be less than the specified chunk size.
203--
204-- /Internal/
205{-# INLINE_NORMAL compactLEParserD #-}
206compactLEParserD ::
207 forall m a. (MonadIO m, Unbox a)
208 => Int -> ParserD.Parser (MutArray a) m (MutArray a)
209compactLEParserD n = ParserD.Parser step initial extract
211 where
213 nBytes = n * SIZE_OF(a)
215 initial =
216 return
217 $ if n <= 0
218 then error
219 $ functionPath
220 ++ ": the size of arrays ["
221 ++ show n ++ "] must be a natural number"
222 else ParserD.IPartial Nothing
224 step Nothing arr =
225 return
226 $ let len = MArray.byteLength arr
227 in if len >= nBytes
228 then ParserD.Done 0 arr
229 else ParserD.Partial 0 (Just arr)
230 step (Just buf) arr =
231 let len = MArray.byteLength buf + MArray.byteLength arr
232 in if len > nBytes
233 then return $ ParserD.Done 1 buf
234 else do
235 buf1 <-
236 if MArray.byteCapacity buf < nBytes
237 then liftIO $ MArray.realloc nBytes buf
238 else return buf
239 buf2 <- MArray.splice buf1 arr
240 return $ ParserD.Partial 0 (Just buf2)
242 extract Nothing = return $ ParserD.Done 0 MArray.nil
243 extract (Just buf) = return $ ParserD.Done 0 buf
245 functionPath =
246 "Streamly.Internal.Data.Array.Mut.Stream.compactLEParserD"
248-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
249-- minimum specified size. Note that if all the arrays in the stream together
250-- are smaller than the specified size the resulting array will be smaller than
251-- the specified size. When we coalesce multiple arrays if the size would exceed
252-- the specified size we stop coalescing further.
253--
254-- /Internal/
255{-# INLINE_NORMAL compactGEFold #-}
256compactGEFold ::
257 forall m a. (MonadIO m, Unbox a)
258 => Int -> FL.Fold m (MutArray a) (MutArray a)
259compactGEFold n = Fold step initial extract
261 where
263 nBytes = n * SIZE_OF(a)
265 initial =
266 return
267 $ if n < 0
268 then error
269 $ functionPath
270 ++ ": the size of arrays ["
271 ++ show n ++ "] must be a natural number"
272 else FL.Partial Nothing
274 step Nothing arr =
275 return
276 $ let len = MArray.byteLength arr
277 in if len >= nBytes
278 then FL.Done arr
279 else FL.Partial (Just arr)
280 step (Just buf) arr = do
281 let len = MArray.byteLength buf + MArray.byteLength arr
282 buf1 <-
283 if MArray.byteCapacity buf < len
284 then liftIO $ MArray.realloc (max len nBytes) buf
285 else return buf
286 buf2 <- MArray.splice buf1 arr
287 if len >= n
288 then return $ FL.Done buf2
289 else return $ FL.Partial (Just buf2)
291 extract Nothing = return MArray.nil
292 extract (Just buf) = return buf
294 functionPath =
295 "Streamly.Internal.Data.Array.Mut.Stream.compactGEFold"
297-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
298-- maximum specified size in bytes.
299--
300-- /Internal/
301compactLE :: (MonadIO m, Unbox a) =>
302 Int -> Stream m (MutArray a) -> Stream m (Either ParseError (MutArray a))
303compactLE n = D.parseManyD (compactLEParserD n)
305-- | Like 'compactLE' but generates arrays of exactly equal to the size
306-- specified except for the last array in the stream which could be shorter.
307--
308-- /Unimplemented/
309{-# INLINE compactEQ #-}
310compactEQ :: -- (MonadIO m, Unbox a) =>
311 Int -> Stream m (MutArray a) -> Stream m (MutArray a)
312compactEQ _n _xs = undefined
313 -- IsStream.fromStreamD $ D.foldMany (compactEQFold n) (IsStream.toStreamD xs)
315-- | Like 'compactLE' but generates arrays of size greater than or equal to the
316-- specified except for the last array in the stream which could be shorter.
317--
318-- /Internal/
319{-# INLINE compactGE #-}
320compactGE ::
321 (MonadIO m, Unbox a)
322 => Int -> Stream m (MutArray a) -> Stream m (MutArray a)
323compactGE n = D.foldMany (compactGEFold n)