@@ -789,30 +789,134 @@ zipWithM f = go
789789mergeByM :: Monad m =>
790790 (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
791791mergeByM cmp = go
792+
793+ where
794+
795+ go mx my = mkStream $ \st yld sng stp -> do
796+ let stop = foldStream st yld sng stp my
797+ single x = foldStream st yld sng stp (goX0 x my)
798+ yield x rx = foldStream st yld sng stp (goX x rx my)
799+ foldStream st yield single stop mx
800+
801+ goX0 x my = mkStream $ \st yld sng _ -> do
802+ let stop = sng x
803+ single y = do
804+ r <- cmp x y
805+ case r of
806+ GT -> yld y (fromPure x)
807+ _ -> yld x (fromPure y)
808+ yield y ry = do
809+ r <- cmp x y
810+ case r of
811+ GT -> yld y (goX0 x ry)
812+ _ -> yld x (y `cons` ry)
813+ in foldStream st yield single stop my
814+
815+ goX x mx my = mkStream $ \st yld _ _ -> do
816+ let stop = yld x mx
817+ single y = do
818+ r <- cmp x y
819+ case r of
820+ GT -> yld y (x `cons` mx)
821+ _ -> yld x (goY0 mx y)
822+ yield y ry = do
823+ r <- cmp x y
824+ case r of
825+ GT -> yld y (goX x mx ry)
826+ _ -> yld x (goY mx y ry)
827+ in foldStream st yield single stop my
828+
829+ goY0 mx y = mkStream $ \st yld sng _ -> do
830+ let stop = sng y
831+ single x = do
832+ r <- cmp x y
833+ case r of
834+ GT -> yld y (fromPure x)
835+ _ -> yld x (fromPure y)
836+ yield x rx = do
837+ r <- cmp x y
838+ case r of
839+ GT -> yld y (x `cons` rx)
840+ _ -> yld x (goY0 rx y)
841+ in foldStream st yield single stop mx
842+
843+ goY mx y my = mkStream $ \st yld _ _ -> do
844+ let stop = yld y my
845+ single x = do
846+ r <- cmp x y
847+ case r of
848+ GT -> yld y (goX0 x my)
849+ _ -> yld x (y `cons` my)
850+ yield x rx = do
851+ r <- cmp x y
852+ case r of
853+ GT -> yld y (goX x rx my)
854+ _ -> yld x (goY rx y my)
855+ in foldStream st yield single stop mx
856+
857+{-# INLINE mergeBy #-}
858+mergeBy :: (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
859+-- XXX GHC: This has slightly worse performance than replacing "r <- cmp x y"
860+-- with "let r = cmp x y" in the monadic version. The definition below is
861+-- exactly the same as mergeByM except this change.
862+-- mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
863+mergeBy cmp = go
864+
792865 where
866+
793867 go mx my = mkStream $ \st yld sng stp -> do
794− let mergeWithY a ra =
795− let stop2 = foldStream st yld sng stp mx
796− single2 b = do
797− r <- cmp a b
798− case r of
799− GT -> yld b (go (a `cons` ra) nil)
800− _ -> yld a (go ra (b `cons` nil))
801− yield2 b rb = do
802− r <- cmp a b
803− case r of
804− GT -> yld b (go (a `cons` ra) rb)
805− _ -> yld a (go ra (b `cons` rb))
806− in foldStream st yield2 single2 stop2 my
807− let stopX = foldStream st yld sng stp my
808− singleX a = mergeWithY a nil
809− yieldX = mergeWithY
810− foldStream st yieldX singleX stopX mx
811−
812−{-# INLINABLE mergeBy #-}
813−mergeBy :: Monad m =>
814− (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
815−mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
868+ let stop = foldStream st yld sng stp my
869+ single x = foldStream st yld sng stp (goX0 x my)
870+ yield x rx = foldStream st yld sng stp (goX x rx my)
871+ foldStream st yield single stop mx
872+
873+ goX0 x my = mkStream $ \st yld sng _ -> do
874+ let stop = sng x
875+ single y = do
876+ case cmp x y of
877+ GT -> yld y (fromPure x)
878+ _ -> yld x (fromPure y)
879+ yield y ry = do
880+ case cmp x y of
881+ GT -> yld y (goX0 x ry)
882+ _ -> yld x (y `cons` ry)
883+ in foldStream st yield single stop my
884+
885+ goX x mx my = mkStream $ \st yld _ _ -> do
886+ let stop = yld x mx
887+ single y = do
888+ case cmp x y of
889+ GT -> yld y (x `cons` mx)
890+ _ -> yld x (goY0 mx y)
891+ yield y ry = do
892+ case cmp x y of
893+ GT -> yld y (goX x mx ry)
894+ _ -> yld x (goY mx y ry)
895+ in foldStream st yield single stop my
896+
897+ goY0 mx y = mkStream $ \st yld sng _ -> do
898+ let stop = sng y
899+ single x = do
900+ case cmp x y of
901+ GT -> yld y (fromPure x)
902+ _ -> yld x (fromPure y)
903+ yield x rx = do
904+ case cmp x y of
905+ GT -> yld y (x `cons` rx)
906+ _ -> yld x (goY0 rx y)
907+ in foldStream st yield single stop mx
908+
909+ goY mx y my = mkStream $ \st yld _ _ -> do
910+ let stop = yld y my
911+ single x = do
912+ case cmp x y of
913+ GT -> yld y (goX0 x my)
914+ _ -> yld x (y `cons` my)
915+ yield x rx = do
916+ case cmp x y of
917+ GT -> yld y (goX x rx my)
918+ _ -> yld x (goY rx y my)
919+ in foldStream st yield single stop mx
816920
817921------------------------------------------------------------------------------
818922-- Transformation comprehensions