git.iko.soy/streamly
on master help sign in
Optimize StreamK mergeBy
70fbefeHarendra Kumar
src/Streamly/Internal/Data/Stream/StreamK.hs+126−22
compare with master
unifiedsplit
@@ -789,30 +789,134 @@ zipWithM f = go
789mergeByM :: Monad m =>
789mergeByM :: Monad m =>
790 (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
790 (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
791mergeByM cmp = go
791mergeByM 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+
792 where
865 where
866+
793 go mx my = mkStream $ \st yld sng stp -> do
867 go mx my = mkStream $ \st yld sng stp -> do
794− let mergeWithY a ra =
868+ let stop = foldStream st yld sng stp my
869+ single x = foldStream st yld sng stp (goX0 x my)
795− let stop2 = foldStream st yld sng stp mx
870+ yield x rx = foldStream st yld sng stp (goX x rx my)
796− single2 b = do
871+ foldStream st yield single stop mx
797− r <- cmp a b
872+
873+ goX0 x my = mkStream $ \st yld sng _ -> do
874+ let stop = sng x
875+ single y = do
798− case r of
876+ case cmp x y of
799− GT -> yld b (go (a `cons` ra) nil)
877+ GT -> yld y (fromPure x)
800− _ -> yld a (go ra (b `cons` nil))
878+ _ -> yld x (fromPure y)
801− yield2 b rb = do
879+ yield y ry = do
802− r <- cmp a b
880+ case cmp x y of
803− case r of
881+ GT -> yld y (goX0 x ry)
804− GT -> yld b (go (a `cons` ra) rb)
805− _ -> yld a (go ra (b `cons` rb))
882+ _ -> yld x (y `cons` ry)
806− in foldStream st yield2 single2 stop2 my
883+ in foldStream st yield single stop my
807− let stopX = foldStream st yld sng stp my
884+
808− singleX a = mergeWithY a nil
885+ goX x mx my = mkStream $ \st yld _ _ -> do
809− yieldX = mergeWithY
886+ let stop = yld x mx
810− foldStream st yieldX singleX stopX mx
887+ single y = do
811−
888+ case cmp x y of
812−{-# INLINABLE mergeBy #-}
889+ GT -> yld y (x `cons` mx)
813−mergeBy :: Monad m =>
890+ _ -> yld x (goY0 mx y)
814− (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
891+ yield y ry = do
815−mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
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
816
920
817------------------------------------------------------------------------------
921------------------------------------------------------------------------------
818-- Transformation comprehensions
922-- Transformation comprehensions