1 files changed, 61 insertions, 30 deletions

diff --git a/remuxer/src/lib.rs b/remuxer/src/lib.rs
index bfcb914..d91bfe5 100644
--- a/remuxer/src/lib.rs
+++ b/remuxer/src/lib.rs
@@ -6,7 +6,7 @@
 pub mod import;
 
 use anyhow::{anyhow, Context, Result};
-use jellycommon::{ItemInfo, SeekIndex, SourceTrack, SourceTrackKind};
+use jellycommon::{BlockIndex, ItemInfo, SeekIndex, SourceTrack, SourceTrackKind};
 use jellymatroska::{
     block::Block,
     read::EbmlReader,
@@ -42,6 +42,7 @@ impl RemuxerContext {
             reader: EbmlReader,
             mapped: u64,
             index: SeekIndex,
+            temp_index: usize,
         }
 
         let mut inputs = selection
@@ -75,6 +76,7 @@ impl RemuxerContext {
                     reader,
                     info,
                     mapped,
+                    temp_index: 0,
                 })
             })
             .into_iter()
@@ -112,6 +114,41 @@ impl RemuxerContext {
             .collect();
         output.write_tag(&MatroskaTag::Tracks(Master::Collected(tracks_header)))?;
 
+        let segment_layout = {
+            let mut pts = 0;
+            let mut cluster_pts = pts + 2_000;
+            let mut clusters = vec![];
+            let mut cluster = vec![];
+            loop {
+                let mut best_block = BlockIndex {
+                    pts: u64::MAX,
+                    size: 0,
+                    source_off: 0,
+                };
+                let mut best_index = 0;
+                for (i, r) in inputs.iter().enumerate() {
+                    if let Some(v) = r.index.blocks.get(r.temp_index) {
+                        if v.pts < best_block.pts {
+                            best_block = v.to_owned();
+                            best_index = i;
+                        }
+                    };
+                }
+                inputs[best_index].temp_index += 1;
+                pts = best_block.pts;
+                if pts > cluster_pts + 2_000 {
+                    clusters.push(std::mem::replace(&mut cluster, vec![]));
+                    cluster_pts = pts;
+                }
+                if pts == u64::MAX {
+                    break;
+                }
+                cluster.push((best_index, best_block))
+            }
+            info!("segment layout computed ({} clusters)", clusters.len());
+            clusters
+        };
+
         struct ReaderD<'a> {
             _info: SourceTrack,
             peek: Option<AbsoluteBlock>,
@@ -144,36 +181,30 @@ impl RemuxerContext {
             });
         }
 
-        loop {
-            let next_index = ks
-                .iter()
-                .enumerate()
-                .fold((0, u64::MAX), |(bi, bpts), (i, r)| {
-                    if let Some(peek) = &r.peek {
-                        let pts = peek.pts();
-                        if pts < bpts {
-                            return (i, pts);
-                        }
-                    }
-                    (bi, bpts)
-                })
-                .0;
+        for (cluster_index, cluster) in segment_layout.into_iter().enumerate() {
+            info!(
+                "writing cluster {cluster_index} with {} blocks",
+                cluster.len()
+            );
+            for (block_index, block) in cluster {
+                let kn = &mut ks[block_index];
+                let mut next = kn.peek.replace(kn.stream.next()?).unwrap(); //.ok_or(anyhow!("eof?"));
 
-            let kn = &mut ks[next_index];
-            let mut next = kn.peek.replace(kn.stream.next()?).unwrap(); //.ok_or(anyhow!("eof?"));
+                assert_eq!(block.size, next.block.data.len(), "the seek index is wrong");
+                assert_eq!(block.pts, next.pts(), "the seek index is wrong");
 
-            let pts = next.pts();
-            next.block.track = kn.mapped;
-            next.block.timestamp_off = 0;
-            let buf = next.block.dump();
-            output.write_tag(&MatroskaTag::Cluster(Master::Collected(vec![
-                MatroskaTag::Timestamp(pts),
-                MatroskaTag::SimpleBlock(buf),
-            ])))?;
+                let pts = next.pts();
+                next.block.track = kn.mapped;
+                next.block.timestamp_off = 0;
+                let buf = next.block.dump();
+                output.write_tag(&MatroskaTag::Cluster(Master::Collected(vec![
+                    MatroskaTag::Timestamp(pts),
+                    MatroskaTag::SimpleBlock(buf),
+                ])))?;
+            }
         }
-
-        // output.write_tag(&MatroskaTag::Segment(Master::End))?;
-        // Ok(())
+        output.write_tag(&MatroskaTag::Segment(Master::End))?;
+        Ok(())
     }
 }
 
@@ -216,11 +247,11 @@ impl SegmentExtractIter<'_> {
                     match item {
                         MatroskaTag::Crc32(_) => (),
                         MatroskaTag::Timestamp(ts) => {
-                            info!("ts={ts}");
+                            trace!("read pts={ts}");
                             pts_base = ts;
                         }
                         MatroskaTag::BlockGroup(_) => {
-                            debug!("group");
+                            trace!("group");
                             let mut children = children.unwrap();
 
                             let mut duration = None;