git.iko.soy/streamly
on 90be387 help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Fold/Foreign.hs
338 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.fold (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.ForeignPtr (touchForeignPtr)
59import Foreign.Ptr (minusPtr, plusPtr)
60import Foreign.Storable (Storable(..))
61import GHC.ForeignPtr (ForeignPtr(..))
62import GHC.Ptr (Ptr(..))
63import GHC.Types (SPEC(..))
64import Streamly.Internal.Data.Array.Foreign.Type (Array(..))
65import Streamly.Internal.Data.Parser.ParserD (Initial(..), Step(..))
66import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
68import qualified Streamly.Internal.Data.Array.Foreign as Array
69import qualified Streamly.Internal.Data.Fold as Fold
70import qualified Streamly.Internal.Data.Parser.ParserD as ParserD
71import qualified Streamly.Internal.Data.Parser.ParserD.Type as ParserD
73import Prelude hiding (concatMap, take)
75-- | Array stream fold.
76--
77-- An array stream fold is basically an array stream "Parser" that does not
78-- fail. In case of array stream folds the count in 'Partial', 'Continue' and
79-- 'Done' is a count of elements that includes the leftover element count in
80-- the array that is currently being processed by the parser. If none of the
81-- elements is consumed by the parser the count is at least the whole array
82-- length. If the whole array is consumed by the parser then the count will be
83-- 0.
84--
85-- /Pre-release/
86--
87newtype Fold m a b = Fold (ParserD.Parser m (Array a) b)
89-------------------------------------------------------------------------------
90-- Constructing array stream folds from element folds and parsers
91-------------------------------------------------------------------------------
93-- | Convert an element 'Fold' into an array stream fold.
94--
95-- /Pre-release/
96{-# INLINE fromFold #-}
97fromFold :: forall m a b. (MonadIO m, Storable a) =>
98 Fold.Fold m a b -> Fold m a b
99fromFold (Fold.Fold fstep finitial fextract) =
100 Fold (ParserD.Parser step initial fextract)
102 where
104 initial = do
105 res <- finitial
106 return
107 $ case res of
108 Fold.Partial s1 -> IPartial s1
109 Fold.Done b -> IDone b
111 step s (Array fp@(ForeignPtr start _) end) = do
112 goArray SPEC (Ptr start) s
114 where
116 goArray !_ !cur !fs | cur >= end = do
117 assert (cur == end) (return ())
118 liftIO $ touchForeignPtr fp
119 return $ Partial 0 fs
120 goArray !_ !cur !fs = do
121 x <- liftIO $ peek cur
122 res <- fstep fs x
123 let elemSize = sizeOf (undefined :: a)
124 next = cur `plusPtr` elemSize
125 case res of
126 Fold.Done b ->
127 return $ Done ((end `minusPtr` next) `div` elemSize) b
128 Fold.Partial fs1 ->
129 goArray SPEC next fs1
131-- | Convert an element 'Parser' into an array stream fold. If the parser fails
132-- the fold would throw an exception.
133--
134-- /Pre-release/
135{-# INLINE fromParser #-}
136fromParser :: forall m a b. (MonadIO m, Storable a) =>
137 ParserD.Parser m a b -> Fold m a b
138fromParser (ParserD.Parser step1 initial1 extract1) =
139 Fold (ParserD.Parser step initial1 extract1)
141 where
143 step s (Array fp@(ForeignPtr start _) end) = do
144 if Ptr start >= end
145 then return $ Continue 0 s
146 else goArray SPEC (Ptr start) s
148 where
150 {-# INLINE partial #-}
151 partial arrRem cur next elemSize st n fs1 = do
152 let next1 = next `plusPtr` negate (n * elemSize)
153 if next1 >= Ptr start && cur < end
154 then goArray SPEC next1 fs1
155 else return $ st (arrRem + n) fs1
157 goArray !_ !cur !fs = do
158 x <- liftIO $ peek cur
159 liftIO $ touchForeignPtr fp
160 res <- step1 fs x
161 let elemSize = sizeOf (undefined :: a)
162 next = cur `plusPtr` elemSize
163 arrRem = (end `minusPtr` next) `div` elemSize
164 case res of
165 ParserD.Done n b -> do
166 return $ Done (arrRem + n) b
167 ParserD.Partial n fs1 ->
168 partial arrRem cur next elemSize Partial n fs1
169 ParserD.Continue n fs1 -> do
170 partial arrRem cur next elemSize Continue n fs1
171 Error err -> return $ Error err
173-- | Adapt an array stream fold.
174--
175-- /Pre-release/
176{-# INLINE fromArrayFold #-}
177fromArrayFold :: forall m a b. (MonadIO m) =>
178 Fold.Fold m (Array a) b -> Fold m a b
179fromArrayFold f = Fold $ ParserD.fromFold f
181-------------------------------------------------------------------------------
182-- Functor
183-------------------------------------------------------------------------------
185-- | Maps a function over the result of fold.
186--
187-- /Pre-release/
188instance Functor m => Functor (Fold m a) where
189 {-# INLINE fmap #-}
190 fmap f (Fold p) = Fold $ fmap f p
192-- | Map a monadic function on the output of a fold.
193--
194-- /Pre-release/
195{-# INLINE rmapM #-}
196rmapM :: Monad m => (b -> m c) -> Fold m a b -> Fold m a c
197rmapM f (Fold p) = Fold $ ParserD.rmapM f p
199-------------------------------------------------------------------------------
200-- Sequential applicative
201-------------------------------------------------------------------------------
203-- | A fold that always yields a pure value without consuming any input.
204--
205-- /Pre-release/
206--
207{-# INLINE fromPure #-}
208fromPure :: Monad m => b -> Fold m a b
209fromPure = Fold . ParserD.yield
211-- | A fold that always yields the result of an effectful action without
212-- consuming any input.
213--
214-- /Pre-release/
215--
216{-# INLINE fromEffect #-}
217fromEffect :: Monad m => m b -> Fold m a b
218fromEffect = Fold . ParserD.yieldM
220-- | Applies two folds sequentially on the input stream and combines their
221-- results using the supplied function.
222--
223-- /Pre-release/
224{-# INLINE serial_ #-}
225serial_ :: MonadThrow m => Fold m x a -> Fold m x b -> Fold m x b
226serial_ (Fold p1) (Fold p2) = Fold $ ParserD.noErrorUnsafeSplit_ p1 p2
228-- | Applies two folds sequentially on the input stream and combines their
229-- results using the supplied function.
230--
231-- /Pre-release/
232{-# INLINE serialWith #-}
233serialWith :: MonadThrow m
234 => (a -> b -> c) -> Fold m x a -> Fold m x b -> Fold m x c
235serialWith f (Fold p1) (Fold p2) =
236 Fold $ ParserD.noErrorUnsafeSplitWith f p1 p2
238-- | 'Applicative' form of 'serialWith'.
239-- > (<*>) = serialWith id
240instance MonadThrow m => Applicative (Fold m a) where
241 {-# INLINE pure #-}
242 pure = fromPure
244 {-# INLINE (<*>) #-}
245 (<*>) = serialWith id
247 {-# INLINE (*>) #-}
248 (*>) = serial_
250#if MIN_VERSION_base(4,10,0)
251 {-# INLINE liftA2 #-}
252 liftA2 f x = (<*>) (fmap f x)
253#endif
255-------------------------------------------------------------------------------
256-- Monad
257-------------------------------------------------------------------------------
259-- | Applies a fold on the input stream, generates the next fold from the
260-- output of the previously applied fold and then applies that fold.
261--
262-- /Pre-release/
263--
264{-# INLINE concatMap #-}
265concatMap :: MonadThrow m =>
266 (b -> Fold m a c) -> Fold m a b -> Fold m a c
267concatMap func (Fold p) =
268 Fold $ ParserD.noErrorUnsafeConcatMap (\x -> let Fold y = func x in y) p
270-- | Monad instance applies folds sequentially. Next fold can depend on the
271-- output of the previous fold. See 'concatMap'.
272--
273-- > (>>=) = flip concatMap
274instance MonadThrow m => Monad (Fold m a) where
275 {-# INLINE return #-}
276 return = pure
278 {-# INLINE (>>=) #-}
279 (>>=) = flip concatMap
281 {-# INLINE (>>) #-}
282 (>>) = (*>)
284-------------------------------------------------------------------------------
285-- Array to Array folds
286-------------------------------------------------------------------------------
288{-# INLINE take #-}
289take :: forall m a b. (Monad m, Storable a) => Int -> Fold m a b -> Fold m a b
290take n (Fold (ParserD.Parser step1 initial1 extract1)) =
291 Fold $ ParserD.Parser step initial extract
293 where
295 initial = do
296 res <- initial1
297 case res of
298 IPartial s ->
299 if n > 0
300 then return $ IPartial $ Tuple' n s
301 else IDone <$> extract1 s
302 IDone b -> return $ IDone b
303 IError err -> return $ IError err
305 {-# INLINE partial #-}
306 partial i1 st j s =
307 let i2 = i1 + j
308 in if i2 > 0
309 then return $ st j (Tuple' i2 s)
310 else Done 0 <$> extract1 s -- i2 == i1 == j == 0
312 step (Tuple' i r) arr = do
313 let len = Array.length arr
314 i1 = i - len
315 if i1 >= 0
316 then do
317 res <- step1 r arr
318 case res of
319 Partial j s -> partial i1 Partial j s
320 Continue j s -> partial i1 Continue j s
321 Done j b -> return $ Done j b
322 Error err -> return $ Error err
323 else do
324 let !(Array (ForeignPtr start contents) _) = arr
325 sz = sizeOf (undefined :: a)
326 end = Ptr start `plusPtr` (i * sz)
327 arr1 = Array (ForeignPtr start contents) end
328 remaining = negate i1
329 res <- step1 r arr1
330 case res of
331 Partial 0 s -> Done remaining <$> extract1 s
332 Partial j s -> return $ Partial (remaining + j) (Tuple' j s)
333 Continue 0 s -> Done remaining <$> extract1 s
334 Continue j s -> return $ Continue (remaining + j) (Tuple' j s)
335 Done j b -> return $ Done (remaining + j) b
336 Error err -> return $ Error err
338 extract (Tuple' _ r) = extract1 r