git.iko.soy/streamly
on 81bfdec help sign in
streamly/core/src/Streamly/Internal/Data/Array/Mut/Stream.hs
326 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 arraysOf
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 (Array(..))
36import Streamly.Internal.Data.Fold.Type (Fold(..))
37import Streamly.Internal.Data.Parser (ParseError)
38import Streamly.Internal.Data.Stream.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
45import qualified Streamly.Internal.Data.Stream.Type as Stream
47-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
48-- @n@ elements each.
49--
50-- Same as the following but may be more efficient:
51--
52-- > arraysOf n = Stream.foldMany (MArray.writeN n)
53--
54-- /Pre-release/
55{-# INLINE arraysOf #-}
56arraysOf :: (MonadIO m, Unbox a)
57 => Int -> Stream m a -> Stream m (Array a)
58arraysOf n = Stream.fromStreamD . MArray.arraysOf n . Stream.toStreamD
60-------------------------------------------------------------------------------
61-- Compact
62-------------------------------------------------------------------------------
64data SpliceState s arr
65 = SpliceInitial s
66 | SpliceBuffering s arr
67 | SpliceYielding arr (SpliceState s arr)
68 | SpliceFinish
70-- XXX This can be removed once compactLEFold/compactLE are implemented.
71--
72-- | This mutates the first array (if it has space) to append values from the
73-- second one. This would work for immutable arrays as well because an
74-- immutable array never has space so a new array is allocated instead of
75-- mutating it.
76--
77-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
78-- maximum specified size. Note that if a single array is bigger than the
79-- specified size we do not split it to fit. When we coalesce multiple arrays
80-- if the size would exceed the specified size we do not coalesce therefore the
81-- actual array size may be less than the specified chunk size.
82--
83-- @since 0.7.0
84{-# INLINE_NORMAL packArraysChunksOf #-}
85packArraysChunksOf :: (MonadIO m, Unbox a)
86 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
87packArraysChunksOf n (D.Stream step state) =
88 D.Stream step' (SpliceInitial state)
90 where
92 {-# INLINE_LATE step' #-}
93 step' gst (SpliceInitial st) = do
94 when (n <= 0) $
95 -- XXX we can pass the module string from the higher level API
96 error $ "Streamly.Internal.Data.Array.Mut.Type.packArraysChunksOf: the size of "
97 ++ "arrays [" ++ show n ++ "] must be a natural number"
98 r <- step gst st
99 case r of
100 D.Yield arr s -> return $
101 let len = MArray.byteLength arr
102 in if len >= n
103 then D.Skip (SpliceYielding arr (SpliceInitial s))
104 else D.Skip (SpliceBuffering s arr)
105 D.Skip s -> return $ D.Skip (SpliceInitial s)
106 D.Stop -> return D.Stop
108 step' gst (SpliceBuffering st buf) = do
109 r <- step gst st
110 case r of
111 D.Yield arr s -> do
112 let len = MArray.byteLength buf + MArray.byteLength arr
113 if len > n
114 then return $
115 D.Skip (SpliceYielding buf (SpliceBuffering s arr))
116 else do
117 buf' <- if MArray.byteCapacity buf < n
118 then liftIO $ MArray.realloc n buf
119 else return buf
120 buf'' <- MArray.splice buf' arr
121 return $ D.Skip (SpliceBuffering s buf'')
122 D.Skip s -> return $ D.Skip (SpliceBuffering s buf)
123 D.Stop -> return $ D.Skip (SpliceYielding buf SpliceFinish)
125 step' _ SpliceFinish = return D.Stop
127 step' _ (SpliceYielding arr next) = return $ D.Yield arr next
129-- XXX Remove this once compactLEFold is implemented
130-- lpackArraysChunksOf = Fold.many compactLEFold
131--
132{-# INLINE_NORMAL lpackArraysChunksOf #-}
133lpackArraysChunksOf :: (MonadIO m, Unbox a)
134 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
135lpackArraysChunksOf n (Fold step1 initial1 extract1) =
136 Fold step initial extract
138 where
140 initial = do
141 when (n <= 0) $
142 -- XXX we can pass the module string from the higher level API
143 error $ "Streamly.Internal.Data.Array.Mut.Type.packArraysChunksOf: the size of "
144 ++ "arrays [" ++ show n ++ "] must be a natural number"
146 r <- initial1
147 return $ first (Tuple' Nothing) r
149 extract (Tuple' Nothing r1) = extract1 r1
150 extract (Tuple' (Just buf) r1) = do
151 r <- step1 r1 buf
152 case r of
153 FL.Partial rr -> extract1 rr
154 FL.Done _ -> return ()
156 step (Tuple' Nothing r1) arr =
157 let len = MArray.byteLength arr
158 in if len >= n
159 then do
160 r <- step1 r1 arr
161 case r of
162 FL.Done _ -> return $ FL.Done ()
163 FL.Partial s -> do
164 extract1 s
165 res <- initial1
166 return $ first (Tuple' Nothing) res
167 else return $ FL.Partial $ Tuple' (Just arr) r1
169 step (Tuple' (Just buf) r1) arr = do
170 let len = MArray.byteLength buf + MArray.byteLength arr
171 buf' <- if MArray.byteCapacity buf < len
172 then liftIO $ MArray.realloc (max n len) buf
173 else return buf
174 buf'' <- MArray.splice buf' arr
176 -- XXX this is common in both the equations of step
177 if len >= n
178 then do
179 r <- step1 r1 buf''
180 case r of
181 FL.Done _ -> return $ FL.Done ()
182 FL.Partial s -> do
183 extract1 s
184 res <- initial1
185 return $ first (Tuple' Nothing) res
186 else return $ FL.Partial $ Tuple' (Just buf'') r1
188-- XXX Same as compactLE, to be removed once that is implemented.
189--
190-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
191-- maximum specified size in bytes.
192--
193-- /Internal/
194{-# INLINE compact #-}
195compact :: (MonadIO m, Unbox a)
196 => Int -> Stream m (Array a) -> Stream m (Array a)
197compact n = Stream.fromStreamD . packArraysChunksOf n . Stream.toStreamD
199-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
200-- maximum specified size. Note that if a single array is bigger than the
201-- specified size we do not split it to fit. When we coalesce multiple arrays
202-- if the size would exceed the specified size we do not coalesce therefore the
203-- actual array size may be less than the specified chunk size.
204--
205-- /Internal/
206{-# INLINE_NORMAL compactLEParserD #-}
207compactLEParserD ::
208 forall m a. (MonadIO m, Unbox a)
209 => Int -> ParserD.Parser (Array a) m (Array a)
210compactLEParserD n = ParserD.Parser step initial extract
212 where
214 nBytes = n * SIZE_OF(a)
216 initial =
217 return
218 $ if n <= 0
219 then error
220 $ functionPath
221 ++ ": the size of arrays ["
222 ++ show n ++ "] must be a natural number"
223 else ParserD.IPartial Nothing
225 step Nothing arr =
226 return
227 $ let len = MArray.byteLength arr
228 in if len >= nBytes
229 then ParserD.Done 0 arr
230 else ParserD.Partial 0 (Just arr)
231 step (Just buf) arr =
232 let len = MArray.byteLength buf + MArray.byteLength arr
233 in if len > nBytes
234 then return $ ParserD.Done 1 buf
235 else do
236 buf1 <-
237 if MArray.byteCapacity buf < nBytes
238 then liftIO $ MArray.realloc nBytes buf
239 else return buf
240 buf2 <- MArray.splice buf1 arr
241 return $ ParserD.Partial 0 (Just buf2)
243 extract Nothing = return $ ParserD.Done 0 MArray.nil
244 extract (Just buf) = return $ ParserD.Done 0 buf
246 functionPath =
247 "Streamly.Internal.Data.Array.Mut.Stream.compactLEParserD"
249-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
250-- minimum specified size. Note that if all the arrays in the stream together
251-- are smaller than the specified size the resulting array will be smaller than
252-- the specified size. When we coalesce multiple arrays if the size would exceed
253-- the specified size we stop coalescing further.
254--
255-- /Internal/
256{-# INLINE_NORMAL compactGEFold #-}
257compactGEFold ::
258 forall m a. (MonadIO m, Unbox a)
259 => Int -> FL.Fold m (Array a) (Array a)
260compactGEFold n = Fold step initial extract
262 where
264 nBytes = n * SIZE_OF(a)
266 initial =
267 return
268 $ if n < 0
269 then error
270 $ functionPath
271 ++ ": the size of arrays ["
272 ++ show n ++ "] must be a natural number"
273 else FL.Partial Nothing
275 step Nothing arr =
276 return
277 $ let len = MArray.byteLength arr
278 in if len >= nBytes
279 then FL.Done arr
280 else FL.Partial (Just arr)
281 step (Just buf) arr = do
282 let len = MArray.byteLength buf + MArray.byteLength arr
283 buf1 <-
284 if MArray.byteCapacity buf < len
285 then liftIO $ MArray.realloc (max len nBytes) buf
286 else return buf
287 buf2 <- MArray.splice buf1 arr
288 if len >= n
289 then return $ FL.Done buf2
290 else return $ FL.Partial (Just buf2)
292 extract Nothing = return MArray.nil
293 extract (Just buf) = return buf
295 functionPath =
296 "Streamly.Internal.Data.Array.Mut.Stream.compactGEFold"
298-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
299-- maximum specified size in bytes.
300--
301-- /Internal/
302compactLE :: (Monad m, MonadIO m, Unbox a) =>
303 Int -> Stream m (Array a) -> Stream m (Either ParseError (Array a))
304compactLE n =
305 Stream.fromStreamD . D.parseMany (compactLEParserD n) . Stream.toStreamD
307-- | Like 'compactLE' but generates arrays of exactly equal to the size
308-- specified except for the last array in the stream which could be shorter.
309--
310-- /Unimplemented/
311{-# INLINE compactEQ #-}
312compactEQ :: -- (MonadIO m, Unbox a) =>
313 Int -> Stream m (Array a) -> Stream m (Array a)
314compactEQ _n _xs = undefined
315 -- IsStream.fromStreamD $ D.foldMany (compactEQFold n) (IsStream.toStreamD xs)
317-- | Like 'compactLE' but generates arrays of size greater than or equal to the
318-- specified except for the last array in the stream which could be shorter.
319--
320-- /Internal/
321{-# INLINE compactGE #-}
322compactGE ::
323 (MonadIO m, Unbox a)
324 => Int -> Stream m (Array a) -> Stream m (Array a)
325compactGE n =
326 Stream.fromStreamD . D.foldMany (compactGEFold n) . Stream.toStreamD