matroska bitstream filter library

1 files changed, 137 insertions, 0 deletions

diff --git a/framework/src/lib.rs b/framework/src/lib.rs
new file mode 100644
index 0000000..84b6de9
--- /dev/null
+++ b/framework/src/lib.rs
@@ -0,0 +1,137 @@
+use ebml_struct::{
+    Size,
+    ids::{
+        EL_ATTACHMENTS, EL_CHAPTERS, EL_CLUSTER, EL_CRC32, EL_CUES, EL_EBML, EL_INFO, EL_SEEKHEAD,
+        EL_SEGMENT, EL_TAGS, EL_TRACKS, EL_VOID,
+    },
+    matroska::{Cluster, Ebml, Info, Tags, Tracks},
+    read::{EbmlReadExt, TagRead},
+    write::{EbmlWriteExt, TagWrite},
+};
+use std::io::{BufReader, BufWriter, ErrorKind, Read, Result, stdin, stdout};
+
+pub trait BitstreamFilter: Sized {
+    const INPUT_CODEC_ID: &str;
+    const OUTPUT_CODEC_ID: &str;
+    fn new(width: u32, height: u32) -> Self;
+    fn process_block(&mut self, a: Vec<u8>) -> Vec<u8>;
+}
+
+pub fn bitstream_filter_main<F: BitstreamFilter>() -> Result<()> {
+    let mut inp = BufReader::new(stdin());
+    let mut out = BufWriter::new(stdout());
+
+    let (id, inner) = inp.read_tag()?;
+    assert_eq!(id, EL_EBML);
+    inner.consume()?;
+
+    Ebml {
+        ebml_version: 1,
+        ebml_read_version: 1,
+        ebml_max_id_length: 4,
+        ebml_max_size_length: 8,
+        doc_type: "matroska".to_string(),
+        doc_type_version: 4,
+        doc_type_read_version: 2,
+        doc_type_extensions: vec![],
+    }
+    .write(&mut out)?;
+
+    let (id, mut segment) = inp.read_tag()?;
+    assert_eq!(id, EL_SEGMENT);
+
+    out.write_tag_id(EL_SEGMENT)?;
+    out.write_size(Size::Unknown)?;
+
+    let (mut info, mut tracks) = read_info_and_tracks(&mut segment)?;
+    info.muxing_app = "ebml-struct".to_string();
+    info.writing_app = "video-codec-experiments".to_string();
+
+    let mut encode_track = None;
+    for t in &mut tracks.entries {
+        if t.track_type == 1 && encode_track.is_none() {
+            assert_eq!(t.codec_id, F::INPUT_CODEC_ID);
+            t.codec_id = F::OUTPUT_CODEC_ID.to_string();
+            t.codec_private = None;
+            t.codec_download_urls =
+                vec!["https://codeberg.org/metamuffin/video-codec-experiments".to_string()];
+
+            encode_track = Some(t.clone());
+        }
+    }
+    let filter_track = encode_track.expect("no video track found");
+
+    tracks.write(&mut out)?;
+    info.write(&mut out)?;
+
+    let mut filter = F::new(
+        filter_track.video.as_ref().unwrap().pixel_width as u32,
+        filter_track.video.as_ref().unwrap().pixel_height as u32,
+    );
+
+    loop {
+        let (id, mut element) = match segment.read_tag() {
+            Ok(x) => x,
+            Err(e) if e.kind() == ErrorKind::UnexpectedEof => break,
+            Err(e) => return Err(e),
+        };
+        match id {
+            EL_CUES | EL_ATTACHMENTS | EL_SEEKHEAD | EL_CHAPTERS | EL_CRC32 | EL_VOID => {
+                element.consume()?;
+            }
+            EL_TAGS => {
+                Tags::read(&mut element)?.write(&mut out)?;
+            }
+            EL_CLUSTER => {
+                let mut cluster = Cluster::read(&mut element)?;
+
+                for b in &mut cluster.simple_blocks {
+                    if b.track == filter_track.track_number {
+                        assert_eq!(b.lacing, None);
+                        assert_eq!(b.discardable, false);
+                        assert_eq!(b.keyframe, true);
+                        b.data = filter.process_block(b.data.clone());
+                    }
+                }
+
+                cluster.write(&mut out)?;
+            }
+            _ => {
+                eprintln!("unhandled element {id:x}");
+                element.consume()?;
+            }
+        }
+    }
+
+    Ok(())
+}
+
+fn read_info_and_tracks(mut segment: impl Read) -> Result<(Info, Tracks)> {
+    let mut info = None;
+    let mut tracks = None;
+    loop {
+        let (id, mut element) = segment.read_tag()?;
+        match id {
+            EL_CUES | EL_TAGS | EL_ATTACHMENTS | EL_SEEKHEAD | EL_CHAPTERS | EL_CRC32 | EL_VOID => {
+                element.consume()?;
+            }
+            EL_INFO => {
+                info = Some(Info::read(&mut element)?);
+            }
+            EL_TRACKS => {
+                tracks = Some(Tracks::read(&mut element)?);
+            }
+            EL_CLUSTER => {
+                eprintln!("not ready for clusters yet");
+                element.consume()?;
+            }
+            _ => {
+                eprintln!("unhandled element {id:x}");
+                element.consume()?;
+            }
+        };
+        if tracks.is_some() && info.is_some() {
+            break Ok((info.unwrap(), tracks.unwrap()));
+        }
+    }
+}