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