1 files changed, 29 insertions, 33 deletions

diff --git a/matroska/src/read.rs b/matroska/src/read.rs
index 3e858f7..68a8efa 100644
--- a/matroska/src/read.rs
+++ b/matroska/src/read.rs
@@ -5,7 +5,7 @@
 */
 use crate::error::Error;
 use crate::Result;
-use crate::{matroska::MatroskaTag, size::EbmlSize, unflatten::IterWithPos, Master};
+use crate::{matroska::MatroskaTag, size::EbmlSize, Master};
 use log::{debug, warn};
 use std::{
     collections::VecDeque,
@@ -17,21 +17,21 @@ impl<T: Read + Seek> ReadAndSeek for T {}
 
 #[derive(Debug, Clone, Copy)]
 pub struct StackTag {
-    end: Option<usize>,
+    end: Option<u64>,
     id: u64,
 }
 
 pub struct EbmlReader {
     inner: Box<dyn ReadAndSeek>,
     stack: Vec<StackTag>,
-    queue: VecDeque<MatroskaTag>,
-    pub position: usize,
+    queue: VecDeque<(Option<u64>, MatroskaTag)>,
+    position: u64,
 }
 
 impl Read for EbmlReader {
     fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
         let r = self.inner.read(buf)?;
-        self.position += r;
+        self.position += r as u64;
         Ok(r)
     }
 }
@@ -53,15 +53,15 @@ impl EbmlReader {
         self.position += 1;
         Ok(b[0])
     }
-    
+
     pub fn read_buf(&mut self, size: impl Into<usize>) -> Result<Vec<u8>> {
         let size = size.into();
         let mut b = vec![0u8; size];
         self.inner.read_exact(&mut b).map_err(Error::Io)?;
-        self.position += size;
+        self.position += size as u64;
         Ok(b)
     }
-    
+
     pub fn read_vint_len(&mut self) -> Result<(u64, usize)> {
         let s = self.read_byte()?;
         let len = s.leading_zeros() + 1;
@@ -76,24 +76,24 @@ impl EbmlReader {
         }
         Ok((value, len as usize))
     }
-    
+
     #[inline]
     pub fn read_vint(&mut self) -> Result<u64> {
         Ok(self.read_vint_len()?.0)
     }
-    
+
     #[inline]
     pub fn read_utf8(&mut self, size: impl Into<usize>) -> Result<String> {
         let b = self.read_buf(size)?;
         Ok(String::from_utf8(b).map_err(|_| Error::InvalidUTF8)?)
     }
-    
+
     #[inline]
     pub fn read_tag_id(&mut self) -> Result<u64> {
         let (value, len) = self.read_vint_len()?;
         Ok(value + (1 << (7 * len)))
     }
-    
+
     #[inline]
     pub fn read_tag_size(&mut self) -> Result<EbmlSize> {
         Ok(EbmlSize::from_vint(self.read_vint_len()?))
@@ -109,7 +109,7 @@ impl EbmlReader {
                     }
                     self.stack.pop();
                     self.queue
-                        .push_back(MatroskaTag::construct_master(e.id, Master::End)?);
+                        .push_back((None, MatroskaTag::construct_master(e.id, Master::End)?));
                 } else {
                     break;
                 }
@@ -118,6 +118,7 @@ impl EbmlReader {
             }
         }
 
+        let start_position = self.position;
         let id = self.read_tag_id()?;
         let size = self.read_tag_size()?;
         let is_master = MatroskaTag::is_master(id)?;
@@ -131,30 +132,29 @@ impl EbmlReader {
         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() {
-                if path.last() == Some(&tag.id) {
+            while let Some(stag @ StackTag { end: None, .. }) = self.stack.last() {
+                if path.last() == Some(&stag.id) {
                     break;
                 } else {
-                    self.queue.push_back(MatroskaTag::construct_master(
-                        self.stack.pop().unwrap().id,
-                        Master::End,
-                    )?);
+                    let end =
+                        MatroskaTag::construct_master(self.stack.pop().unwrap().id, Master::End)?;
+                    self.queue.push_back((None, end));
                 }
             }
         }
 
         if is_master {
             self.stack.push(StackTag {
-                end: size.some().map(|s| s + self.position),
+                end: size.some().map(|s| s as u64 + self.position),
                 id,
             });
         }
-        self.queue.push_back(tag);
+        self.queue.push_back((Some(start_position), tag));
         Ok(())
     }
 
     /// context should be the next expected tag, such that the stack can be derived from its path.
-    pub fn seek(&mut self, position: usize, context: MatroskaTag) -> Result<()> {
+    pub fn seek(&mut self, position: u64, context: MatroskaTag) -> Result<()> {
         let path = context.path().ok_or(Error::GlobalTagsAsContext)?;
         debug!(
             "seeking to {position} with a context restored from path {:x?}",
@@ -162,9 +162,7 @@ impl EbmlReader {
         );
         self.queue.clear();
         self.position = position;
-        self.inner
-            .seek(SeekFrom::Start(position as u64))
-            .map_err(Error::Io)?;
+        self.inner.seek(SeekFrom::Start(position as u64))?;
         self.stack = path
             .iter()
             .map(|id| StackTag { id: *id, end: None })
@@ -173,13 +171,8 @@ impl EbmlReader {
     }
 }
 
-impl IterWithPos for EbmlReader {
-    type Item = Result<MatroskaTag>;
-
-    fn position(&self) -> usize {
-        self.position
-    }
-
+impl Iterator for EbmlReader {
+    type Item = Result<(Option<u64>, MatroskaTag)>;
     fn next(&mut self) -> Option<Self::Item> {
         if let Some(t) = self.queue.pop_front() {
             // match t {
@@ -198,7 +191,10 @@ impl IterWithPos for EbmlReader {
                         match self.queue.pop_front() {
                             Some(q) => Some(Ok(q)),
                             None => match self.stack.pop() {
-                                Some(q) => Some(MatroskaTag::construct_master(q.id, Master::End)),
+                                Some(q) => Some(Ok((
+                                    None,
+                                    MatroskaTag::construct_master(q.id, Master::End).unwrap(),
+                                ))),
                                 None => Some(Err(e)),
                             },
                         }