git.iko.soy/streamly
on 1305145 help sign in
streamly/src/Streamly/Internal/Data/Unfold.hs
541 lines 18 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ExistentialQuantification #-}
4{-# LANGUAGE FlexibleContexts #-}
5{-# LANGUAGE PatternSynonyms #-}
6{-# LANGUAGE RankNTypes #-}
7{-# LANGUAGE RecordWildCards #-}
8{-# LANGUAGE ScopedTypeVariables #-}
10#include "inline.hs"
12-- |
13-- Module : Streamly.Internal.Data.Unfold
14-- Copyright : (c) 2019 Composewell Technologies
15-- License : BSD3
16-- Maintainer : harendra.kumar@gmail.com
17-- Stability : experimental
18-- Portability : GHC
19--
20-- Streams forcing a closed control flow loop can be categorized under
21-- two types, unfolds and folds, both of these are duals of each other.
22--
23-- Unfold streams are really generators of a sequence of elements, we can also
24-- call them pull style streams. These are lazy producers of streams. On each
25-- evaluation the producer generates the next element. A consumer can
26-- therefore pull elements from the stream whenever it wants to. A stream
27-- consumer can multiplex pull streams by pulling elements from the chosen
28-- streams, therefore, pull streams allow merging or multiplexing. On the
29-- other hand, with this representation we cannot split or demultiplex a
30-- stream. So really these are stream sources that can be generated from a
31-- seed and can be merged or zipped into a single stream.
32--
33-- The dual of Unfolds are Folds. Folds can also be called as push style
34-- streams or reducers. These are strict consumers of streams. We keep pushing
35-- elements to a fold and we can extract the result at any point. A driver can
36-- choose which fold to push to and can also push the same element to multiple
37-- folds. Therefore, folds allow splitting or demultiplexing a stream. On the
38-- other hand, we cannot merge streams using this representation. So really
39-- these are stream consumers that reduce the stream to a single value, these
40-- consumers can be composed such that a stream can be split over multiple
41-- consumers.
42--
43-- Performance:
44--
45-- Composing a tree or graph of computations with unfolds can be much more
46-- efficient compared to composing with the Monad instance. The reason is that
47-- unfolds allow the compiler to statically know the state and optimize it
48-- using stream fusion whereas it is not possible with the monad bind because
49-- the state is determined dynamically.
51-- Open control flow style streams can also have two representations. StreamK
52-- is a producer style representation. We can also have a consumer style
53-- representation. We can use that for composable folds in StreamK
54-- representation.
55--
56module Streamly.Internal.Data.Unfold
57 (
58 -- * Unfold Type
59 Unfold
61 -- * Operations on Input
62 , close
63 , first
64 , second
65 , lmap
66 -- coapply
67 -- comonad
69 -- * Operations on Output
70 , fold
71 -- pipe
73 -- * Unfolds
74 , singleton
75 , replicateM
76 , fromList
77 , enumerateFromStepIntegral
78 , enumerateFromToIntegral
79 , enumerateFromIntegral
81 -- * Transformations
82 , map
84 -- * Filtering
85 , takeWhileM
86 , takeWhile
87 , take
88 , filter
89 , filterM
91 -- * Nesting
92 , concat
93 , concatMapM
94 , outerProduct
96 -- * Exceptions
97 , before
98 , after
99 , onException
100 , finally
101 , bracket
102 , handle
103 )
104where
106import Control.Exception (Exception)
107import Data.Void (Void)
108import GHC.Types (SPEC(..))
109import Prelude hiding (concat, map, takeWhile, take, filter)
111import Streamly.Streams.StreamD.Type (Stream(..), Step(..))
112#if __GLASGOW_HASKELL__ < 800
113import Streamly.Streams.StreamD.Type (pattern Stream)
114#endif
115import Streamly.Internal.Data.Unfold.Types (Unfold(..))
116import Streamly.Internal.Data.Fold.Types (Fold(..))
117import Streamly.Internal.Data.SVar (defState)
118import Control.Monad.Catch (MonadCatch)
120import qualified Control.Monad.Catch as MC
122-------------------------------------------------------------------------------
123-- Input operations
124-------------------------------------------------------------------------------
126-- | Supply the seed to an unfold closing the input end of the unfold.
127--
128-- /Internal/
129--
130{-# INLINE_NORMAL close #-}
131close :: Unfold m a b -> a -> Unfold m Void b
132close (Unfold ustep uinject) a = Unfold ustep (const $ uinject a)
134-- XXX rename to closeFst/closeSnd
135--
136-- | Supply the first component of the tuple to an unfold that accepts a tuple
137-- as a seed resulting in a fold that accepts the second component of the tuple
138-- as a seed.
139--
140-- /Internal/
141--
142{-# INLINE_NORMAL first #-}
143first :: Unfold m (a, b) c -> a -> Unfold m b c
144first (Unfold ustep uinject) a = Unfold ustep (\b -> uinject (a, b))
146-- | Supply the second component of the tuple to an unfold that accepts a tuple
147-- as a seed resulting in a fold that accepts the first component of the tuple
148-- as a seed.
149--
150-- /Internal/
151--
152{-# INLINE_NORMAL second #-}
153second :: Unfold m (a, b) c -> b -> Unfold m a c
154second (Unfold ustep uinject) b = Unfold ustep (\a -> uinject (a, b))
156{-# INLINE_NORMAL lmap #-}
157lmap :: (a -> c) -> Unfold m c b -> Unfold m a b
158lmap f (Unfold ustep uinject) = Unfold ustep (\x -> uinject $ f x)
160-------------------------------------------------------------------------------
161-- Output operations
162-------------------------------------------------------------------------------
164-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
165-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
166-- application of the fold on the unfolded stream.
167--
168-- /Internal/
169--
170{-# INLINE_NORMAL fold #-}
171fold :: Monad m => Unfold m a b -> Fold m b c -> a -> m c
172fold (Unfold ustep inject) (Fold fstep initial extract) a =
173 initial >>= \x -> inject a >>= go SPEC x
174 where
175 -- XXX !acc?
176 {-# INLINE_LATE go #-}
177 go !_ acc st = acc `seq` do
178 r <- ustep st
179 case r of
180 Yield x s -> do
181 acc' <- fstep acc x
182 go SPEC acc' s
183 Skip s -> go SPEC acc s
184 Stop -> extract acc
186{-# INLINE_NORMAL map #-}
187map :: Monad m => (b -> c) -> Unfold m a b -> Unfold m a c
188map f (Unfold ustep uinject) = Unfold step uinject
189 where
190 {-# INLINE_LATE step #-}
191 step st = do
192 r <- ustep st
193 return $ case r of
194 Yield x s -> Yield (f x) s
195 Skip s -> Skip s
196 Stop -> Stop
198-------------------------------------------------------------------------------
199-- Unfolds
200-------------------------------------------------------------------------------
202-- | Identity unfold. Generates a singleton stream with the seed as the only
203-- element in the stream.
204--
205-- > singleton = replicateM 1
206--
207{-# INLINE singleton #-}
208singleton :: Monad m => Unfold m a a
209singleton = Unfold step inject
210 where
211 inject x = return $ Just x
212 {-# INLINE_LATE step #-}
213 step (Just x) = return $ Yield x Nothing
214 step Nothing = return Stop
216-- | Generates a stream replicating the seed @n@ times.
217--
218{-# INLINE replicateM #-}
219replicateM :: Monad m => Int -> Unfold m a a
220replicateM n = Unfold step inject
221 where
222 inject x = return (x, n)
223 {-# INLINE_LATE step #-}
224 step (x, i) = return $
225 if i <= 0
226 then Stop
227 else Yield x (x, (i - 1))
229-- | Convert a list of pure values to a 'Stream'
230{-# INLINE_LATE fromList #-}
231fromList :: Monad m => Unfold m [a] a
232fromList = Unfold step inject
233 where
234 inject x = return x
235 {-# INLINE_LATE step #-}
236 step (x:xs) = return $ Yield x xs
237 step [] = return Stop
239-------------------------------------------------------------------------------
240-- Filtering
241-------------------------------------------------------------------------------
243{-# INLINE_NORMAL take #-}
244take :: Monad m => Int -> Unfold m a b -> Unfold m a b
245take n (Unfold step1 inject1) = Unfold step inject
246 where
247 inject x = do
248 s <- inject1 x
249 return (s, 0)
250 {-# INLINE_LATE step #-}
251 step (st, i) | i < n = do
252 r <- step1 st
253 return $ case r of
254 Yield x s -> Yield x (s, i + 1)
255 Skip s -> Skip (s, i)
256 Stop -> Stop
257 step (_, _) = return Stop
259{-# INLINE_NORMAL takeWhileM #-}
260takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
261takeWhileM f (Unfold step1 inject1) = Unfold step inject1
262 where
263 {-# INLINE_LATE step #-}
264 step st = do
265 r <- step1 st
266 case r of
267 Yield x s -> do
268 b <- f x
269 return $ if b then Yield x s else Stop
270 Skip s -> return $ Skip s
271 Stop -> return Stop
273{-# INLINE takeWhile #-}
274takeWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
275takeWhile f = takeWhileM (return . f)
277{-# INLINE_NORMAL filterM #-}
278filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
279filterM f (Unfold step1 inject1) = Unfold step inject1
280 where
281 {-# INLINE_LATE step #-}
282 step st = do
283 r <- step1 st
284 case r of
285 Yield x s -> do
286 b <- f x
287 return $ if b then Yield x s else Skip s
288 Skip s -> return $ Skip s
289 Stop -> return Stop
291{-# INLINE filter #-}
292filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
293filter f = filterM (return . f)
295-------------------------------------------------------------------------------
296-- Enumeration
297-------------------------------------------------------------------------------
299-- | Can be used to enumerate unbounded integrals. This does not check for
300-- overflow or underflow for bounded integrals.
301{-# INLINE_NORMAL enumerateFromStepIntegral #-}
302enumerateFromStepIntegral :: (Integral a, Monad m) => Unfold m (a, a) a
303enumerateFromStepIntegral = Unfold step inject
304 where
305 inject (from, stride) = from `seq` stride `seq` return (from, stride)
306 {-# INLINE_LATE step #-}
307 step !(x, stride) = return $ Yield x $! (x + stride, stride)
309-- We are assuming that "to" is constrained by the type to be within
310-- max/min bounds.
311{-# INLINE enumerateFromToIntegral #-}
312enumerateFromToIntegral :: (Monad m, Integral a) => a -> Unfold m a a
313enumerateFromToIntegral to =
314 takeWhile (<= to) $ second enumerateFromStepIntegral 1
316{-# INLINE enumerateFromIntegral #-}
317enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
318enumerateFromIntegral = enumerateFromToIntegral maxBound
320-------------------------------------------------------------------------------
321-- Nested
322-------------------------------------------------------------------------------
324data ConcatState s1 s2 = ConcatOuter s1 | ConcatInner s1 s2
326{-# INLINE_NORMAL concat #-}
327concat :: Monad m => Unfold m a b -> Unfold m b c -> Unfold m a c
328concat (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
329 where
330 inject x = do
331 s <- inject1 x
332 return $ ConcatOuter s
334 {-# INLINE_LATE step #-}
335 step (ConcatOuter st) = do
336 r <- step1 st
337 case r of
338 Yield x s -> do
339 innerSt <- inject2 x
340 return $ Skip (ConcatInner s innerSt)
341 Skip s -> return $ Skip (ConcatOuter s)
342 Stop -> return Stop
344 step (ConcatInner ost ist) = do
345 r <- step2 ist
346 return $ case r of
347 Yield x s -> Yield x (ConcatInner ost s)
348 Skip s -> Skip (ConcatInner ost s)
349 Stop -> Skip (ConcatOuter ost)
351data OuterProductState s1 s2 sy x y =
352 OuterProductOuter s1 y | OuterProductInner s1 sy s2 x
354{-# INLINE_NORMAL outerProduct #-}
355outerProduct :: Monad m
356 => Unfold m a b -> Unfold m c d -> Unfold m (a, c) (b, d)
357outerProduct (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
358 where
359 inject (x, y) = do
360 s1 <- inject1 x
361 return $ OuterProductOuter s1 y
363 {-# INLINE_LATE step #-}
364 step (OuterProductOuter st1 sy) = do
365 r <- step1 st1
366 case r of
367 Yield x s -> do
368 s2 <- inject2 sy
369 return $ Skip (OuterProductInner s sy s2 x)
370 Skip s -> return $ Skip (OuterProductOuter s sy)
371 Stop -> return Stop
373 step (OuterProductInner ost sy ist x) = do
374 r <- step2 ist
375 return $ case r of
376 Yield y s -> Yield (x, y) (OuterProductInner ost sy s x)
377 Skip s -> Skip (OuterProductInner ost sy s x)
378 Stop -> Skip (OuterProductOuter ost sy)
380-- XXX This can be used to implement a Monad instance for "Unfold m ()".
382data ConcatMapState s1 s2 = ConcatMapOuter s1 | ConcatMapInner s1 s2
384{-# INLINE_NORMAL concatMapM #-}
385concatMapM :: Monad m
386 => (b -> m (Unfold m () c)) -> Unfold m a b -> Unfold m a c
387concatMapM f (Unfold step1 inject1) = Unfold step inject
388 where
389 inject x = do
390 s <- inject1 x
391 return $ ConcatMapOuter s
393 {-# INLINE_LATE step #-}
394 step (ConcatMapOuter st) = do
395 r <- step1 st
396 case r of
397 Yield x s -> do
398 Unfold step2 inject2 <- f x
399 innerSt <- inject2 ()
400 return $ Skip (ConcatMapInner s (Stream (\_ ss -> step2 ss)
401 innerSt))
402 Skip s -> return $ Skip (ConcatMapOuter s)
403 Stop -> return Stop
405 step (ConcatMapInner ost (Stream istep ist)) = do
406 r <- istep defState ist
407 return $ case r of
408 Yield x s -> Yield x (ConcatMapInner ost (Stream istep s))
409 Skip s -> Skip (ConcatMapInner ost (Stream istep s))
410 Stop -> Skip (ConcatMapOuter ost)
412------------------------------------------------------------------------------
413-- Exceptions
414------------------------------------------------------------------------------
416-- | Run a side effect before the unfold yields its first element.
417--
418-- /Internal/
419{-# INLINE_NORMAL before #-}
420before :: Monad m => m c -> Unfold m a b -> Unfold m a b
421before action (Unfold step1 inject1) = Unfold step inject
423 where
425 inject x = do
426 _ <- action
427 st <- inject1 x
428 return st
430 {-# INLINE_LATE step #-}
431 step st = do
432 res <- step1 st
433 case res of
434 Yield x s -> return $ Yield x s
435 Skip s -> return $ Skip s
436 Stop -> return Stop
438-- | Run a side effect whenever the unfold stops normally.
439--
440-- /Internal/
441{-# INLINE_NORMAL after #-}
442after :: Monad m => m c -> Unfold m a b -> Unfold m a b
443after action (Unfold step1 inject1) = Unfold step inject1
445 where
447 {-# INLINE_LATE step #-}
448 step st = do
449 res <- step1 st
450 case res of
451 Yield x s -> return $ Yield x s
452 Skip s -> return $ Skip s
453 Stop -> action >> return Stop
455-- | Run a side effect whenever the unfold aborts due to an exception.
456--
457-- /Internal/
458{-# INLINE_NORMAL onException #-}
459onException :: MonadCatch m => m c -> Unfold m a b -> Unfold m a b
460onException action (Unfold step1 inject1) = Unfold step inject1
462 where
464 {-# INLINE_LATE step #-}
465 step st = do
466 res <- step1 st `MC.onException` action
467 case res of
468 Yield x s -> return $ Yield x s
469 Skip s -> return $ Skip s
470 Stop -> return Stop
472-- | Run a side effect whenever the unfold stops normally or aborts due to an
473-- exception.
474--
475-- /Internal/
476{-# INLINE_NORMAL finally #-}
477finally :: MonadCatch m => m c -> Unfold m a b -> Unfold m a b
478finally action (Unfold step1 inject1) = Unfold step inject1
480 where
482 {-# INLINE_LATE step #-}
483 step st = do
484 res <- step1 st `MC.onException` action
485 case res of
486 Yield x s -> return $ Yield x s
487 Skip s -> return $ Skip s
488 Stop -> action >> return Stop
490-- | @bracket before after between@ runs the @before@ action and then runs the
491-- @between@ function using the output of @before@ as its argument. When the
492-- @between@ unfold is done or if an exception occurs while running @between@
493-- the @after@ action is run with the output of @before@ as argument.
494--
495-- /Internal/
496{-# INLINE_NORMAL bracket #-}
497bracket :: MonadCatch m
498 => m c -> (c -> m d) -> (c -> Unfold m a b) -> Unfold m a b
499bracket bef aft bet = Unfold step inject
501 where
503 inject x = do
504 r <- bef
505 go (bet r) r
506 where
507 go (Unfold step2 inject2) r = do
508 st <- inject2 x
509 return (Stream (\_ -> step2) st, r)
511 {-# INLINE_LATE step #-}
512 step (Stream step2 state, v) = do
513 res <- step2 defState state `MC.onException` aft v
514 case res of
515 Yield x s -> return $ Yield x (Stream step2 s, v)
516 Skip s -> return $ Skip (Stream step2 s, v)
517 Stop -> aft v >> return Stop
519-- | When unfolding an unfold if an exception occurs, unfold aborts
520-- and the specified exception handler is run with the exception as argument.
521--
522-- /Internal/
523{-# INLINE_NORMAL handle #-}
524handle :: (MonadCatch m, Exception e)
525 => (e -> m b) -> Unfold m a b -> Unfold m a b
526handle f (Unfold step1 inject1) = Unfold step inject
528 where
530 inject x = inject1 x >>= return . Just
532 {-# INLINE_LATE step #-}
533 step (Just st) = do
534 res <- MC.try $ step1 st
535 case res of
536 Left e -> f e >>= \x -> return (Yield x Nothing)
537 Right r -> case r of
538 Yield x s -> return $ Yield x (Just s)
539 Skip s -> return $ Skip (Just s)
540 Stop -> return Stop
541 step Nothing = return Stop