@@ -10,11 +10,12 @@ module Network.GRPC.Util.Parser (
1010 , consumeExactly
1111 , getExactly
1212 -- * Execution
13+ , IsFinal
14+ , Leftover
15+ , ProcessResult (.. )
1316 , processAll
14- , processAllIO
1517 ) where
1618
17- import Control.Exception
1819import Control.Monad
1920import Data.Bifunctor
2021import Data.Binary (Get )
@@ -158,33 +159,74 @@ getExactly len get =
158159 Execution
159160-------------------------------------------------------------------------------}
160161
162+ type IsFinal = Bool
163+ type Leftover = Lazy. ByteString
164+
165+ data ProcessResult e b =
166+ -- | Parse error during processing
167+ ProcessError e
168+
169+ -- | Parsing succeeded (compare to 'ProcessedWithoutFinal')
170+ | ProcessedWithFinal b Leftover
171+
172+ -- | Parsing succeeded, but we did not recognize the final message on time
173+ --
174+ -- There are two ways that parsing can terminate: the final few chunks may
175+ -- look like this:
176+ --
177+ -- > chunk1 -- not marked final
178+ -- > chunk2 -- not marked final
179+ -- > chunk3 -- marked final
180+ --
181+ -- or like this:
182+ --
183+ -- > chunk1 -- not marked final
184+ -- > chunk2 -- not marked final
185+ -- > chunk3 -- not marked final
186+ -- > empty chunk -- marked final
187+ --
188+ -- In the former case, we know that we are processing the final message /as/
189+ -- we are processing it ('ProcessedFinal'); in the latter case, we realize
190+ -- this only after we receive the final empty chunk.
191+ | ProcessedWithoutFinal Leftover
192+
161193-- | Process all incoming data
162194--
163195-- Returns any unprocessed data.
164- processAll :: forall m e a .
196+ -- Also returns if we knew that the final result
197+ -- was in fact the final result when we received it (this may or may not be the
198+ -- case, depending on
199+ processAll :: forall m e a b .
165200 Monad m
166- => m Strict. ByteString-- ^ Get next chunk (empty indicates end of input)
167- -> (a -> m () ) -- ^ Process single value
168- -> Parser e a -- ^ Parser
169- -> m (Either e Lazy. ByteString
170- processAll getChunk processOne parser =
201+ => m (Strict. ByteStringIsFinal ) -- ^ Get next chunk
202+ -> (a -> m () ) -- ^ Process value
203+ -> (a -> m b ) -- ^ Process final value
204+ -> Parser e a -- ^ Parser
205+ -> m (ProcessResult e b )
206+ processAll getChunk processOne processFinal parser =
171207 go $ runParser parser nil
172208 where
173- go :: Result e a -> m (Either e Lazy. ByteString
174- go (Failed err) = return $ Left
175- go (Done result bs') = processOne result >> go (runParser parser bs')
176- go (NeedData parser' acc) = do
177- bs <- getChunk
178- if not (BS.Strict. null bs)
179- then go $ runParser parser' (snoc acc bs)
180- else return $ Right
181-
182- -- | Wrapper around 'processAll' that throws errors as exceptions
183- processAllIO :: forall e a .
184- Exception e
185- => IO Strict. ByteString
186- -> (a -> IO () )
187- -> Parser e a
188- -> IO Lazy. ByteString
189- processAllIO getChunk processOne parser =
190- processAll getChunk processOne parser >>= either throwIO return
209+ go :: Result e a -> m (ProcessResult e b )
210+ go (Failed err) = return $ ProcessError err
211+ go (Done a left) = processOne a >> go (runParser parser left)
212+ go (NeedData parser' left) = do
213+ (bs, isFinal) <- getChunk
214+ if not isFinal
215+ then go $ runParser parser' (left `snoc` bs)
216+ else goFinal [] $ runParser parser' (left `snoc` bs)
217+
218+ -- We have received the final chunk; extract all messages until we are done
219+ goFinal :: [a ] -> Result e a -> m (ProcessResult e b )
220+ goFinal _ (Failed err) = return $ ProcessError err
221+ goFinal acc (Done a left) = goFinal (a: acc) $ runParser parser left
222+ goFinal acc (NeedData _ left) = do
223+ mb <- processLastFew (reverse acc)
224+ return $ case mb of
225+ Just b -> ProcessedWithFinal b $ toLazy left
226+ Nothing -> ProcessedWithoutFinal $ toLazy left
227+
228+ processLastFew :: [a ] -> m (Maybe b )
229+ processLastFew [] = return Nothing
230+ processLastFew [a] = Just <$> processFinal a
231+ processLastFew (a: as) = processOne a >> processLastFew as
232+
0 commit comments