git.iko.soy/streamly
on 956faed help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Mut/Foreign.hs
326 lines 12 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Array.Stream.Mut.Foreign
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.Stream.Mut.Foreign
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 Control.Monad.Catch (MonadThrow)
33import Data.Bifunctor (first)
34import Foreign.Storable (Storable(..))
35import Streamly.Internal.Data.Array.Foreign.Mut.Type (Array(..))
36import Streamly.Internal.Data.Fold.Type (Fold(..))
37import Streamly.Internal.Data.Stream.Serial (SerialT(..))
38import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
40import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MArray
41import qualified Streamly.Internal.Data.Fold.Type as FL
42import qualified Streamly.Internal.Data.Stream.StreamD as D
43import qualified Streamly.Internal.Data.Parser.ParserD as ParserD
45-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
46-- @n@ elements each.
47--
48-- Same as the following but may be more efficient:
49--
50-- > arraysOf n = Stream.foldMany (MArray.writeN n)
51--
52-- /Pre-release/
53{-# INLINE arraysOf #-}
54arraysOf :: (MonadIO m, Storable a)
55 => Int -> SerialT m a -> SerialT m (Array a)
56arraysOf n (SerialT xs) =
57 SerialT $ D.toStreamK $ MArray.arraysOf n $ D.fromStreamK xs
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, Storable a)
85 => Int -> D.Stream m (Array a) -> D.Stream m (Array 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.Foreign.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, Storable a)
133 => Int -> Fold m (Array a) () -> Fold m (Array 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.Foreign.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, Storable a)
195 => Int -> SerialT m (Array a) -> SerialT m (Array a)
196compact n (SerialT xs) =
197 SerialT $ D.toStreamK $ packArraysChunksOf n (D.fromStreamK xs)
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. (MonadThrow m, MonadIO m, Storable a)
209 => Int -> ParserD.Parser m (Array a) (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 MArray.nil
244 extract (Just buf) = return buf
246 functionPath =
247 "Streamly.Internal.Data.Array.Stream.Mut.Foreign.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, Storable 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.Stream.Mut.Foreign.compactGEFold"
298-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
299-- maximum specified size in bytes.
300--
301-- /Internal/
302compactLE :: (MonadThrow m, MonadIO m, Storable a) =>
303 Int -> SerialT m (Array a) -> SerialT m (Array a)
304compactLE n (SerialT xs) =
305 SerialT $ D.toStreamK $ D.parseMany (compactLEParserD n) (D.fromStreamK xs)
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, Storable a) =>
313 Int -> SerialT m (Array a) -> SerialT 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, Storable a)
324 => Int -> SerialT m (Array a) -> SerialT m (Array a)
325compactGE n (SerialT xs) =
326 SerialT $ D.toStreamK $ D.foldMany (compactGEFold n) (D.fromStreamK xs)