git.iko.soy/streamly
on master help sign in
Fix the performance of takeGE
3daa138Harendra Kumar
benchmark/Streamly/Benchmark/Data/Parser.hs+5
compare with master
unifiedsplit
@@ -118,6 +118,10 @@ takeBetween value = Stream.parse (PR.takeBetween 0 value Fold.drain)
118takeEQ :: Monad m => Int -> Stream m a -> m (Either ParseError ())
118takeEQ :: Monad m => Int -> Stream m a -> m (Either ParseError ())
119takeEQ value = Stream.parse (PR.takeEQ value Fold.drain)
119takeEQ value = Stream.parse (PR.takeEQ value Fold.drain)
120
120
121+{-# INLINE takeGE #-}
122+takeGE :: Monad m => Int -> Stream m a -> m (Either ParseError ())
123+takeGE value = Stream.parse (PR.takeGE value Fold.drain)
124+
121{-# INLINE dropWhile #-}
125{-# INLINE dropWhile #-}
122dropWhile :: Monad m => Int -> Stream m Int -> m (Either ParseError ())
126dropWhile :: Monad m => Int -> Stream m Int -> m (Either ParseError ())
123dropWhile value = Stream.parse (PR.dropWhile (<= value))
127dropWhile value = Stream.parse (PR.dropWhile (<= value))
@@ -441,6 +445,7 @@ o_1_space_serial :: Int -> [Benchmark]
441o_1_space_serial value =
445o_1_space_serial value =
442 [ benchIOSink value "takeBetween" $ takeBetween value
446 [ benchIOSink value "takeBetween" $ takeBetween value
443 , benchIOSink value "takeEQ" $ takeEQ value
447 , benchIOSink value "takeEQ" $ takeEQ value
448+ , benchIOSink value "takeGE" $ takeGE value
444 , benchIOSink value "takeWhile" $ takeWhile value
449 , benchIOSink value "takeWhile" $ takeWhile value
445 , benchIOSink value "takeWhileP" $ takeWhileP value
450 , benchIOSink value "takeWhileP" $ takeWhileP value
446 , benchIOSink value "takeP" $ takeP value
451 , benchIOSink value "takeP" $ takeP value
core/src/Streamly/Internal/Data/Parser/ParserD.hs+36−28
@@ -817,6 +817,11 @@ takeEQ n (Fold fstep finitial fextract) = Parser step initial extract
817 $ "takeEQ: Expecting exactly " ++ show n
817 $ "takeEQ: Expecting exactly " ++ show n
818 ++ " elements, input terminated on " ++ show (i - 1)
818 ++ " elements, input terminated on " ++ show (i - 1)
819
819
820+{-# ANN type TakeGEState Fuse #-}
821+data TakeGEState s =
822+ TakeGELT !Int !s
823+ | TakeGEGE !s
824+
820-- | Take at least @n@ input elements, but can collect more.
825-- | Take at least @n@ input elements, but can collect more.
821--
826--
822-- * Stops - when the collecting fold stops.
827-- * Stops - when the collecting fold stops.
@@ -837,48 +842,51 @@ takeGE n (Fold fstep finitial fextract) = Parser step initial extract
837
842
838 where
843 where
839
844
840− cnt = max n 0
841−
842 initial = do
845 initial = do
843 res <- finitial
846 res <- finitial
844− return $ case res of
847+ case res of
845− FL.Partial s -> IPartial $ Tuple' 0 s
848+ FL.Partial s ->
849+ if n > 0
850+ then return $ IPartial $ TakeGELT 1 s
851+ else return $ IPartial $ TakeGEGE s
846− FL.Done b ->
852+ FL.Done b -> return $
847− if cnt == 0
853+ if n > 0
848− then IDone b
849− else IError
854+ then IError
850− $ "takeGE: Expecting at least " ++ show cnt
855+ $ "takeGE: Expecting at least " ++ show n
851 ++ " elements, fold terminated without"
856 ++ " elements, fold terminated without"
852 ++ " consuming any elements"
857 ++ " consuming any elements"
858+ else IDone b
853
859
854− step (Tuple' i r) a
860+ step (TakeGELT i1 r) a = do
855− | i1 < cnt = do
856− res <- fstep r a
861+ res <- fstep r a
862+ if n > i1
863+ then
857 return
864 return
858 $ case res of
865 $ case res of
859− FL.Partial s -> Continue 0 $ Tuple' i1 s
866+ FL.Partial s -> Continue 0 $ TakeGELT (i1 + 1) s
860 FL.Done _ ->
867 FL.Done _ ->
861 Error
868 Error
862− $ "takeGE: Expecting at least " ++ show cnt
869+ $ "takeGE: Expecting at least " ++ show n
863 ++ " elements, fold terminated on " ++ show i1
870 ++ " elements, fold terminated on " ++ show i1
864− | otherwise = do
871+ else
865− res <- fstep r a
872+ -- assert (n <= i1)
866 return
873 return
867 $ case res of
874 $ case res of
868− FL.Partial s -> Partial 0 $ Tuple' i1 s
875+ FL.Partial s -> Partial 0 $ TakeGEGE s
869 FL.Done b -> Done 0 b
876 FL.Done b -> Done 0 b
877+ step (TakeGEGE r) a = do
878+ res <- fstep r a
879+ return
880+ $ case res of
881+ FL.Partial s -> Partial 0 $ TakeGEGE s
882+ FL.Done b -> Done 0 b
870
883
871− where
872−
873− i1 = i + 1
874−
875− extract (Tuple' i r)
884+ extract (TakeGELT i _) =
876− | i >= cnt = fmap (Done 0) $ fextract r
877− | otherwise =
878− return
885+ return
879− $ Error
886+ $ Error
880− $ "takeGE: Expecting at least " ++ show cnt
887+ $ "takeGE: Expecting at least " ++ show n
881− ++ " elements, input terminated on " ++ show i
888+ ++ " elements, input terminated on " ++ show (i - 1)
889+ extract (TakeGEGE r) = fmap (Done 0) $ fextract r
882
890
883-------------------------------------------------------------------------------
891-------------------------------------------------------------------------------
884-- Conditional splitting
892-- Conditional splitting