stream wrapper for remux mkv/webm

1 files changed, 241 insertions, 254 deletions

diff --git a/remuxer/src/lib.rs b/remuxer/src/lib.rs
index 21b7fd6..a9702fb 100644
--- a/remuxer/src/lib.rs
+++ b/remuxer/src/lib.rs
@@ -25,191 +25,180 @@ use std::{
     time::Instant,
 };
 
-#[derive(Debug, Clone)]
-pub struct RemuxerContext {}
+pub fn remux_stream_into(
+    writer: impl Write + 'static,
+    range: Range<usize>,
+    path_base: PathBuf,
+    item: NodePublic,
+    track_sources: Vec<LocalTrack>,
+    selection: Vec<usize>,
+    webm: bool,
+) -> anyhow::Result<()> {
+    info!("remuxing {:?} to have tracks {selection:?}", item.title);
+    let writer = TrimWriter::new(writer, range.clone());
+    let mut output = EbmlWriter::new(writer, 0);
 
-impl RemuxerContext {
-    #[allow(clippy::new_without_default)]
-    pub fn new() -> Self {
-        Self {}
+    struct ReaderC {
+        info: SourceTrack,
+        reader: EbmlReader,
+        mapped: u64,
+        index: SeekIndex,
+        source_track_index: usize,
+        codec_private: Option<Vec<u8>>,
+        layouting_progress_index: usize,
     }
 
-    pub fn generate_into(
-        &self,
-        writer: impl Write + 'static,
-        range: Range<usize>,
-        path_base: PathBuf,
-        item: NodePublic,
-        track_sources: Vec<LocalTrack>,
-        selection: Vec<usize>,
-        webm: bool,
-    ) -> anyhow::Result<()> {
-        info!("remuxing {:?} to have tracks {selection:?}", item.title);
-        let writer = TrimWriter::new(writer, range.clone());
-        let mut output = EbmlWriter::new(writer, 0);
+    let timing_cp = Instant::now();
 
-        struct ReaderC {
-            info: SourceTrack,
-            reader: EbmlReader,
-            mapped: u64,
-            index: SeekIndex,
-            source_track_index: usize,
-            codec_private: Option<Vec<u8>>,
-            layouting_progress_index: usize,
-        }
-
-        let timing_cp = Instant::now();
-
-        let mut inputs = selection
-            .iter()
-            .enumerate()
-            .map(|(index, sel)| {
-                let info = item
-                    .media
-                    .as_ref()
-                    .unwrap()
-                    .tracks
-                    .get(*sel)
-                    .ok_or(anyhow!("track not available"))?
-                    .to_owned();
-                let private = &track_sources[*sel];
-                let source_path = path_base.join(&private.path);
-                let mapped = index as u64 + 1;
-                info!("\t- {sel} {source_path:?} ({} => {mapped})", private.track);
-                info!("\t    {}", info);
-                let file = File::open(&source_path).context("opening source file")?;
-                let mut index =
-                    File::open(source_path.with_extension(format!("si.{}", private.track)))
-                        .context("opening seek index file")?;
-                let index = bincode::decode_from_std_read::<SeekIndex, _, _>(
-                    &mut index,
-                    bincode::config::standard(),
-                )?;
-                debug!("\t    seek index: {} blocks loaded", index.blocks.len());
-                let reader = EbmlReader::new(file);
-                Ok(ReaderC {
-                    index,
-                    reader,
-                    info,
-                    mapped,
-                    source_track_index: private.track,
-                    codec_private: private.codec_private.clone(),
-                    layouting_progress_index: 0,
-                })
+    let mut inputs = selection
+        .iter()
+        .enumerate()
+        .map(|(index, sel)| {
+            let info = item
+                .media
+                .as_ref()
+                .unwrap()
+                .tracks
+                .get(*sel)
+                .ok_or(anyhow!("track not available"))?
+                .to_owned();
+            let private = &track_sources[*sel];
+            let source_path = path_base.join(&private.path);
+            let mapped = index as u64 + 1;
+            info!("\t- {sel} {source_path:?} ({} => {mapped})", private.track);
+            info!("\t    {}", info);
+            let file = File::open(&source_path).context("opening source file")?;
+            let mut index = File::open(source_path.with_extension(format!("si.{}", private.track)))
+                .context("opening seek index file")?;
+            let index = bincode::decode_from_std_read::<SeekIndex, _, _>(
+                &mut index,
+                bincode::config::standard(),
+            )?;
+            debug!("\t    seek index: {} blocks loaded", index.blocks.len());
+            let reader = EbmlReader::new(file);
+            Ok(ReaderC {
+                index,
+                reader,
+                info,
+                mapped,
+                source_track_index: private.track,
+                codec_private: private.codec_private.clone(),
+                layouting_progress_index: 0,
             })
-            .collect::<anyhow::Result<Vec<_>>>()?;
+        })
+        .collect::<anyhow::Result<Vec<_>>>()?;
 
-        info!("(perf) prepare inputs: {:?}", Instant::now() - timing_cp);
-        let timing_cp = Instant::now();
+    info!("(perf) prepare inputs: {:?}", Instant::now() - timing_cp);
+    let timing_cp = Instant::now();
 
-        output.write_tag(&MatroskaTag::Ebml(Master::Collected(vec![
-            MatroskaTag::EbmlVersion(1),
-            MatroskaTag::EbmlReadVersion(1),
-            MatroskaTag::EbmlMaxIdLength(4),
-            MatroskaTag::EbmlMaxSizeLength(8),
-            MatroskaTag::DocType(if webm {
-                "webm".to_string()
-            } else {
-                "matroska".to_string()
-            }),
-            MatroskaTag::DocTypeVersion(4),
-            MatroskaTag::DocTypeReadVersion(2),
-        ])))?;
+    output.write_tag(&MatroskaTag::Ebml(Master::Collected(vec![
+        MatroskaTag::EbmlVersion(1),
+        MatroskaTag::EbmlReadVersion(1),
+        MatroskaTag::EbmlMaxIdLength(4),
+        MatroskaTag::EbmlMaxSizeLength(8),
+        MatroskaTag::DocType(if webm {
+            "webm".to_string()
+        } else {
+            "matroska".to_string()
+        }),
+        MatroskaTag::DocTypeVersion(4),
+        MatroskaTag::DocTypeReadVersion(2),
+    ])))?;
 
-        output.write_tag(&MatroskaTag::Segment(Master::Start))?;
-        let segment_offset = output.position();
+    output.write_tag(&MatroskaTag::Segment(Master::Start))?;
+    let segment_offset = output.position();
 
-        output.write_tag(&MatroskaTag::Info(Master::Collected(vec![
-            MatroskaTag::TimestampScale(1_000_000),
-            MatroskaTag::Duration(item.media.unwrap().duration * 1000.0),
-            MatroskaTag::Title(item.title.clone()),
-            MatroskaTag::MuxingApp("jellyremux".to_string()),
-            MatroskaTag::WritingApp("jellything".to_string()),
-        ])))?;
-        output.write_tag(&MatroskaTag::Tags(Master::Collected(vec![])))?;
+    output.write_tag(&MatroskaTag::Info(Master::Collected(vec![
+        MatroskaTag::TimestampScale(1_000_000),
+        MatroskaTag::Duration(item.media.unwrap().duration * 1000.0),
+        MatroskaTag::Title(item.title.clone()),
+        MatroskaTag::MuxingApp("jellyremux".to_string()),
+        MatroskaTag::WritingApp("jellything".to_string()),
+    ])))?;
+    output.write_tag(&MatroskaTag::Tags(Master::Collected(vec![])))?;
 
-        let tracks_header = inputs
-            .iter_mut()
-            .map(|rc| track_to_ebml(rc.mapped, &rc.info, rc.codec_private.take()))
-            .collect();
-        output.write_tag(&MatroskaTag::Tracks(Master::Collected(tracks_header)))?;
+    let tracks_header = inputs
+        .iter_mut()
+        .map(|rc| track_to_ebml(rc.mapped, &rc.info, rc.codec_private.take()))
+        .collect();
+    output.write_tag(&MatroskaTag::Tracks(Master::Collected(tracks_header)))?;
 
-        struct ClusterLayout {
-            position: usize,
-            timestamp: u64,
-            source_offsets: Vec<Option<usize>>,
-            blocks: Vec<(usize, BlockIndex)>,
-        }
+    struct ClusterLayout {
+        position: usize,
+        timestamp: u64,
+        source_offsets: Vec<Option<usize>>,
+        blocks: Vec<(usize, BlockIndex)>,
+    }
 
-        let mut segment_layout: Vec<ClusterLayout> = {
-            let mut cluster_pts = 0;
-            let mut clusters = vec![];
-            let mut cluster = vec![];
-            let mut source_offsets = vec![None; inputs.len()];
-            let mut gp = 0usize; // cluster position (in the segment)
-            let mut p = 0usize; // block position (in the cluster)
-            loop {
-                let (track, block) = {
-                    let mut best_block = BlockIndex {
-                        pts: u64::MAX,
-                        size: 0,
-                        source_off: 0,
-                    };
-                    let mut best_track = 0;
-                    for (i, r) in inputs.iter().enumerate() {
-                        if let Some(v) = r.index.blocks.get(r.layouting_progress_index) {
-                            if v.pts < best_block.pts {
-                                best_block = v.to_owned();
-                                best_track = i;
-                            }
-                        };
-                    }
-                    (best_track, best_block)
+    let mut segment_layout: Vec<ClusterLayout> = {
+        let mut cluster_pts = 0;
+        let mut clusters = vec![];
+        let mut cluster = vec![];
+        let mut source_offsets = vec![None; inputs.len()];
+        let mut gp = 0usize; // cluster position (in the segment)
+        let mut p = 0usize; // block position (in the cluster)
+        loop {
+            let (track, block) = {
+                let mut best_block = BlockIndex {
+                    pts: u64::MAX,
+                    size: 0,
+                    source_off: 0,
                 };
-                inputs[track].layouting_progress_index += 1;
-                source_offsets[track].get_or_insert(block.source_off);
-                if block.pts > cluster_pts + 1_000 {
-                    let cluster_content_size = 1 + 1 // timestamp {tag, size}
+                let mut best_track = 0;
+                for (i, r) in inputs.iter().enumerate() {
+                    if let Some(v) = r.index.blocks.get(r.layouting_progress_index) {
+                        if v.pts < best_block.pts {
+                            best_block = v.to_owned();
+                            best_track = i;
+                        }
+                    };
+                }
+                (best_track, best_block)
+            };
+            inputs[track].layouting_progress_index += 1;
+            source_offsets[track].get_or_insert(block.source_off);
+            if block.pts > cluster_pts + 1_000 {
+                let cluster_content_size = 1 + 1 // timestamp {tag, size}
                      + bad_vint_length(cluster_pts) // timestamp tag value
                      + p;
-                    let cluster_size = 4 // tag length
+                let cluster_size = 4 // tag length
                       + vint_length(cluster_content_size as u64) // size varint
                       + cluster_content_size;
-                    clusters.push(ClusterLayout {
-                        position: gp, // relative to the first cluster
-                        timestamp: cluster_pts,
-                        source_offsets,
-                        blocks: std::mem::take(&mut cluster),
-                    });
+                clusters.push(ClusterLayout {
+                    position: gp, // relative to the first cluster
+                    timestamp: cluster_pts,
+                    source_offsets,
+                    blocks: std::mem::take(&mut cluster),
+                });
 
-                    cluster_pts = block.pts;
-                    source_offsets = vec![None; inputs.len()];
-                    gp += cluster_size;
-                    p = 0;
-                }
-                if block.pts == u64::MAX {
-                    break;
-                }
+                cluster_pts = block.pts;
+                source_offsets = vec![None; inputs.len()];
+                gp += cluster_size;
+                p = 0;
+            }
+            if block.pts == u64::MAX {
+                break;
+            }
 
-                let simpleblock_size = 1 + 2 + 1 // block {tracknum, pts_off, flags}
+            let simpleblock_size = 1 + 2 + 1 // block {tracknum, pts_off, flags}
                                             // TODO does not work, if more than 127 tracks are present
                                             + block.size; // block payload
-                p += 1; // simpleblock tag
-                p += vint_length(simpleblock_size as u64); // simpleblock size vint
-                p += simpleblock_size;
+            p += 1; // simpleblock tag
+            p += vint_length(simpleblock_size as u64); // simpleblock size vint
+            p += simpleblock_size;
 
-                cluster.push((track, block))
-            }
-            info!("segment layout computed ({} clusters)", clusters.len());
-            clusters
-        };
-        info!(
-            "(perf) compute segment layout: {:?}",
-            Instant::now() - timing_cp
-        );
-        let timing_cp = Instant::now();
+            cluster.push((track, block))
+        }
+        info!("segment layout computed ({} clusters)", clusters.len());
+        clusters
+    };
+    info!(
+        "(perf) compute segment layout: {:?}",
+        Instant::now() - timing_cp
+    );
+    let timing_cp = Instant::now();
 
-        let max_cue_size = 4 // cues id
+    let max_cue_size = 4 // cues id
             + 8 // cues len
             + ( // cues content
                 1 // cp id 
@@ -236,115 +225,113 @@ impl RemuxerContext {
             + 1 // void id
             + 8; // void len
 
-        let first_cluster_offset_predict = max_cue_size + output.position();
+    let first_cluster_offset_predict = max_cue_size + output.position();
 
-        // make the cluster position relative to the segment start as they should
-        segment_layout
-            .iter_mut()
-            .for_each(|e| e.position += first_cluster_offset_predict - segment_offset);
+    // make the cluster position relative to the segment start as they should
+    segment_layout
+        .iter_mut()
+        .for_each(|e| e.position += first_cluster_offset_predict - segment_offset);
 
-        output.write_tag(&MatroskaTag::Cues(Master::Collected(
-            segment_layout
-                .iter()
-                .map(|cluster| {
-                    MatroskaTag::CuePoint(Master::Collected(
-                        Some(MatroskaTag::CueTime(cluster.timestamp))
-                            .into_iter()
-                            // TODO: Subtitles should not have cues for every cluster
-                            .chain(inputs.iter().map(|i| {
-                                MatroskaTag::CueTrackPositions(Master::Collected(vec![
-                                    MatroskaTag::CueTrack(i.mapped),
-                                    MatroskaTag::CueClusterPosition(cluster.position as u64),
-                                ]))
-                            }))
-                            .collect(),
-                    ))
-                })
-                .collect(),
-        )))?;
-        output.write_padding(first_cluster_offset_predict)?;
-        let first_cluster_offset = output.position();
-        assert_eq!(first_cluster_offset, first_cluster_offset_predict);
+    output.write_tag(&MatroskaTag::Cues(Master::Collected(
+        segment_layout
+            .iter()
+            .map(|cluster| {
+                MatroskaTag::CuePoint(Master::Collected(
+                    Some(MatroskaTag::CueTime(cluster.timestamp))
+                        .into_iter()
+                        // TODO: Subtitles should not have cues for every cluster
+                        .chain(inputs.iter().map(|i| {
+                            MatroskaTag::CueTrackPositions(Master::Collected(vec![
+                                MatroskaTag::CueTrack(i.mapped),
+                                MatroskaTag::CueClusterPosition(cluster.position as u64),
+                            ]))
+                        }))
+                        .collect(),
+                ))
+            })
+            .collect(),
+    )))?;
+    output.write_padding(first_cluster_offset_predict)?;
+    let first_cluster_offset = output.position();
+    assert_eq!(first_cluster_offset, first_cluster_offset_predict);
 
-        let mut skip = 0;
-        for (i, cluster) in segment_layout.iter().enumerate() {
-            if (cluster.position + segment_offset) >= range.start {
-                break;
-            }
-            skip = i;
+    let mut skip = 0;
+    for (i, cluster) in segment_layout.iter().enumerate() {
+        if (cluster.position + segment_offset) >= range.start {
+            break;
         }
+        skip = i;
+    }
 
-        if skip != 0 {
-            info!("skipping {skip} clusters");
-            output.seek(SeekFrom::Start(
-                (segment_layout[skip].position + segment_offset) as u64,
-            ))?;
-        }
+    if skip != 0 {
+        info!("skipping {skip} clusters");
+        output.seek(SeekFrom::Start(
+            (segment_layout[skip].position + segment_offset) as u64,
+        ))?;
+    }
 
-        struct ReaderD<'a> {
-            peek: Option<Block>,
-            stream: SegmentExtractIter<'a>,
-            mapped: u64,
-        }
+    struct ReaderD<'a> {
+        peek: Option<Block>,
+        stream: SegmentExtractIter<'a>,
+        mapped: u64,
+    }
 
-        let mut track_readers = inputs
-            .iter_mut()
-            .enumerate()
-            .map(|(i, inp)| {
-                inp.reader
-                    .seek(
-                        segment_layout[skip].source_offsets[i].unwrap(), // TODO will crash if there is a "hole"
-                        MatroskaTag::Cluster(Master::Start),
-                    )
-                    .context("seeking in input")?;
-                let mut stream =
-                    SegmentExtractIter::new(&mut inp.reader, inp.source_track_index as u64);
+    let mut track_readers = inputs
+        .iter_mut()
+        .enumerate()
+        .map(|(i, inp)| {
+            inp.reader
+                .seek(
+                    segment_layout[skip].source_offsets[i].unwrap(), // TODO will crash if there is a "hole"
+                    MatroskaTag::Cluster(Master::Start),
+                )
+                .context("seeking in input")?;
+            let mut stream =
+                SegmentExtractIter::new(&mut inp.reader, inp.source_track_index as u64);
 
-                Ok(ReaderD {
-                    mapped: inp.mapped,
-                    peek: Some(stream.next()?),
-                    stream,
-                })
+            Ok(ReaderD {
+                mapped: inp.mapped,
+                peek: Some(stream.next()?),
+                stream,
             })
-            .collect::<anyhow::Result<Vec<_>>>()?;
+        })
+        .collect::<anyhow::Result<Vec<_>>>()?;
 
-        info!("(perf) seek inputs: {:?}", Instant::now() - timing_cp);
+    info!("(perf) seek inputs: {:?}", Instant::now() - timing_cp);
 
-        for (cluster_index, cluster) in segment_layout.into_iter().enumerate().skip(skip) {
-            debug!(
-                "writing cluster {cluster_index} (pts_base={}) with {} blocks",
-                cluster.timestamp,
-                cluster.blocks.len()
-            );
-            {
-                let cue_error =
-                    cluster.position as i64 - (output.position() - segment_offset) as i64;
-                if cue_error != 0 {
-                    warn!("calculation was {} bytes off", cue_error);
-                }
+    for (cluster_index, cluster) in segment_layout.into_iter().enumerate().skip(skip) {
+        debug!(
+            "writing cluster {cluster_index} (pts_base={}) with {} blocks",
+            cluster.timestamp,
+            cluster.blocks.len()
+        );
+        {
+            let cue_error = cluster.position as i64 - (output.position() - segment_offset) as i64;
+            if cue_error != 0 {
+                warn!("calculation was {} bytes off", cue_error);
             }
+        }
 
-            let mut cluster_blocks = vec![MatroskaTag::Timestamp(cluster.timestamp)];
-            for (block_track, index_block) in cluster.blocks {
-                let track_reader = &mut track_readers[block_track];
-                let mut block = track_reader
-                    .peek
-                    .replace(track_reader.stream.next()?)
-                    .expect("source file too short");
+        let mut cluster_blocks = vec![MatroskaTag::Timestamp(cluster.timestamp)];
+        for (block_track, index_block) in cluster.blocks {
+            let track_reader = &mut track_readers[block_track];
+            let mut block = track_reader
+                .peek
+                .replace(track_reader.stream.next()?)
+                .expect("source file too short");
 
-                assert_eq!(index_block.size, block.data.len(), "seek index is wrong");
+            assert_eq!(index_block.size, block.data.len(), "seek index is wrong");
 
-                block.track = track_reader.mapped;
-                block.timestamp_off = (index_block.pts - cluster.timestamp).try_into().unwrap();
-                trace!("n={} tso={}", block.track, block.timestamp_off);
-                let buf = block.dump();
-                cluster_blocks.push(MatroskaTag::SimpleBlock(buf))
-            }
-            output.write_tag(&MatroskaTag::Cluster(Master::Collected(cluster_blocks)))?;
+            block.track = track_reader.mapped;
+            block.timestamp_off = (index_block.pts - cluster.timestamp).try_into().unwrap();
+            trace!("n={} tso={}", block.track, block.timestamp_off);
+            let buf = block.dump();
+            cluster_blocks.push(MatroskaTag::SimpleBlock(buf))
         }
-        output.write_tag(&MatroskaTag::Segment(Master::End))?;
-        Ok(())
+        output.write_tag(&MatroskaTag::Cluster(Master::Collected(cluster_blocks)))?;
     }
+    output.write_tag(&MatroskaTag::Segment(Master::End))?;
+    Ok(())
 }
 
 pub fn track_to_ebml(