1 files changed, 17 insertions, 79 deletions

diff --git a/remuxer/src/segment_extractor.rs b/remuxer/src/segment_extractor.rs
index acbec53..095bdfe 100644
--- a/remuxer/src/segment_extractor.rs
+++ b/remuxer/src/segment_extractor.rs
@@ -1,10 +1,6 @@
 use anyhow::{anyhow, Result};
-use jellymatroska::{
-    block::Block,
-    unflatten::{Unflat, Unflatten},
-    MatroskaTag,
-};
-use log::{debug, trace, warn};
+use jellymatroska::{block::Block, read::EbmlReader, unflatten::IterWithPos, MatroskaTag};
+use log::{debug, trace};
 use std::collections::VecDeque;
 
 pub struct AbsoluteBlock {
@@ -13,9 +9,8 @@ pub struct AbsoluteBlock {
 }
 
 pub struct SegmentExtractIter<'a> {
-    segment: Unflatten<'a>,
+    segment: &'a mut EbmlReader,
     extract: u64,
-    emission_queue: VecDeque<AbsoluteBlock>,
 }
 
 impl AbsoluteBlock {
@@ -25,83 +20,26 @@ impl AbsoluteBlock {
 }
 
 impl<'a> SegmentExtractIter<'a> {
-    pub fn new(segment: Unflatten<'a>, extract: u64) -> Self {
-        Self {
-            segment,
-            extract,
-            emission_queue: Default::default(),
-        }
+    pub fn new(segment: &'a mut EbmlReader, extract: u64) -> Self {
+        Self { segment, extract }
     }
 
-    pub fn next(&mut self) -> Result<AbsoluteBlock> {
+    pub fn next(&mut self) -> Result<Block> {
         loop {
-            if let Some(b) = self.emission_queue.pop_front() {
-                break Ok(b);
-            }
-            self.read()?;
-        }
-    }
-
-    pub fn read(&mut self) -> Result<()> {
-        let Unflat { children, item } = self.segment.n().ok_or(anyhow!("eof"))??;
-        let mut pts_base = 0;
-        match item {
-            MatroskaTag::SeekHead(_) => {}
-            MatroskaTag::Info(_) => {}
-            MatroskaTag::Cluster(_) => {
-                let mut children = children.unwrap();
-                while let Some(Ok(Unflat { children, item })) = children.n() {
-                    match item {
-                        MatroskaTag::Crc32(_) => (),
-                        MatroskaTag::Timestamp(ts) => {
-                            trace!("read pts={ts}");
-                            pts_base = ts;
-                        }
-                        MatroskaTag::BlockGroup(_) => {
-                            trace!("group");
-                            let mut children = children.unwrap();
-
-                            // let mut duration = None;
-                            let mut block = None;
-
-                            while let Some(Ok(Unflat { children: _, item })) = children.n() {
-                                match item {
-                                    MatroskaTag::Block(buf) => block = Some(buf),
-                                    // MatroskaTag::BlockDuration(v) => duration = Some(v),
-                                    _ => debug!("ignored {item:?}"),
-                                }
-                            }
-                            // TODO duration
-                            let block = Block::parse(&block.unwrap())?;
-                            if block.track == self.extract {
-                                trace!("block: track={} tso={}", block.track, block.timestamp_off);
-                                self.emission_queue.push_back(AbsoluteBlock {
-                                    pts_base,
-                                    inner: block,
-                                });
-                            }
-                        }
-                        MatroskaTag::SimpleBlock(buf) => {
-                            let block = Block::parse(&buf)?;
-                            if block.track == self.extract {
-                                trace!("block: track={} tso={}", block.track, block.timestamp_off);
-                                self.emission_queue.push_back(AbsoluteBlock {
-                                    pts_base,
-                                    inner: block,
-                                });
-                            }
-                        }
-                        _ => warn!("(rsc) tag ignored: {item:?}"),
+            let item = self.segment.next().ok_or(anyhow!("eof"))??;
+            match item {
+                MatroskaTag::Void(_) => (),
+                MatroskaTag::Crc32(_) => (),
+                MatroskaTag::Cluster(_) => (),
+                MatroskaTag::SimpleBlock(buf) | MatroskaTag::Block(buf) => {
+                    let block = Block::parse(&buf)?;
+                    if block.track == self.extract {
+                        trace!("block: track={} tso={}", block.track, block.timestamp_off);
+                        return Ok(block);
                     }
                 }
+                _ => debug!("(rs) tag ignored: {item:?}"),
             }
-            MatroskaTag::Tags(_) => {}
-            MatroskaTag::Cues(_) => {}
-            MatroskaTag::Chapters(_) => {}
-            MatroskaTag::Tracks(_) => {}
-            MatroskaTag::Void(_) => {}
-            _ => debug!("(rs) tag ignored: {item:?}"),
         }
-        Ok(())
     }
 }