git.iko.soy/streamly
on f8d80cc help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Fold/Foreign.hs
336 lines 11 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Array.Stream.Fold.Foreign
3-- Copyright : (c) 2021 Composewell Technologies
4-- License : BSD-3-Clause
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
8--
9-- Fold a stream of foreign arrays. @Fold m a b@ in this module works
10-- on a stream of "Array a" and produces an output of type @b@.
11--
12-- Though @Fold m a b@ in this module works on a stream of @Array a@ it is
13-- different from @Data.Fold m (Array a) b@. While the latter works on arrays
14-- as a whole treating them as atomic elements, the folds in this module can
15-- work on the stream of arrays as if it is an element stream with all the
16-- arrays coalesced together. This module allows adapting the element stream
17-- folds in Data.Fold to correctly work on an array stream as if it is an
18-- element stream. For example:
19--
20-- >>> import qualified Streamly.Data.Fold as Fold
21-- >>> import qualified Streamly.Internal.Data.Array.Stream.Foreign as ArrayStream
22-- >>> import qualified Streamly.Internal.Data.Array.Stream.Fold.Foreign as ArrayFold
23-- >>> import qualified Streamly.Internal.Data.Stream.IsStream as Stream (arraysOf)
24-- >>> import qualified Streamly.Prelude as Stream
25--
26-- >>> ArrayStream.foldArr (ArrayFold.fromFold (Fold.take 7 Fold.toList)) $ Stream.arraysOf 5 $ Stream.fromList "hello world"
27-- "hello w"
28--
29module Streamly.Internal.Data.Array.Stream.Fold.Foreign
30 (
31 Fold (..)
33 -- * Construction
34 , fromFold
35 , fromParser
36 , fromArrayFold
38 -- * Mapping
39 , rmapM
41 -- * Applicative
42 , fromPure
43 , fromEffect
44 , serialWith
46 -- * Monad
47 , concatMap
49 -- * Transformation
50 , take
51 )
52where
54import Control.Applicative (liftA2)
55import Control.Exception (assert)
56import Control.Monad.Catch (MonadThrow)
57import Control.Monad.IO.Class (MonadIO(..))
58import Foreign.Ptr (minusPtr, plusPtr)
59import Foreign.Storable (Storable(..))
60import GHC.Types (SPEC(..))
61import Streamly.Internal.Data.Array.Foreign.Mut.Type (touch)
62import Streamly.Internal.Data.Array.Foreign.Type (Array(..))
63import Streamly.Internal.Data.Parser.ParserD (Initial(..), Step(..))
64import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
66import qualified Streamly.Internal.Data.Array.Foreign as Array
67import qualified Streamly.Internal.Data.Fold as Fold
68import qualified Streamly.Internal.Data.Parser.ParserD as ParserD
69import qualified Streamly.Internal.Data.Parser.ParserD.Type as ParserD
71import Prelude hiding (concatMap, take)
73-- | Array stream fold.
74--
75-- An array stream fold is basically an array stream "Parser" that does not
76-- fail. In case of array stream folds the count in 'Partial', 'Continue' and
77-- 'Done' is a count of elements that includes the leftover element count in
78-- the array that is currently being processed by the parser. If none of the
79-- elements is consumed by the parser the count is at least the whole array
80-- length. If the whole array is consumed by the parser then the count will be
81-- 0.
82--
83-- /Pre-release/
84--
85newtype Fold m a b = Fold (ParserD.Parser m (Array a) b)
87-------------------------------------------------------------------------------
88-- Constructing array stream folds from element folds and parsers
89-------------------------------------------------------------------------------
91-- | Convert an element 'Fold' into an array stream fold.
92--
93-- /Pre-release/
94{-# INLINE fromFold #-}
95fromFold :: forall m a b. (MonadIO m, Storable a) =>
96 Fold.Fold m a b -> Fold m a b
97fromFold (Fold.Fold fstep finitial fextract) =
98 Fold (ParserD.Parser step initial fextract)
100 where
102 initial = do
103 res <- finitial
104 return
105 $ case res of
106 Fold.Partial s1 -> IPartial s1
107 Fold.Done b -> IDone b
109 step s (Array contents start end) = do
110 goArray SPEC start s
112 where
114 goArray !_ !cur !fs | cur >= end = do
115 assert (cur == end) (return ())
116 liftIO $ touch contents
117 return $ Partial 0 fs
118 goArray !_ !cur !fs = do
119 x <- liftIO $ peek cur
120 res <- fstep fs x
121 let elemSize = sizeOf (undefined :: a)
122 next = cur `plusPtr` elemSize
123 case res of
124 Fold.Done b ->
125 return $ Done ((end `minusPtr` next) `div` elemSize) b
126 Fold.Partial fs1 ->
127 goArray SPEC next fs1
129-- | Convert an element 'Parser' into an array stream fold. If the parser fails
130-- the fold would throw an exception.
131--
132-- /Pre-release/
133{-# INLINE fromParser #-}
134fromParser :: forall m a b. (MonadIO m, Storable a) =>
135 ParserD.Parser m a b -> Fold m a b
136fromParser (ParserD.Parser step1 initial1 extract1) =
137 Fold (ParserD.Parser step initial1 extract1)
139 where
141 step s (Array contents start end) = do
142 if start >= end
143 then return $ Continue 0 s
144 else goArray SPEC start s
146 where
148 {-# INLINE partial #-}
149 partial arrRem cur next elemSize st n fs1 = do
150 let next1 = next `plusPtr` negate (n * elemSize)
151 if next1 >= start && cur < end
152 then goArray SPEC next1 fs1
153 else return $ st (arrRem + n) fs1
155 goArray !_ !cur !fs = do
156 x <- liftIO $ peek cur
157 liftIO $ touch contents
158 res <- step1 fs x
159 let elemSize = sizeOf (undefined :: a)
160 next = cur `plusPtr` elemSize
161 arrRem = (end `minusPtr` next) `div` elemSize
162 case res of
163 ParserD.Done n b -> do
164 return $ Done (arrRem + n) b
165 ParserD.Partial n fs1 ->
166 partial arrRem cur next elemSize Partial n fs1
167 ParserD.Continue n fs1 -> do
168 partial arrRem cur next elemSize Continue n fs1
169 Error err -> return $ Error err
171-- | Adapt an array stream fold.
172--
173-- /Pre-release/
174{-# INLINE fromArrayFold #-}
175fromArrayFold :: forall m a b. (MonadIO m) =>
176 Fold.Fold m (Array a) b -> Fold m a b
177fromArrayFold f = Fold $ ParserD.fromFold f
179-------------------------------------------------------------------------------
180-- Functor
181-------------------------------------------------------------------------------
183-- | Maps a function over the result of fold.
184--
185-- /Pre-release/
186instance Functor m => Functor (Fold m a) where
187 {-# INLINE fmap #-}
188 fmap f (Fold p) = Fold $ fmap f p
190-- | Map a monadic function on the output of a fold.
191--
192-- /Pre-release/
193{-# INLINE rmapM #-}
194rmapM :: Monad m => (b -> m c) -> Fold m a b -> Fold m a c
195rmapM f (Fold p) = Fold $ ParserD.rmapM f p
197-------------------------------------------------------------------------------
198-- Sequential applicative
199-------------------------------------------------------------------------------
201-- | A fold that always yields a pure value without consuming any input.
202--
203-- /Pre-release/
204--
205{-# INLINE fromPure #-}
206fromPure :: Monad m => b -> Fold m a b
207fromPure = Fold . ParserD.fromPure
209-- | A fold that always yields the result of an effectful action without
210-- consuming any input.
211--
212-- /Pre-release/
213--
214{-# INLINE fromEffect #-}
215fromEffect :: Monad m => m b -> Fold m a b
216fromEffect = Fold . ParserD.fromEffect
218-- | Applies two folds sequentially on the input stream and combines their
219-- results using the supplied function.
220--
221-- /Pre-release/
222{-# INLINE serial_ #-}
223serial_ :: MonadThrow m => Fold m x a -> Fold m x b -> Fold m x b
224serial_ (Fold p1) (Fold p2) = Fold $ ParserD.noErrorUnsafeSplit_ p1 p2
226-- | Applies two folds sequentially on the input stream and combines their
227-- results using the supplied function.
228--
229-- /Pre-release/
230{-# INLINE serialWith #-}
231serialWith :: MonadThrow m
232 => (a -> b -> c) -> Fold m x a -> Fold m x b -> Fold m x c
233serialWith f (Fold p1) (Fold p2) =
234 Fold $ ParserD.noErrorUnsafeSplitWith f p1 p2
236-- | 'Applicative' form of 'serialWith'.
237-- > (<*>) = serialWith id
238instance MonadThrow m => Applicative (Fold m a) where
239 {-# INLINE pure #-}
240 pure = fromPure
242 {-# INLINE (<*>) #-}
243 (<*>) = serialWith id
245 {-# INLINE (*>) #-}
246 (*>) = serial_
248#if MIN_VERSION_base(4,10,0)
249 {-# INLINE liftA2 #-}
250 liftA2 f x = (<*>) (fmap f x)
251#endif
253-------------------------------------------------------------------------------
254-- Monad
255-------------------------------------------------------------------------------
257-- | Applies a fold on the input stream, generates the next fold from the
258-- output of the previously applied fold and then applies that fold.
259--
260-- /Pre-release/
261--
262{-# INLINE concatMap #-}
263concatMap :: MonadThrow m =>
264 (b -> Fold m a c) -> Fold m a b -> Fold m a c
265concatMap func (Fold p) =
266 Fold $ ParserD.noErrorUnsafeConcatMap (\x -> let Fold y = func x in y) p
268-- | Monad instance applies folds sequentially. Next fold can depend on the
269-- output of the previous fold. See 'concatMap'.
270--
271-- > (>>=) = flip concatMap
272instance MonadThrow m => Monad (Fold m a) where
273 {-# INLINE return #-}
274 return = pure
276 {-# INLINE (>>=) #-}
277 (>>=) = flip concatMap
279 {-# INLINE (>>) #-}
280 (>>) = (*>)
282-------------------------------------------------------------------------------
283-- Array to Array folds
284-------------------------------------------------------------------------------
286{-# INLINE take #-}
287take :: forall m a b. (Monad m, Storable a) => Int -> Fold m a b -> Fold m a b
288take n (Fold (ParserD.Parser step1 initial1 extract1)) =
289 Fold $ ParserD.Parser step initial extract
291 where
293 initial = do
294 res <- initial1
295 case res of
296 IPartial s ->
297 if n > 0
298 then return $ IPartial $ Tuple' n s
299 else IDone <$> extract1 s
300 IDone b -> return $ IDone b
301 IError err -> return $ IError err
303 {-# INLINE partial #-}
304 partial i1 st j s =
305 let i2 = i1 + j
306 in if i2 > 0
307 then return $ st j (Tuple' i2 s)
308 else Done 0 <$> extract1 s -- i2 == i1 == j == 0
310 step (Tuple' i r) arr = do
311 let len = Array.length arr
312 i1 = i - len
313 if i1 >= 0
314 then do
315 res <- step1 r arr
316 case res of
317 Partial j s -> partial i1 Partial j s
318 Continue j s -> partial i1 Continue j s
319 Done j b -> return $ Done j b
320 Error err -> return $ Error err
321 else do
322 let !(Array contents start _) = arr
323 sz = sizeOf (undefined :: a)
324 end = start `plusPtr` (i * sz)
325 arr1 = Array contents start end
326 remaining = negate i1
327 res <- step1 r arr1
328 case res of
329 Partial 0 s -> Done remaining <$> extract1 s
330 Partial j s -> return $ Partial (remaining + j) (Tuple' j s)
331 Continue 0 s -> Done remaining <$> extract1 s
332 Continue j s -> return $ Continue (remaining + j) (Tuple' j s)
333 Done j b -> return $ Done (remaining + j) b
334 Error err -> return $ Error err
336 extract (Tuple' _ r) = extract1 r