parsing should work

1 files changed, 81 insertions, 5 deletions

diff --git a/ebml/src/read.rs b/ebml/src/read.rs
index 23b4e21..3e02e18 100644
--- a/ebml/src/read.rs
+++ b/ebml/src/read.rs
@@ -1,13 +1,25 @@
 use anyhow::{bail, Result};
-use std::io::{Read, Seek};
+use log::warn;
+use std::{
+    collections::VecDeque,
+    io::{Read, Seek},
+};
 
 use crate::matroska::MatroskaTag;
 
 trait ReadAndSeek: Read + Seek {}
 impl<T: Read + Seek> ReadAndSeek for T {}
 
+#[derive(Debug, Clone, Copy)]
+pub struct StackTag {
+    end: Option<usize>,
+    id: u64,
+}
+
 pub struct EbmlReader {
     inner: Box<dyn ReadAndSeek>,
+    stack: Vec<StackTag>,
+    queue: VecDeque<MatroskaTag>,
     pub position: usize,
 }
 
@@ -18,7 +30,9 @@ pub trait EbmlRead: Sized {
 impl EbmlReader {
     pub fn new<T: Seek + Read + 'static>(inner: T) -> Self {
         Self {
+            queue: VecDeque::new(),
             inner: Box::new(inner),
+            stack: vec![],
             position: 0,
         }
     }
@@ -64,14 +78,70 @@ impl EbmlReader {
     pub fn read_tag_size(&mut self) -> Result<EbmlSize> {
         Ok(EbmlSize::from_vint(self.read_vint_len()?))
     }
-    pub fn read_tag(&mut self) -> Result<MatroskaTag> {
+    pub fn read_stuff(&mut self) -> Result<()> {
+        if let Some(e) = self.stack.last().map(|e| *e) {
+            if let Some(end) = e.end {
+                if self.position >= end {
+                    if self.position != end {
+                        warn!("we missed the end")
+                    }
+                    self.stack.pop();
+                    self.queue.push_back(MatroskaTag::parse(e.id, &[0])?);
+                }
+            }
+        }
+
         let id = self.read_tag_id()?;
         let size = self.read_tag_size()?;
-        if MatroskaTag::is_master(id)? {
-            Ok(MatroskaTag::parse(id, &[])?)
+        let is_master = MatroskaTag::is_master(id)?;
+        let tag = if is_master {
+            MatroskaTag::parse(id, &[])?
         } else {
             let data = self.read_buf(size)?;
-            Ok(MatroskaTag::parse(id, &data)?)
+            MatroskaTag::parse(id, &data)?
+        };
+
+        if let Some(path) = tag.path() {
+            // we have slightly different rules for closing tags implicitly
+            // this closes as many tags as needed to make the next tag a valid child
+            while let Some(tag @ StackTag { end: None, .. }) = self.stack.last() {
+                let mut valid_child = false;
+                for p in path {
+                    if *p == tag.id {
+                        valid_child = true;
+                    }
+                }
+                if valid_child {
+                    break;
+                } else {
+                    self.queue
+                        .push_back(MatroskaTag::parse(self.stack.pop().unwrap().id, &[0])?);
+                }
+            }
+        }
+
+        if is_master {
+            self.stack.push(StackTag {
+                end: size.some().map(|s| s + self.position),
+                id,
+            });
+        }
+        self.queue.push_back(tag);
+        Ok(())
+    }
+}
+
+impl Iterator for EbmlReader {
+    type Item = Result<MatroskaTag>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if let Some(t) = self.queue.pop_front() {
+            Some(Ok(t))
+        } else {
+            match self.read_stuff() {
+                Ok(()) => self.next(),
+                Err(e) => Some(Err(e)),
+            }
         }
     }
 }
@@ -89,6 +159,12 @@ impl EbmlSize {
             Self::Exact(value as usize)
         }
     }
+    pub fn some(self) -> Option<usize> {
+        match self {
+            EbmlSize::Exact(s) => Some(s),
+            EbmlSize::Unknown => None,
+        }
+    }
 }
 impl Into<usize> for EbmlSize {
     fn into(self) -> usize {