From 4a42b7ae820ab320060f4d8b8d2ed148c05f611c Mon Sep 17 00:00:00 2001
From: Tangent 128 <Tangent128@gmail.com>
Date: Tue, 24 Jan 2017 03:20:49 -0500
Subject: [PATCH] implement WebM iterator

---
 src/lib.rs | 54 +++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 53 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index 679d27b..f3a2652 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -94,7 +94,8 @@ pub fn decode_tag(bytes: &[u8]) -> Result<Option<(u64, Varint, usize)>, Error> {
 }
 
 pub trait Schema<'a> {
-    type Element;
+    type Element: 'a;
+
     fn should_unwrap(element_id: u64) -> bool;
     fn decode<'b: 'a>(element_id: u64, bytes: &'b[u8]) -> Result<Self::Element, Error>;
 }
@@ -146,6 +147,35 @@ pub fn decode_element<'a, 'b: 'a, T: Schema<'a>>(bytes: &'b[u8]) -> Result<Optio
     }
 }
 
+pub struct WebmIterator<'b> {
+    slice: &'b[u8],
+    position: usize,
+}
+
+impl<'b> WebmIterator<'b> {
+    pub fn new(bytes: &'b[u8]) -> Self {
+        WebmIterator {
+            slice: bytes,
+            position: 0
+        }
+    }
+}
+
+impl<'b> Iterator for WebmIterator<'b> {
+    type Item = WebmElement<'b>;
+
+    fn next(&mut self) -> Option<WebmElement<'b>> {
+        match decode_element::<Webm>(&self.slice[self.position..]) {
+            Err(_) => None,
+            Ok(None) => None,
+            Ok(Some((element, element_size))) => {
+                self.position += element_size;
+                Some(element)
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
 
@@ -227,4 +257,26 @@ mod tests {
         }
     }
 
+    fn assert_webm_blob(test: Option<WebmElement>, tag: u64, payload_size: usize) {
+        match test {
+            Some(WebmElement::Unknown(element_tag, bytes)) => {
+                assert_eq!(element_tag, tag);
+                assert_eq!(bytes.len(), payload_size);
+            },
+            None => {
+                panic!("Did not parse expected WebM element; result: {:?}", test);
+            }
+        }
+    }
+
+    #[test]
+    fn decode_webm_test1() {
+        let mut iter = WebmIterator::new(TEST_FILE);
+        // EBML Header
+        assert_webm_blob(iter.next(), 0x0A45DFA3, 31);
+
+        // Segment
+        assert_webm_blob(iter.next(), 0x08538067, 56124);
+    }
+
 }