Use another if let in stream parser

This commit is contained in:
Tangent Wantwight 2023-02-03 23:10:50 -05:00
parent 884aa37888
commit 4db64df278
1 changed files with 18 additions and 20 deletions

View File

@ -1,8 +1,11 @@
use bytes::{Buf, BufMut, Bytes, BytesMut}; use bytes::{Buf, BufMut, Bytes, BytesMut};
use futures::{TryStreamExt, stream::{Stream, StreamExt}}; use futures::{
stream::{Stream, StreamExt},
TryStreamExt,
};
use std::task::{Context, Poll}; use std::task::{Context, Poll};
use crate::ebml::FromEbml; use crate::ebml::{EbmlLayout, FromEbml};
use crate::error::WebmetroError; use crate::error::WebmetroError;
pub struct EbmlStreamingParser<S> { pub struct EbmlStreamingParser<S> {
@ -45,20 +48,16 @@ where
cx: &mut Context, cx: &mut Context,
) -> Poll<Option<Result<T, WebmetroError>>> { ) -> Poll<Option<Result<T, WebmetroError>>> {
loop { loop {
match T::check_space(&self.buffer)? { if let Some(info) = T::check_space(&self.buffer)? {
None => { self.borrowed = self.buffer.split_to(info.element_len).freeze();
// need to refill buffer, below self.borrowed.advance(info.body_offset);
} return Poll::Ready(Some(
Some(info) => { T::decode(info.element_id, &self.borrowed).map_err(Into::into),
let mut bytes = self.buffer.split_to(info.element_len).freeze(); ));
bytes.advance(info.body_offset);
self.borrowed = bytes;
return Poll::Ready(Some(
T::decode(info.element_id, &self.borrowed).map_err(Into::into),
));
}
} }
// need to refill buffer
if let Some(limit) = self.buffer_size_limit { if let Some(limit) = self.buffer_size_limit {
if limit <= self.buffer.len() { if limit <= self.buffer.len() {
return Poll::Ready(Some(Err(WebmetroError::ResourcesExceeded))); return Poll::Ready(Some(Err(WebmetroError::ResourcesExceeded)));
@ -80,9 +79,8 @@ where
pub async fn next<'a, T: FromEbml<'a>>(&'a mut self) -> Result<Option<T>, WebmetroError> { pub async fn next<'a, T: FromEbml<'a>>(&'a mut self) -> Result<Option<T>, WebmetroError> {
loop { loop {
if let Some(info) = T::check_space(&self.buffer)? { if let Some(info) = T::check_space(&self.buffer)? {
let mut bytes = self.buffer.split_to(info.element_len).freeze(); self.borrowed = self.buffer.split_to(info.element_len).freeze();
bytes.advance(info.body_offset); self.borrowed.advance(info.body_offset);
self.borrowed = bytes;
return Ok(Some(T::decode(info.element_id, &self.borrowed)?)); return Ok(Some(T::decode(info.element_id, &self.borrowed)?));
} }
@ -193,8 +191,8 @@ mod tests {
Result::<(), WebmetroError>::Ok(()) Result::<(), WebmetroError>::Ok(())
} }
.now_or_never() .now_or_never()
.expect("Test tried to block on I/O") .expect("Test tried to block on I/O")
.expect("Parse failed"); .expect("Parse failed");
} }
} }