on da62ada help sign in
503 lines 16 kB history raw
1{-# OPTIONS_GHC -Wno-orphans #-}
3-- |
4-- Module : Streamly.Internal.Data.Stream.IsStream.Common
5-- Copyright : (c) 2017 Composewell Technologies
6-- License : BSD-3-Clause
7-- Maintainer : streamly@composewell.com
8-- Stability : experimental
9-- Portability : GHC
10--
12module Streamly.Internal.Data.Stream.IsStream.Common
13 (
14 -- * Generation
15 yield
16 , yieldM
17 , repeatM
18 , timesWith
19 , absTimesWith
20 , relTimesWith
22 -- * Elimination
23 , fold
25 -- * Transformation
26 , scanlMAfter'
27 , postscanlM'
28 , smapM
29 -- $smapM_Notes
30 , take
31 , takeWhile
32 , drop
33 , findIndices
34 , intersperseM
35 , interjectSuffix
36 , reverse
37 , reverse'
39 -- * Nesting
40 , concatM
41 , concatMapM
42 , splitOnSeq
43 )
44where
46#include "inline.hs"
48import Control.Concurrent (threadDelay)
49import Control.Monad.IO.Class (MonadIO(..))
50import Foreign.Storable (Storable)
51import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array)
52import Streamly.Internal.Data.Fold.Types (Fold (..))
53import Streamly.Internal.Data.Stream.IsStream.Combinators (maxYields)
54import Streamly.Internal.Data.Stream.Prelude (fromStreamS, toStreamS)
55import Streamly.Internal.Data.Stream.Serial (SerialT)
56import Streamly.Internal.Data.Stream.StreamD.Type (fromStreamD, toStreamD)
57import Streamly.Internal.Data.Stream.StreamK.Type (IsStream())
58import Streamly.Internal.Data.SVar (MonadAsync)
59import Streamly.Internal.Data.Time.Units (AbsTime, RelTime64, addToAbsTime64)
61import qualified Streamly.Internal.Data.Stream.Parallel as Par
62import qualified Streamly.Internal.Data.Stream.Prelude as P
63import qualified Streamly.Internal.Data.Stream.Serial as Serial
64import qualified Streamly.Internal.Data.Stream.StreamK as K (repeatM)
65import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
66import qualified Streamly.Internal.Data.Stream.StreamD as D
67#ifdef USE_STREAMK_ONLY
68import qualified Streamly.Internal.Data.Stream.StreamK as S
69#else
70import qualified Streamly.Internal.Data.Stream.StreamD as S
71#endif
73import Prelude hiding (take, takeWhile, drop, reverse)
75------------------------------------------------------------------------------
76-- Generation
77------------------------------------------------------------------------------
79-- Faster than yieldM because there is no bind.
80--
81-- |
82-- @
83-- yield a = a \`cons` nil
84-- @
85--
86-- Create a singleton stream from a pure value.
87--
88-- The following holds in monadic streams, but not in Zip streams:
89--
90-- @
91-- yield = pure
92-- yield = yieldM . pure
93-- @
94--
95-- In Zip applicative streams 'yield' is not the same as 'pure' because in that
96-- case 'pure' is equivalent to 'repeat' instead. 'yield' and 'pure' are
97-- equally efficient, in other cases 'yield' may be slightly more efficient
98-- than the other equivalent definitions.
99--
100-- @since 0.4.0
101{-# INLINE yield #-}
102yield :: IsStream t => a -> t m a
103yield = K.yield
105-- |
106-- @
107-- yieldM m = m \`consM` nil
108-- @
109--
110-- Create a singleton stream from a monadic action.
111--
112-- @
113-- > toList $ yieldM getLine
114-- hello
115-- ["hello"]
116-- @
117--
118-- @since 0.4.0
119{-# INLINE yieldM #-}
120yieldM :: (Monad m, IsStream t) => m a -> t m a
121yieldM = K.yieldM
123-- |
124-- @
125-- repeatM = fix . consM
126-- repeatM = cycle1 . yieldM
127-- @
128--
129-- Generate a stream by repeatedly executing a monadic action forever.
130--
131-- @
132-- drain $ serially $ S.take 10 $ S.repeatM $ (threadDelay 1000000 >> print 1)
133-- drain $ asyncly $ S.take 10 $ S.repeatM $ (threadDelay 1000000 >> print 1)
134-- @
135--
136-- /Concurrent, infinite (do not use with 'parallely')/
137--
138-- @since 0.2.0
139{-# INLINE_EARLY repeatM #-}
140repeatM :: (IsStream t, MonadAsync m) => m a -> t m a
141repeatM = K.repeatM
143{-# RULES "repeatM serial" repeatM = repeatMSerial #-}
144{-# INLINE repeatMSerial #-}
145repeatMSerial :: MonadAsync m => m a -> SerialT m a
146repeatMSerial = fromStreamS . S.repeatM
148------------------------------------------------------------------------------
149-- Generation - Time related
150------------------------------------------------------------------------------
152-- | @timesWith g@ returns a stream of time value tuples. The first component
153-- of the tuple is an absolute time reference (epoch) denoting the start of the
154-- stream and the second component is a time relative to the reference.
155--
156-- The argument @g@ specifies the granularity of the relative time in seconds.
157-- A lower granularity clock gives higher precision but is more expensive in
158-- terms of CPU usage. Any granularity lower than 1 ms is treated as 1 ms.
159--
160-- @
161-- >>> S.mapM_ (\x -> print x >> threadDelay 1000000) $ S.timesWith 0.01
162-- > (AbsTime (TimeSpec {sec = 2496295, nsec = 536223000}),RelTime64 (NanoSecond64 0))
163-- > (AbsTime (TimeSpec {sec = 2496295, nsec = 536223000}),RelTime64 (NanoSecond64 1002028000))
164-- > (AbsTime (TimeSpec {sec = 2496295, nsec = 536223000}),RelTime64 (NanoSecond64 1996656000))
165-- @
166--
167-- Note: This API is not safe on 32-bit machines.
168--
169-- /Internal/
170--
171{-# INLINE timesWith #-}
172timesWith :: (IsStream t, MonadAsync m) => Double -> t m (AbsTime, RelTime64)
173timesWith g = fromStreamD $ D.times g
175-- | @absTimesWith g@ returns a stream of absolute timestamps using a clock of
176-- granularity @g@ specified in seconds. A low granularity clock is more
177-- expensive in terms of CPU usage. Any granularity lower than 1 ms is treated
178-- as 1 ms.
179--
180-- @
181-- >>> S.mapM_ print $ S.delayPre 1 $ S.absTimesWith 0.01
182-- @
183--
184-- Note: This API is not safe on 32-bit machines.
185--
186-- /Internal/
187--
188{-# INLINE absTimesWith #-}
189absTimesWith :: (IsStream t, MonadAsync m, Functor (t m))
190 => Double -> t m AbsTime
191absTimesWith = fmap (uncurry addToAbsTime64) . timesWith
193-- | @relTimesWith g@ returns a stream of relative time values starting from 0,
194-- using a clock of granularity @g@ specified in seconds. A low granularity
195-- clock is more expensive in terms of CPU usage. Any granularity lower than 1
196-- ms is treated as 1 ms.
197--
198-- @
199-- >>> S.mapM_ print $ S.delayPre 1 $ S.relTimesWith 0.01
200-- > RelTime64 (NanoSecond64 0)
201-- > RelTime64 (NanoSecond64 91139000)
202-- > RelTime64 (NanoSecond64 204052000)
203-- @
204--
205-- Note: This API is not safe on 32-bit machines.
206--
207-- /Internal/
208--
209{-# INLINE relTimesWith #-}
210relTimesWith :: (IsStream t, MonadAsync m, Functor (t m))
211 => Double -> t m RelTime64
212relTimesWith = fmap snd . timesWith
214------------------------------------------------------------------------------
215-- Elimination - Running a Fold
216------------------------------------------------------------------------------
218-- | Fold a stream using the supplied left 'Fold' and reducing the resulting
219-- expression strictly at each step. The behavior is similar to 'foldl''. A
220-- 'Fold' can terminate early without consuming the full stream. See the
221-- documentation of individual 'Fold's for termination behavior.
222--
223-- >>> S.fold FL.sum (S.enumerateFromTo 1 100)
224-- 5050
225--
226-- @since 0.7.0
227{-# INLINE fold #-}
228fold :: Monad m => Fold m a b -> SerialT m a -> m b
229fold = P.foldOnce
231------------------------------------------------------------------------------
232-- Transformation
233------------------------------------------------------------------------------
235-- | @scanlMAfter' accumulate initial done stream@ is like 'scanlM'' except
236-- that it provides an additional @done@ function to be applied on the
237-- accumulator when the stream stops. The result of @done@ is also emitted in
238-- the stream.
239--
240-- This function can be used to allocate a resource in the beginning of the
241-- scan and release it when the stream ends or to flush the internal state of
242-- the scan at the end.
243--
244-- /Internal/
245--
246{-# INLINE scanlMAfter' #-}
247scanlMAfter' :: (IsStream t, Monad m)
248 => (b -> a -> m b) -> m b -> (b -> m b) -> t m a -> t m b
249scanlMAfter' step initial done stream =
250 fromStreamD $ D.scanlMAfter' step initial done $ toStreamD stream
252-- XXX this needs to be concurrent
253-- | Like 'postscanl'' but with a monadic step function and a monadic seed.
254--
255-- /Since: 0.7.0/
256--
257-- /Since: 0.8.0 (signature change)/
258{-# INLINE postscanlM' #-}
259postscanlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> m b -> t m a -> t m b
260postscanlM' step z m = fromStreamD $ D.postscanlM' step z $ toStreamD m
262-- $smapM_Notes
263--
264-- The stateful step function can be simplified to @(s -> a -> m b)@ to provide
265-- a read-only environment. However, that would just be 'mapM'.
266--
267-- The initial action could be @m (s, Maybe b)@, and we can also add a final
268-- action @s -> m (Maybe b)@. This can be used to get pre/post scan like
269-- functionality and also to flush the state in the end like scanlMAfter'.
270-- We can also use it along with a fusible version of bracket to get
271-- scanlMAfter' like functionality. See issue #677.
272--
273-- This can be further generalized to a type similar to Fold/Parser, giving it
274-- filtering and parsing capability as well (this is in fact equivalent to
275-- parseMany):
276--
277-- smapM :: (s -> a -> m (Step s b)) -> m s -> t m a -> t m b
278--
280-- | A stateful 'mapM', equivalent to a left scan, more like mapAccumL.
281-- Hopefully, this is a better alternative to @scan@. Separation of state from
282-- the output makes it easier to think in terms of a shared state, and also
283-- makes it easier to keep the state fully strict and the output lazy.
284--
285-- See also: 'scanlM''
286--
287-- /Internal/
288--
289{-# INLINE smapM #-}
290smapM :: (IsStream t, Monad m) =>
291 (s -> a -> m (s, b))
292 -> m s
293 -> t m a
294 -> t m b
295smapM step initial stream =
296 -- XXX implement this directly instead of using scanlM'
297 -- Once we have postscanlM' with monadic initial we can use this code
298 -- let r = postscanlM'
299 -- (\(s, _) a -> step s a)
300 -- (fmap (,undefined) initial)
301 -- stream
302 let r = postscanlM'
303 (\(s, _) a -> step s a)
304 (fmap (,undefined) initial)
305 stream
306 in Serial.map snd r
308------------------------------------------------------------------------------
309-- Transformation - Trimming
310------------------------------------------------------------------------------
312-- | Take first 'n' elements from the stream and discard the rest.
313--
314-- @since 0.1.0
315{-# INLINE take #-}
316take :: (IsStream t, Monad m) => Int -> t m a -> t m a
317take n m = fromStreamS $ S.take n $ toStreamS
318 (maxYields (Just (fromIntegral n)) m)
320-- | End the stream as soon as the predicate fails on an element.
321--
322-- @since 0.1.0
323{-# INLINE takeWhile #-}
324takeWhile :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> t m a
325takeWhile p m = fromStreamS $ S.takeWhile p $ toStreamS m
327-- | Discard first 'n' elements from the stream and take the rest.
328--
329-- @since 0.1.0
330{-# INLINE drop #-}
331drop :: (IsStream t, Monad m) => Int -> t m a -> t m a
332drop n m = fromStreamS $ S.drop n $ toStreamS m
334------------------------------------------------------------------------------
335-- Searching
336------------------------------------------------------------------------------
338-- | Find all the indices where the element in the stream satisfies the given
339-- predicate.
340--
341-- > findIndices = fold Fold.findIndices
342--
343-- @since 0.5.0
344{-# INLINE findIndices #-}
345findIndices :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> t m Int
346findIndices p m = fromStreamS $ S.findIndices p (toStreamS m)
348------------------------------------------------------------------------------
349-- Transformation by Inserting
350------------------------------------------------------------------------------
352-- intersperseM = intersperseBySpan 1
354-- | Insert an effect and its output before consuming an element of a stream
355-- except the first one.
356--
357-- @
358-- >>> S.toList $ S.trace putChar $ S.intersperseM (putChar '.' >> return ',') $ S.fromList "hello"
359-- > h.,e.,l.,l.,o"h,e,l,l,o"
360-- @
361--
362-- @since 0.5.0
363{-# INLINE intersperseM #-}
364intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a
365intersperseM m = fromStreamS . S.intersperseM m . toStreamS
367-- | Intersperse a monadic action into the input stream after every @n@
368-- seconds.
369--
370-- @
371-- > S.drain $ S.interjectSuffix 1 (putChar ',') $ S.mapM (\\x -> threadDelay 1000000 >> putChar x) $ S.fromList "hello"
372-- "h,e,l,l,o"
373-- @
374--
375-- /Internal/
376{-# INLINE interjectSuffix #-}
377interjectSuffix
378 :: (IsStream t, MonadAsync m)
379 => Double -> m a -> t m a -> t m a
380interjectSuffix n f xs = xs `Par.parallelFst` repeatM timed
381 where timed = liftIO (threadDelay (round $ n * 1000000)) >> f
383------------------------------------------------------------------------------
384-- Transformation by Reordering
385------------------------------------------------------------------------------
387-- XXX Use a compact region list to temporarily store the list, in both reverse
388-- as well as in reverse'.
389--
390-- /Note:/ 'reverse'' is much faster than this, use that when performance
391-- matters.
392--
393-- > reverse = S.foldlT (flip S.cons) S.nil
394--
395-- | Returns the elements of the stream in reverse order. The stream must be
396-- finite. Note that this necessarily buffers the entire stream in memory.
397--
398-- /Since 0.7.0 (Monad m constraint)/
399--
400-- /Since: 0.1.1/
401{-# INLINE reverse #-}
402reverse :: (IsStream t, Monad m) => t m a -> t m a
403reverse s = fromStreamS $ S.reverse $ toStreamS s
405-- | Like 'reverse' but several times faster, requires a 'Storable' instance.
406--
407-- /Internal/
408{-# INLINE reverse' #-}
409reverse' :: (IsStream t, MonadIO m, Storable a) => t m a -> t m a
410reverse' s = fromStreamD $ D.reverse' $ toStreamD s
412------------------------------------------------------------------------------
413-- Combine streams and flatten
414------------------------------------------------------------------------------
416-- | Map a stream producing monadic function on each element of the stream
417-- and then flatten the results into a single stream. Since the stream
418-- generation function is monadic, unlike 'concatMap', it can produce an
419-- effect at the beginning of each iteration of the inner loop.
420--
421-- @since 0.6.0
422{-# INLINE concatMapM #-}
423concatMapM :: (IsStream t, Monad m) => (a -> m (t m b)) -> t m a -> t m b
424concatMapM f m = fromStreamD $ D.concatMapM (fmap toStreamD . f) (toStreamD m)
426-- | Given a stream value in the underlying monad, lift and join the underlying
427-- monad with the stream monad.
428--
429-- @
430-- concatM = concat . yieldM
431-- concatM = concat . lift -- requires @(MonadTrans t)@
432-- concatM = join . lift -- requires @(MonadTrans t@, @Monad (t m))@
433-- @
434--
435-- See also: 'concat', 'sequence'
436--
437-- /Internal/
438--
439{-# INLINE concatM #-}
440concatM :: (IsStream t, Monad m) => m (t m a) -> t m a
441concatM generator = concatMapM (\() -> generator) (yield ())
443------------------------------------------------------------------------------
444-- Split stream and fold
445------------------------------------------------------------------------------
447-- | Like 'splitOn' but the separator is a sequence of elements instead of a
448-- single element.
449--
450-- For illustration, let's define a function that operates on pure lists:
451--
452-- @
453-- splitOnSeq' pat xs = S.toList $ S.splitOnSeq (A.fromList pat) (FL.toList) (S.fromList xs)
454-- @
455--
456-- >>> splitOnSeq' "" "hello"
457-- > ["h","e","l","l","o"]
458--
459-- >>> splitOnSeq' "hello" ""
460-- > [""]
461--
462-- >>> splitOnSeq' "hello" "hello"
463-- > ["",""]
464--
465-- >>> splitOnSeq' "x" "hello"
466-- > ["hello"]
467--
468-- >>> splitOnSeq' "h" "hello"
469-- > ["","ello"]
470--
471-- >>> splitOnSeq' "o" "hello"
472-- > ["hell",""]
473--
474-- >>> splitOnSeq' "e" "hello"
475-- > ["h","llo"]
476--
477-- >>> splitOnSeq' "l" "hello"
478-- > ["he","","o"]
479--
480-- >>> splitOnSeq' "ll" "hello"
481-- > ["he","o"]
482--
483-- 'splitOnSeq' is an inverse of 'intercalate'. The following law always holds:
484--
485-- > intercalate . splitOn == id
486--
487-- The following law holds when the separator is non-empty and contains none of
488-- the elements present in the input lists:
489--
490-- > splitOn . intercalate == id
491--
492-- /Internal/
494-- XXX We can use a polymorphic vector implemented by Array# to represent the
495-- sequence, that way we can avoid the Storable constraint. If we still need
496-- Storable Array for performance, we can use a separate splitOnArray API for
497-- that. We can also have an API where the sequence itself is a lazy stream, so
498-- that we can search files in files for example.
499{-# INLINE splitOnSeq #-}
500splitOnSeq
501 :: (IsStream t, MonadIO m, Storable a, Enum a, Eq a)
502 => Array a -> Fold m a b -> t m a -> t m b
503splitOnSeq patt f m = D.fromStreamD $ D.splitOnSeq patt f (D.toStreamD m)