git.iko.soy/streamly
on cb4a2bb help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Foreign.hs
387 lines 14 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Array.Stream.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 immutable arrays.
10--
11module Streamly.Internal.Data.Array.Stream.Foreign
12 (
13 -- * Creation
14 arraysOf
16 -- * Flattening to elements
17 , concat
18 , concatRev
19 , interpose
20 , interposeSuffix
21 , intercalateSuffix
22 , unlines
24 -- * Elimination
25 , toArray
27 -- * Compaction
28 , lpackArraysChunksOf
29#if !defined(mingw32_HOST_OS)
30 , groupIOVecsOf
31#endif
32 , compact
34 -- * Splitting
35 , splitOn
36 , splitOnSuffix
37 )
38where
40#include "inline.hs"
42import Control.Monad.IO.Class (MonadIO(..))
43import Data.Word (Word8)
44import Foreign.ForeignPtr (withForeignPtr, touchForeignPtr)
45import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
46import Foreign.Ptr (minusPtr, plusPtr, castPtr)
47import Foreign.Storable (Storable(..))
48import GHC.ForeignPtr (ForeignPtr(..))
49import GHC.Ptr (Ptr(..))
50import Prelude hiding (length, null, last, map, (!!), read, concat, unlines)
52#if !defined(mingw32_HOST_OS)
53import Streamly.Internal.FileSystem.FDIO (IOVec(..))
54#endif
55import Streamly.Internal.Data.Array.Foreign.Types (Array(..), length)
56import Streamly.Internal.Data.Fold.Types (Fold(..))
57import Streamly.Internal.Data.Stream.Serial (SerialT)
58import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
59import Streamly.Internal.Data.SVar (adaptState)
61import qualified Streamly.Internal.Data.Array.Foreign as A
62import qualified Streamly.Internal.Data.Array.Foreign.Types as A
63import qualified Streamly.Internal.Data.Array.Foreign.Mut.Types as MA
64import qualified Streamly.Internal.Data.Array.Stream.Mut.Foreign as AS
65import qualified Streamly.Internal.Data.Fold.Types as FL
66import qualified Streamly.Internal.Data.Stream.IsStream as S
67import qualified Streamly.Internal.Data.Stream.StreamD as D
69-------------------------------------------------------------------------------
70-- Generation
71-------------------------------------------------------------------------------
73-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
74-- @n@ elements each.
75--
76-- > arraysOf n = Stream.chunksOf n (Array.writeN n)
77--
78-- /Pre-release/
79{-# INLINE arraysOf #-}
80arraysOf :: (IsStream t, MonadIO m, Storable a)
81 => Int -> t m a -> t m (Array a)
82arraysOf n str = D.fromStreamD $ A.arraysOf n (D.toStreamD str)
84-------------------------------------------------------------------------------
85-- Append
86-------------------------------------------------------------------------------
88-- XXX efficiently compare two streams of arrays. Two streams can have chunks
89-- of different sizes, we can handle that in the stream comparison abstraction.
90-- This could be useful e.g. to fast compare whether two files differ.
92-- | Convert a stream of arrays into a stream of their elements.
93--
94-- Same as the following but more efficient:
95--
96-- > concat = Stream.concatUnfold Array.read
97--
98-- @since 0.7.0
99{-# INLINE concat #-}
100concat :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
101-- concat m = D.fromStreamD $ A.flattenArrays (D.toStreamD m)
102-- concat m = D.fromStreamD $ D.concatMap A.toStreamD (D.toStreamD m)
103concat m = D.fromStreamD $ D.concatUnfold A.read (D.toStreamD m)
105-- | Convert a stream of arrays into a stream of their elements reversing the
106-- contents of each array before flattening.
107--
108-- > concatRev = Stream.concatUnfold Array.readRev
109--
110-- @since 0.7.0
111{-# INLINE concatRev #-}
112concatRev :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
113concatRev m = D.fromStreamD $ A.flattenArraysRev (D.toStreamD m)
115-------------------------------------------------------------------------------
116-- Intersperse and append
117-------------------------------------------------------------------------------
119-- | Flatten a stream of arrays after inserting the given element between
120-- arrays.
121--
122-- /Pre-release/
123{-# INLINE interpose #-}
124interpose :: (MonadIO m, IsStream t, Storable a) => a -> t m (Array a) -> t m a
125interpose x = S.interpose x A.read
127{-# INLINE intercalateSuffix #-}
128intercalateSuffix :: (MonadIO m, IsStream t, Storable a)
129 => Array a -> t m (Array a) -> t m a
130intercalateSuffix arr = S.intercalateSuffix arr A.read
132-- | Flatten a stream of arrays appending the given element after each
133-- array.
134--
135-- @since 0.7.0
136{-# INLINE interposeSuffix #-}
137interposeSuffix :: (MonadIO m, IsStream t, Storable a)
138 => a -> t m (Array a) -> t m a
139-- interposeSuffix x = D.fromStreamD . A.unlines x . D.toStreamD
140interposeSuffix x = S.interposeSuffix x A.read
142data FlattenState s a =
143 OuterLoop s
144 | InnerLoop s !(ForeignPtr a) !(Ptr a) !(Ptr a)
146{-# INLINE_NORMAL unlines #-}
147unlines :: forall m a. (MonadIO m, Storable a)
148 => a -> D.Stream m (Array a) -> D.Stream m a
149unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
150 where
151 {-# INLINE_LATE step' #-}
152 step' gst (OuterLoop st) = do
153 r <- step (adaptState gst) st
154 return $ case r of
155 D.Yield Array{..} s ->
156 let p = unsafeForeignPtrToPtr aStart
157 in D.Skip (InnerLoop s aStart p aEnd)
158 D.Skip s -> D.Skip (OuterLoop s)
159 D.Stop -> D.Stop
161 step' _ (InnerLoop st _ p end) | p == end =
162 return $ D.Yield sep $ OuterLoop st
164 step' _ (InnerLoop st startf p end) = do
165 x <- liftIO $ do
166 r <- peek p
167 touchForeignPtr startf
168 return r
169 return $ D.Yield x (InnerLoop st startf
170 (p `plusPtr` sizeOf (undefined :: a)) end)
172-------------------------------------------------------------------------------
173-- Compact
174-------------------------------------------------------------------------------
176{-# INLINE_NORMAL packArraysChunksOf #-}
177packArraysChunksOf :: (MonadIO m, Storable a)
178 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
179packArraysChunksOf n str =
180 D.map A.unsafeFreeze $ AS.packArraysChunksOf n $ D.map A.unsafeThaw str
182-- XXX instead of writing two different versions of this operation, we should
183-- write it as a pipe.
184{-# INLINE_NORMAL lpackArraysChunksOf #-}
185lpackArraysChunksOf :: (MonadIO m, Storable a)
186 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
187lpackArraysChunksOf n fld =
188 FL.map A.unsafeThaw $ AS.lpackArraysChunksOf n (FL.map A.unsafeFreeze fld)
190#if !defined(mingw32_HOST_OS)
192-- | @groupIOVecsOf maxBytes maxEntries@ groups arrays in the incoming stream
193-- to create a stream of 'IOVec' arrays with a maximum of @maxBytes@ bytes in
194-- each array and a maximum of @maxEntries@ entries in each array.
195--
196-- @since 0.7.0
197{-# INLINE_NORMAL groupIOVecsOf #-}
198groupIOVecsOf :: MonadIO m
199 => Int -> Int -> D.Stream m (Array a) -> D.Stream m (Array IOVec)
200groupIOVecsOf n maxIOVLen str =
201 D.map A.unsafeFreeze
202 $ AS.groupIOVecsOf n maxIOVLen
203 $ D.map A.unsafeThaw str
204#endif
206-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
207-- maximum specified size in bytes.
208--
209-- @since 0.7.0
210{-# INLINE compact #-}
211compact :: (MonadIO m, Storable a)
212 => Int -> SerialT m (Array a) -> SerialT m (Array a)
213compact n xs = D.fromStreamD $ packArraysChunksOf n (D.toStreamD xs)
215-------------------------------------------------------------------------------
216-- Split
217-------------------------------------------------------------------------------
219data SplitState s arr
220 = Initial s
221 | Buffering s arr
222 | Splitting s arr
223 | Yielding arr (SplitState s arr)
224 | Finishing
226-- | Split a stream of arrays on a given separator byte, dropping the separator
227-- and coalescing all the arrays between two separators into a single array.
228--
229-- @since 0.7.0
230{-# INLINE_NORMAL _splitOn #-}
231_splitOn
232 :: MonadIO m
233 => Word8
234 -> D.Stream m (Array Word8)
235 -> D.Stream m (Array Word8)
236_splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
238 where
240 {-# INLINE_LATE step' #-}
241 step' gst (Initial st) = do
242 r <- step gst st
243 case r of
244 D.Yield arr s -> do
245 (arr1, marr2) <- A.breakOn byte arr
246 return $ case marr2 of
247 Nothing -> D.Skip (Buffering s arr1)
248 Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
249 D.Skip s -> return $ D.Skip (Initial s)
250 D.Stop -> return D.Stop
252 step' gst (Buffering st buf) = do
253 r <- step gst st
254 case r of
255 D.Yield arr s -> do
256 (arr1, marr2) <- A.breakOn byte arr
257 buf' <- A.spliceTwo buf arr1
258 return $ case marr2 of
259 Nothing -> D.Skip (Buffering s buf')
260 Just x -> D.Skip (Yielding buf' (Splitting s x))
261 D.Skip s -> return $ D.Skip (Buffering s buf)
262 D.Stop -> return $
263 if A.byteLength buf == 0
264 then D.Stop
265 else D.Skip (Yielding buf Finishing)
267 step' _ (Splitting st buf) = do
268 (arr1, marr2) <- A.breakOn byte buf
269 return $ case marr2 of
270 Nothing -> D.Skip $ Buffering st arr1
271 Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
273 step' _ (Yielding arr next) = return $ D.Yield arr next
274 step' _ Finishing = return D.Stop
276-- | Split a stream of arrays on a given separator byte, dropping the separator
277-- and coalescing all the arrays between two separators into a single array.
278--
279-- @since 0.7.0
280{-# INLINE splitOn #-}
281splitOn
282 :: (IsStream t, MonadIO m)
283 => Word8
284 -> t m (Array Word8)
285 -> t m (Array Word8)
286splitOn byte s =
287 D.fromStreamD $ D.splitInnerBy (A.breakOn byte) A.spliceTwo $ D.toStreamD s
289{-# INLINE splitOnSuffix #-}
290splitOnSuffix
291 :: (IsStream t, MonadIO m)
292 => Word8
293 -> t m (Array Word8)
294 -> t m (Array Word8)
295-- splitOn byte s = D.fromStreamD $ A.splitOn byte $ D.toStreamD s
296splitOnSuffix byte s =
297 D.fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.spliceTwo $ D.toStreamD s
299-------------------------------------------------------------------------------
300-- Elimination
301-------------------------------------------------------------------------------
303-- XXX Both of these implementations of splicing seem to perform equally well.
304-- We need to perform benchmarks over a range of sizes though.
306-- CAUTION! length must more than equal to lengths of all the arrays in the
307-- stream.
308{-# INLINE spliceArraysLenUnsafe #-}
309spliceArraysLenUnsafe :: (MonadIO m, Storable a)
310 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
311spliceArraysLenUnsafe len buffered = do
312 arr <- liftIO $ MA.newArray len
313 end <- S.foldlM' writeArr (return $ MA.aEnd arr) buffered
314 return $ arr {MA.aEnd = end}
316 where
318 writeArr dst (MA.Array as ae _) =
319 liftIO $ withForeignPtr as $ \src -> do
320 let count = ae `minusPtr` src
321 A.memcpy (castPtr dst) (castPtr src) count
322 return $ dst `plusPtr` count
324{-# INLINE _spliceArrays #-}
325_spliceArrays :: (MonadIO m, Storable a)
326 => SerialT m (Array a) -> m (Array a)
327_spliceArrays s = do
328 buffered <- S.foldr S.cons S.nil s
329 len <- S.sum (S.map length buffered)
330 arr <- liftIO $ MA.newArray len
331 end <- S.foldlM' writeArr (return $ MA.aEnd arr) s
332 return $ A.unsafeFreeze $ arr {MA.aEnd = end}
334 where
336 writeArr dst (Array as ae) =
337 liftIO $ withForeignPtr as $ \src -> do
338 let count = ae `minusPtr` src
339 A.memcpy (castPtr dst) (castPtr src) count
340 return $ dst `plusPtr` count
342{-# INLINE _spliceArraysBuffered #-}
343_spliceArraysBuffered :: (MonadIO m, Storable a)
344 => SerialT m (Array a) -> m (Array a)
345_spliceArraysBuffered s = do
346 buffered <- S.foldr S.cons S.nil s
347 len <- S.sum (S.map length buffered)
348 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
350{-# INLINE spliceArraysRealloced #-}
351spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
352 => SerialT m (Array a) -> m (Array a)
353spliceArraysRealloced s = do
354 let idst = liftIO $ MA.newArray (A.bytesToElemCount (undefined :: a)
355 (A.mkChunkSizeKB 4))
357 arr <- S.foldlM' MA.spliceWithDoubling idst (S.map A.unsafeThaw s)
358 liftIO $ A.unsafeFreeze <$> MA.shrinkToFit arr
360-- | Given a stream of arrays, splice them all together to generate a single
361-- array. The stream must be /finite/.
362--
363-- @since 0.7.0
364{-# INLINE toArray #-}
365toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
366toArray = spliceArraysRealloced
367-- spliceArrays = _spliceArraysBuffered
369-- exponentially increasing sizes of the chunks upto the max limit.
370-- XXX this will be easier to implement with parsers/terminating folds
371-- With this we should be able to reduce the number of chunks/allocations.
372-- The reallocation/copy based toArray can also be implemented using this.
373--
374{-
375{-# INLINE toArraysInRange #-}
376toArraysInRange :: (IsStream t, MonadIO m, Storable a)
377 => Int -> Int -> Fold m (Array a) b -> Fold m a b
378toArraysInRange low high (Fold step initial extract) =
379-}
381{-
382-- | Fold the input to a pure buffered stream (List) of arrays.
383{-# INLINE _toArraysOf #-}
384_toArraysOf :: (MonadIO m, Storable a)
385 => Int -> Fold m a (SerialT Identity (Array a))
386_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
387-}