diff --git a/src/iterator.rs b/src/iterator.rs index d746b3c..ed5c2fd 100644 --- a/src/iterator.rs +++ b/src/iterator.rs @@ -1,7 +1,6 @@ use std::marker::PhantomData; use ebml::FromEbml; -use webm::WebmElement; pub struct EbmlIterator<'a, T: FromEbml<'a>>(&'a [u8], PhantomData T>); @@ -10,17 +9,12 @@ pub fn ebml_iter<'a, T: FromEbml<'a>>(source: &'a [u8])-> EbmlIterator<'a, T> { } impl<'a, T: FromEbml<'a>> Iterator for EbmlIterator<'a, T> { - type Item = WebmElement<'a>; + type Item = T; - fn next(&mut self) -> Option> { - WebmElement::check_space(self.0).unwrap_or(None).and_then(|element_size| { - let (element_data, rest) = self.0.split_at(element_size); - self.0 = rest; - match WebmElement::decode_element(element_data) { - Err(_) => None, - Ok(None) => panic!("Buffer was supposed to have enough data to parse element, somehow did not."), - Ok(Some((element, _))) => Some(element) - } + fn next(&mut self) -> Option { + T::decode_element(self.0).unwrap_or(None).and_then(|(element, element_size)| { + self.0 = &self.0[element_size..]; + Some(element) }) } } diff --git a/src/slice.rs b/src/slice.rs index 8e126b7..32733ce 100644 --- a/src/slice.rs +++ b/src/slice.rs @@ -9,19 +9,9 @@ impl<'b> EbmlEventSource for EbmlSlice<'b> { type Error = EbmlError; fn poll_event<'a, T: FromEbml<'a>>(&'a mut self) -> Result>, EbmlError> { - T::check_space(self.0).and_then(|size_option| { - match size_option { - None => Ok(None), - Some(element_size) => { - let (element_data, rest) = self.0.split_at(element_size); - self.0 = rest; - match T::decode_element(element_data) { - Err(err) => Err(err), - Ok(None) => panic!("Buffer was supposed to have enough data to parse element, somehow did not."), - Ok(Some((element, _))) => Ok(Some(element)) - } - } - } - }).map(Async::Ready) + T::decode_element(self.0).map(|option| option.map(|(element, element_size)| { + self.0 = &self.0[element_size..]; + element + })).map(Async::Ready) } } diff --git a/src/webm.rs b/src/webm.rs index c0f1d71..907e0f2 100644 --- a/src/webm.rs +++ b/src/webm.rs @@ -134,7 +134,7 @@ mod tests { #[test] fn decode_webm_test1() { - let mut iter = parse_webm(TEST_FILE).into_iter(); + let mut iter = parse_webm(TEST_FILE); // test that we match the structure of the test file assert_eq!(iter.next(), Some(WebmElement::EbmlHead));