git.iko.soy/streamly
on 810ba5b help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Mut/Foreign.hs
235 lines 8.8 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"
29import Control.Monad.IO.Class (MonadIO(..))
30import Control.Monad (when)
31import Data.Bifunctor (first)
32import Foreign.Storable (Storable(..))
33import Streamly.Internal.Data.Array.Foreign.Mut.Type (Array(..))
34import Streamly.Internal.Data.Fold.Type (Fold(..))
35import Streamly.Internal.Data.Stream.Serial (SerialT)
36import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
37import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
39import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MArray
40import qualified Streamly.Internal.Data.Fold.Type as FL
41import qualified Streamly.Internal.Data.Stream.StreamD as D
43-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
44-- @n@ elements each.
45--
46-- Same as the following but may be more efficient:
47--
48-- > arraysOf n = Stream.foldMany (MArray.writeN n)
49--
50-- /Pre-release/
51{-# INLINE arraysOf #-}
52arraysOf :: (IsStream t, MonadIO m, Storable a)
53 => Int -> t m a -> t m (Array a)
54arraysOf n = D.fromStreamD . MArray.arraysOf n . D.toStreamD
56-------------------------------------------------------------------------------
57-- Compact
58-------------------------------------------------------------------------------
60data SpliceState s arr
61 = SpliceInitial s
62 | SpliceBuffering s arr
63 | SpliceYielding arr (SpliceState s arr)
64 | SpliceFinish
66-- XXX This can be removed once compactLEFold/compactLE are implemented.
67--
68-- | This mutates the first array (if it has space) to append values from the
69-- second one. This would work for immutable arrays as well because an
70-- immutable array never has space so a new array is allocated instead of
71-- mutating it.
72--
73-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
74-- maximum specified size. Note that if a single array is bigger than the
75-- specified size we do not split it to fit. When we coalesce multiple arrays
76-- if the size would exceed the specified size we do not coalesce therefore the
77-- actual array size may be less than the specified chunk size.
78--
79-- @since 0.7.0
80{-# INLINE_NORMAL packArraysChunksOf #-}
81packArraysChunksOf :: (MonadIO m, Storable a)
82 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
83packArraysChunksOf n (D.Stream step state) =
84 D.Stream step' (SpliceInitial state)
86 where
88 {-# INLINE_LATE step' #-}
89 step' gst (SpliceInitial st) = do
90 when (n <= 0) $
91 -- XXX we can pass the module string from the higher level API
92 error $ "Streamly.Internal.Data.Array.Foreign.Mut.Type.packArraysChunksOf: the size of "
93 ++ "arrays [" ++ show n ++ "] must be a natural number"
94 r <- step gst st
95 case r of
96 D.Yield arr s -> return $
97 let len = MArray.byteLength arr
98 in if len >= n
99 then D.Skip (SpliceYielding arr (SpliceInitial s))
100 else D.Skip (SpliceBuffering s arr)
101 D.Skip s -> return $ D.Skip (SpliceInitial s)
102 D.Stop -> return D.Stop
104 step' gst (SpliceBuffering st buf) = do
105 r <- step gst st
106 case r of
107 D.Yield arr s -> do
108 let len = MArray.byteLength buf + MArray.byteLength arr
109 if len > n
110 then return $
111 D.Skip (SpliceYielding buf (SpliceBuffering s arr))
112 else do
113 buf' <- if MArray.byteCapacity buf < n
114 then liftIO $ MArray.realloc n buf
115 else return buf
116 buf'' <- MArray.spliceWith buf' arr
117 return $ D.Skip (SpliceBuffering s buf'')
118 D.Skip s -> return $ D.Skip (SpliceBuffering s buf)
119 D.Stop -> return $ D.Skip (SpliceYielding buf SpliceFinish)
121 step' _ SpliceFinish = return D.Stop
123 step' _ (SpliceYielding arr next) = return $ D.Yield arr next
125-- XXX Remove this once compactLEFold is implemented
126-- lpackArraysChunksOf = Fold.many compactLEFold
127--
128{-# INLINE_NORMAL lpackArraysChunksOf #-}
129lpackArraysChunksOf :: (MonadIO m, Storable a)
130 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
131lpackArraysChunksOf n (Fold step1 initial1 extract1) =
132 Fold step initial extract
134 where
136 initial = do
137 when (n <= 0) $
138 -- XXX we can pass the module string from the higher level API
139 error $ "Streamly.Internal.Data.Array.Foreign.Mut.Type.packArraysChunksOf: the size of "
140 ++ "arrays [" ++ show n ++ "] must be a natural number"
142 r <- initial1
143 return $ first (Tuple' Nothing) r
145 extract (Tuple' Nothing r1) = extract1 r1
146 extract (Tuple' (Just buf) r1) = do
147 r <- step1 r1 buf
148 case r of
149 FL.Partial rr -> extract1 rr
150 FL.Done _ -> return ()
152 step (Tuple' Nothing r1) arr =
153 let len = MArray.byteLength arr
154 in if len >= n
155 then do
156 r <- step1 r1 arr
157 case r of
158 FL.Done _ -> return $ FL.Done ()
159 FL.Partial s -> do
160 extract1 s
161 res <- initial1
162 return $ first (Tuple' Nothing) res
163 else return $ FL.Partial $ Tuple' (Just arr) r1
165 step (Tuple' (Just buf) r1) arr = do
166 let len = MArray.byteLength buf + MArray.byteLength arr
167 buf' <- if MArray.byteCapacity buf < len
168 then liftIO $ MArray.realloc (max n len) buf
169 else return buf
170 buf'' <- MArray.spliceWith buf' arr
172 -- XXX this is common in both the equations of step
173 if len >= n
174 then do
175 r <- step1 r1 buf''
176 case r of
177 FL.Done _ -> return $ FL.Done ()
178 FL.Partial s -> do
179 extract1 s
180 res <- initial1
181 return $ first (Tuple' Nothing) res
182 else return $ FL.Partial $ Tuple' (Just buf'') r1
184-- XXX Same as compactLE, to be removed once that is implemented.
185--
186-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
187-- maximum specified size in bytes.
188--
189-- /Internal/
190{-# INLINE compact #-}
191compact :: (MonadIO m, Storable a)
192 => Int -> SerialT m (Array a) -> SerialT m (Array a)
193compact n xs = D.fromStreamD $ packArraysChunksOf n (D.toStreamD xs)
195-- See lpackArraysChunksOf/packArraysChunksOf to implement this.
196--
197-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
198-- maximum specified size. Note that if a single array is bigger than the
199-- specified size we do not split it to fit. When we coalesce multiple arrays
200-- if the size would exceed the specified size we do not coalesce therefore the
201-- actual array size may be less than the specified chunk size.
202--
203-- /Unimplemented/
204{-# INLINE_NORMAL compactLEFold #-}
205compactLEFold :: -- (MonadIO m, Storable a) =>
206 Int -> Fold m (Array a) (Array a)
207compactLEFold = undefined
209-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
210-- maximum specified size in bytes.
211--
212-- /Internal/
213compactLE :: (MonadIO m {-, Storable a-}) =>
214 Int -> SerialT m (Array a) -> SerialT m (Array a)
215compactLE n xs = D.fromStreamD $ D.foldMany (compactLEFold n) (D.toStreamD xs)
217-- | Like 'compactLE' but generates arrays of exactly equal to the size
218-- specified except for the last array in the stream which could be shorter.
219--
220-- /Unimplemented/
221{-# INLINE compactEQ #-}
222compactEQ :: -- (MonadIO m, Storable a) =>
223 Int -> SerialT m (Array a) -> SerialT m (Array a)
224compactEQ _n _xs = undefined
225 -- D.fromStreamD $ D.foldMany (compactEQFold n) (D.toStreamD xs)
227-- | Like 'compactLE' but generates arrays of size greater than or equal to the
228-- specified except for the last array in the stream which could be shorter.
229--
230-- /Unimplemented/
231{-# INLINE compactGE #-}
232compactGE :: -- (MonadIO m, Storable a) =>
233 Int -> SerialT m (Array a) -> SerialT m (Array a)
234compactGE _n _xs = undefined
235 -- D.fromStreamD $ D.foldMany (compactGEFold n) (D.toStreamD xs)