git.iko.soy/streamly
on master help sign in
Moved nested StreamD ops to StreamD.Flatten.
ab5648cPranay Sashank
src/Streamly/Internal/Data/Stream/StreamD.hs+1−717
compare with master
unifiedsplit
@@ -314,12 +314,10 @@ import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_, control)
314import Data.Functor.Identity (Identity(..))
314import Data.Functor.Identity (Identity(..))
315import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
315import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
316import Data.Maybe (fromJust, isJust)
316import Data.Maybe (fromJust, isJust)
317−import Fusion.Plugin.Types (Fuse(..))
318import Streamly.Internal.Data.Time.Units
317import Streamly.Internal.Data.Time.Units
319 (TimeUnit64, toRelTime64, diffAbsTime64)
318 (TimeUnit64, toRelTime64, diffAbsTime64)
320import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
319import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
321import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
320import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
322−import Streamly.Internal.Data.Unfold.Types (Unfold(..))
323
321
324import qualified Streamly.Internal.Data.Pipe.Types as Pipe
322import qualified Streamly.Internal.Data.Pipe.Types as Pipe
325import qualified Streamly.Internal.Data.Fold as FL
323import qualified Streamly.Internal.Data.Fold as FL
@@ -337,6 +335,7 @@ import Prelude hiding
337import Streamly.Internal.Data.Stream.StreamD.Type
335import Streamly.Internal.Data.Stream.StreamD.Type
338import Streamly.Internal.Data.Stream.StreamD.Generate
336import Streamly.Internal.Data.Stream.StreamD.Generate
339import Streamly.Internal.Data.Stream.StreamD.Eliminate
337import Streamly.Internal.Data.Stream.StreamD.Eliminate
338+import Streamly.Internal.Data.Stream.StreamD.Flatten
340import Streamly.Internal.Data.Stream.StreamD.SplitGroup
339import Streamly.Internal.Data.Stream.StreamD.SplitGroup
341import Streamly.Internal.Data.SVar
340import Streamly.Internal.Data.SVar
342
341
@@ -600,721 +599,6 @@ stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
600mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
599mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
601mapM_ m = drain . mapM m
600mapM_ m = drain . mapM m
602
601
603−-------------------------------------------------------------------------------
604−-- Stream transformations using Unfolds
605−-------------------------------------------------------------------------------
606−
607−-- Define a unique structure to use in inspection testing
608−data ConcatMapUState o i =
609− ConcatMapUOuter o
610− | ConcatMapUInner o i
611−
612−-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
613−-- to streams and then flattens the generated streams into a single output
614−-- stream.
615−
616−-- This is like 'concatMap' but uses an unfold with an explicit state to
617−-- generate the stream instead of a 'Stream' type generator. This allows better
618−-- optimization via fusion. This can be many times more efficient than
619−-- 'concatMap'.
620−
621−{-# INLINE_NORMAL concatMapU #-}
622−concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
623−concatMapU (Unfold istep inject) (Stream ostep ost) =
624− Stream step (ConcatMapUOuter ost)
625− where
626− {-# INLINE_LATE step #-}
627− step gst (ConcatMapUOuter o) = do
628− r <- ostep (adaptState gst) o
629− case r of
630− Yield a o' -> do
631− i <- inject a
632− i `seq` return (Skip (ConcatMapUInner o' i))
633− Skip o' -> return $ Skip (ConcatMapUOuter o')
634− Stop -> return $ Stop
635−
636− step _ (ConcatMapUInner o i) = do
637− r <- istep i
638− return $ case r of
639− Yield x i' -> Yield x (ConcatMapUInner o i')
640− Skip i' -> Skip (ConcatMapUInner o i')
641− Stop -> Skip (ConcatMapUOuter o)
642−
643−data ConcatUnfoldInterleaveState o i =
644− ConcatUnfoldInterleaveOuter o [i]
645− | ConcatUnfoldInterleaveInner o [i]
646− | ConcatUnfoldInterleaveInnerL [i] [i]
647− | ConcatUnfoldInterleaveInnerR [i] [i]
648−
649−-- XXX use arrays to store state instead of lists.
650−-- XXX In general we can use different scheduling strategies e.g. how to
651−-- schedule the outer vs inner loop or assigning weights to different streams
652−-- or outer and inner loops.
653−
654−-- After a yield, switch to the next stream. Do not switch streams on Skip.
655−-- Yield from outer stream switches to the inner stream.
656−--
657−-- There are two choices here, (1) exhaust the outer stream first and then
658−-- start yielding from the inner streams, this is much simpler to implement,
659−-- (2) yield at least one element from an inner stream before going back to
660−-- outer stream and opening the next stream from it.
661−--
662−-- Ideally, we need some scheduling bias to inner streams vs outer stream.
663−-- Maybe we can configure the behavior.
664−--
665−{-# INLINE_NORMAL concatUnfoldInterleave #-}
666−concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
667−concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
668− Stream step (ConcatUnfoldInterleaveOuter ost [])
669− where
670− {-# INLINE_LATE step #-}
671− step gst (ConcatUnfoldInterleaveOuter o ls) = do
672− r <- ostep (adaptState gst) o
673− case r of
674− Yield a o' -> do
675− i <- inject a
676− i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
677− Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
678− Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
679−
680− step _ (ConcatUnfoldInterleaveInner _ []) = undefined
681− step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
682− r <- istep st
683− return $ case r of
684− Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
685− Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
686− Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
687−
688− step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
689− step _ (ConcatUnfoldInterleaveInnerL [] rs) =
690− return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
691−
692− step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
693− r <- istep st
694− return $ case r of
695− Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
696− Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
697− Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
698−
699− step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
700− step _ (ConcatUnfoldInterleaveInnerR ls []) =
701− return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
702−
703− step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
704− r <- istep st
705− return $ case r of
706− Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
707− Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
708− Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
709−
710−-- XXX In general we can use different scheduling strategies e.g. how to
711−-- schedule the outer vs inner loop or assigning weights to different streams
712−-- or outer and inner loops.
713−--
714−-- This could be inefficient if the tasks are too small.
715−--
716−-- Compared to concatUnfoldInterleave this one switches streams on Skips.
717−--
718−{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
719−concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
720−concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
721− Stream step (ConcatUnfoldInterleaveOuter ost [])
722− where
723− {-# INLINE_LATE step #-}
724− step gst (ConcatUnfoldInterleaveOuter o ls) = do
725− r <- ostep (adaptState gst) o
726− case r of
727− Yield a o' -> do
728− i <- inject a
729− i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
730− Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
731− Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
732−
733− step _ (ConcatUnfoldInterleaveInner o []) =
734− return $ Skip (ConcatUnfoldInterleaveOuter o [])
735−
736− step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
737− r <- istep st
738− return $ case r of
739− Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
740− Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
741− Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
742−
743− step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
744− step _ (ConcatUnfoldInterleaveInnerL [] rs) =
745− return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
746−
747− step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
748− r <- istep st
749− return $ case r of
750− Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
751− Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
752− Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
753−
754− step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
755− step _ (ConcatUnfoldInterleaveInnerR ls []) =
756− return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
757−
758− step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
759− r <- istep st
760− return $ case r of
761− Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
762− Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
763− Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
764−
765−data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
766−
767−-- Note that this could be much faster compared to the CPS stream. However, as
768−-- the number of streams being composed increases this may become expensive.
769−-- Need to see where the breaking point is between the two.
770−--
771−{-# INLINE_NORMAL append #-}
772−append :: Monad m => Stream m a -> Stream m a -> Stream m a
773−append (Stream step1 state1) (Stream step2 state2) =
774− Stream step (AppendFirst state1)
775−
776− where
777−
778− {-# INLINE_LATE step #-}
779− step gst (AppendFirst st) = do
780− r <- step1 gst st
781− return $ case r of
782− Yield a s -> Yield a (AppendFirst s)
783− Skip s -> Skip (AppendFirst s)
784− Stop -> Skip (AppendSecond state2)
785−
786− step gst (AppendSecond st) = do
787− r <- step2 gst st
788− return $ case r of
789− Yield a s -> Yield a (AppendSecond s)
790− Skip s -> Skip (AppendSecond s)
791− Stop -> Stop
792−
793−data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
794− | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
795−
796−{-# INLINE_NORMAL interleave #-}
797−interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
798−interleave (Stream step1 state1) (Stream step2 state2) =
799− Stream step (InterleaveFirst state1 state2)
800−
801− where
802−
803− {-# INLINE_LATE step #-}
804− step gst (InterleaveFirst st1 st2) = do
805− r <- step1 gst st1
806− return $ case r of
807− Yield a s -> Yield a (InterleaveSecond s st2)
808− Skip s -> Skip (InterleaveFirst s st2)
809− Stop -> Skip (InterleaveSecondOnly st2)
810−
811− step gst (InterleaveSecond st1 st2) = do
812− r <- step2 gst st2
813− return $ case r of
814− Yield a s -> Yield a (InterleaveFirst st1 s)
815− Skip s -> Skip (InterleaveSecond st1 s)
816− Stop -> Skip (InterleaveFirstOnly st1)
817−
818− step gst (InterleaveFirstOnly st1) = do
819− r <- step1 gst st1
820− return $ case r of
821− Yield a s -> Yield a (InterleaveFirstOnly s)
822− Skip s -> Skip (InterleaveFirstOnly s)
823− Stop -> Stop
824−
825− step gst (InterleaveSecondOnly st2) = do
826− r <- step2 gst st2
827− return $ case r of
828− Yield a s -> Yield a (InterleaveSecondOnly s)
829− Skip s -> Skip (InterleaveSecondOnly s)
830− Stop -> Stop
831−
832−{-# INLINE_NORMAL interleaveMin #-}
833−interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
834−interleaveMin (Stream step1 state1) (Stream step2 state2) =
835− Stream step (InterleaveFirst state1 state2)
836−
837− where
838−
839− {-# INLINE_LATE step #-}
840− step gst (InterleaveFirst st1 st2) = do
841− r <- step1 gst st1
842− return $ case r of
843− Yield a s -> Yield a (InterleaveSecond s st2)
844− Skip s -> Skip (InterleaveFirst s st2)
845− Stop -> Stop
846−
847− step gst (InterleaveSecond st1 st2) = do
848− r <- step2 gst st2
849− return $ case r of
850− Yield a s -> Yield a (InterleaveFirst st1 s)
851− Skip s -> Skip (InterleaveSecond st1 s)
852− Stop -> Stop
853−
854− step _ (InterleaveFirstOnly _) = undefined
855− step _ (InterleaveSecondOnly _) = undefined
856−
857−{-# INLINE_NORMAL interleaveSuffix #-}
858−interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
859−interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
860− Stream step (InterleaveFirst state1 state2)
861−
862− where
863−
864− {-# INLINE_LATE step #-}
865− step gst (InterleaveFirst st1 st2) = do
866− r <- step1 gst st1
867− return $ case r of
868− Yield a s -> Yield a (InterleaveSecond s st2)
869− Skip s -> Skip (InterleaveFirst s st2)
870− Stop -> Stop
871−
872− step gst (InterleaveSecond st1 st2) = do
873− r <- step2 gst st2
874− return $ case r of
875− Yield a s -> Yield a (InterleaveFirst st1 s)
876− Skip s -> Skip (InterleaveSecond st1 s)
877− Stop -> Skip (InterleaveFirstOnly st1)
878−
879− step gst (InterleaveFirstOnly st1) = do
880− r <- step1 gst st1
881− return $ case r of
882− Yield a s -> Yield a (InterleaveFirstOnly s)
883− Skip s -> Skip (InterleaveFirstOnly s)
884− Stop -> Stop
885−
886− step _ (InterleaveSecondOnly _) = undefined
887−
888−data InterleaveInfixState s1 s2 a
889− = InterleaveInfixFirst s1 s2
890− | InterleaveInfixSecondBuf s1 s2
891− | InterleaveInfixSecondYield s1 s2 a
892− | InterleaveInfixFirstYield s1 s2 a
893− | InterleaveInfixFirstOnly s1
894−
895−{-# INLINE_NORMAL interleaveInfix #-}
896−interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
897−interleaveInfix (Stream step1 state1) (Stream step2 state2) =
898− Stream step (InterleaveInfixFirst state1 state2)
899−
900− where
901−
902− {-# INLINE_LATE step #-}
903− step gst (InterleaveInfixFirst st1 st2) = do
904− r <- step1 gst st1
905− return $ case r of
906− Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
907− Skip s -> Skip (InterleaveInfixFirst s st2)
908− Stop -> Stop
909−
910− step gst (InterleaveInfixSecondBuf st1 st2) = do
911− r <- step2 gst st2
912− return $ case r of
913− Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
914− Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
915− Stop -> Skip (InterleaveInfixFirstOnly st1)
916−
917− step gst (InterleaveInfixSecondYield st1 st2 x) = do
918− r <- step1 gst st1
919− return $ case r of
920− Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
921− Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
922− Stop -> Stop
923−
924− step _ (InterleaveInfixFirstYield st1 st2 x) = do
925− return $ Yield x (InterleaveInfixSecondBuf st1 st2)
926−
927− step gst (InterleaveInfixFirstOnly st1) = do
928− r <- step1 gst st1
929− return $ case r of
930− Yield a s -> Yield a (InterleaveInfixFirstOnly s)
931− Skip s -> Skip (InterleaveInfixFirstOnly s)
932− Stop -> Stop
933−
934−{-# INLINE_NORMAL roundRobin #-}
935−roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
936−roundRobin (Stream step1 state1) (Stream step2 state2) =
937− Stream step (InterleaveFirst state1 state2)
938−
939− where
940−
941− {-# INLINE_LATE step #-}
942− step gst (InterleaveFirst st1 st2) = do
943− r <- step1 gst st1
944− return $ case r of
945− Yield a s -> Yield a (InterleaveSecond s st2)
946− Skip s -> Skip (InterleaveSecond s st2)
947− Stop -> Skip (InterleaveSecondOnly st2)
948−
949− step gst (InterleaveSecond st1 st2) = do
950− r <- step2 gst st2
951− return $ case r of
952− Yield a s -> Yield a (InterleaveFirst st1 s)
953− Skip s -> Skip (InterleaveFirst st1 s)
954− Stop -> Skip (InterleaveFirstOnly st1)
955−
956− step gst (InterleaveSecondOnly st2) = do
957− r <- step2 gst st2
958− return $ case r of
959− Yield a s -> Yield a (InterleaveSecondOnly s)
960− Skip s -> Skip (InterleaveSecondOnly s)
961− Stop -> Stop
962−
963− step gst (InterleaveFirstOnly st1) = do
964− r <- step1 gst st1
965− return $ case r of
966− Yield a s -> Yield a (InterleaveFirstOnly s)
967− Skip s -> Skip (InterleaveFirstOnly s)
968− Stop -> Stop
969−
970−data ICUState s1 s2 i1 i2 =
971− ICUFirst s1 s2
972− | ICUSecond s1 s2
973− | ICUSecondOnly s2
974− | ICUFirstOnly s1
975− | ICUFirstInner s1 s2 i1
976− | ICUSecondInner s1 s2 i2
977− | ICUFirstOnlyInner s1 i1
978− | ICUSecondOnlyInner s2 i2
979−
980−-- | Interleave streams (full streams, not the elements) unfolded from two
981−-- input streams and concat. Stop when the first stream stops. If the second
982−-- stream ends before the first one then first stream still keeps running alone
983−-- without any interleaving with the second stream.
984−--
985−-- [a1, a2, ... an] [b1, b2 ...]
986−-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
987−-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
988−-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
989−--
990−{-# INLINE_NORMAL gintercalateSuffix #-}
991−gintercalateSuffix
992− :: Monad m
993− => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
994−gintercalateSuffix
995− (Unfold istep1 inject1) (Stream step1 state1)
996− (Unfold istep2 inject2) (Stream step2 state2) =
997− Stream step (ICUFirst state1 state2)
998−
999− where
1000−
1001− {-# INLINE_LATE step #-}
1002− step gst (ICUFirst s1 s2) = do
1003− r <- step1 (adaptState gst) s1
1004− case r of
1005− Yield a s -> do
1006− i <- inject1 a
1007− i `seq` return (Skip (ICUFirstInner s s2 i))
1008− Skip s -> return $ Skip (ICUFirst s s2)
1009− Stop -> return Stop
1010−
1011− step gst (ICUFirstOnly s1) = do
1012− r <- step1 (adaptState gst) s1
1013− case r of
1014− Yield a s -> do
1015− i <- inject1 a
1016− i `seq` return (Skip (ICUFirstOnlyInner s i))
1017− Skip s -> return $ Skip (ICUFirstOnly s)
1018− Stop -> return Stop
1019−
1020− step _ (ICUFirstInner s1 s2 i1) = do
1021− r <- istep1 i1
1022− return $ case r of
1023− Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
1024− Skip i' -> Skip (ICUFirstInner s1 s2 i')
1025− Stop -> Skip (ICUSecond s1 s2)
1026−
1027− step _ (ICUFirstOnlyInner s1 i1) = do
1028− r <- istep1 i1
1029− return $ case r of
1030− Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
1031− Skip i' -> Skip (ICUFirstOnlyInner s1 i')
1032− Stop -> Skip (ICUFirstOnly s1)
1033−
1034− step gst (ICUSecond s1 s2) = do
1035− r <- step2 (adaptState gst) s2
1036− case r of
1037− Yield a s -> do
1038− i <- inject2 a
1039− i `seq` return (Skip (ICUSecondInner s1 s i))
1040− Skip s -> return $ Skip (ICUSecond s1 s)
1041− Stop -> return $ Skip (ICUFirstOnly s1)
1042−
1043− step _ (ICUSecondInner s1 s2 i2) = do
1044− r <- istep2 i2
1045− return $ case r of
1046− Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
1047− Skip i' -> Skip (ICUSecondInner s1 s2 i')
1048− Stop -> Skip (ICUFirst s1 s2)
1049−
1050− step _ (ICUSecondOnly _s2) = undefined
1051− step _ (ICUSecondOnlyInner _s2 _i2) = undefined
1052−
1053−{-# ANN type InterposeSuffixState Fuse #-}
1054−data InterposeSuffixState s1 i1 =
1055− InterposeSuffixFirst s1
1056− -- | InterposeSuffixFirstYield s1 i1
1057− | InterposeSuffixFirstInner s1 i1
1058− | InterposeSuffixSecond s1
1059−
1060−-- Note that if an unfolded layer turns out to be nil we still emit the
1061−-- separator effect. An alternate behavior could be to emit the separator
1062−-- effect only if at least one element has been yielded by the unfolding.
1063−-- However, that becomes a bit complicated, so we have chosen the former
1064−-- behvaior for now.
1065−{-# INLINE_NORMAL interposeSuffix #-}
1066−interposeSuffix
1067− :: Monad m
1068− => m c -> Unfold m b c -> Stream m b -> Stream m c
1069−interposeSuffix
1070− action
1071− (Unfold istep1 inject1) (Stream step1 state1) =
1072− Stream step (InterposeSuffixFirst state1)
1073−
1074− where
1075−
1076− {-# INLINE_LATE step #-}
1077− step gst (InterposeSuffixFirst s1) = do
1078− r <- step1 (adaptState gst) s1
1079− case r of
1080− Yield a s -> do
1081− i <- inject1 a
1082− i `seq` return (Skip (InterposeSuffixFirstInner s i))
1083− -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
1084− Skip s -> return $ Skip (InterposeSuffixFirst s)
1085− Stop -> return Stop
1086−
1087− {-
1088− step _ (InterposeSuffixFirstYield s1 i1) = do
1089− r <- istep1 i1
1090− return $ case r of
1091− Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
1092− Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
1093− Stop -> Skip (InterposeSuffixFirst s1)
1094− -}
1095−
1096− step _ (InterposeSuffixFirstInner s1 i1) = do
1097− r <- istep1 i1
1098− return $ case r of
1099− Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
1100− Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
1101− Stop -> Skip (InterposeSuffixSecond s1)
1102−
1103− step _ (InterposeSuffixSecond s1) = do
1104− r <- action
1105− return $ Yield r (InterposeSuffixFirst s1)
1106−
1107−data ICALState s1 s2 i1 i2 a =
1108− ICALFirst s1 s2
1109− -- | ICALFirstYield s1 s2 i1
1110− | ICALFirstInner s1 s2 i1
1111− | ICALFirstOnly s1
1112− | ICALFirstOnlyInner s1 i1
1113− | ICALSecondInject s1 s2
1114− | ICALFirstInject s1 s2 i2
1115− -- | ICALFirstBuf s1 s2 i1 i2
1116− | ICALSecondInner s1 s2 i1 i2
1117− -- -- | ICALSecondInner s1 s2 i1 i2 a
1118− -- -- | ICALFirstResume s1 s2 i1 i2 a
1119−
1120−-- | Interleave streams (full streams, not the elements) unfolded from two
1121−-- input streams and concat. Stop when the first stream stops. If the second
1122−-- stream ends before the first one then first stream still keeps running alone
1123−-- without any interleaving with the second stream.
1124−--
1125−-- [a1, a2, ... an] [b1, b2 ...]
1126−-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
1127−-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
1128−-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
1129−--
1130−{-# INLINE_NORMAL gintercalate #-}
1131−gintercalate
1132− :: Monad m
1133− => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
1134−gintercalate
1135− (Unfold istep1 inject1) (Stream step1 state1)
1136− (Unfold istep2 inject2) (Stream step2 state2) =
1137− Stream step (ICALFirst state1 state2)
1138−
1139− where
1140−
1141− {-# INLINE_LATE step #-}
1142− step gst (ICALFirst s1 s2) = do
1143− r <- step1 (adaptState gst) s1
1144− case r of
1145− Yield a s -> do
1146− i <- inject1 a
1147− i `seq` return (Skip (ICALFirstInner s s2 i))
1148− -- i `seq` return (Skip (ICALFirstYield s s2 i))
1149− Skip s -> return $ Skip (ICALFirst s s2)
1150− Stop -> return Stop
1151−
1152− {-
1153− step _ (ICALFirstYield s1 s2 i1) = do
1154− r <- istep1 i1
1155− return $ case r of
1156− Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
1157− Skip i' -> Skip (ICALFirstYield s1 s2 i')
1158− Stop -> Skip (ICALFirst s1 s2)
1159− -}
1160−
1161− step _ (ICALFirstInner s1 s2 i1) = do
1162− r <- istep1 i1
1163− return $ case r of
1164− Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
1165− Skip i' -> Skip (ICALFirstInner s1 s2 i')
1166− Stop -> Skip (ICALSecondInject s1 s2)
1167−
1168− step gst (ICALFirstOnly s1) = do
1169− r <- step1 (adaptState gst) s1
1170− case r of
1171− Yield a s -> do
1172− i <- inject1 a
1173− i `seq` return (Skip (ICALFirstOnlyInner s i))
1174− Skip s -> return $ Skip (ICALFirstOnly s)
1175− Stop -> return Stop
1176−
1177− step _ (ICALFirstOnlyInner s1 i1) = do
1178− r <- istep1 i1
1179− return $ case r of
1180− Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
1181− Skip i' -> Skip (ICALFirstOnlyInner s1 i')
1182− Stop -> Skip (ICALFirstOnly s1)
1183−
1184− -- We inject the second stream even before checking if the first stream
1185− -- would yield any more elements. There is no clear choice whether we
1186− -- should do this before or after that. Doing it after may make the state
1187− -- machine a bit simpler though.
1188− step gst (ICALSecondInject s1 s2) = do
1189− r <- step2 (adaptState gst) s2
1190− case r of
1191− Yield a s -> do
1192− i <- inject2 a
1193− i `seq` return (Skip (ICALFirstInject s1 s i))
1194− Skip s -> return $ Skip (ICALSecondInject s1 s)
1195− Stop -> return $ Skip (ICALFirstOnly s1)
1196−
1197− step gst (ICALFirstInject s1 s2 i2) = do
1198− r <- step1 (adaptState gst) s1
1199− case r of
1200− Yield a s -> do
1201− i <- inject1 a
1202− i `seq` return (Skip (ICALSecondInner s s2 i i2))
1203− -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
1204− Skip s -> return $ Skip (ICALFirstInject s s2 i2)
1205− Stop -> return Stop
1206−
1207− {-
1208− step _ (ICALFirstBuf s1 s2 i1 i2) = do
1209− r <- istep1 i1
1210− return $ case r of
1211− Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
1212− Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
1213− Stop -> Stop
1214−
1215− step _ (ICALSecondInner s1 s2 i1 i2 v) = do
1216− r <- istep2 i2
1217− return $ case r of
1218− Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
1219− Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
1220− Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
1221− -}
1222−
1223− step _ (ICALSecondInner s1 s2 i1 i2) = do
1224− r <- istep2 i2
1225− return $ case r of
1226− Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
1227− Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
1228− Stop -> Skip (ICALFirstInner s1 s2 i1)
1229− -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
1230−
1231− {-
1232− step _ (ICALFirstResume s1 s2 i1 i2 x) = do
1233− return $ Yield x (ICALFirstInner s1 s2 i1 i2)
1234− -}
1235−
1236−{-# ANN type InterposeState Fuse #-}
1237−data InterposeState s1 i1 a =
1238− InterposeFirst s1
1239− -- | InterposeFirstYield s1 i1
1240− | InterposeFirstInner s1 i1
1241− | InterposeFirstInject s1
1242− -- | InterposeFirstBuf s1 i1
1243− | InterposeSecondYield s1 i1
1244− -- -- | InterposeSecondYield s1 i1 a
1245− -- -- | InterposeFirstResume s1 i1 a
1246−
1247−-- Note that this only interposes the pure values, we may run many effects to
1248−-- generate those values as some effects may not generate anything (Skip).
1249−{-# INLINE_NORMAL interpose #-}
1250−interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
1251−interpose
1252− action
1253− (Unfold istep1 inject1) (Stream step1 state1) =
1254− Stream step (InterposeFirst state1)
1255−
1256− where
1257−
1258− {-# INLINE_LATE step #-}
1259− step gst (InterposeFirst s1) = do
1260− r <- step1 (adaptState gst) s1
1261− case r of
1262− Yield a s -> do
1263− i <- inject1 a
1264− i `seq` return (Skip (InterposeFirstInner s i))
1265− -- i `seq` return (Skip (InterposeFirstYield s i))
1266− Skip s -> return $ Skip (InterposeFirst s)
1267− Stop -> return Stop
1268−
1269− {-
1270− step _ (InterposeFirstYield s1 i1) = do
1271− r <- istep1 i1
1272− return $ case r of
1273− Yield x i' -> Yield x (InterposeFirstInner s1 i')
1274− Skip i' -> Skip (InterposeFirstYield s1 i')
1275− Stop -> Skip (InterposeFirst s1)
1276− -}
1277−
1278− step _ (InterposeFirstInner s1 i1) = do
1279− r <- istep1 i1
1280− return $ case r of
1281− Yield x i' -> Yield x (InterposeFirstInner s1 i')
1282− Skip i' -> Skip (InterposeFirstInner s1 i')
1283− Stop -> Skip (InterposeFirstInject s1)
1284−
1285− step gst (InterposeFirstInject s1) = do
1286− r <- step1 (adaptState gst) s1
1287− case r of
1288− Yield a s -> do
1289− i <- inject1 a
1290− -- i `seq` return (Skip (InterposeFirstBuf s i))
1291− i `seq` return (Skip (InterposeSecondYield s i))
1292− Skip s -> return $ Skip (InterposeFirstInject s)
1293− Stop -> return Stop
1294−
1295− {-
1296− step _ (InterposeFirstBuf s1 i1) = do
1297− r <- istep1 i1
1298− return $ case r of
1299− Yield x i' -> Skip (InterposeSecondYield s1 i' x)
1300− Skip i' -> Skip (InterposeFirstBuf s1 i')
1301− Stop -> Stop
1302− -}
1303−
1304− {-
1305− step _ (InterposeSecondYield s1 i1 v) = do
1306− r <- action
1307− return $ Yield r (InterposeFirstResume s1 i1 v)
1308− -}
1309− step _ (InterposeSecondYield s1 i1) = do
1310− r <- action
1311− return $ Yield r (InterposeFirstInner s1 i1)
1312−
1313− {-
1314− step _ (InterposeFirstResume s1 i1 v) = do
1315− return $ Yield v (InterposeFirstInner s1 i1)
1316− -}
1317−
1318------------------------------------------------------------------------------
602------------------------------------------------------------------------------
1319-- Exceptions
603-- Exceptions
1320------------------------------------------------------------------------------
604------------------------------------------------------------------------------
src/Streamly/Internal/Data/Stream/StreamD/Flatten.hsnew file+764
@@ -0,0 +1,764 @@
1+#include "inline.hs"
2+
3+-- |
4+-- Module : Streamly.Internal.Data.Stream.StreamD.Flatten
5+-- Copyright : (c) 2018 Composewell Technologies
6+-- (c) Roman Leshchinskiy 2008-2010
7+-- (c) The University of Glasgow, 2009
8+-- License : BSD3
9+-- Maintainer : streamly@composewell.com
10+-- Stability : experimental
11+-- Portability : GHC
12+
13+module Streamly.Internal.Data.Stream.StreamD.Flatten
14+ (
15+ -- ** Flattening nested streams
16+ concatMapM
17+ , concatMap
18+ , ConcatMapUState (..)
19+ , concatMapU
20+ , ConcatUnfoldInterleaveState (..)
21+ , concatUnfoldInterleave
22+ , concatUnfoldRoundrobin
23+ , AppendState(..)
24+ , append
25+ , InterleaveState(..)
26+ , interleave
27+ , interleaveMin
28+ , interleaveSuffix
29+ , interleaveInfix
30+ , roundRobin -- interleaveFair?/ParallelFair
31+ , gintercalateSuffix
32+ , interposeSuffix
33+ , gintercalate
34+ , interpose
35+ )
36+where
37+
38+import Fusion.Plugin.Types (Fuse(..))
39+import Streamly.Internal.Data.Unfold.Types (Unfold(..))
40+
41+
42+import Prelude hiding
43+ ( map, mapM, mapM_, repeat, foldr, last, take, filter
44+ , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem
45+ , notElem, null, head, tail, zipWith, lookup, foldr1, sequence
46+ , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
47+ , reverse, iterate, splitAt)
48+import Streamly.Internal.Data.Stream.StreamD.Type
49+import Streamly.Internal.Data.SVar
50+
51+-------------------------------------------------------------------------------
52+-- Stream transformations using Unfolds
53+-------------------------------------------------------------------------------
54+
55+-- Define a unique structure to use in inspection testing
56+data ConcatMapUState o i =
57+ ConcatMapUOuter o
58+ | ConcatMapUInner o i
59+
60+-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
61+-- to streams and then flattens the generated streams into a single output
62+-- stream.
63+
64+-- This is like 'concatMap' but uses an unfold with an explicit state to
65+-- generate the stream instead of a 'Stream' type generator. This allows better
66+-- optimization via fusion. This can be many times more efficient than
67+-- 'concatMap'.
68+
69+{-# INLINE_NORMAL concatMapU #-}
70+concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
71+concatMapU (Unfold istep inject) (Stream ostep ost) =
72+ Stream step (ConcatMapUOuter ost)
73+ where
74+ {-# INLINE_LATE step #-}
75+ step gst (ConcatMapUOuter o) = do
76+ r <- ostep (adaptState gst) o
77+ case r of
78+ Yield a o' -> do
79+ i <- inject a
80+ i `seq` return (Skip (ConcatMapUInner o' i))
81+ Skip o' -> return $ Skip (ConcatMapUOuter o')
82+ Stop -> return $ Stop
83+
84+ step _ (ConcatMapUInner o i) = do
85+ r <- istep i
86+ return $ case r of
87+ Yield x i' -> Yield x (ConcatMapUInner o i')
88+ Skip i' -> Skip (ConcatMapUInner o i')
89+ Stop -> Skip (ConcatMapUOuter o)
90+
91+data ConcatUnfoldInterleaveState o i =
92+ ConcatUnfoldInterleaveOuter o [i]
93+ | ConcatUnfoldInterleaveInner o [i]
94+ | ConcatUnfoldInterleaveInnerL [i] [i]
95+ | ConcatUnfoldInterleaveInnerR [i] [i]
96+
97+-- XXX use arrays to store state instead of lists.
98+-- XXX In general we can use different scheduling strategies e.g. how to
99+-- schedule the outer vs inner loop or assigning weights to different streams
100+-- or outer and inner loops.
101+
102+-- After a yield, switch to the next stream. Do not switch streams on Skip.
103+-- Yield from outer stream switches to the inner stream.
104+--
105+-- There are two choices here, (1) exhaust the outer stream first and then
106+-- start yielding from the inner streams, this is much simpler to implement,
107+-- (2) yield at least one element from an inner stream before going back to
108+-- outer stream and opening the next stream from it.
109+--
110+-- Ideally, we need some scheduling bias to inner streams vs outer stream.
111+-- Maybe we can configure the behavior.
112+--
113+{-# INLINE_NORMAL concatUnfoldInterleave #-}
114+concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
115+concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
116+ Stream step (ConcatUnfoldInterleaveOuter ost [])
117+ where
118+ {-# INLINE_LATE step #-}
119+ step gst (ConcatUnfoldInterleaveOuter o ls) = do
120+ r <- ostep (adaptState gst) o
121+ case r of
122+ Yield a o' -> do
123+ i <- inject a
124+ i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
125+ Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
126+ Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
127+
128+ step _ (ConcatUnfoldInterleaveInner _ []) = undefined
129+ step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
130+ r <- istep st
131+ return $ case r of
132+ Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
133+ Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
134+ Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
135+
136+ step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
137+ step _ (ConcatUnfoldInterleaveInnerL [] rs) =
138+ return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
139+
140+ step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
141+ r <- istep st
142+ return $ case r of
143+ Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
144+ Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
145+ Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
146+
147+ step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
148+ step _ (ConcatUnfoldInterleaveInnerR ls []) =
149+ return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
150+
151+ step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
152+ r <- istep st
153+ return $ case r of
154+ Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
155+ Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
156+ Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
157+
158+-- XXX In general we can use different scheduling strategies e.g. how to
159+-- schedule the outer vs inner loop or assigning weights to different streams
160+-- or outer and inner loops.
161+--
162+-- This could be inefficient if the tasks are too small.
163+--
164+-- Compared to concatUnfoldInterleave this one switches streams on Skips.
165+--
166+{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
167+concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
168+concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
169+ Stream step (ConcatUnfoldInterleaveOuter ost [])
170+ where
171+ {-# INLINE_LATE step #-}
172+ step gst (ConcatUnfoldInterleaveOuter o ls) = do
173+ r <- ostep (adaptState gst) o
174+ case r of
175+ Yield a o' -> do
176+ i <- inject a
177+ i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
178+ Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
179+ Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
180+
181+ step _ (ConcatUnfoldInterleaveInner o []) =
182+ return $ Skip (ConcatUnfoldInterleaveOuter o [])
183+
184+ step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
185+ r <- istep st
186+ return $ case r of
187+ Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
188+ Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
189+ Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
190+
191+ step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
192+ step _ (ConcatUnfoldInterleaveInnerL [] rs) =
193+ return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
194+
195+ step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
196+ r <- istep st
197+ return $ case r of
198+ Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
199+ Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
200+ Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
201+
202+ step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
203+ step _ (ConcatUnfoldInterleaveInnerR ls []) =
204+ return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
205+
206+ step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
207+ r <- istep st
208+ return $ case r of
209+ Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
210+ Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
211+ Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
212+
213+data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
214+
215+-- Note that this could be much faster compared to the CPS stream. However, as
216+-- the number of streams being composed increases this may become expensive.
217+-- Need to see where the breaking point is between the two.
218+--
219+{-# INLINE_NORMAL append #-}
220+append :: Monad m => Stream m a -> Stream m a -> Stream m a
221+append (Stream step1 state1) (Stream step2 state2) =
222+ Stream step (AppendFirst state1)
223+
224+ where
225+
226+ {-# INLINE_LATE step #-}
227+ step gst (AppendFirst st) = do
228+ r <- step1 gst st
229+ return $ case r of
230+ Yield a s -> Yield a (AppendFirst s)
231+ Skip s -> Skip (AppendFirst s)
232+ Stop -> Skip (AppendSecond state2)
233+
234+ step gst (AppendSecond st) = do
235+ r <- step2 gst st
236+ return $ case r of
237+ Yield a s -> Yield a (AppendSecond s)
238+ Skip s -> Skip (AppendSecond s)
239+ Stop -> Stop
240+
241+data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
242+ | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
243+
244+{-# INLINE_NORMAL interleave #-}
245+interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
246+interleave (Stream step1 state1) (Stream step2 state2) =
247+ Stream step (InterleaveFirst state1 state2)
248+
249+ where
250+
251+ {-# INLINE_LATE step #-}
252+ step gst (InterleaveFirst st1 st2) = do
253+ r <- step1 gst st1
254+ return $ case r of
255+ Yield a s -> Yield a (InterleaveSecond s st2)
256+ Skip s -> Skip (InterleaveFirst s st2)
257+ Stop -> Skip (InterleaveSecondOnly st2)
258+
259+ step gst (InterleaveSecond st1 st2) = do
260+ r <- step2 gst st2
261+ return $ case r of
262+ Yield a s -> Yield a (InterleaveFirst st1 s)
263+ Skip s -> Skip (InterleaveSecond st1 s)
264+ Stop -> Skip (InterleaveFirstOnly st1)
265+
266+ step gst (InterleaveFirstOnly st1) = do
267+ r <- step1 gst st1
268+ return $ case r of
269+ Yield a s -> Yield a (InterleaveFirstOnly s)
270+ Skip s -> Skip (InterleaveFirstOnly s)
271+ Stop -> Stop
272+
273+ step gst (InterleaveSecondOnly st2) = do
274+ r <- step2 gst st2
275+ return $ case r of
276+ Yield a s -> Yield a (InterleaveSecondOnly s)
277+ Skip s -> Skip (InterleaveSecondOnly s)
278+ Stop -> Stop
279+
280+{-# INLINE_NORMAL interleaveMin #-}
281+interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
282+interleaveMin (Stream step1 state1) (Stream step2 state2) =
283+ Stream step (InterleaveFirst state1 state2)
284+
285+ where
286+
287+ {-# INLINE_LATE step #-}
288+ step gst (InterleaveFirst st1 st2) = do
289+ r <- step1 gst st1
290+ return $ case r of
291+ Yield a s -> Yield a (InterleaveSecond s st2)
292+ Skip s -> Skip (InterleaveFirst s st2)
293+ Stop -> Stop
294+
295+ step gst (InterleaveSecond st1 st2) = do
296+ r <- step2 gst st2
297+ return $ case r of
298+ Yield a s -> Yield a (InterleaveFirst st1 s)
299+ Skip s -> Skip (InterleaveSecond st1 s)
300+ Stop -> Stop
301+
302+ step _ (InterleaveFirstOnly _) = undefined
303+ step _ (InterleaveSecondOnly _) = undefined
304+
305+{-# INLINE_NORMAL interleaveSuffix #-}
306+interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
307+interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
308+ Stream step (InterleaveFirst state1 state2)
309+
310+ where
311+
312+ {-# INLINE_LATE step #-}
313+ step gst (InterleaveFirst st1 st2) = do
314+ r <- step1 gst st1
315+ return $ case r of
316+ Yield a s -> Yield a (InterleaveSecond s st2)
317+ Skip s -> Skip (InterleaveFirst s st2)
318+ Stop -> Stop
319+
320+ step gst (InterleaveSecond st1 st2) = do
321+ r <- step2 gst st2
322+ return $ case r of
323+ Yield a s -> Yield a (InterleaveFirst st1 s)
324+ Skip s -> Skip (InterleaveSecond st1 s)
325+ Stop -> Skip (InterleaveFirstOnly st1)
326+
327+ step gst (InterleaveFirstOnly st1) = do
328+ r <- step1 gst st1
329+ return $ case r of
330+ Yield a s -> Yield a (InterleaveFirstOnly s)
331+ Skip s -> Skip (InterleaveFirstOnly s)
332+ Stop -> Stop
333+
334+ step _ (InterleaveSecondOnly _) = undefined
335+
336+data InterleaveInfixState s1 s2 a
337+ = InterleaveInfixFirst s1 s2
338+ | InterleaveInfixSecondBuf s1 s2
339+ | InterleaveInfixSecondYield s1 s2 a
340+ | InterleaveInfixFirstYield s1 s2 a
341+ | InterleaveInfixFirstOnly s1
342+
343+{-# INLINE_NORMAL interleaveInfix #-}
344+interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
345+interleaveInfix (Stream step1 state1) (Stream step2 state2) =
346+ Stream step (InterleaveInfixFirst state1 state2)
347+
348+ where
349+
350+ {-# INLINE_LATE step #-}
351+ step gst (InterleaveInfixFirst st1 st2) = do
352+ r <- step1 gst st1
353+ return $ case r of
354+ Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
355+ Skip s -> Skip (InterleaveInfixFirst s st2)
356+ Stop -> Stop
357+
358+ step gst (InterleaveInfixSecondBuf st1 st2) = do
359+ r <- step2 gst st2
360+ return $ case r of
361+ Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
362+ Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
363+ Stop -> Skip (InterleaveInfixFirstOnly st1)
364+
365+ step gst (InterleaveInfixSecondYield st1 st2 x) = do
366+ r <- step1 gst st1
367+ return $ case r of
368+ Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
369+ Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
370+ Stop -> Stop
371+
372+ step _ (InterleaveInfixFirstYield st1 st2 x) = do
373+ return $ Yield x (InterleaveInfixSecondBuf st1 st2)
374+
375+ step gst (InterleaveInfixFirstOnly st1) = do
376+ r <- step1 gst st1
377+ return $ case r of
378+ Yield a s -> Yield a (InterleaveInfixFirstOnly s)
379+ Skip s -> Skip (InterleaveInfixFirstOnly s)
380+ Stop -> Stop
381+
382+{-# INLINE_NORMAL roundRobin #-}
383+roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
384+roundRobin (Stream step1 state1) (Stream step2 state2) =
385+ Stream step (InterleaveFirst state1 state2)
386+
387+ where
388+
389+ {-# INLINE_LATE step #-}
390+ step gst (InterleaveFirst st1 st2) = do
391+ r <- step1 gst st1
392+ return $ case r of
393+ Yield a s -> Yield a (InterleaveSecond s st2)
394+ Skip s -> Skip (InterleaveSecond s st2)
395+ Stop -> Skip (InterleaveSecondOnly st2)
396+
397+ step gst (InterleaveSecond st1 st2) = do
398+ r <- step2 gst st2
399+ return $ case r of
400+ Yield a s -> Yield a (InterleaveFirst st1 s)
401+ Skip s -> Skip (InterleaveFirst st1 s)
402+ Stop -> Skip (InterleaveFirstOnly st1)
403+
404+ step gst (InterleaveSecondOnly st2) = do
405+ r <- step2 gst st2
406+ return $ case r of
407+ Yield a s -> Yield a (InterleaveSecondOnly s)
408+ Skip s -> Skip (InterleaveSecondOnly s)
409+ Stop -> Stop
410+
411+ step gst (InterleaveFirstOnly st1) = do
412+ r <- step1 gst st1
413+ return $ case r of
414+ Yield a s -> Yield a (InterleaveFirstOnly s)
415+ Skip s -> Skip (InterleaveFirstOnly s)
416+ Stop -> Stop
417+
418+data ICUState s1 s2 i1 i2 =
419+ ICUFirst s1 s2
420+ | ICUSecond s1 s2
421+ | ICUSecondOnly s2
422+ | ICUFirstOnly s1
423+ | ICUFirstInner s1 s2 i1
424+ | ICUSecondInner s1 s2 i2
425+ | ICUFirstOnlyInner s1 i1
426+ | ICUSecondOnlyInner s2 i2
427+
428+-- | Interleave streams (full streams, not the elements) unfolded from two
429+-- input streams and concat. Stop when the first stream stops. If the second
430+-- stream ends before the first one then first stream still keeps running alone
431+-- without any interleaving with the second stream.
432+--
433+-- [a1, a2, ... an] [b1, b2 ...]
434+-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
435+-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
436+-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
437+--
438+{-# INLINE_NORMAL gintercalateSuffix #-}
439+gintercalateSuffix
440+ :: Monad m
441+ => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
442+gintercalateSuffix
443+ (Unfold istep1 inject1) (Stream step1 state1)
444+ (Unfold istep2 inject2) (Stream step2 state2) =
445+ Stream step (ICUFirst state1 state2)
446+
447+ where
448+
449+ {-# INLINE_LATE step #-}
450+ step gst (ICUFirst s1 s2) = do
451+ r <- step1 (adaptState gst) s1
452+ case r of
453+ Yield a s -> do
454+ i <- inject1 a
455+ i `seq` return (Skip (ICUFirstInner s s2 i))
456+ Skip s -> return $ Skip (ICUFirst s s2)
457+ Stop -> return Stop
458+
459+ step gst (ICUFirstOnly s1) = do
460+ r <- step1 (adaptState gst) s1
461+ case r of
462+ Yield a s -> do
463+ i <- inject1 a
464+ i `seq` return (Skip (ICUFirstOnlyInner s i))
465+ Skip s -> return $ Skip (ICUFirstOnly s)
466+ Stop -> return Stop
467+
468+ step _ (ICUFirstInner s1 s2 i1) = do
469+ r <- istep1 i1
470+ return $ case r of
471+ Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
472+ Skip i' -> Skip (ICUFirstInner s1 s2 i')
473+ Stop -> Skip (ICUSecond s1 s2)
474+
475+ step _ (ICUFirstOnlyInner s1 i1) = do
476+ r <- istep1 i1
477+ return $ case r of
478+ Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
479+ Skip i' -> Skip (ICUFirstOnlyInner s1 i')
480+ Stop -> Skip (ICUFirstOnly s1)
481+
482+ step gst (ICUSecond s1 s2) = do
483+ r <- step2 (adaptState gst) s2
484+ case r of
485+ Yield a s -> do
486+ i <- inject2 a
487+ i `seq` return (Skip (ICUSecondInner s1 s i))
488+ Skip s -> return $ Skip (ICUSecond s1 s)
489+ Stop -> return $ Skip (ICUFirstOnly s1)
490+
491+ step _ (ICUSecondInner s1 s2 i2) = do
492+ r <- istep2 i2
493+ return $ case r of
494+ Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
495+ Skip i' -> Skip (ICUSecondInner s1 s2 i')
496+ Stop -> Skip (ICUFirst s1 s2)
497+
498+ step _ (ICUSecondOnly _s2) = undefined
499+ step _ (ICUSecondOnlyInner _s2 _i2) = undefined
500+
501+{-# ANN type InterposeSuffixState Fuse #-}
502+data InterposeSuffixState s1 i1 =
503+ InterposeSuffixFirst s1
504+ -- | InterposeSuffixFirstYield s1 i1
505+ | InterposeSuffixFirstInner s1 i1
506+ | InterposeSuffixSecond s1
507+
508+-- Note that if an unfolded layer turns out to be nil we still emit the
509+-- separator effect. An alternate behavior could be to emit the separator
510+-- effect only if at least one element has been yielded by the unfolding.
511+-- However, that becomes a bit complicated, so we have chosen the former
512+-- behvaior for now.
513+{-# INLINE_NORMAL interposeSuffix #-}
514+interposeSuffix
515+ :: Monad m
516+ => m c -> Unfold m b c -> Stream m b -> Stream m c
517+interposeSuffix
518+ action
519+ (Unfold istep1 inject1) (Stream step1 state1) =
520+ Stream step (InterposeSuffixFirst state1)
521+
522+ where
523+
524+ {-# INLINE_LATE step #-}
525+ step gst (InterposeSuffixFirst s1) = do
526+ r <- step1 (adaptState gst) s1
527+ case r of
528+ Yield a s -> do
529+ i <- inject1 a
530+ i `seq` return (Skip (InterposeSuffixFirstInner s i))
531+ -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
532+ Skip s -> return $ Skip (InterposeSuffixFirst s)
533+ Stop -> return Stop
534+
535+ {-
536+ step _ (InterposeSuffixFirstYield s1 i1) = do
537+ r <- istep1 i1
538+ return $ case r of
539+ Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
540+ Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
541+ Stop -> Skip (InterposeSuffixFirst s1)
542+ -}
543+
544+ step _ (InterposeSuffixFirstInner s1 i1) = do
545+ r <- istep1 i1
546+ return $ case r of
547+ Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
548+ Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
549+ Stop -> Skip (InterposeSuffixSecond s1)
550+
551+ step _ (InterposeSuffixSecond s1) = do
552+ r <- action
553+ return $ Yield r (InterposeSuffixFirst s1)
554+
555+data ICALState s1 s2 i1 i2 a =
556+ ICALFirst s1 s2
557+ -- | ICALFirstYield s1 s2 i1
558+ | ICALFirstInner s1 s2 i1
559+ | ICALFirstOnly s1
560+ | ICALFirstOnlyInner s1 i1
561+ | ICALSecondInject s1 s2
562+ | ICALFirstInject s1 s2 i2
563+ -- | ICALFirstBuf s1 s2 i1 i2
564+ | ICALSecondInner s1 s2 i1 i2
565+ -- -- | ICALSecondInner s1 s2 i1 i2 a
566+ -- -- | ICALFirstResume s1 s2 i1 i2 a
567+
568+-- | Interleave streams (full streams, not the elements) unfolded from two
569+-- input streams and concat. Stop when the first stream stops. If the second
570+-- stream ends before the first one then first stream still keeps running alone
571+-- without any interleaving with the second stream.
572+--
573+-- [a1, a2, ... an] [b1, b2 ...]
574+-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
575+-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
576+-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
577+--
578+{-# INLINE_NORMAL gintercalate #-}
579+gintercalate
580+ :: Monad m
581+ => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
582+gintercalate
583+ (Unfold istep1 inject1) (Stream step1 state1)
584+ (Unfold istep2 inject2) (Stream step2 state2) =
585+ Stream step (ICALFirst state1 state2)
586+
587+ where
588+
589+ {-# INLINE_LATE step #-}
590+ step gst (ICALFirst s1 s2) = do
591+ r <- step1 (adaptState gst) s1
592+ case r of
593+ Yield a s -> do
594+ i <- inject1 a
595+ i `seq` return (Skip (ICALFirstInner s s2 i))
596+ -- i `seq` return (Skip (ICALFirstYield s s2 i))
597+ Skip s -> return $ Skip (ICALFirst s s2)
598+ Stop -> return Stop
599+
600+ {-
601+ step _ (ICALFirstYield s1 s2 i1) = do
602+ r <- istep1 i1
603+ return $ case r of
604+ Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
605+ Skip i' -> Skip (ICALFirstYield s1 s2 i')
606+ Stop -> Skip (ICALFirst s1 s2)
607+ -}
608+
609+ step _ (ICALFirstInner s1 s2 i1) = do
610+ r <- istep1 i1
611+ return $ case r of
612+ Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
613+ Skip i' -> Skip (ICALFirstInner s1 s2 i')
614+ Stop -> Skip (ICALSecondInject s1 s2)
615+
616+ step gst (ICALFirstOnly s1) = do
617+ r <- step1 (adaptState gst) s1
618+ case r of
619+ Yield a s -> do
620+ i <- inject1 a
621+ i `seq` return (Skip (ICALFirstOnlyInner s i))
622+ Skip s -> return $ Skip (ICALFirstOnly s)
623+ Stop -> return Stop
624+
625+ step _ (ICALFirstOnlyInner s1 i1) = do
626+ r <- istep1 i1
627+ return $ case r of
628+ Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
629+ Skip i' -> Skip (ICALFirstOnlyInner s1 i')
630+ Stop -> Skip (ICALFirstOnly s1)
631+
632+ -- We inject the second stream even before checking if the first stream
633+ -- would yield any more elements. There is no clear choice whether we
634+ -- should do this before or after that. Doing it after may make the state
635+ -- machine a bit simpler though.
636+ step gst (ICALSecondInject s1 s2) = do
637+ r <- step2 (adaptState gst) s2
638+ case r of
639+ Yield a s -> do
640+ i <- inject2 a
641+ i `seq` return (Skip (ICALFirstInject s1 s i))
642+ Skip s -> return $ Skip (ICALSecondInject s1 s)
643+ Stop -> return $ Skip (ICALFirstOnly s1)
644+
645+ step gst (ICALFirstInject s1 s2 i2) = do
646+ r <- step1 (adaptState gst) s1
647+ case r of
648+ Yield a s -> do
649+ i <- inject1 a
650+ i `seq` return (Skip (ICALSecondInner s s2 i i2))
651+ -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
652+ Skip s -> return $ Skip (ICALFirstInject s s2 i2)
653+ Stop -> return Stop
654+
655+ {-
656+ step _ (ICALFirstBuf s1 s2 i1 i2) = do
657+ r <- istep1 i1
658+ return $ case r of
659+ Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
660+ Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
661+ Stop -> Stop
662+
663+ step _ (ICALSecondInner s1 s2 i1 i2 v) = do
664+ r <- istep2 i2
665+ return $ case r of
666+ Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
667+ Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
668+ Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
669+ -}
670+
671+ step _ (ICALSecondInner s1 s2 i1 i2) = do
672+ r <- istep2 i2
673+ return $ case r of
674+ Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
675+ Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
676+ Stop -> Skip (ICALFirstInner s1 s2 i1)
677+ -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
678+
679+ {-
680+ step _ (ICALFirstResume s1 s2 i1 i2 x) = do
681+ return $ Yield x (ICALFirstInner s1 s2 i1 i2)
682+ -}
683+
684+{-# ANN type InterposeState Fuse #-}
685+data InterposeState s1 i1 a =
686+ InterposeFirst s1
687+ -- | InterposeFirstYield s1 i1
688+ | InterposeFirstInner s1 i1
689+ | InterposeFirstInject s1
690+ -- | InterposeFirstBuf s1 i1
691+ | InterposeSecondYield s1 i1
692+ -- -- | InterposeSecondYield s1 i1 a
693+ -- -- | InterposeFirstResume s1 i1 a
694+
695+-- Note that this only interposes the pure values, we may run many effects to
696+-- generate those values as some effects may not generate anything (Skip).
697+{-# INLINE_NORMAL interpose #-}
698+interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
699+interpose
700+ action
701+ (Unfold istep1 inject1) (Stream step1 state1) =
702+ Stream step (InterposeFirst state1)
703+
704+ where
705+
706+ {-# INLINE_LATE step #-}
707+ step gst (InterposeFirst s1) = do
708+ r <- step1 (adaptState gst) s1
709+ case r of
710+ Yield a s -> do
711+ i <- inject1 a
712+ i `seq` return (Skip (InterposeFirstInner s i))
713+ -- i `seq` return (Skip (InterposeFirstYield s i))
714+ Skip s -> return $ Skip (InterposeFirst s)
715+ Stop -> return Stop
716+
717+ {-
718+ step _ (InterposeFirstYield s1 i1) = do
719+ r <- istep1 i1
720+ return $ case r of
721+ Yield x i' -> Yield x (InterposeFirstInner s1 i')
722+ Skip i' -> Skip (InterposeFirstYield s1 i')
723+ Stop -> Skip (InterposeFirst s1)
724+ -}
725+
726+ step _ (InterposeFirstInner s1 i1) = do
727+ r <- istep1 i1
728+ return $ case r of
729+ Yield x i' -> Yield x (InterposeFirstInner s1 i')
730+ Skip i' -> Skip (InterposeFirstInner s1 i')
731+ Stop -> Skip (InterposeFirstInject s1)
732+
733+ step gst (InterposeFirstInject s1) = do
734+ r <- step1 (adaptState gst) s1
735+ case r of
736+ Yield a s -> do
737+ i <- inject1 a
738+ -- i `seq` return (Skip (InterposeFirstBuf s i))
739+ i `seq` return (Skip (InterposeSecondYield s i))
740+ Skip s -> return $ Skip (InterposeFirstInject s)
741+ Stop -> return Stop
742+
743+ {-
744+ step _ (InterposeFirstBuf s1 i1) = do
745+ r <- istep1 i1
746+ return $ case r of
747+ Yield x i' -> Skip (InterposeSecondYield s1 i' x)
748+ Skip i' -> Skip (InterposeFirstBuf s1 i')
749+ Stop -> Stop
750+ -}
751+
752+ {-
753+ step _ (InterposeSecondYield s1 i1 v) = do
754+ r <- action
755+ return $ Yield r (InterposeFirstResume s1 i1 v)
756+ -}
757+ step _ (InterposeSecondYield s1 i1) = do
758+ r <- action
759+ return $ Yield r (InterposeFirstInner s1 i1)
760+
761+ {-
762+ step _ (InterposeFirstResume s1 i1 v) = do
763+ return $ Yield v (InterposeFirstInner s1 i1)
764+ -}
streamly.cabal+1
@@ -456,6 +456,7 @@ library
456 , Streamly.Internal.Data.Stream.StreamD.Generate
456 , Streamly.Internal.Data.Stream.StreamD.Generate
457 , Streamly.Internal.Data.Stream.StreamD.SplitGroup
457 , Streamly.Internal.Data.Stream.StreamD.SplitGroup
458 , Streamly.Internal.Data.Stream.StreamD.Eliminate
458 , Streamly.Internal.Data.Stream.StreamD.Eliminate
459+ , Streamly.Internal.Data.Stream.StreamD.Flatten
459 , Streamly.Internal.Data.Stream.StreamDK.Type
460 , Streamly.Internal.Data.Stream.StreamDK.Type
460 , Streamly.Internal.Data.Stream.StreamDK
461 , Streamly.Internal.Data.Stream.StreamDK
461 , Streamly.Internal.Data.Stream.Enumeration
462 , Streamly.Internal.Data.Stream.Enumeration