@@ -878,65 +878,42 @@ cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
878878-- Merging
879879------------------------------------------------------------------------------
880880
881−data MergeWithState l r o
882− = MergeBoth l r
883− | MergeWithLeft l r o
884− | MergeWithRight l r o
885− | MergeNoMore (Either l r)
886−
887881{-# INLINE_NORMAL mergeBy #-}
888882mergeBy :: (Monad m) => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
889883mergeBy f (Stream stepa ta) (Stream stepb tb) =
890− Stream step (MergeBoth ta tb)
884+ Stream step (Just ta, Just tb, Nothing, Nothing)
885+ where
891886 {-# INLINE_LATE step #-}
892− step gst (MergeBoth sa sb) = do
893− ra <- stepa (rstState gst) sa
894− rb <- stepb (rstState gst) sb
895− case (ra, rb) of
896− (Yield xa sa', Yield xb sb') -> do
897− case f xa xb of
898− LT ->
899− return $ Yield xa (MergeWithRight sa' sb' xb)
900− _ ->
901− return $ Yield xb (MergeWithLeft sa' sb' xa)
902− (Yield xa sa', Stop) -> return $ Yield xa (MergeNoMore (Left sa'))
903− (Stop, Yield xb sb') -> return $ Yield xb (MergeNoMore (Right sb'))
904− (Yield xa sa', Skip sb') -> return $ Skip (MergeWithLeft sa' sb' xa)
905− (Skip sa', Yield xb sb') -> return $ Skip (MergeWithRight sa' sb' xb)
906− (Skip sa', Skip sb') -> return $ Skip (MergeBoth sa' sb')
907− (Skip sa', Stop) -> return $ Skip (MergeNoMore (Left sa'))
908− (Stop, Skip sb') -> return $ Skip (MergeNoMore (Right sb'))
909− (Stop, Stop) -> return Stop
910− step gst (MergeWithLeft sa sb x) = do
911− rb <- stepb (rstState gst) sb
912− case rb of
913− Yield xb sb' -> do
914− case f x xb of
915− LT -> return $ Yield x (MergeWithRight sa sb' xb)
916− _ -> return $ Yield xb (MergeWithLeft sa sb' x)
917− Skip sb' -> return $ Skip (MergeWithLeft sa sb' x)
918− Stop -> return $ Yield x (MergeNoMore (Left sa))
919− step gst (MergeWithRight sa sb x) = do
920− ra <- stepa (rstState gst) sa
921− case ra of
922− Yield xa sa' -> do
923− case f xa x of
924− LT -> return $ Yield xa (MergeWithRight sa' sb x)
925− _ -> return $ Yield x (MergeWithLeft sa' sb xa)
926− Skip sa' -> return $ Skip (MergeWithRight sa' sb x)
927− Stop -> return $ Yield x (MergeNoMore (Right sb))
928− step gst (MergeNoMore (Left sa)) = do
929− ra <- stepa (rstState gst) sa
930− case ra of
931− Yield xa sa' -> return $ Yield xa (MergeNoMore (Left sa'))
932− Skip sa' -> return $ Skip (MergeNoMore (Left sa'))
933− Stop -> return Stop
934− step gst (MergeNoMore (Right sb)) = do
935− rb <- stepb (rstState gst) sb
936− case rb of
937− Yield xb sb' -> return $ Yield xb (MergeNoMore (Right sb'))
938− Skip sb' -> return $ Skip (MergeNoMore (Right sb'))
939− Stop -> return Stop
887+
888+ -- one of the values is missing, and the corresponding stream is running
889+ step gst (Just sa, sb, Nothing, b) = do
890+ r <- stepa (rstState gst) sa
891+ return $ case r of
892+ Yield a sa' -> Skip (Just sa', sb, Just a, b)
893+ Skip sa' -> Skip (Just sa', sb, Nothing, b)
894+ Stop -> Skip (Nothing, sb, Nothing, b)
895+
896+ step gst (sa, Just sb, a, Nothing) = do
897+ r <- stepb (rstState gst) sb
898+ return $ case r of
899+ Yield b sb' -> Skip (sa, Just sb', a, Just b)
900+ Skip sb' -> Skip (sa, Just sb', a, Nothing)
901+ Stop -> Skip (sa, Nothing, a, Nothing)
902+
903+ -- both the values are available
904+ step _ (sa, sb, Just a, Just b) =
905+ return $ case f a b of
906+ LT -> Yield a (sa, sb, Nothing, Just b)
907+ _ -> Yield b (sa, sb, Just a, Nothing)
908+
909+ -- one of the values is missing, corresponding stream is done
910+ step _ (Nothing, sb, Nothing, Just b) = do
911+ return $ Yield b (Nothing, sb, Nothing, Nothing)
912+
913+ step _ (sa, Nothing, Just a, Nothing) = do
914+ return $ Yield a (sa, Nothing, Nothing, Nothing)
915+
916+ step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
940917
941918------------------------------------------------------------------------------
942919-- Transformation comprehensions