git.iko.soy/streamly
on fa9fa09 help sign in
streamly/core/src/Streamly/Internal/Data/MutArray/Stream.hs
329 lines 12 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.MutArray.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--
11-- We can either push these in the MutArray module with a "chunks" prefix or
12-- keep this as a separate module and release it.
13--
14module Streamly.Internal.Data.MutArray.Stream
15 (
16 -- * Generation
17 MArray.chunksOf
18 , MArray.pinnedChunksOf
19 , MArray.writeChunks -- chunksWrite?
20 , MArray.splitOn -- chunksSplitOn
22 -- * Compaction
23 , packArraysChunksOf
24 , SpliceState (..)
25 , lpackArraysChunksOf
26 , compact -- chunksCompact
27 , compactLE
28 , compactEQ
29 , compactGE
31 -- * Elimination
32 , MArray.flattenArrays -- chunksConcat
33 , MArray.flattenArraysRev -- chunksConcatRev
34 , MArray.fromArrayStreamK -- chunksCoalesce
35 )
36where
38#include "inline.hs"
39#include "ArrayMacros.h"
41import Control.Monad.IO.Class (MonadIO(..))
42import Control.Monad (when)
43import Data.Bifunctor (first)
44import Data.Proxy (Proxy(..))
45import Streamly.Internal.Data.Unbox (Unbox, sizeOf)
46import Streamly.Internal.Data.MutArray.Type (MutArray(..))
47import Streamly.Internal.Data.Fold.Type (Fold(..))
48import Streamly.Internal.Data.Parser (ParseError)
49import Streamly.Internal.Data.Stream.Type (Stream)
50import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
52import qualified Streamly.Internal.Data.MutArray as MArray
53import qualified Streamly.Internal.Data.Fold.Type as FL
54import qualified Streamly.Internal.Data.Parser as ParserD
55import qualified Streamly.Internal.Data.Stream as D
57-------------------------------------------------------------------------------
58-- Compact
59-------------------------------------------------------------------------------
61data SpliceState s arr
62 = SpliceInitial s
63 | SpliceBuffering s arr
64 | SpliceYielding arr (SpliceState s arr)
65 | SpliceFinish
67-- XXX This can be removed once compactLEFold/compactLE are implemented.
68--
69-- | This mutates the first array (if it has space) to append values from the
70-- second one. This would work for immutable arrays as well because an
71-- immutable array never has space so a new array is allocated instead of
72-- mutating it.
73--
74-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
75-- maximum specified size. Note that if a single array is bigger than the
76-- specified size we do not split it to fit. When we coalesce multiple arrays
77-- if the size would exceed the specified size we do not coalesce therefore the
78-- actual array size may be less than the specified chunk size.
79--
80-- @since 0.7.0
81{-# INLINE_NORMAL packArraysChunksOf #-}
82packArraysChunksOf :: (MonadIO m, Unbox a)
83 => Int -> D.Stream m (MutArray a) -> D.Stream m (MutArray a)
84packArraysChunksOf n (D.Stream step state) =
85 D.Stream step' (SpliceInitial state)
87 where
89 {-# INLINE_LATE step' #-}
90 step' gst (SpliceInitial st) = do
91 when (n <= 0) $
92 -- XXX we can pass the module string from the higher level API
93 error $ "Streamly.Internal.Data.MutArray.Type.packArraysChunksOf: the size of "
94 ++ "arrays [" ++ show n ++ "] must be a natural number"
95 r <- step gst st
96 case r of
97 D.Yield arr s -> return $
98 let len = MArray.byteLength arr
99 in if len >= n
100 then D.Skip (SpliceYielding arr (SpliceInitial s))
101 else D.Skip (SpliceBuffering s arr)
102 D.Skip s -> return $ D.Skip (SpliceInitial s)
103 D.Stop -> return D.Stop
105 step' gst (SpliceBuffering st buf) = do
106 r <- step gst st
107 case r of
108 D.Yield arr s -> do
109 let len = MArray.byteLength buf + MArray.byteLength arr
110 if len > n
111 then return $
112 D.Skip (SpliceYielding buf (SpliceBuffering s arr))
113 else do
114 buf' <- if MArray.byteCapacity buf < n
115 then liftIO $ MArray.realloc n buf
116 else return buf
117 buf'' <- MArray.splice buf' arr
118 return $ D.Skip (SpliceBuffering s buf'')
119 D.Skip s -> return $ D.Skip (SpliceBuffering s buf)
120 D.Stop -> return $ D.Skip (SpliceYielding buf SpliceFinish)
122 step' _ SpliceFinish = return D.Stop
124 step' _ (SpliceYielding arr next) = return $ D.Yield arr next
126-- XXX Remove this once compactLEFold is implemented
127-- lpackArraysChunksOf = Fold.many compactLEFold
128--
129{-# INLINE_NORMAL lpackArraysChunksOf #-}
130lpackArraysChunksOf :: (MonadIO m, Unbox a)
131 => Int -> Fold m (MutArray a) () -> Fold m (MutArray a) ()
132lpackArraysChunksOf n (Fold step1 initial1 _ final1) =
133 Fold step initial extract final
135 where
137 initial = do
138 when (n <= 0) $
139 -- XXX we can pass the module string from the higher level API
140 error $ "Streamly.Internal.Data.MutArray.Type.packArraysChunksOf: the size of "
141 ++ "arrays [" ++ show n ++ "] must be a natural number"
143 r <- initial1
144 return $ first (Tuple' Nothing) r
146 step (Tuple' Nothing r1) arr =
147 let len = MArray.byteLength arr
148 in if len >= n
149 then do
150 r <- step1 r1 arr
151 case r of
152 FL.Done _ -> return $ FL.Done ()
153 FL.Partial s -> do
154 _ <- final1 s
155 res <- initial1
156 return $ first (Tuple' Nothing) res
157 else return $ FL.Partial $ Tuple' (Just arr) r1
159 step (Tuple' (Just buf) r1) arr = do
160 let len = MArray.byteLength buf + MArray.byteLength arr
161 buf' <- if MArray.byteCapacity buf < len
162 then liftIO $ MArray.realloc (max n len) buf
163 else return buf
164 buf'' <- MArray.splice buf' arr
166 -- XXX this is common in both the equations of step
167 if len >= n
168 then do
169 r <- step1 r1 buf''
170 case r of
171 FL.Done _ -> return $ FL.Done ()
172 FL.Partial s -> do
173 _ <- final1 s
174 res <- initial1
175 return $ first (Tuple' Nothing) res
176 else return $ FL.Partial $ Tuple' (Just buf'') r1
178 -- XXX Several folds do extract >=> final, therefore, we need to make final
179 -- return "m b" rather than using extract post it if we want extract to be
180 -- partial.
181 --
182 -- extract forces the pending buffer to be sent to the fold which is not
183 -- what we want.
184 extract _ = error "lpackArraysChunksOf: not designed for scanning"
186 final (Tuple' Nothing r1) = final1 r1
187 final (Tuple' (Just buf) r1) = do
188 r <- step1 r1 buf
189 case r of
190 FL.Partial rr -> final1 rr
191 FL.Done _ -> return ()
193-- XXX Same as compactLE, to be removed once that is implemented.
194--
195-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
196-- maximum specified size in bytes.
197--
198-- /Internal/
199{-# INLINE compact #-}
200compact :: (MonadIO m, Unbox a)
201 => Int -> Stream m (MutArray a) -> Stream m (MutArray a)
202compact = packArraysChunksOf
204-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
205-- maximum specified size. Note that if a single array is bigger than the
206-- specified size we do not split it to fit. When we coalesce multiple arrays
207-- if the size would exceed the specified size we do not coalesce therefore the
208-- actual array size may be less than the specified chunk size.
209--
210-- /Internal/
211{-# INLINE_NORMAL compactLEParserD #-}
212compactLEParserD ::
213 forall m a. (MonadIO m, Unbox a)
214 => Int -> ParserD.Parser (MutArray a) m (MutArray a)
215compactLEParserD n = ParserD.Parser step initial extract
217 where
219 nBytes = n * SIZE_OF(a)
221 initial =
222 return
223 $ if n <= 0
224 then error
225 $ functionPath
226 ++ ": the size of arrays ["
227 ++ show n ++ "] must be a natural number"
228 else ParserD.IPartial Nothing
230 step Nothing arr =
231 return
232 $ let len = MArray.byteLength arr
233 in if len >= nBytes
234 then ParserD.Done 0 arr
235 else ParserD.Partial 0 (Just arr)
236 step (Just buf) arr =
237 let len = MArray.byteLength buf + MArray.byteLength arr
238 in if len > nBytes
239 then return $ ParserD.Done 1 buf
240 else do
241 buf1 <-
242 if MArray.byteCapacity buf < nBytes
243 then liftIO $ MArray.realloc nBytes buf
244 else return buf
245 buf2 <- MArray.splice buf1 arr
246 return $ ParserD.Partial 0 (Just buf2)
248 extract Nothing = return $ ParserD.Done 0 MArray.nil
249 extract (Just buf) = return $ ParserD.Done 0 buf
251 functionPath =
252 "Streamly.Internal.Data.Stream.MutChunked.compactLEParserD"
254-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
255-- minimum specified size. Note that if all the arrays in the stream together
256-- are smaller than the specified size the resulting array will be smaller than
257-- the specified size. When we coalesce multiple arrays if the size would exceed
258-- the specified size we stop coalescing further.
259--
260-- /Internal/
261{-# INLINE_NORMAL compactGEFold #-}
262compactGEFold ::
263 forall m a. (MonadIO m, Unbox a)
264 => Int -> FL.Fold m (MutArray a) (MutArray a)
265compactGEFold n = Fold step initial extract extract
267 where
269 nBytes = n * SIZE_OF(a)
271 initial =
272 return
273 $ if n < 0
274 then error
275 $ functionPath
276 ++ ": the size of arrays ["
277 ++ show n ++ "] must be a natural number"
278 else FL.Partial Nothing
280 step Nothing arr =
281 return
282 $ let len = MArray.byteLength arr
283 in if len >= nBytes
284 then FL.Done arr
285 else FL.Partial (Just arr)
286 step (Just buf) arr = do
287 let len = MArray.byteLength buf + MArray.byteLength arr
288 buf1 <-
289 if MArray.byteCapacity buf < len
290 then liftIO $ MArray.realloc (max len nBytes) buf
291 else return buf
292 buf2 <- MArray.splice buf1 arr
293 if len >= n
294 then return $ FL.Done buf2
295 else return $ FL.Partial (Just buf2)
297 extract Nothing = return MArray.nil
298 extract (Just buf) = return buf
300 functionPath =
301 "Streamly.Internal.Data.Stream.MutChunked.compactGEFold"
303-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
304-- maximum specified size in bytes.
305--
306-- /Internal/
307compactLE :: (MonadIO m, Unbox a) =>
308 Int -> Stream m (MutArray a) -> Stream m (Either ParseError (MutArray a))
309compactLE n = D.parseManyD (compactLEParserD n)
311-- | Like 'compactLE' but generates arrays of exactly equal to the size
312-- specified except for the last array in the stream which could be shorter.
313--
314-- /Unimplemented/
315{-# INLINE compactEQ #-}
316compactEQ :: -- (MonadIO m, Unbox a) =>
317 Int -> Stream m (MutArray a) -> Stream m (MutArray a)
318compactEQ _n _xs = undefined
319 -- IsStream.fromStreamD $ D.foldMany (compactEQFold n) (IsStream.toStreamD xs)
321-- | Like 'compactLE' but generates arrays of size greater than or equal to the
322-- specified except for the last array in the stream which could be shorter.
323--
324-- /Internal/
325{-# INLINE compactGE #-}
326compactGE ::
327 (MonadIO m, Unbox a)
328 => Int -> Stream m (MutArray a) -> Stream m (MutArray a)
329compactGE n = D.foldMany (compactGEFold n)