git.iko.soy/streamly
on master help sign in
Add takeP combinator
commitc375be2944100f5575c04201a4a75c2797727d8c authorAdithya Kumar <adithya@composewell.com> dateSun 1 Aug 2021, 16:42 (5 years ago) parentd9abfa1
2 files changed+68
compare with master
unifiedsplit
src/Streamly/Internal/Data/Parser.hs+17 src/Streamly/Internal/Data/Parser/ParserD.hs+51
src/Streamly/Internal/Data/Parser.hs+17
@@ -78,6 +78,7 @@ module Streamly.Internal.Data.Parser
7878 , takeGE -- takeBetween n maxBound
7979
8080 -- Grab a sequence of input elements by inspecting them
81+ , takeP
8182 , lookAhead
8283 , takeWhileP
8384 , takeWhile
@@ -852,6 +853,22 @@ longest p1 p2 = K.toParserK $ D.longest (K.fromParserK p1) (K.fromParserK p2)
852853lookAhead :: MonadCatch m => Parser m a b -> Parser m a b
853854lookAhead p = K.toParserK $ D.lookAhead $ K.fromParserK p
854855
856+-- | Takes at-most @n@ input elements.
857+--
858+-- * Stops - when the collecting parser stops.
859+-- * Fails - when the collecting parser fails.
860+--
861+-- >>> Stream.parse (Parser.takeP 4 (Parser.takeEQ 2 Fold.toList)) $ Stream.fromList [1, 2, 3, 4, 5]
862+-- [1, 2]
863+--
864+-- >>> Stream.parse (Parser.takeP 4 (Parser.takeEQ 5 Fold.toList)) $ Stream.fromList [1, 2, 3, 4, 5]
865+-- *** Exception: ParseError "takeEQ: Expecting exactly 5 elements, input terminated on 4"
866+--
867+-- /Internal/
868+{-# INLINE takeP #-}
869+takeP :: MonadCatch m => Int -> Parser m a b -> Parser m a b
870+takeP i p = K.toParserK $ D.takeP i $ K.fromParserK p
871+
855872-------------------------------------------------------------------------------
856873-- Interleaving
857874-------------------------------------------------------------------------------
src/Streamly/Internal/Data/Parser/ParserD.hs+51
@@ -52,6 +52,7 @@ module Streamly.Internal.Data.Parser.ParserD
5252 , takeGE -- takeBetween n maxBound
5353
5454 -- Grab a sequence of input elements by inspecting them
55+ , takeP
5556 , lookAhead
5657 , takeWhile
5758 , takeWhile1
@@ -814,6 +815,56 @@ spanByRolling eq f1 f2 =
814815-- nested parsers
815816-------------------------------------------------------------------------------
816817
818+-- | See 'Streamly.Internal.Data.Parser.takeP'.
819+--
820+-- /Internal/
821+{-# INLINE takeP #-}
822+takeP :: Monad m => Int -> Parser m a b -> Parser m a b
823+takeP lim (Parser pstep pinitial pextract) = Parser step initial extract
824+
825+ where
826+
827+ initial = do
828+ res <- pinitial
829+ case res of
830+ IPartial s ->
831+ if lim > 0
832+ then return $ IPartial $ Tuple' 0 s
833+ else IDone <$> pextract s
834+ IDone b -> return $ IDone b
835+ IError e -> return $ IError e
836+
837+ step (Tuple' cnt r) a = do
838+ assert (cnt < lim) (return ())
839+ res <- pstep r a
840+ let cnt1 = cnt + 1
841+ case res of
842+ Partial 0 s -> do
843+ assert (cnt1 >= 0) (return ())
844+ if cnt1 < lim
845+ then return $ Partial 0 $ Tuple' cnt1 s
846+ else Done 0 <$> pextract s
847+ Continue 0 s -> do
848+ assert (cnt1 >= 0) (return ())
849+ if cnt1 < lim
850+ then return $ Continue 0 $ Tuple' cnt1 s
851+ -- XXX This should error out?
852+ -- If designed properly, this will probably error out.
853+ -- "pextract" should error out
854+ else Done 0 <$> pextract s
855+ Partial n s -> do
856+ let taken = cnt1 - n
857+ assert (taken >= 0) (return ())
858+ return $ Partial n $ Tuple' taken s
859+ Continue n s -> do
860+ let taken = cnt1 - n
861+ assert (taken >= 0) (return ())
862+ return $ Continue n $ Tuple' taken s
863+ Done n b -> return $ Done n b
864+ Error str -> return $ Error str
865+
866+ extract (Tuple' _ r) = pextract r
867+
817868-- | See 'Streamly.Internal.Data.Parser.lookahead'.
818869--
819870-- /Pre-release/