@@ -817,6 +817,11 @@ takeEQ n (Fold fstep finitial fextract) = Parser step initial extract
817817 $ "takeEQ: Expecting exactly " ++ show n
818818 ++ " elements, input terminated on " ++ show (i - 1)
819819
820+{-# ANN type TakeGEState Fuse #-}
821+data TakeGEState s =
822+ TakeGELT !Int !s
823+ | TakeGEGE !s
824+
820825-- | Take at least @n@ input elements, but can collect more.
821826--
822827-- * Stops - when the collecting fold stops.
@@ -837,48 +842,51 @@ takeGE n (Fold fstep finitial fextract) = Parser step initial extract
837842
838843 where
839844
840− cnt = max n 0
841−
842845 initial = do
843846 res <- finitial
844− return $ case res of
845− FL.Partial s -> IPartial $ Tuple' 0 s
846− FL.Done b ->
847− if cnt == 0
848− then IDone b
849− else IError
850− $ "takeGE: Expecting at least " ++ show cnt
847+ case res of
848+ FL.Partial s ->
849+ if n > 0
850+ then return $ IPartial $ TakeGELT 1 s
851+ else return $ IPartial $ TakeGEGE s
852+ FL.Done b -> return $
853+ if n > 0
854+ then IError
855+ $ "takeGE: Expecting at least " ++ show n
851856 ++ " elements, fold terminated without"
852857 ++ " consuming any elements"
858+ else IDone b
853859
854− step (Tuple' i r) a
855− | i1 < cnt = do
856− res <- fstep r a
860+ step (TakeGELT i1 r) a = do
861+ res <- fstep r a
862+ if n > i1
863+ then
857864 return
858865 $ case res of
859− FL.Partial s -> Continue 0 $ Tuple' i1 s
866+ FL.Partial s -> Continue 0 $ TakeGELT (i1 + 1) s
860867 FL.Done _ ->
861868 Error
862− $ "takeGE: Expecting at least " ++ show cnt
869+ $ "takeGE: Expecting at least " ++ show n
863870 ++ " elements, fold terminated on " ++ show i1
864− | otherwise = do
865− res <- fstep r a
871+ else
872+ -- assert (n <= i1)
866873 return
867874 $ case res of
868− FL.Partial s -> Partial 0 $ Tuple' i1 s
875+ FL.Partial s -> Partial 0 $ TakeGEGE s
869876 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
870883
871− where
872−
873− i1 = i + 1
874−
875− extract (Tuple' i r)
876− | i >= cnt = fmap (Done 0) $ fextract r
877− | otherwise =
878− return
879− $ Error
880− $ "takeGE: Expecting at least " ++ show cnt
881− ++ " elements, input terminated on " ++ show i
884+ extract (TakeGELT i _) =
885+ return
886+ $ Error
887+ $ "takeGE: Expecting at least " ++ show n
888+ ++ " elements, input terminated on " ++ show (i - 1)
889+ extract (TakeGEGE r) = fmap (Done 0) $ fextract r
882890
883891-------------------------------------------------------------------------------
884892-- Conditional splitting