git.iko.soy/streamly
on 03218ad help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Mut/Foreign.hs
311 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#if !defined(mingw32_HOST_OS)
21 , groupIOVecsOf
22#endif
23 , compact
24 , compactLE
25 , compactEQ
26 , compactGE
27 )
28where
30#include "inline.hs"
32import Control.Monad.IO.Class (MonadIO(..))
33import Control.Monad (when)
34import Data.Bifunctor (first)
35import Foreign.Storable (Storable(..))
36#if !defined(mingw32_HOST_OS)
37import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
38import Foreign.Ptr (castPtr)
39import Streamly.Internal.FileSystem.FDIO (IOVec(..))
40import Streamly.Internal.Data.Array.Foreign.Mut.Type (length)
41import Streamly.Internal.Data.SVar (adaptState)
42#endif
43import Streamly.Internal.Data.Array.Foreign.Mut.Type (Array(..))
44import Streamly.Internal.Data.Fold.Type (Fold(..))
45import Streamly.Internal.Data.Stream.Serial (SerialT)
46import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
47import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
49import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MArray
50import qualified Streamly.Internal.Data.Fold.Type as FL
51import qualified Streamly.Internal.Data.Stream.StreamD as D
53import Prelude hiding (length)
55-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
56-- @n@ elements each.
57--
58-- Same as the following but may be more efficient:
59--
60-- > arraysOf n = Stream.foldMany (MArray.writeN n)
61--
62-- /Pre-release/
63{-# INLINE arraysOf #-}
64arraysOf :: (IsStream t, MonadIO m, Storable a)
65 => Int -> t m a -> t m (Array a)
66arraysOf n = D.fromStreamD . MArray.arraysOf n . D.toStreamD
68-------------------------------------------------------------------------------
69-- Compact
70-------------------------------------------------------------------------------
72data SpliceState s arr
73 = SpliceInitial s
74 | SpliceBuffering s arr
75 | SpliceYielding arr (SpliceState s arr)
76 | SpliceFinish
78-- | This mutates the first array (if it has space) to append values from the
79-- second one. This would work for immutable arrays as well because an
80-- immutable array never has space so a new array is allocated instead of
81-- mutating it.
82--
83-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
84-- maximum specified size. Note that if a single array is bigger than the
85-- specified size we do not split it to fit. When we coalesce multiple arrays
86-- if the size would exceed the specified size we do not coalesce therefore the
87-- actual array size may be less than the specified chunk size.
88--
89-- @since 0.7.0
90{-# INLINE_NORMAL packArraysChunksOf #-}
91packArraysChunksOf :: (MonadIO m, Storable a)
92 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
93packArraysChunksOf n (D.Stream step state) =
94 D.Stream step' (SpliceInitial state)
96 where
98 {-# INLINE_LATE step' #-}
99 step' gst (SpliceInitial st) = do
100 when (n <= 0) $
101 -- XXX we can pass the module string from the higher level API
102 error $ "Streamly.Internal.Data.Array.Foreign.Mut.Type.packArraysChunksOf: the size of "
103 ++ "arrays [" ++ show n ++ "] must be a natural number"
104 r <- step gst st
105 case r of
106 D.Yield arr s -> return $
107 let len = MArray.byteLength arr
108 in if len >= n
109 then D.Skip (SpliceYielding arr (SpliceInitial s))
110 else D.Skip (SpliceBuffering s arr)
111 D.Skip s -> return $ D.Skip (SpliceInitial s)
112 D.Stop -> return D.Stop
114 step' gst (SpliceBuffering st buf) = do
115 r <- step gst st
116 case r of
117 D.Yield arr s -> do
118 let len = MArray.byteLength buf + MArray.byteLength arr
119 if len > n
120 then return $
121 D.Skip (SpliceYielding buf (SpliceBuffering s arr))
122 else do
123 buf' <- if MArray.byteCapacity buf < n
124 then liftIO $ MArray.realloc n buf
125 else return buf
126 buf'' <- MArray.spliceWith buf' arr
127 return $ D.Skip (SpliceBuffering s buf'')
128 D.Skip s -> return $ D.Skip (SpliceBuffering s buf)
129 D.Stop -> return $ D.Skip (SpliceYielding buf SpliceFinish)
131 step' _ SpliceFinish = return D.Stop
133 step' _ (SpliceYielding arr next) = return $ D.Yield arr next
135{-# INLINE_NORMAL lpackArraysChunksOf #-}
136lpackArraysChunksOf :: (MonadIO m, Storable a)
137 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
138lpackArraysChunksOf n (Fold step1 initial1 extract1) =
139 Fold step initial extract
141 where
143 initial = do
144 when (n <= 0) $
145 -- XXX we can pass the module string from the higher level API
146 error $ "Streamly.Internal.Data.Array.Foreign.Mut.Type.packArraysChunksOf: the size of "
147 ++ "arrays [" ++ show n ++ "] must be a natural number"
149 r <- initial1
150 return $ first (Tuple' Nothing) r
152 extract (Tuple' Nothing r1) = extract1 r1
153 extract (Tuple' (Just buf) r1) = do
154 r <- step1 r1 buf
155 case r of
156 FL.Partial rr -> extract1 rr
157 FL.Done _ -> return ()
159 step (Tuple' Nothing r1) arr =
160 let len = MArray.byteLength arr
161 in if len >= n
162 then do
163 r <- step1 r1 arr
164 case r of
165 FL.Done _ -> return $ FL.Done ()
166 FL.Partial s -> do
167 extract1 s
168 res <- initial1
169 return $ first (Tuple' Nothing) res
170 else return $ FL.Partial $ Tuple' (Just arr) r1
172 step (Tuple' (Just buf) r1) arr = do
173 let len = MArray.byteLength buf + MArray.byteLength arr
174 buf' <- if MArray.byteCapacity buf < len
175 then liftIO $ MArray.realloc (max n len) buf
176 else return buf
177 buf'' <- MArray.spliceWith buf' arr
179 -- XXX this is common in both the equations of step
180 if len >= n
181 then do
182 r <- step1 r1 buf''
183 case r of
184 FL.Done _ -> return $ FL.Done ()
185 FL.Partial s -> do
186 extract1 s
187 res <- initial1
188 return $ first (Tuple' Nothing) res
189 else return $ FL.Partial $ Tuple' (Just buf'') r1
191-- XXX replace by compactLE
192--
193-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
194-- maximum specified size in bytes.
195--
196-- @since 0.7.0
197{-# INLINE compact #-}
198compact :: (MonadIO m, Storable a)
199 => Int -> SerialT m (Array a) -> SerialT m (Array a)
200compact n xs = D.fromStreamD $ packArraysChunksOf n (D.toStreamD xs)
202-- XXX Replace the above functions by a compactLEFold
203-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
204-- maximum specified size. Note that if a single array is bigger than the
205-- specified size we do not split it to fit. When we coalesce multiple arrays
206-- if the size would exceed the specified size we do not coalesce therefore the
207-- actual array size may be less than the specified chunk size.
208--
209-- /Unimplemented/
210{-# INLINE_NORMAL compactLEFold #-}
211compactLEFold :: -- (MonadIO m, Storable a) =>
212 Int -> Fold m (Array a) (Array a)
213compactLEFold = undefined
215compactLE :: (MonadIO m {-, Storable a-}) =>
216 Int -> SerialT m (Array a) -> SerialT m (Array a)
217compactLE n xs = D.fromStreamD $ D.foldMany (compactLEFold n) (D.toStreamD xs)
219-- | Like 'compact' but generates arrays of exactly equal to the size specified
220-- except for the last array in the stream which could be shorter.
221--
222-- /Unimplemented/
223{-# INLINE compactEQ #-}
224compactEQ :: -- (MonadIO m, Storable a) =>
225 Int -> SerialT m (Array a) -> SerialT m (Array a)
226compactEQ _n _xs = undefined
227 -- D.fromStreamD $ D.foldMany (compactEQFold n) (D.toStreamD xs)
229-- | Like 'compact' but generates arrays of size greater than or equal to the
230-- specified except for the last array in the stream which could be shorter.
231--
232-- /Unimplemented/
233{-# INLINE compactGE #-}
234compactGE :: -- (MonadIO m, Storable a) =>
235 Int -> SerialT m (Array a) -> SerialT m (Array a)
236compactGE _n _xs = undefined
237 -- D.fromStreamD $ D.foldMany (compactGEFold n) (D.toStreamD xs)
239-------------------------------------------------------------------------------
240-- IOVec
241-------------------------------------------------------------------------------
243#if !defined(mingw32_HOST_OS)
244data GatherState s arr
245 = GatherInitial s
246 | GatherBuffering s arr Int
247 | GatherYielding arr (GatherState s arr)
248 | GatherFinish
250-- | @groupIOVecsOf maxBytes maxEntries@ groups arrays in the incoming stream
251-- to create a stream of 'IOVec' arrays with a maximum of @maxBytes@ bytes in
252-- each array and a maximum of @maxEntries@ entries in each array.
253--
254-- @since 0.7.0
255{-# INLINE_NORMAL groupIOVecsOf #-}
256groupIOVecsOf :: MonadIO m
257 => Int -> Int -> D.Stream m (Array a) -> D.Stream m (Array IOVec)
258groupIOVecsOf n maxIOVLen (D.Stream step state) =
259 D.Stream step' (GatherInitial state)
261 where
263 {-# INLINE_LATE step' #-}
264 step' gst (GatherInitial st) = do
265 when (n <= 0) $
266 -- XXX we can pass the module string from the higher level API
267 error $ "Streamly.Internal.Data.Array.Foreign.Mut.Type.groupIOVecsOf: the size of "
268 ++ "groups [" ++ show n ++ "] must be a natural number"
269 when (maxIOVLen <= 0) $
270 -- XXX we can pass the module string from the higher level API
271 error $ "Streamly.Internal.Data.Array.Foreign.Mut.Type.groupIOVecsOf: the number of "
272 ++ "IOVec entries [" ++ show n ++ "] must be a natural number"
273 r <- step (adaptState gst) st
274 case r of
275 D.Yield arr s -> do
276 let p = unsafeForeignPtrToPtr (aStart arr)
277 len = MArray.byteLength arr
278 iov <- liftIO $ MArray.newArray maxIOVLen
279 iov' <- liftIO $ MArray.unsafeSnoc iov (IOVec (castPtr p)
280 (fromIntegral len))
281 if len >= n
282 then return $ D.Skip (GatherYielding iov' (GatherInitial s))
283 else return $ D.Skip (GatherBuffering s iov' len)
284 D.Skip s -> return $ D.Skip (GatherInitial s)
285 D.Stop -> return D.Stop
287 step' gst (GatherBuffering st iov len) = do
288 r <- step (adaptState gst) st
289 case r of
290 D.Yield arr s -> do
291 let p = unsafeForeignPtrToPtr (aStart arr)
292 alen = MArray.byteLength arr
293 len' = len + alen
294 if len' > n || length iov >= maxIOVLen
295 then do
296 iov' <- liftIO $ MArray.newArray maxIOVLen
297 iov'' <- liftIO $ MArray.unsafeSnoc iov' (IOVec (castPtr p)
298 (fromIntegral alen))
299 return $ D.Skip (GatherYielding iov
300 (GatherBuffering s iov'' alen))
301 else do
302 iov' <- liftIO $ MArray.unsafeSnoc iov (IOVec (castPtr p)
303 (fromIntegral alen))
304 return $ D.Skip (GatherBuffering s iov' len')
305 D.Skip s -> return $ D.Skip (GatherBuffering s iov len)
306 D.Stop -> return $ D.Skip (GatherYielding iov GatherFinish)
308 step' _ GatherFinish = return D.Stop
310 step' _ (GatherYielding iov next) = return $ D.Yield iov next
311#endif