- commit
- 61922d88d480137198353d07d3f3b1c37a501aca
- parent
- 404b41529394701f5fa02708ab284660cbe219e2
- author
- Dami <dami.etoile@gmail.com>
| M | crates/ | 70 | ++--------------------------- |
| A | crates/ | 74 | ++++++++++++++++++++++++++++++ |
@@ -19,6 +19,10 @@ use crate::config::{
1919 GamescopeEncoderBackend, GamescopePipewireConfig, HostConfig, VaapiRateControl,2020 };2121 22+mod annexb;23+24+use annexb::AnnexBAccessUnitParser;25+2226 2327 use gstreamer::prelude::ObjectExt;2428 @@ -2010,56 +2014,6 @@ fn diagnostic(bytes: &[u8]) -> String {
20102014 String::from_utf8_lossy(bytes).trim().to_owned()20112015 }20122016 2013-2014-struct AnnexBAccessUnitParser {2015- buffer: Vec<u8>,2016-}2017-2018-impl AnnexBAccessUnitParser {2019- fn push(&mut self, bytes: &[u8]) -> Result<Vec<Vec<u8>>> {2020- self.buffer.extend_from_slice(bytes);2021- ensure!(2022- self.buffer.len() <= MAX_ENCODE_BUFFER,2023- "encoded stream exceeded {MAX_ENCODE_BUFFER} bytes without a frame boundary"2024- );2025-2026- let mut frames = Vec::new();2027- while let Some(first) = find_h264_aud(&self.buffer, 0) {2028- if first > 0 {2029- self.buffer.drain(..first);2030- }2031- let Some(next) = find_h264_aud(&self.buffer, 4) else {2032- break;2033- };2034- let frame: Vec<u8> = self.buffer.drain(..next).collect();2035- if !frame.is_empty() {2036- frames.push(frame);2037- }2038- }2039- Ok(frames)2040- }2041-}2042-2043-fn find_h264_aud(data: &[u8], start: usize) -> Option<usize> {2044- let mut index = start;2045- while index + 4 <= data.len() {2046- let nal_offset = if data.get(index..index + 4) == Some(&[0, 0, 0, 1]) {2047- Some(index + 4)2048- } else if data.get(index..index + 3) == Some(&[0, 0, 1]) {2049- Some(index + 3)2050- } else {2051- None2052- };2053- if let Some(nal_offset) = nal_offset2054- && data.get(nal_offset).is_some_and(|byte| byte & 0x1f == 9)2055- {2056- return Some(index);2057- }2058- index += 1;2059- }2060- None2061-}2062-20632017 fn is_h264_keyframe(data: &[u8]) -> bool {20642018 h264_nal_types(data).any(|nal_type| nal_type == 5)20652019 }@@ -2452,22 +2406,6 @@ mod tests {
24522406 assert!(!description.contains("appsink name=sigil_sink max-buffers=1 drop=true"));24532407 }24542408 2455- 2456- fn splits_access_units_on_aud() {2457- let first = [0, 0, 0, 1, 9, 0x10, 0, 0, 1, 5, 1, 2, 3];2458- let second = [0, 0, 0, 1, 9, 0x10, 0, 0, 1, 1, 4, 5];2459- let third = [0, 0, 1, 9, 0x10, 0, 0, 1, 1, 6, 7];2460- let mut parser = AnnexBAccessUnitParser::default();2461-2462- let mut bytes = Vec::new();2463- bytes.extend_from_slice(&first);2464- bytes.extend_from_slice(&second);2465- bytes.extend_from_slice(&third);2466- let frames = parser.push(&bytes).unwrap();2467-2468- assert_eq!(frames, vec![first.to_vec(), second.to_vec()]);2469- }2470-24712409 24722410 fn detects_h264_keyframes() {24732411 assert!(is_h264_keyframe(&[0, 0, 0, 1, 5, 1, 2]));@@ -0,0 +1,74 @@
1+use anyhow::{Result, ensure};2+3+use super::MAX_ENCODE_BUFFER;4+5+6+pub(super) struct AnnexBAccessUnitParser {7+ buffer: Vec<u8>,8+}9+10+impl AnnexBAccessUnitParser {11+ pub(super) fn push(&mut self, bytes: &[u8]) -> Result<Vec<Vec<u8>>> {12+ self.buffer.extend_from_slice(bytes);13+ ensure!(14+ self.buffer.len() <= MAX_ENCODE_BUFFER,15+ "encoded stream exceeded {MAX_ENCODE_BUFFER} bytes without a frame boundary"16+ );17+18+ let mut frames = Vec::new();19+ while let Some(first) = find_h264_aud(&self.buffer, 0) {20+ if first > 0 {21+ self.buffer.drain(..first);22+ }23+ let Some(next) = find_h264_aud(&self.buffer, 4) else {24+ break;25+ };26+ let frame: Vec<u8> = self.buffer.drain(..next).collect();27+ if !frame.is_empty() {28+ frames.push(frame);29+ }30+ }31+ Ok(frames)32+ }33+}34+35+fn find_h264_aud(data: &[u8], start: usize) -> Option<usize> {36+ let mut index = start;37+ while index + 4 <= data.len() {38+ let nal_offset = if data.get(index..index + 4) == Some(&[0, 0, 0, 1]) {39+ Some(index + 4)40+ } else if data.get(index..index + 3) == Some(&[0, 0, 1]) {41+ Some(index + 3)42+ } else {43+ None44+ };45+ if let Some(nal_offset) = nal_offset46+ && data.get(nal_offset).is_some_and(|byte| byte & 0x1f == 9)47+ {48+ return Some(index);49+ }50+ index += 1;51+ }52+ None53+}54+55+56+mod tests {57+ use super::*;58+59+ 60+ fn splits_access_units_on_aud() {61+ let first = [0, 0, 0, 1, 9, 0x10, 0, 0, 1, 5, 1, 2, 3];62+ let second = [0, 0, 0, 1, 9, 0x10, 0, 0, 1, 1, 4, 5];63+ let third = [0, 0, 1, 9, 0x10, 0, 0, 1, 1, 6, 7];64+ let mut parser = AnnexBAccessUnitParser::default();65+66+ let mut bytes = Vec::new();67+ bytes.extend_from_slice(&first);68+ bytes.extend_from_slice(&second);69+ bytes.extend_from_slice(&third);70+ let frames = parser.push(&bytes).unwrap();71+72+ assert_eq!(frames, vec![first.to_vec(), second.to_vec()]);73+ }74+}