validate matroska kf flag with bitstream parser

6 files changed, 211 insertions, 9 deletions

diff --git a/remuxer/src/bin/analyze_kf_placement.rs b/remuxer/src/bin/analyze_kf_placement.rs
index 519d3b9..0a2fae9 100644
--- a/remuxer/src/bin/analyze_kf_placement.rs
+++ b/remuxer/src/bin/analyze_kf_placement.rs
@@ -5,7 +5,10 @@
 */
 
 use anyhow::{Result, anyhow};
-use jellyremuxer::demuxers::{Demuxer, DemuxerNew, matroska::MatroskaDemuxer};
+use jellyremuxer::{
+    demuxers::{Demuxer, DemuxerNew, matroska::MatroskaDemuxer},
+    kf_detect::get_is_keyframe_fn,
+};
 use std::{env::args, fs::File};
 use winter_matroska::TrackType;
 
@@ -16,34 +19,54 @@ fn main() -> Result<()> {
     let mut reader = MatroskaDemuxer::new(Box::new(file));
 
     let tracks = reader.tracks()?.unwrap();
-    let video_track = tracks
+    let (video_track, codec) = tracks
         .entries
         .iter()
         .find(|t| matches!(t.track_type, TrackType::Video))
-        .map(|t| t.track_number)
+        .map(|t| (t.track_number, &t.codec_id))
         .unwrap();
 
+    let is_keyframe = get_is_keyframe_fn(codec)?;
+
     reader.seek_cluster(None)?;
-    let mut num_kf_first = 0;
-    let mut num_kf_later = 0;
+    let mut num_kf_flag_first = 0;
+    let mut num_kf_bs_first = 0;
+    let mut num_kf_flag_later = 0;
+    let mut num_kf_bs_later = 0;
+    let mut num_kf_flag_missing = 0;
     while let Some((_, cluster)) = reader.read_cluster()? {
         let mut first = true;
+        let mut has_kf = false;
         for block in cluster.simple_blocks {
             if block.track != video_track {
                 continue;
             }
+            if is_keyframe(&block).unwrap_or(false) {
+                if first {
+                    num_kf_bs_first += 1
+                } else {
+                    num_kf_bs_later += 1;
+                }
+            }
             if block.flags.keyframe() {
+                has_kf = true;
                 if first {
-                    num_kf_first += 1
+                    num_kf_flag_first += 1
                 } else {
-                    num_kf_later += 1;
+                    num_kf_flag_later += 1;
                 }
             }
             first = false
         }
+        if !has_kf {
+            num_kf_flag_missing += 1;
+        }
     }
 
-    println!("{num_kf_first:>4} kf first");
-    println!("{num_kf_later:>4} kf later");
+    println!("{num_kf_flag_missing:>4} kf missing (flag)");
+    println!("{num_kf_flag_first:>4} kf first (flag)");
+    println!("{num_kf_bs_first:>4} kf first (bitstream)");
+    println!("{num_kf_flag_later:>4} kf later (flag)");
+    println!("{num_kf_bs_later:>4} kf later (bitstream)");
     Ok(())
 }
diff --git a/remuxer/src/codec_param/vp9.rs b/remuxer/src/codec_param/vp9.rs
index 59c0331..f310350 100644
--- a/remuxer/src/codec_param/vp9.rs
+++ b/remuxer/src/codec_param/vp9.rs
@@ -1,3 +1,9 @@
+/*
+    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
+    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
+    Copyright (C) 2026 metamuffin <metamuffin.org>
+*/
+
 use winter_matroska::TrackEntry;
 
 enum VPLevel {
diff --git a/remuxer/src/kf_detect/av1.rs b/remuxer/src/kf_detect/av1.rs
new file mode 100644
index 0000000..62bf6ea
--- /dev/null
+++ b/remuxer/src/kf_detect/av1.rs
@@ -0,0 +1,123 @@
+/*
+    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
+    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
+    Copyright (C) 2026 metamuffin <metamuffin.org>
+*/
+
+use log::debug;
+
+#[allow(unused)]
+mod obu_type {
+    pub const SEQUENCE_HEADER: u8 = 1;
+    pub const TEMPORAL_DELIMITER: u8 = 2;
+    pub const FRAME_HEADER: u8 = 3;
+    pub const TILE_GROUP: u8 = 4;
+    pub const METADATA: u8 = 5;
+    pub const FRAME: u8 = 6;
+    pub const REDUNDANT_FRAME_HEADER: u8 = 7;
+    pub const TILE_LIST: u8 = 8;
+    pub const PADDING: u8 = 15;
+}
+#[allow(unused)]
+mod frame_type {
+    pub const KEY_FRAME: u8 = 0;
+    pub const INTER_FRAME: u8 = 1;
+    pub const INTRA_ONLY_FRAME: u8 = 2;
+    pub const SWITCH_FRAME: u8 = 3;
+}
+
+struct ParserState {
+    seen_frame_header: bool,
+}
+
+pub fn is_keyframe(mut s: &[u8]) -> Option<bool> {
+    let mut ps = ParserState {
+        seen_frame_header: false,
+    };
+    while !s.is_empty() {
+        if let Some(kf) = read_obu(&mut s, &mut ps) {
+            return Some(kf);
+        }
+    }
+    None
+}
+
+fn read_obu(s: &mut &[u8], ps: &mut ParserState) -> Option<bool> {
+    let obu_forbidden_bit = (s[0] >> 7) != 0;
+    let obu_type = (s[0] >> 3) & 0b1111;
+    let obu_extension_flag = (s[0] >> 2) & 0b1 != 0;
+    let obu_has_size_field = (s[0] >> 1) & 0b1 != 0;
+    let _obu_reserved_1bit = (s[0] >> 0) & 0b1 != 0;
+    assert!(!obu_forbidden_bit);
+    if obu_extension_flag {
+        todo!()
+    }
+    *s = &s[1..];
+    let obu_size = if obu_has_size_field {
+        read_leb128(s)
+    } else {
+        todo!()
+    };
+    debug!("obu type {obu_type} (size={obu_size})");
+    if obu_type != obu_type::SEQUENCE_HEADER
+        && obu_type != obu_type::TEMPORAL_DELIMITER
+        && obu_extension_flag
+    {
+        todo!()
+    }
+    let obu_payload_start = *s;
+    match obu_type {
+        obu_type::FRAME_HEADER => {
+            assert!(!ps.seen_frame_header);
+            if let Some(kf) = read_uncompressed_header(s, ps, false) {
+                return Some(kf);
+            }
+        }
+        obu_type::FRAME => {
+            assert!(!ps.seen_frame_header);
+            if let Some(kf) = read_uncompressed_header(s, ps, true) {
+                return Some(kf);
+            }
+            ps.seen_frame_header = false; // tile_group_obu might reset this
+        }
+        obu_type::TEMPORAL_DELIMITER => {
+            ps.seen_frame_header = false;
+        }
+        obu_type::REDUNDANT_FRAME_HEADER => {
+            assert!(ps.seen_frame_header);
+        }
+        _ => {}
+    }
+    *s = &obu_payload_start[obu_size as usize..];
+    None
+}
+
+fn read_uncompressed_header(s: &mut &[u8], ps: &mut ParserState, is_frame: bool) -> Option<bool> {
+    // TODO handle frame_id_numbers_present_flag
+    let show_existing_frame = (s[0] >> 7) != 0;
+    if is_frame {
+        assert!(
+            !show_existing_frame,
+            "OBU_FRAME requires show_existing_frame = 0"
+        );
+    }
+    if show_existing_frame {
+        return None;
+    }
+    ps.seen_frame_header = true;
+    let frame_type = (s[0] >> 5) & 0b11;
+    Some(frame_type == frame_type::KEY_FRAME)
+}
+
+fn read_leb128(s: &mut &[u8]) -> u32 {
+    let mut value = 0;
+    for i in 0..8 {
+        let byte = s[0];
+        value |= ((byte & 0x7f) as u32) << (i * 7);
+        *s = &s[1..];
+        if (byte & 0x80) == 0 {
+            break;
+        }
+    }
+    value
+}
diff --git a/remuxer/src/kf_detect/mod.rs b/remuxer/src/kf_detect/mod.rs
new file mode 100644
index 0000000..1782d93
--- /dev/null
+++ b/remuxer/src/kf_detect/mod.rs
@@ -0,0 +1,21 @@
+/*
+    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
+    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
+    Copyright (C) 2026 metamuffin <metamuffin.org>
+*/
+
+use anyhow::{Result, bail};
+use winter_matroska::block::Block;
+
+mod av1;
+mod vp9;
+
+pub type IsKeyframeFn = fn(&Block) -> Option<bool>;
+
+pub fn get_is_keyframe_fn(codec: &str) -> Result<IsKeyframeFn> {
+    Ok(match codec {
+        "V_AV1" => |b| av1::is_keyframe(&b.data),
+        "V_VP9" => |b| vp9::is_keyframe(&b.data),
+        _ => bail!("unsupported codec {codec:?}"),
+    })
+}
diff --git a/remuxer/src/kf_detect/vp9.rs b/remuxer/src/kf_detect/vp9.rs
new file mode 100644
index 0000000..47c56ba
--- /dev/null
+++ b/remuxer/src/kf_detect/vp9.rs
@@ -0,0 +1,28 @@
+/*
+    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
+    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
+    Copyright (C) 2026 metamuffin <metamuffin.org>
+*/
+
+#[allow(unused)]
+mod frame_type {
+    pub const KEY_FRAME: u8 = 0;
+    pub const NON_KEY_FRAME: u8 = 0;
+}
+
+pub fn is_keyframe(s: &[u8]) -> Option<bool> {
+    let frame_marker = (s[0] >> 6) & 0b11;
+    assert_eq!(frame_marker, 2);
+    let profile_low_bit = (s[0] >> 5) & 0b1;
+    let profile_high_bit = (s[0] >> 4) & 0b1;
+    let profile = (profile_high_bit << 1) | profile_low_bit;
+    if profile == 3 {
+        todo!()
+    }
+    let show_existing_frame = (s[0] >> 3) & 0b1 != 0;
+    if show_existing_frame {
+        todo!()
+    }
+    let frame_type = (s[0] >> 2) & 0b1;
+    Some(frame_type == frame_type::KEY_FRAME)
+}
diff --git a/remuxer/src/lib.rs b/remuxer/src/lib.rs
index 6afe7cb..8a78fec 100644
--- a/remuxer/src/lib.rs
+++ b/remuxer/src/lib.rs
@@ -9,6 +9,7 @@ mod codec_param;
 pub mod demuxers;
 pub mod magic;
 pub mod muxers;
+pub mod kf_detect;
 
 pub use codec_param::codec_param;
 pub use winter_matroska as matroska;