4d78dd5 · streamly · git.iko.soy

git.iko.soy/streamly

Add generic chunked version of continuation parsing

commit4d78dd581fb589413065081e2fa22a2b130765d6 authorAdithya Kumar <adithya@composewell.com> dateSun 6 Aug 2023, 19:25 (3 years ago) parentdc7a543

3 files changed+307−8

compare with master

unifiedsplit

core/src/Streamly/Internal/Data/MutArray/Generic.hs+15−6 core/src/Streamly/Internal/Data/Parser/ParserK/Type.hs+142 core/src/Streamly/Internal/Data/StreamK.hs+150−2

core/src/Streamly/Internal/Data/MutArray/Generic.hs+15−6

@@ -101,6 +101,7 @@ module Streamly.Internal.Data.MutArray.Generic

101101 -- ** Random reads

102102 , getIndex

103103 , getIndexUnsafe

104+ , getIndexUnsafeWith

104105 -- , getIndices

105106 -- , getFromThenTo

106107 -- , getIndexRev

@@ -466,17 +467,25 @@ uninit arr@MutArray{..} len =

466467-- Random reads

467468-------------------------------------------------------------------------------

468469

469−-- | Return the element at the specified index without checking the bounds.

470+-- | Return the element at the specified index without checking the bounds from

471+-- a @MutableArray# RealWorld@.

470472--

471473-- Unsafe because it does not check the bounds of the array.

472−{-# INLINE_NORMAL getIndexUnsafe #-}

473−getIndexUnsafe :: MonadIO m => Int -> MutArray a -> m a

474−getIndexUnsafe n MutArray {..} =

474+{-# INLINE getIndexUnsafeWith #-}

475+getIndexUnsafeWith :: MonadIO m => MutableArray# RealWorld a -> Int -> m a

476+getIndexUnsafeWith _arrContents# n =

475477 liftIO

476478 $ IO

477479 $ \s# ->

478− let !(I# i#) = arrStart + n

479− in readArray# arrContents# i# s#

480+ let !(I# i#) = n

481+ in readArray# _arrContents# i# s#

482+

483+-- | Return the element at the specified index without checking the bounds.

484+--

485+-- Unsafe because it does not check the bounds of the array.

486+{-# INLINE_NORMAL getIndexUnsafe #-}

487+getIndexUnsafe :: MonadIO m => Int -> MutArray a -> m a

488+getIndexUnsafe n MutArray {..} = getIndexUnsafeWith arrContents# (n + arrStart)

480489

481490-- | /O(1)/ Lookup the element at the given index. Index starts from 0.

482491--

core/src/Streamly/Internal/Data/Parser/ParserK/Type.hs+142

@@ -29,6 +29,7 @@ module Streamly.Internal.Data.Parser.ParserK.Type

2929 , ParserK (..)

3030 , fromParser

3131 , fromParserSingular

32+ , fromParserGeneric

3233 -- , toParser

3334 , fromPure

3435 , fromEffect

@@ -52,6 +53,10 @@ import Streamly.Internal.System.IO (unsafeInlineIO)

5253

5354import qualified Control.Monad.Fail as Fail

5455import qualified Streamly.Internal.Data.Array.Type as Array

56+import qualified Streamly.Internal.Data.MutArray.Generic as GenArr

57+ ( getIndexUnsafeWith

58+ )

59+import qualified Streamly.Internal.Data.Array.Generic as GenArr

5560import qualified Streamly.Internal.Data.Parser.ParserD.Type as ParserD

5661

5762-- Note: We cannot use an Array directly as input because we need to identify

@@ -371,6 +376,10 @@ instance MonadTrans (ParserK a) where

371376-- Convert ParserD to ParserK

372377-------------------------------------------------------------------------------

373378

379+--------------------------------------------------------------------------------

380+-- Chunked

381+--------------------------------------------------------------------------------

382+

374383{-# INLINE parseDToK #-}

375384parseDToK

376385 :: forall m a s b r. (Monad m, Unbox a)

@@ -504,6 +513,10 @@ fromParser :: (Monad m, Unbox a) => ParserD.Parser a m b -> ParserK (Array a) m

504513fromParser (ParserD.Parser step initial extract) =

505514 MkParser $ parseDToK step initial extract

506515

516+--------------------------------------------------------------------------------

517+-- Singular

518+--------------------------------------------------------------------------------

519+

507520{-# INLINE parseDToKSingular #-}

508521parseDToKSingular

509522 :: forall m a s b r. (Monad m)

@@ -600,6 +613,135 @@ fromParserSingular :: Monad m => ParserD.Parser a m b -> ParserK a m b

600613fromParserSingular (ParserD.Parser step initial extract) =

601614 MkParser $ parseDToKSingular step initial extract

602615

616+--------------------------------------------------------------------------------

617+-- Chunked Generic

618+--------------------------------------------------------------------------------

619+

620+{-# INLINE parseDToKGeneric #-}

621+parseDToKGeneric

622+ :: forall m a s b r. (Monad m)

623+ => (s -> a -> m (ParserD.Step s b))

624+ -> m (ParserD.Initial s b)

625+ -> (s -> m (ParserD.Step s b))

626+ -> (ParseResult b -> Int -> Input (GenArr.Array a) -> m (Step (GenArr.Array a) m r))

627+ -> Int

628+ -> Int

629+ -> Input (GenArr.Array a)

630+ -> m (Step (GenArr.Array a) m r)

631+parseDToKGeneric pstep initial extract cont !offset0 !usedCount !input = do

632+ res <- initial

633+ case res of

634+ ParserD.IPartial pst -> do

635+ case input of

636+ Chunk arr -> parseContChunk usedCount offset0 pst arr

637+ None -> parseContNothing usedCount pst

638+ ParserD.IDone b -> cont (Success offset0 b) usedCount input

639+ ParserD.IError err -> cont (Failure offset0 err) usedCount input

640+

641+ where

642+

643+ {-# NOINLINE parseContChunk #-}

644+ parseContChunk !count !offset !state arr@(GenArr.Array contents start len) = do

645+ if offset >= 0

646+ then go SPEC (start + offset) state

647+ else return $ Continue offset (parseCont count state)

648+

649+ where

650+

651+ {-# INLINE end #-}

652+ end = start + len

653+

654+ {-# INLINE onDone #-}

655+ onDone n b =

656+ assert (n <= GenArr.length arr)

657+ (cont (Success n b) (count + n - offset) (Chunk arr))

658+

659+ {-# INLINE callParseCont #-}

660+ callParseCont constr n pst1 =

661+ assert (n < 0 || n >= GenArr.length arr)

662+ (return $ constr n (parseCont (count + n - offset) pst1))

663+

664+ {-# INLINE onPartial #-}

665+ onPartial = callParseCont Partial

666+

667+ {-# INLINE onContinue #-}

668+ onContinue = callParseCont Continue

669+

670+ {-# INLINE onError #-}

671+ onError n err =

672+ cont (Failure n err) (count + n - offset) (Chunk arr)

673+

674+ {-# INLINE onBack #-}

675+ onBack offset1 constr pst = do

676+ let pos = offset1 - start

677+ in if pos >= 0

678+ then go SPEC offset1 pst

679+ else constr pos pst

680+

681+ go !_ !cur !pst | cur >= end =

682+ onContinue len pst

683+ go !_ !cur !pst = do

684+ let !x = unsafeInlineIO $ GenArr.getIndexUnsafeWith contents cur

685+ pRes <- pstep pst x

686+ let next = cur + 1

687+ back n = next - n

688+ curOff = cur - start

689+ nextOff = next - start

690+ -- The "n" here is stream position index wrt the array start, and

691+ -- not the backtrack count as returned by byte stream parsers.

692+ case pRes of

693+ ParserD.Done 0 b ->

694+ onDone nextOff b

695+ ParserD.Done 1 b ->

696+ onDone curOff b

697+ ParserD.Done n b ->

698+ onDone (back n - start) b

699+ ParserD.Partial 0 pst1 ->

700+ go SPEC next pst1

701+ ParserD.Partial 1 pst1 ->

702+ go SPEC cur pst1

703+ ParserD.Partial n pst1 ->

704+ onBack (back n) onPartial pst1

705+ ParserD.Continue 0 pst1 ->

706+ go SPEC next pst1

707+ ParserD.Continue 1 pst1 ->

708+ go SPEC cur pst1

709+ ParserD.Continue n pst1 ->

710+ onBack (back n) onContinue pst1

711+ ParserD.Error err ->

712+ onError curOff err

713+

714+ {-# NOINLINE parseContNothing #-}

715+ parseContNothing !count !pst = do

716+ r <- extract pst

717+ case r of

718+ -- IMPORTANT: the n here is from the byte stream parser, that means

719+ -- it is the backtrack element count and not the stream position

720+ -- index into the current input array.

721+ ParserD.Done n b ->

722+ assert (n >= 0)

723+ (cont (Success (- n) b) (count - n) None)

724+ ParserD.Continue n pst1 ->

725+ assert (n >= 0)

726+ (return $ Continue (- n) (parseCont (count - n) pst1))

727+ ParserD.Error err ->

728+ -- XXX It is called only when there is no input arr. So using 0

729+ -- as the position is correct?

730+ cont (Failure 0 err) count None

731+ ParserD.Partial _ _ -> error "Bug: parseDToKGeneric Partial unreachable"

732+

733+ {-# INLINE parseCont #-}

734+ parseCont !cnt !pst (Chunk arr) = parseContChunk cnt 0 pst arr

735+ parseCont !cnt !pst None = parseContNothing cnt pst

736+

737+-- | Similar to "fromParser" but is not constrained.

738+--

739+{-# INLINE_LATE fromParserGeneric #-}

740+fromParserGeneric ::

741+ Monad m => ParserD.Parser a m b -> ParserK (GenArr.Array a) m b

742+fromParserGeneric (ParserD.Parser step initial extract) =

743+ MkParser $ parseDToKGeneric step initial extract

744+

603745{-

604746-------------------------------------------------------------------------------

605747-- Convert CPS style 'Parser' to direct style 'D.Parser'

core/src/Streamly/Internal/Data/StreamK.hs+150−2

@@ -49,6 +49,8 @@ module Streamly.Internal.Data.StreamK

4949 , parseChunks

5050 , parseBreakSingular

5151 , parseSingular

52+ , parseBreakGenericChunks

53+ , parseGenericChunks

5254

5355 -- ** Specialized Folds

5456 , head

@@ -143,6 +145,7 @@ import Streamly.Internal.Data.Unbox (sizeOf, Unbox)

143145

144146import qualified Control.Monad.Catch as MC

145147import qualified Streamly.Internal.Data.Array.Type as Array

148+import qualified Streamly.Internal.Data.Array.Generic as GenArr

146149import qualified Streamly.Internal.Data.Fold.Type as FL

147150import qualified Streamly.Internal.Data.Parser as Parser

148151import qualified Streamly.Internal.Data.Parser.ParserD.Type as PR

@@ -1284,7 +1287,7 @@ parseD :: Monad m =>

12841287parseD f = fmap fst . parseDBreak f

12851288

12861289-------------------------------------------------------------------------------

1287−-- Chunked parsing using ParserK

1290+-- ParserK Chunked

12881291-------------------------------------------------------------------------------

12891292

12901293-- The backracking buffer consists of arrays in the most-recent-first order. We

@@ -1427,7 +1430,7 @@ parseChunks :: (Monad m, Unbox a) =>

14271430parseChunks f = fmap fst . parseBreakChunks f

14281431

14291432-------------------------------------------------------------------------------

1430−-- Parsing using ParserK.Generic

1433+-- ParserK Singular

14311434-------------------------------------------------------------------------------

14321435

14331436{-# INLINE backTrackSingular #-}

@@ -1555,6 +1558,151 @@ parseSingular :: Monad m =>

15551558 ParserK.ParserK a m b -> StreamK m a -> m (Either ParseError b)

15561559parseSingular f = fmap fst . parseBreakSingular f

15571560

1561+-------------------------------------------------------------------------------

1562+-- ParserK Chunked Generic

1563+-------------------------------------------------------------------------------

1564+

1565+{-# INLINE backTrackGenericChunks #-}

1566+backTrackGenericChunks ::

1567+ Int

1568+ -> [GenArr.Array a]

1569+ -> StreamK m (GenArr.Array a)

1570+ -> (StreamK m (GenArr.Array a), [GenArr.Array a])

1571+backTrackGenericChunks = go

1572+

1573+ where

1574+

1575+ go _ [] stream = (stream, [])

1576+ go n xs stream | n <= 0 = (stream, xs)

1577+ go n (x:xs) stream =

1578+ let len = GenArr.length x

1579+ in if n > len

1580+ then go (n - len) xs (cons x stream)

1581+ else if n == len

1582+ then (cons x stream, xs)

1583+ else let arr1 = GenArr.getSliceUnsafe (len - n) n x

1584+ arr2 = GenArr.getSliceUnsafe 0 (len - n) x

1585+ in (cons arr1 stream, arr2:xs)

1586+

1587+-- | Similar to 'parseBreak' but works on generic arrays

1588+--

1589+{-# INLINE_NORMAL parseBreakGenericChunks #-}

1590+parseBreakGenericChunks

1591+ :: forall m a b. Monad m

1592+ => ParserK.ParserK (GenArr.Array a) m b

1593+ -> StreamK m (GenArr.Array a)

1594+ -> m (Either ParseError b, StreamK m (GenArr.Array a))

1595+parseBreakGenericChunks parser input = do

1596+ let parserk = ParserK.runParser parser parserDone 0 0

1597+ in go [] parserk input

1598+

1599+ where

1600+

1601+ {-# INLINE goStop #-}

1602+ goStop

1603+ :: [GenArr.Array a]

1604+ -> (ParserK.Input (GenArr.Array a)

1605+ -> m (ParserK.Step (GenArr.Array a) m b))

1606+ -> m (Either ParseError b, StreamK m (GenArr.Array a))

1607+ goStop backBuf parserk = do

1608+ pRes <- parserk ParserK.None

1609+ case pRes of

1610+ -- If we stop in an alternative, it will try calling the next

1611+ -- parser, the next parser may call initial returning Partial and

1612+ -- then immediately we have to call extract on it.

1613+ ParserK.Partial 0 cont1 ->

1614+ go [] cont1 nil

1615+ ParserK.Partial n cont1 -> do

1616+ let n1 = negate n

1617+ assertM(n1 >= 0 && n1 <= sum (Prelude.map GenArr.length backBuf))

1618+ let (s1, backBuf1) = backTrackGenericChunks n1 backBuf nil

1619+ in go backBuf1 cont1 s1

1620+ ParserK.Continue 0 cont1 ->

1621+ go backBuf cont1 nil

1622+ ParserK.Continue n cont1 -> do

1623+ let n1 = negate n

1624+ assertM(n1 >= 0 && n1 <= sum (Prelude.map GenArr.length backBuf))

1625+ let (s1, backBuf1) = backTrackGenericChunks n1 backBuf nil

1626+ in go backBuf1 cont1 s1

1627+ ParserK.Done 0 b ->

1628+ return (Right b, nil)

1629+ ParserK.Done n b -> do

1630+ let n1 = negate n

1631+ assertM(n1 >= 0 && n1 <= sum (Prelude.map GenArr.length backBuf))

1632+ let (s1, _) = backTrackGenericChunks n1 backBuf nil

1633+ in return (Right b, s1)

1634+ ParserK.Error _ err -> return (Left (ParseError err), nil)

1635+

1636+ seekErr n len =

1637+ error $ "parseBreak: Partial: forward seek not implemented n = "

1638+ ++ show n ++ " len = " ++ show len

1639+

1640+ yieldk

1641+ :: [GenArr.Array a]

1642+ -> (ParserK.Input (GenArr.Array a)

1643+ -> m (ParserK.Step (GenArr.Array a) m b))

1644+ -> GenArr.Array a

1645+ -> StreamK m (GenArr.Array a)

1646+ -> m (Either ParseError b, StreamK m (GenArr.Array a))

1647+ yieldk backBuf parserk arr stream = do

1648+ pRes <- parserk (ParserK.Chunk arr)

1649+ let len = GenArr.length arr

1650+ case pRes of

1651+ ParserK.Partial n cont1 ->

1652+ case compare n len of

1653+ EQ -> go [] cont1 stream

1654+ LT -> do

1655+ if n >= 0

1656+ then yieldk [] cont1 arr stream

1657+ else do

1658+ let n1 = negate n

1659+ bufLen = sum (Prelude.map GenArr.length backBuf)

1660+ s = cons arr stream

1661+ assertM(n1 >= 0 && n1 <= bufLen)

1662+ let (s1, _) = backTrackGenericChunks n1 backBuf s

1663+ go [] cont1 s1

1664+ GT -> seekErr n len

1665+ ParserK.Continue n cont1 ->

1666+ case compare n len of

1667+ EQ -> go (arr:backBuf) cont1 stream

1668+ LT -> do

1669+ if n >= 0

1670+ then yieldk backBuf cont1 arr stream

1671+ else do

1672+ let n1 = negate n

1673+ bufLen = sum (Prelude.map GenArr.length backBuf)

1674+ s = cons arr stream

1675+ assertM(n1 >= 0 && n1 <= bufLen)

1676+ let (s1, backBuf1) = backTrackGenericChunks n1 backBuf s

1677+ go backBuf1 cont1 s1

1678+ GT -> seekErr n len

1679+ ParserK.Done n b -> do

1680+ let n1 = len - n

1681+ assertM(n1 <= sum (Prelude.map GenArr.length (arr:backBuf)))

1682+ let (s1, _) = backTrackGenericChunks n1 (arr:backBuf) stream

1683+ in return (Right b, s1)

1684+ ParserK.Error _ err -> return (Left (ParseError err), nil)

1685+

1686+ go

1687+ :: [GenArr.Array a]

1688+ -> (ParserK.Input (GenArr.Array a)

1689+ -> m (ParserK.Step (GenArr.Array a) m b))

1690+ -> StreamK m (GenArr.Array a)

1691+ -> m (Either ParseError b, StreamK m (GenArr.Array a))

1692+ go backBuf parserk stream = do

1693+ let stop = goStop backBuf parserk

1694+ single a = yieldk backBuf parserk a nil

1695+ in foldStream

1696+ defState (yieldk backBuf parserk) single stop stream

1697+

1698+{-# INLINE parseGenericChunks #-}

1699+parseGenericChunks ::

1700+ (Monad m)

1701+ => ParserK.ParserK (GenArr.Array a) m b

1702+ -> StreamK m (GenArr.Array a)

1703+ -> m (Either ParseError b)

1704+parseGenericChunks f = fmap fst . parseBreakGenericChunks f

1705+

15581706-------------------------------------------------------------------------------

15591707-- Sorting

15601708-------------------------------------------------------------------------------