git.iko.soy/streamly
on cc44d97 help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Mut/Foreign.hs
327 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.Stream.IsStream.Type
39 (IsStream, fromStreamD, toStreamD)
40import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
42import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MArray
43import qualified Streamly.Internal.Data.Fold.Type as FL
44import qualified Streamly.Internal.Data.Stream.StreamD as D
45import qualified Streamly.Internal.Data.Parser.ParserD as ParserD
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 :: (IsStream t, MonadIO m, Storable a)
57 => Int -> t m a -> t m (Array a)
58arraysOf n = fromStreamD . MArray.arraysOf n . 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, Storable 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.Foreign.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, Storable 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.Foreign.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, Storable a)
196 => Int -> SerialT m (Array a) -> SerialT m (Array a)
197compact n (SerialT xs) =
198 SerialT $ D.toStreamK $ packArraysChunksOf n (D.fromStreamK xs)
200-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
201-- maximum specified size. Note that if a single array is bigger than the
202-- specified size we do not split it to fit. When we coalesce multiple arrays
203-- if the size would exceed the specified size we do not coalesce therefore the
204-- actual array size may be less than the specified chunk size.
205--
206-- /Internal/
207{-# INLINE_NORMAL compactLEParserD #-}
208compactLEParserD ::
209 forall m a. (MonadThrow m, MonadIO m, Storable a)
210 => Int -> ParserD.Parser m (Array a) (Array a)
211compactLEParserD n = ParserD.Parser step initial extract
213 where
215 nBytes = n * SIZE_OF(a)
217 initial =
218 return
219 $ if n <= 0
220 then error
221 $ functionPath
222 ++ ": the size of arrays ["
223 ++ show n ++ "] must be a natural number"
224 else ParserD.IPartial Nothing
226 step Nothing arr =
227 return
228 $ let len = MArray.byteLength arr
229 in if len >= nBytes
230 then ParserD.Done 0 arr
231 else ParserD.Partial 0 (Just arr)
232 step (Just buf) arr =
233 let len = MArray.byteLength buf + MArray.byteLength arr
234 in if len > nBytes
235 then return $ ParserD.Done 1 buf
236 else do
237 buf1 <-
238 if MArray.byteCapacity buf < nBytes
239 then liftIO $ MArray.realloc nBytes buf
240 else return buf
241 buf2 <- MArray.splice buf1 arr
242 return $ ParserD.Partial 0 (Just buf2)
244 extract Nothing = return MArray.nil
245 extract (Just buf) = return buf
247 functionPath =
248 "Streamly.Internal.Data.Array.Stream.Mut.Foreign.compactLEParserD"
250-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
251-- minimum specified size. Note that if all the arrays in the stream together
252-- are smaller than the specified size the resulting array will be smaller than
253-- the specified size. When we coalesce multiple arrays if the size would exceed
254-- the specified size we stop coalescing further.
255--
256-- /Internal/
257{-# INLINE_NORMAL compactGEFold #-}
258compactGEFold ::
259 forall m a. (MonadIO m, Storable a)
260 => Int -> FL.Fold m (Array a) (Array a)
261compactGEFold n = Fold step initial extract
263 where
265 nBytes = n * SIZE_OF(a)
267 initial =
268 return
269 $ if n < 0
270 then error
271 $ functionPath
272 ++ ": the size of arrays ["
273 ++ show n ++ "] must be a natural number"
274 else FL.Partial Nothing
276 step Nothing arr =
277 return
278 $ let len = MArray.byteLength arr
279 in if len >= nBytes
280 then FL.Done arr
281 else FL.Partial (Just arr)
282 step (Just buf) arr = do
283 let len = MArray.byteLength buf + MArray.byteLength arr
284 buf1 <-
285 if MArray.byteCapacity buf < len
286 then liftIO $ MArray.realloc (max len nBytes) buf
287 else return buf
288 buf2 <- MArray.splice buf1 arr
289 if len >= n
290 then return $ FL.Done buf2
291 else return $ FL.Partial (Just buf2)
293 extract Nothing = return MArray.nil
294 extract (Just buf) = return buf
296 functionPath =
297 "Streamly.Internal.Data.Array.Stream.Mut.Foreign.compactGEFold"
299-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
300-- maximum specified size in bytes.
301--
302-- /Internal/
303compactLE :: (MonadThrow m, MonadIO m, Storable a) =>
304 Int -> SerialT m (Array a) -> SerialT m (Array a)
305compactLE n (SerialT xs) =
306 SerialT $ D.toStreamK $ D.parseMany (compactLEParserD n) (D.fromStreamK xs)
308-- | Like 'compactLE' but generates arrays of exactly equal to the size
309-- specified except for the last array in the stream which could be shorter.
310--
311-- /Unimplemented/
312{-# INLINE compactEQ #-}
313compactEQ :: -- (MonadIO m, Storable a) =>
314 Int -> SerialT m (Array a) -> SerialT m (Array a)
315compactEQ _n _xs = undefined
316 -- IsStream.fromStreamD $ D.foldMany (compactEQFold n) (IsStream.toStreamD xs)
318-- | Like 'compactLE' but generates arrays of size greater than or equal to the
319-- specified except for the last array in the stream which could be shorter.
320--
321-- /Internal/
322{-# INLINE compactGE #-}
323compactGE ::
324 (MonadIO m, Storable a)
325 => Int -> SerialT m (Array a) -> SerialT m (Array a)
326compactGE n (SerialT xs) =
327 SerialT $ D.toStreamK $ D.foldMany (compactGEFold n) (D.fromStreamK xs)