@@ -0,0 +1,229 @@
1+-- |
2+-- Module : Streamly.Internal.Data.Unfold.Resume
3+-- Copyright : (c) 2021 Composewell Technologies
4+-- License : BSD-3-Clause
5+-- Maintainer : streamly@composewell.com
6+-- Stability : experimental
7+-- Portability : GHC
8+--
9+-- Unfolds in "Streamly.Internal.Data.Unfold" module run to completion, they
10+-- cannot be stopped and then resumed. In contrast, the unfolds in this module
11+-- allow extraction of the intermediate state which can be used to resume the
12+-- unfold later on.
13+--
14+-- For simple cases the seed itself can be used to represent the intermediate
15+-- state as well. For example, if we are iterating through a list or an array
16+-- we can return the remaining list or the remaining array slice as the
17+-- state and resume from it later.
18+--
19+-- In general, the intermediate state could be different than the seed. Also,
20+-- when unfolds are nested, for example using 'concat', then each nesting layer
21+-- would have its own state.
22+--
23+-- In imperative terms, an unfold is like a for loop (or an iterator), the
24+-- state of the unfold is the state of the loop. The state of the 'concat' of
25+-- unfolds is like the state of each loop in the nesting of loops.
26+--
27+
28+module Streamly.Internal.Data.Unfold.Resume
29+ ( Unfold (..)
30+
31+ -- * Converting
32+ , simplify
33+
34+ -- * Unfolds
35+ , nil
36+ , nilM
37+ , unfoldrM
38+ , fromStreamD
39+ , fromList
40+
41+ -- * Combinators
42+ , concat
43+ , parse
44+ , parseMany
45+ , parseManyD
46+ )
47+where
48+
49+#include "inline.hs"
50+
51+import Control.Exception (assert)
52+import Control.Monad.Catch (MonadThrow, throwM)
53+import GHC.Exts (SpecConstrAnnotation(..))
54+import GHC.Types (SPEC(..))
55+import Streamly.Internal.Data.Parser.ParserD (ParseError(..), Step(..))
56+import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))
57+import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..))
58+import Streamly.Internal.Data.SVar (defState)
59+import Streamly.Internal.Data.Unfold.Source (Source)
60+
61+import qualified Streamly.Internal.Data.Parser.ParserD as ParserD
62+import qualified Streamly.Internal.Data.Parser.ParserK.Types as ParserK
63+import qualified Streamly.Internal.Data.Unfold.Types as Unfold
64+import qualified Streamly.Internal.Data.Unfold.Source as Source
65+
66+import Streamly.Internal.Data.Unfold.Resume.Type
67+import Prelude hiding (concat)
68+
69+-- XXX We should write unfolds as resumable unfolds where possible and define
70+-- the regular unfolds using "simplify".
71+--
72+-------------------------------------------------------------------------------
73+-- Converting to regular unfolds
74+-------------------------------------------------------------------------------
75+
76+-- | Simplify a resumable unfold to a regular unfold.
77+--
78+-- /Internal/
79+{-# INLINE simplify #-}
80+simplify :: Unfold m a b -> Unfold.Unfold m a b
81+simplify (Unfold step inject _) = Unfold.Unfold step inject
82+
83+-------------------------------------------------------------------------------
84+-- Unfolds
85+-------------------------------------------------------------------------------
86+
87+-- | Convert a StreamD stream into an Unfold.
88+--
89+-- /Internal/
90+{-# INLINE_NORMAL fromStreamD #-}
91+fromStreamD :: Monad m => Unfold m (Stream m a) a
92+fromStreamD = Unfold step return (return . Just)
93+ where
94+
95+ {-# INLINE_LATE step #-}
96+ step (UnStream step1 state1) = do
97+ r <- step1 defState state1
98+ return $ case r of
99+ Yield x s -> Yield x (Stream step1 s)
100+ Skip s -> Skip (Stream step1 s)
101+ Stop -> Stop
102+
103+-------------------------------------------------------------------------------
104+-- Parsing
105+-------------------------------------------------------------------------------
106+
107+-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
108+-- to make a newtype.
109+{-# ANN type List NoSpecConstr #-}
110+newtype List a = List {getList :: [a]}
111+
112+{-# INLINE_NORMAL parseD #-}
113+parseD
114+ :: MonadThrow m
115+ => ParserD.Parser m a b
116+ -> Unfold m (Source s a) a
117+ -> Source s a
118+ -> m (b, Maybe (Source s a))
119+parseD
120+ (ParserD.Parser pstep initial extract)
121+ (Unfold ustep uinject uextract)
122+ seed = do
123+
124+ state <- uinject seed
125+ initial >>= go SPEC state (List [])
126+
127+ where
128+
129+ -- XXX currently we are using a dumb list based approach for backtracking
130+ -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
131+ -- That will allow us more efficient random back and forth movement.
132+ {-# INLINE go #-}
133+ go !_ st buf !pst = do
134+ r <- ustep st
135+ case r of
136+ Yield x s -> do
137+ pRes <- pstep pst x
138+ case pRes of
139+ Partial 0 pst1 -> go SPEC s (List []) pst1
140+ Partial n pst1 -> do
141+ assert (n <= length (x:getList buf)) (return ())
142+ let src0 = Prelude.take n (x:getList buf)
143+ src = Prelude.reverse src0
144+ gobuf SPEC s (List []) (List src) pst1
145+ Continue 0 pst1 -> go SPEC s (List (x:getList buf)) pst1
146+ Continue n pst1 -> do
147+ assert (n <= length (x:getList buf)) (return ())
148+ let (src0, buf1) = splitAt n (x:getList buf)
149+ src = Prelude.reverse src0
150+ gobuf SPEC s (List buf1) (List src) pst1
151+ Done n b -> do
152+ assert (n <= length (x:getList buf)) (return ())
153+ let src0 = Prelude.take n (x:getList buf)
154+ src = Prelude.reverse src0
155+ s1 <- uextract s
156+ return (b, fmap (Source.unread src) s1)
157+ Error err -> throwM $ ParseError err
158+ Skip s -> go SPEC s buf pst
159+ Stop -> do
160+ b <- extract pst
161+ -- XXX we should return the remaining buffer
162+ return (b, Nothing)
163+
164+ gobuf !_ s buf (List []) !pst = go SPEC s buf pst
165+ gobuf !_ s buf (List (x:xs)) !pst = do
166+ pRes <- pstep pst x
167+ case pRes of
168+ Partial 0 pst1 ->
169+ gobuf SPEC s (List []) (List xs) pst1
170+ Partial n pst1 -> do
171+ assert (n <= length (x:getList buf)) (return ())
172+ let src0 = Prelude.take n (x:getList buf)
173+ src = Prelude.reverse src0 ++ xs
174+ gobuf SPEC s (List []) (List src) pst1
175+ Continue 0 pst1 ->
176+ gobuf SPEC s (List (x:getList buf)) (List xs) pst1
177+ Continue n pst1 -> do
178+ assert (n <= length (x:getList buf)) (return ())
179+ let (src0, buf1) = splitAt n (x:getList buf)
180+ src = Prelude.reverse src0 ++ xs
181+ gobuf SPEC s (List buf1) (List src) pst1
182+ Done n b -> do
183+ assert (n <= length (x:getList buf)) (return ())
184+ let src0 = Prelude.take n (x:getList buf)
185+ src = Prelude.reverse src0
186+ s1 <- uextract s
187+ return (b, fmap (Source.unread src) s1)
188+ Error err -> throwM $ ParseError err
189+
190+-- | Parse a resumable unfold returning the parsed value and the state of the
191+-- unfold which can be resumed later to unfold the remaining values.
192+--
193+-- /Internal/
194+{-# INLINE [3] parse #-}
195+parse
196+ :: MonadThrow m
197+ => ParserK.Parser m a b
198+ -> Unfold m (Source s a) a
199+ -> Source s a
200+ -> m (b, Maybe (Source s a))
201+parse = parseD . ParserK.fromParserK
202+
203+-------------------------------------------------------------------------------
204+-- Nested parsing
205+-------------------------------------------------------------------------------
206+
207+{-# INLINE parseManyD #-}
208+parseManyD :: MonadThrow m =>
209+ ParserD.Parser m a b -> Unfold m (Source x a) a -> Unfold m (Source x a) b
210+parseManyD parser reader = Unfold step inject return
211+
212+ where
213+
214+ inject = return . Just
215+
216+ {-# INLINE_LATE step #-}
217+ step (Just src) = do
218+ (b, s1) <- parseD parser reader src
219+ return $ Yield b s1
220+ step Nothing = return Stop
221+
222+-- | Apply a parser repeatedly on an unfold to generate an unfold of parsed
223+-- values.
224+--
225+-- /Internal/
226+{-# INLINE parseMany #-}
227+parseMany :: MonadThrow m =>
228+ ParserK.Parser m a b -> Unfold m (Source x a) a -> Unfold m (Source x a) b
229+parseMany parser = parseManyD (ParserK.fromParserK parser)