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