diff --git a/CHANGELOG.md b/CHANGELOG.md index 89272f1..7d7af62 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,8 +35,13 @@ reconnect screen instead, for up to the time you pick (30 s to 5 min). gets its own silence, and OBS reconnecting is handed the same Twitch session, so the destination carries on instead of restarting. Twitch VOD-audio sessions are kept the same way. -- **HEVC and AV1 hold the last frame.** The screen is H.264, so those - tracks re-send their last keyframe once a second instead. +- **2K (HEVC) channels get the screen too.** Twitch streams Partners and + Affiliates at 1440p in HEVC under Enhanced Broadcasting, so the screen + is encoded in HEVC for HEVC tracks, at each track's own resolution and + as light as the H.264 one. The log names each track's codec when it + goes on air, for example `reconnect screen on air (2560x1440 HEVC, + 1920x1080 HEVC)`. 10-bit HEVC and AV1 hold the last frame instead, + re-sent once a second. - **Stopping in OBS never opens the screen,** even while it is up for a frozen OBS: the stream ends once your delay has aired. Only a drop without a goodbye opens it. **Protect every disconnect** (under @@ -486,6 +491,10 @@ stream settings: h264_texture_amf, 6000 kbps`. ### Fixes +- **Integrations: a clear error when OBS wants a WebSocket password you + haven't saved.** It used to try an empty one, and OBS closed the + connection with an unclear code. The same goes for a malformed OBS + greeting that asks for a password without saying how. - **Stopping in OBS no longer cuts the end of a delayed stream.** With a delay on and no crash protection hold (protection off, or a stop it lets through), the last seconds of the delay used to vanish: viewers diff --git a/src/crash_hold.rs b/src/crash_hold.rs index 346b13b..fca863c 100644 --- a/src/crash_hold.rs +++ b/src/crash_hold.rs @@ -9,11 +9,13 @@ //! pump's reconnect path takes over from there), the hold time runs out, //! or the streamer ends it early: //! -//! - H.264 tracks loop the reconnect screen at the track's own resolution, -//! framed exactly like the stream's own tags (legacy, Enhanced RTMP, or -//! an Enhanced Broadcasting multitrack tag with its track id). -//! - Other tracks (HEVC, AV1) hold their last keyframe, re-sent once a -//! second, because the screen's encoder only speaks H.264. +//! - H.264 and 8-bit 4:2:0 HEVC tracks loop the reconnect screen, encoded +//! in the track's own codec at its own resolution and framed exactly +//! like the stream's own tags (legacy, Enhanced RTMP, or an Enhanced +//! Broadcasting multitrack tag with its track id). +//! - Other tracks (10-bit HEVC, AV1) hold their last keyframe, re-sent +//! once a second, because the screen's encoders only speak 8-bit H.264 +//! and HEVC. //! - Every AAC track gets digital silence in its own format. //! //! Timestamps continue from the last real frame, so the platform sees one @@ -23,7 +25,7 @@ use crate::controller::{Controller, DestinationState}; use crate::crash_protection::CrashProtection; use crate::h264::{AudioEgress, VideoEgress}; use crate::rtmp::client::EgressSink; -use crate::slate::{SlateLoop, StreamShape}; +use crate::slate::{SlateCodec, SlateLoop, StreamShape}; use std::borrow::Cow; use std::collections::{HashMap, HashSet}; use std::io; @@ -47,8 +49,8 @@ const HELD_FRAME_EVERY: u64 = SLATE_FPS as u64; /// it, or a delay raised meanwhile) and rejoin where the video is. const REBUILD_GRACE: Duration = Duration::from_secs(5); -/// Used when a destination's resolution can't be read from its SPS. -const FALLBACK_SHAPE: (usize, usize) = (1280, 720); +/// Used when an H.264 track's resolution can't be read from its SPS. +const FALLBACK_SIZE: (u32, u32) = (1280, 720); // FLV and Enhanced RTMP tag header fields. const FLV_VIDEO_AVC: u8 = 7; @@ -65,6 +67,8 @@ const ONE_TRACK: u8 = 0; const FRAME_KEY: u8 = 1; const FRAME_INTER: u8 = 2; const FOURCC_AVC1: [u8; 4] = *b"avc1"; +const FOURCC_HVC1: [u8; 4] = *b"hvc1"; +const FOURCC_HEV1: [u8; 4] = *b"hev1"; const FOURCC_MP4A: [u8; 4] = *b"mp4a"; /// Raw AAC-LC frames that decode to digital silence (verified with @@ -155,20 +159,28 @@ struct HoldPlan { } enum VideoSource { - /// The reconnect screen, encoded at the track's own resolution. - Screen { - shape: StreamShape, - framing: Framing, - }, + /// The reconnect screen, encoded like the track it replaces. + Screen(ScreenTarget), /// The track's last keyframe as the destination receives it. HeldFrame(Arc<[u8]>), } +/// How the reconnect screen stands in for one video track. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct ScreenTarget { + framing: Framing, + /// The track's own FourCC, which Enhanced RTMP tags repeat. + fourcc: [u8; 4], + codec: SlateCodec, + shape: StreamShape, +} + /// A `VideoSource` ready to play, as the tags the destination gets. enum VideoTrack { /// The loop wrapped once in the track's framing, not once per frame. Screen { shape: StreamShape, + codec: SlateCodec, /// The two IDR variants, then the P frames, in loop order. keyframes: [Vec; 2], deltas: Vec>, @@ -177,14 +189,18 @@ enum VideoTrack { } impl VideoTrack { - fn screen(slate: &SlateLoop, framing: Framing) -> Self { + fn screen(slate: &SlateLoop, target: ScreenTarget) -> Self { + let tag = |keyframe: bool, sample: &[u8]| { + video_tag(target.framing, target.fourcc, keyframe, sample) + }; VideoTrack::Screen { shape: slate.shape, - keyframes: [0, 1].map(|replay| video_tag(framing, true, &slate.keyframes[replay])), + codec: target.codec, + keyframes: [0, 1].map(|replay| tag(true, &slate.keyframes[replay])), deltas: slate .deltas .iter() - .map(|sample| video_tag(framing, false, sample)) + .map(|sample| tag(false, sample)) .collect(), } } @@ -214,16 +230,19 @@ pub fn prebuild_for(ctrl: &Controller, dest: &DestinationState, settings: &Crash return; }; for (_, header) in video_headers_for(ctrl, egress) { - if avc_framing(&header).is_some() { - ctrl.slate_cache.prebuild(settings, screen_shape(&header)); + if let Some(target) = screen_target(&header) { + ctrl.slate_cache + .prebuild(settings, target.shape, target.codec); } } } -/// Whether the hold plays this video tag's track by holding its last +/// Whether the hold may play this video tag's track by holding its last /// keyframe: an Enhanced RTMP keyframe that isn't H.264 in a framing the -/// reconnect screen matches (HEVC, AV1, bundled multitrack layouts). The -/// controller keeps the latest one per track. +/// reconnect screen matches (HEVC, AV1, bundled multitrack layouts). HEVC +/// usually gets the screen instead; its keyframe is kept for the HEVC the +/// screen can't encode (10-bit). The controller keeps the latest one per +/// track. pub fn is_held_keyframe(tag: &[u8]) -> bool { let Some(&first) = tag.first() else { return false; @@ -286,17 +305,16 @@ fn video_headers_for(ctrl: &Controller, egress: VideoEgress) -> Vec<(u8, Vec .collect() } -/// The reconnect screen for an H.264 track, or its last keyframe for -/// anything else. +/// The reconnect screen for a track it can be encoded for, or its last +/// keyframe for anything else. fn video_source( ctrl: &Controller, egress: VideoEgress, track: u8, header: &[u8], ) -> Result { - if let Some(framing) = avc_framing(header) { - let shape = screen_shape(header); - return Ok(VideoSource::Screen { shape, framing }); + if let Some(target) = screen_target(header) { + return Ok(VideoSource::Screen(target)); } let codec = crate::h264::seq_header_codec(header).label(); let Some(keyframe) = ctrl.held_keyframe(track) else { @@ -312,29 +330,46 @@ fn video_source( Ok(VideoSource::HeldFrame(sent)) } -/// The screen's shape for an H.264 track: its resolution, rounded down -/// to even sides, at the screen's frame rate. -fn screen_shape(header: &[u8]) -> StreamShape { - let (width, height) = crate::h264::sps_dimensions(header) - .map(|(w, h)| (w as usize & !1, h as usize & !1)) - .unwrap_or(FALLBACK_SHAPE); +/// The screen's shape for a track of `width` x `height`: rounded down to +/// even sides, at the screen's frame rate. +fn screen_shape((width, height): (u32, u32)) -> StreamShape { StreamShape { - width, - height, + width: width as usize & !1, + height: height as usize & !1, fps: SLATE_FPS, } } -/// The framing of an H.264 sequence header the reconnect screen can -/// match, or `None` for other codecs and bundled multitrack layouts. -fn avc_framing(header: &[u8]) -> Option { +/// How the reconnect screen replaces the track `header` describes, or +/// `None` when it can't: codecs other than H.264 and 8-bit 4:2:0 HEVC, +/// HEVC whose picture size can't be read, and bundled multitrack layouts. +fn screen_target(header: &[u8]) -> Option { let &first = header.first()?; + let avc_shape = || screen_shape(crate::h264::sps_dimensions(header).unwrap_or(FALLBACK_SIZE)); if first & EX_VIDEO == 0 { let is_avc_config = first & 0x0F == FLV_VIDEO_AVC && header.get(1) == Some(&0); - return is_avc_config.then_some(Framing::Legacy); + return is_avc_config.then(|| ScreenTarget { + framing: Framing::Legacy, + fourcc: FOURCC_AVC1, + codec: SlateCodec::H264, + shape: avc_shape(), + }); } - let (framing, fourcc, _) = enhanced_sequence_start(header, EX_VIDEO_MULTITRACK)?; - (fourcc == FOURCC_AVC1).then_some(framing) + let (framing, fourcc, config) = enhanced_sequence_start(header, EX_VIDEO_MULTITRACK)?; + let (codec, shape) = match fourcc { + FOURCC_AVC1 => (SlateCodec::H264, avc_shape()), + FOURCC_HVC1 | FOURCC_HEV1 => ( + SlateCodec::Hevc, + screen_shape(crate::h264::hevc_8bit_420_dimensions(config)?), + ), + _ => return None, + }; + Some(ScreenTarget { + framing, + fourcc, + codec, + shape, + }) } /// Parse an Enhanced RTMP sequence-start tag (video or audio): its @@ -443,15 +478,19 @@ pub async fn play( let mut video = Vec::with_capacity(plan.video.len()); for source in plan.video { video.push(match source { - VideoSource::Screen { shape, framing } => { - let Some(slate) = ctrl.slate_cache.get(&settings, shape).await else { + VideoSource::Screen(target) => { + let built = ctrl + .slate_cache + .get(&settings, target.shape, target.codec) + .await; + let Some(slate) = built else { ctrl.log(format!( "[{}] crash protection couldn't build the reconnect screen", dest.id )); return Ok(ended); }; - VideoTrack::screen(&slate, framing) + VideoTrack::screen(&slate, target) } VideoSource::HeldFrame(tag) => VideoTrack::HeldFrame(tag), }); @@ -461,19 +500,22 @@ pub async fn play( dest.id, describe(&video) )); - // No sequence header: the loop's keyframes carry their own SPS/PPS - // in-band under ids the stream doesn't use (see slate::encoder), so the - // destination's decoder keeps the stream's config for the way back. + // No sequence header: the loop's keyframes carry their own parameter + // sets in-band under ids the stream doesn't use (see slate::encoder), + // so the destination's decoder keeps the stream's config for the way + // back. let base = after_ts.wrapping_add(1); stream_loop(ctrl, dest, sink, &video, &plan.audio, base).await } -/// "1920x1080, 1280x720, last keyframe held" for the log. +/// "2560x1440 HEVC, 1280x720 H.264, last keyframe held" for the log. fn describe(video: &[VideoTrack]) -> String { video .iter() .map(|track| match track { - VideoTrack::Screen { shape, .. } => format!("{}x{}", shape.width, shape.height), + VideoTrack::Screen { shape, codec, .. } => { + format!("{}x{} {}", shape.width, shape.height, codec.label()) + } VideoTrack::HeldFrame(_) => "last keyframe held".to_string(), }) .collect::>() @@ -612,20 +654,21 @@ fn video_frame(track: &VideoTrack, index: u64) -> Option<&[u8]> { } } -/// An H.264 coded-frame tag in `framing`, zero composition time. -fn video_tag(framing: Framing, keyframe: bool, sample: &[u8]) -> Vec { +/// A coded-frame tag in `framing`, zero composition time. Enhanced RTMP +/// framings name the codec by `fourcc`; legacy framing is H.264 only. +fn video_tag(framing: Framing, fourcc: [u8; 4], keyframe: bool, sample: &[u8]) -> Vec { let frame_type = if keyframe { FRAME_KEY } else { FRAME_INTER }; let mut tag = Vec::with_capacity(sample.len() + 12); match framing { Framing::Legacy => tag.extend_from_slice(&[frame_type << 4 | FLV_VIDEO_AVC, AVC_NALU]), Framing::Enhanced => { tag.push(EX_VIDEO | frame_type << 4 | EX_CODED_FRAMES); - tag.extend_from_slice(&FOURCC_AVC1); + tag.extend_from_slice(&fourcc); } Framing::OneTrack(track) => { tag.push(EX_VIDEO | frame_type << 4 | EX_VIDEO_MULTITRACK); tag.push(ONE_TRACK << 4 | EX_CODED_FRAMES); - tag.extend_from_slice(&FOURCC_AVC1); + tag.extend_from_slice(&fourcc); tag.push(track); } } @@ -654,42 +697,48 @@ fn audio_tag(audio: &SilentAudio) -> Vec { tag } -/// Encoded reconnect-screen loops, one per destination shape, rebuilt when -/// the screen's look changes. Encoding a 1080p loop takes a few hundred -/// ms of one core, so loops are built ahead of time while OBS is live and -/// a hold usually starts with one ready. +/// Encoded reconnect-screen loops, one per destination shape and codec, +/// rebuilt when the screen's look changes. Encoding a 1080p loop takes a +/// few hundred ms of one core, so loops are built ahead of time while OBS +/// is live and a hold usually starts with one ready. #[derive(Default)] pub struct SlateCache { - /// Per (width, height): the loop and the settings it was built with. + /// Per (width, height, codec): the loop and the settings it was built + /// with. ready: crate::sync::Mutex)>>, building: crate::sync::Mutex>, } -/// A slate loop's (width, height). -type SlateKey = (usize, usize); +/// A slate loop's (width, height, codec). +type SlateKey = (usize, usize, SlateCodec); + +fn slate_key(shape: StreamShape, codec: SlateCodec) -> SlateKey { + (shape.width, shape.height, codec) +} impl SlateCache { - fn cached(&self, settings: &CrashProtection, shape: StreamShape) -> Option> { + fn cached(&self, settings: &CrashProtection, key: SlateKey) -> Option> { self.ready .lock() - .get(&(shape.width, shape.height)) + .get(&key) .filter(|(built_with, _)| built_with.same_screen(settings)) .map(|(_, slate)| slate.clone()) } - /// The loop for `shape`, building it now if it isn't ready. When a - /// build of it is already running (a prebuild, or another destination - /// at the same resolution) this waits for that one rather than encode - /// the same loop twice at the moment OBS drops. + /// The loop for `shape` in `codec`, building it now if it isn't + /// ready. When a build of it is already running (a prebuild, or + /// another destination at the same resolution) this waits for that one + /// rather than encode the same loop twice at the moment OBS drops. pub async fn get( self: &Arc, settings: &CrashProtection, shape: StreamShape, + codec: SlateCodec, ) -> Option> { const BUILD_POLL: Duration = Duration::from_millis(50); - let key = (shape.width, shape.height); + let key = slate_key(shape, codec); loop { - if let Some(slate) = self.cached(settings, shape) { + if let Some(slate) = self.cached(settings, key) { return Some(slate); } if self.building.lock().insert(key) { @@ -703,7 +752,7 @@ impl SlateCache { let cache = self.clone(); let settings = settings.clone(); let task = tokio::spawn(async move { - let slate = build(settings.clone(), shape).await; + let slate = build(settings.clone(), shape, codec).await; if let Some(slate) = &slate { cache.ready.lock().insert(key, (settings, slate.clone())); } @@ -713,17 +762,22 @@ impl SlateCache { task.await.ok().flatten() } - /// Start building the loop for `shape` in the background, unless it is - /// ready or already being built. - pub fn prebuild(self: &Arc, settings: &CrashProtection, shape: StreamShape) { - let key = (shape.width, shape.height); - if self.cached(settings, shape).is_some() || self.building.lock().contains(&key) { + /// Start building the loop for `shape` in `codec` in the background, + /// unless it is ready or already being built. + pub fn prebuild( + self: &Arc, + settings: &CrashProtection, + shape: StreamShape, + codec: SlateCodec, + ) { + let key = slate_key(shape, codec); + if self.cached(settings, key).is_some() || self.building.lock().contains(&key) { return; } let cache = self.clone(); let settings = settings.clone(); tokio::spawn(async move { - cache.get(&settings, shape).await; + cache.get(&settings, shape, codec).await; }); } @@ -734,9 +788,13 @@ impl SlateCache { } /// Encode a loop off the async runtime. -async fn build(settings: CrashProtection, shape: StreamShape) -> Option> { +async fn build( + settings: CrashProtection, + shape: StreamShape, + codec: SlateCodec, +) -> Option> { tokio::task::spawn_blocking(move || { - crate::slate::build_loop(&settings, shape) + crate::slate::build_loop(&settings, shape, codec) .ok() .map(Arc::new) }) @@ -806,35 +864,78 @@ mod tests { assert!(silent_audio_from_header(&one_track_header(0x95, b"Opus", 1, &[0; 8])).is_none()); } + fn framing_and_codec(header: &[u8]) -> Option<(Framing, SlateCodec)> { + screen_target(header).map(|target| (target.framing, target.codec)) + } + #[test] - fn the_screen_matches_each_h264_framing_and_nothing_else() { - assert_eq!(avc_framing(&[0x17, 0, 0, 0, 0, 1]), Some(Framing::Legacy)); + fn the_screen_matches_each_h264_framing() { + let legacy = Some((Framing::Legacy, SlateCodec::H264)); + assert_eq!(framing_and_codec(&[0x17, 0, 0, 0, 0, 1]), legacy); let mut enhanced = vec![EX_VIDEO | FRAME_KEY << 4 | EX_SEQUENCE_START]; enhanced.extend_from_slice(b"avc1"); - assert_eq!(avc_framing(&enhanced), Some(Framing::Enhanced)); + let enhanced_h264 = Some((Framing::Enhanced, SlateCodec::H264)); + assert_eq!(framing_and_codec(&enhanced), enhanced_h264); let rung = one_track_header(0x96, b"avc1", 2, &[1]); - assert_eq!(avc_framing(&rung), Some(Framing::OneTrack(2))); - assert_eq!(avc_framing(&one_track_header(0x96, b"hvc1", 2, &[1])), None); + let rung_h264 = Some((Framing::OneTrack(2), SlateCodec::H264)); + assert_eq!(framing_and_codec(&rung), rung_h264); // ManyTracks bundles several tracks' configs in one tag. let mut bundle = one_track_header(0x96, b"avc1", 0, &[1]); bundle[1] = 1 << 4; - assert_eq!(avc_framing(&bundle), None); + assert_eq!(framing_and_codec(&bundle), None); + assert_eq!( + framing_and_codec(&one_track_header(0x96, b"av01", 2, &[1])), + None + ); + } + + /// Twitch 2K channels send Enhanced Broadcasting tracks as 8-bit HEVC: + /// they get the screen in HEVC at their own size, under their own + /// FourCC and track id. 10-bit HEVC, or a config whose size can't be + /// read, falls back to the held keyframe. + #[test] + fn eight_bit_hevc_tracks_get_the_screen_in_hevc() { + let config = crate::slate::test_hevc_config(2560, 1440); + let rung = one_track_header(0x96, b"hvc1", 1, &config); + let target = screen_target(&rung).unwrap(); + assert_eq!(target.framing, Framing::OneTrack(1)); + assert_eq!(target.codec, SlateCodec::Hevc); + assert_eq!(target.fourcc, *b"hvc1"); + assert_eq!((target.shape.width, target.shape.height), (2560, 1440)); + + let mut ten_bit = config.clone(); + ten_bit[17] = 0xF8 | 2; // bitDepthLumaMinus8 + assert_eq!( + screen_target(&one_track_header(0x96, b"hvc1", 1, &ten_bit)), + None + ); + assert_eq!( + screen_target(&one_track_header(0x96, b"hvc1", 1, &[1])), + None + ); } #[test] fn screen_frames_are_framed_like_the_track_they_replace() { assert_eq!( - video_tag(Framing::Legacy, true, &[9, 9]), + video_tag(Framing::Legacy, FOURCC_AVC1, true, &[9, 9]), vec![0x17, 1, 0, 0, 0, 9, 9] ); - assert_eq!(video_tag(Framing::Legacy, false, &[9])[0], 0x27); - let rung = video_tag(Framing::OneTrack(3), true, &[9, 9]); + assert_eq!( + video_tag(Framing::Legacy, FOURCC_AVC1, false, &[9])[0], + 0x27 + ); + let rung = video_tag(Framing::OneTrack(3), FOURCC_AVC1, true, &[9, 9]); assert_eq!(rung[..7], [0x96, 0x01, b'a', b'v', b'c', b'1', 3]); // A vertical destination fed track 3 sees the plain legacy frame. let flat = select_video_bytes(&rung, VideoEgress::Track(3)).unwrap(); assert_eq!(&flat[..], &[0x17, 1, 0, 0, 0, 9, 9]); - let enhanced = video_tag(Framing::Enhanced, false, &[9]); + let enhanced = video_tag(Framing::Enhanced, FOURCC_AVC1, false, &[9]); assert_eq!(enhanced[..5], [0xA1, b'a', b'v', b'c', b'1']); + let hevc = video_tag(Framing::OneTrack(1), FOURCC_HVC1, true, &[9]); + assert_eq!(hevc[..10], [0x96, 0x01, b'h', b'v', b'c', b'1', 1, 0, 0, 0]); + // An HEVC keyframe is one the controller recognises as a keyframe. + assert!(crate::h264::classify_video_tag(&hevc).is_idr); } #[test] @@ -887,11 +988,12 @@ mod tests { let first = { let cache = cache.clone(); let settings = settings.clone(); - tokio::spawn(async move { cache.get(&settings, shape).await }) + tokio::spawn(async move { cache.get(&settings, shape, SlateCodec::H264).await }) }; tokio::task::yield_now().await; first.abort(); - let later = tokio::time::timeout(Duration::from_secs(10), cache.get(&settings, shape)); + let later = cache.get(&settings, shape, SlateCodec::H264); + let later = tokio::time::timeout(Duration::from_secs(10), later); assert!(later.await.expect("not wedged").is_some()); } diff --git a/src/h264.rs b/src/h264.rs index 67fb46d..9d00e93 100644 --- a/src/h264.rs +++ b/src/h264.rs @@ -667,6 +667,98 @@ pub fn sps_dimensions(seq_header: &[u8]) -> Option<(u32, u32)> { Some((width, height)) } +/// Picture size of an 8-bit 4:2:0 HEVC stream with 4-byte NAL lengths, +/// read from its HEVCDecoderConfigurationRecord (hvcC) with the +/// conformance window applied: the HEVC the reconnect screen can be +/// encoded for. `None` for any other HEVC, or a malformed or truncated +/// record. Never panics. +pub fn hevc_8bit_420_dimensions(hvcc: &[u8]) -> Option<(u32, u32)> { + const CHROMA_420: u32 = 1; + // The record repeats the SPS's chroma format and bit depths in + // bytes 16..19, and its NAL length size in byte 21. + let record_is_8bit_420 = u32::from(hvcc.get(16)? & 0x03) == CHROMA_420 + && hvcc.get(17)? & 0x07 == 0 + && hvcc.get(18)? & 0x07 == 0; + if !record_is_8bit_420 || hvcc.get(21)? & 0x03 != 3 { + return None; + } + // Drop the 2-byte NAL header, as `sps_dimensions` drops H.264's one. + let rbsp = strip_emulation_prevention(hevc_first_sps(hvcc)?.get(2..)?); + let mut r = BitReader::new(&rbsp); + let _vps_id = r.read_bits(4)?; + let max_sub_layers_minus1 = r.read_bits(3)?; + let _temporal_id_nesting = r.read_bit()?; + skip_hevc_profile_tier_level(&mut r, max_sub_layers_minus1)?; + let _sps_id = r.read_ue()?; + if r.read_ue()? != CHROMA_420 { + return None; + } + let mut width = r.read_ue()?; + let mut height = r.read_ue()?; + if r.read_bit()? == 1 { + // conformance_window offsets count chroma samples: 2 luma in 4:2:0. + let crop = |a: u32, b: u32| a.checked_add(b)?.checked_mul(2); + let (left, right, top, bottom) = (r.read_ue()?, r.read_ue()?, r.read_ue()?, r.read_ue()?); + width = width.checked_sub(crop(left, right)?)?; + height = height.checked_sub(crop(top, bottom)?)?; + } + let is_8bit = r.read_ue()? == 0 && r.read_ue()? == 0; + (is_8bit && width > 0 && height > 0).then_some((width, height)) +} + +/// The first SPS NAL unit (header included) in an hvcC's arrays. +fn hevc_first_sps(hvcc: &[u8]) -> Option<&[u8]> { + const NAL_SPS: u8 = 33; + let read_u16 = |at: usize| Some(u16::from_be_bytes([*hvcc.get(at)?, *hvcc.get(at + 1)?])); + let array_count = *hvcc.get(22)?; + let mut at = 23; + for _ in 0..array_count { + let nal_type = hvcc.get(at)? & 0x3F; + let nal_count = read_u16(at + 1)?; + at += 3; + for _ in 0..nal_count { + let len = read_u16(at)? as usize; + let nal = hvcc.get(at + 2..at + 2 + len)?; + if nal_type == NAL_SPS { + return Some(nal); + } + at += 2 + len; + } + } + None +} + +/// Skip profile_tier_level(1, max_sub_layers_minus1) (H.265 7.3.3). +fn skip_hevc_profile_tier_level(r: &mut BitReader, max_sub_layers_minus1: u32) -> Option<()> { + // general_profile_space through general_level_idc: 96 bits. + for _ in 0..3 { + r.read_bits(32)?; + } + let mut present = Vec::new(); + for _ in 0..max_sub_layers_minus1 { + let profile = r.read_bit()? == 1; + let level = r.read_bit()? == 1; + present.push((profile, level)); + } + if max_sub_layers_minus1 > 0 { + for _ in max_sub_layers_minus1..8 { + r.read_bits(2)?; // reserved_zero_2bits + } + } + for (profile, level) in present { + if profile { + // sub_layer profile fields: 88 bits. + r.read_bits(32)?; + r.read_bits(32)?; + r.read_bits(24)?; + } + if level { + r.read_bits(8)?; + } + } + Some(()) +} + /// Best-effort codec of a sequence-header tag, across the framings OBS /// emits (legacy AVC, Enhanced single-track, Enhanced OneTrack multi-track). /// Used for the per-destination "1080x1920 ยท H.264" readout. Returns @@ -1776,6 +1868,42 @@ mod tests { assert_eq!(sps_dimensions(&hevc), None); } + /// Read back the size from the reconnect screen's own HEVC parameter + /// sets, cropping included (1080 and 360 aren't whole CTBs). + #[test] + fn hevc_dimensions_read_the_visible_size() { + for (width, height) in [(1920, 1080), (1080, 1920), (2560, 1440), (200, 360)] { + let config = crate::slate::test_hevc_config(width, height); + assert_eq!( + hevc_8bit_420_dimensions(&config), + Some((width as u32, height as u32)), + "{width}x{height}" + ); + } + } + + /// 10-bit, 4:4:4 or 2-byte NAL lengths are HEVC the screen can't stand + /// in for; every truncation of a good record fails cleanly. + #[test] + fn hevc_dimensions_refuse_what_the_screen_cannot_encode() { + let config = crate::slate::test_hevc_config(1920, 1080); + for (byte, value) in [(17, 0xFA), (18, 0xFA), (16, 0xFF), (21, 0x0D)] { + let mut other = config.clone(); + other[byte] = value; + assert_eq!( + hevc_8bit_420_dimensions(&other), + None, + "byte {byte} = {value:#x}" + ); + } + // Any truncation returns without panicking (these are wire bytes), + // and one cut before the SPS's size fields reads nothing. + for len in 0..config.len() { + let _ = hevc_8bit_420_dimensions(&config[..len]); + } + assert_eq!(hevc_8bit_420_dimensions(&config[..30]), None); + } + /// The SPS fields that change how many bits sit in front of the picture /// size. Everything else is fixed by `build_rich_sps_rbsp`. struct SpsFields { diff --git a/src/slate/bitstream.rs b/src/slate/bitstream.rs index 5822116..0c7bcb6 100644 --- a/src/slate/bitstream.rs +++ b/src/slate/bitstream.rs @@ -1,5 +1,5 @@ -//! H.264 bit-level writing: fixed-width and Exp-Golomb fields, RBSP -//! trailing bits, and NAL unit framing with emulation prevention. +//! H.264 and HEVC bit-level writing: fixed-width and Exp-Golomb fields, +//! RBSP trailing bits, and NAL unit framing with emulation prevention. /// MSB-first bit writer for one RBSP. #[derive(Default)] @@ -69,6 +69,12 @@ impl BitWriter { /// Close the RBSP with rbsp_stop_one_bit and alignment zeros. pub fn finish(mut self) -> Vec { self.bit(true); + self.into_aligned_bytes() + } + + /// Pad with zeros to a byte boundary and return the bytes, for an RBSP + /// whose stop bit is already written (a CABAC flush writes its own). + pub fn into_aligned_bytes(mut self) -> Vec { self.align_with_zeros(); self.bytes } @@ -79,11 +85,28 @@ pub const NAL_IDR_SLICE: u8 = 5; pub const NAL_SPS: u8 = 7; pub const NAL_PPS: u8 = 8; -/// NAL header plus the RBSP with emulation-prevention bytes inserted, so -/// no 00 00 0x (x <= 3) start-code pattern appears inside the payload. +pub const HEVC_NAL_TRAIL_R: u8 = 1; +pub const HEVC_NAL_IDR_N_LP: u8 = 20; +pub const HEVC_NAL_VPS: u8 = 32; +pub const HEVC_NAL_SPS: u8 = 33; +pub const HEVC_NAL_PPS: u8 = 34; + +/// H.264 NAL unit: header byte plus the escaped RBSP. pub fn nal_unit(nal_ref_idc: u8, nal_type: u8, rbsp: &[u8]) -> Vec { - let mut out = Vec::with_capacity(rbsp.len() + rbsp.len() / 64 + 1); - out.push((nal_ref_idc << 5) | nal_type); + escaped(&[(nal_ref_idc << 5) | nal_type], rbsp) +} + +/// HEVC NAL unit: two header bytes (layer 0, temporal id 0) plus the +/// escaped RBSP. +pub fn hevc_nal_unit(nal_type: u8, rbsp: &[u8]) -> Vec { + escaped(&[nal_type << 1, 1], rbsp) +} + +/// `header` plus the RBSP with emulation-prevention bytes inserted, so no +/// 00 00 0x (x <= 3) start-code pattern appears inside the payload. +fn escaped(header: &[u8], rbsp: &[u8]) -> Vec { + let mut out = Vec::with_capacity(header.len() + rbsp.len() + rbsp.len() / 64 + 1); + out.extend_from_slice(header); let mut zeros = 0; for &byte in rbsp { if zeros >= 2 && byte <= 3 { @@ -145,5 +168,9 @@ mod tests { vec![0x67, 0, 0, 3, 1, 0, 0, 3, 0, 0] ); assert_eq!(nal_unit(0, NAL_SLICE, &[0, 0, 4]), vec![0x01, 0, 0, 4]); + assert_eq!( + hevc_nal_unit(HEVC_NAL_SPS, &[0, 0, 2]), + vec![0x42, 0x01, 0, 0, 3, 2] + ); } } diff --git a/src/slate/cabac.rs b/src/slate/cabac.rs new file mode 100644 index 0000000..9c37f1c --- /dev/null +++ b/src/slate/cabac.rs @@ -0,0 +1,261 @@ +//! CABAC arithmetic encoder for the reconnect screen's HEVC slices +//! (H.265 9.3.4.3 and the informative encoder of 9.3.5): context-coded, +//! bypass and terminate bins, plus the restart after PCM samples. +//! +//! The probability tables are the ones H.264 shares with HEVC. + +use super::bitstream::BitWriter; + +/// rangeTabLps (Table 9-52), indexed by pStateIdx then qRangeIdx. +const RANGE_TAB_LPS: [[u8; 4]; 64] = [ + [128, 176, 208, 240], + [128, 167, 197, 227], + [128, 158, 187, 216], + [123, 150, 178, 205], + [116, 142, 169, 195], + [111, 135, 160, 185], + [105, 128, 152, 175], + [100, 122, 144, 166], + [95, 116, 137, 158], + [90, 110, 130, 150], + [85, 104, 123, 142], + [81, 99, 117, 135], + [77, 94, 111, 128], + [73, 89, 105, 122], + [69, 85, 100, 116], + [66, 80, 95, 110], + [62, 76, 90, 104], + [59, 72, 86, 99], + [56, 69, 81, 94], + [53, 65, 77, 89], + [51, 62, 73, 85], + [48, 59, 69, 80], + [46, 56, 66, 76], + [43, 53, 63, 72], + [41, 50, 59, 69], + [39, 48, 56, 65], + [37, 45, 54, 62], + [35, 43, 51, 59], + [33, 41, 48, 56], + [32, 39, 46, 53], + [30, 37, 43, 50], + [29, 35, 41, 48], + [27, 33, 39, 45], + [26, 31, 37, 43], + [24, 30, 35, 41], + [23, 28, 33, 39], + [22, 27, 32, 37], + [21, 26, 30, 35], + [20, 24, 29, 33], + [19, 23, 27, 31], + [18, 22, 26, 30], + [17, 21, 25, 28], + [16, 20, 23, 27], + [15, 19, 22, 25], + [14, 18, 21, 24], + [14, 17, 20, 23], + [13, 16, 19, 22], + [12, 15, 18, 21], + [12, 14, 17, 20], + [11, 14, 16, 19], + [11, 13, 15, 18], + [10, 12, 15, 17], + [10, 12, 14, 16], + [9, 11, 13, 15], + [9, 11, 12, 14], + [8, 10, 12, 14], + [8, 9, 11, 13], + [7, 9, 11, 12], + [7, 9, 10, 12], + [7, 8, 10, 11], + [6, 8, 9, 11], + [6, 7, 9, 10], + [6, 7, 8, 9], + [2, 2, 2, 2], +]; + +/// transIdxLps (Table 9-53): the state after coding the less probable +/// symbol. After the more probable one the state just counts up to 62. +const TRANS_IDX_LPS: [u8; 64] = [ + 0, 0, 1, 2, 2, 4, 4, 5, 6, 7, 8, 9, 9, 11, 11, 12, 13, 13, 15, 15, 16, 16, 18, 18, 19, 19, 21, + 21, 22, 22, 23, 24, 24, 25, 26, 26, 27, 27, 28, 29, 29, 30, 30, 30, 31, 32, 32, 33, 33, 33, 34, + 34, 35, 35, 35, 36, 36, 36, 37, 37, 37, 38, 38, 63, +]; +const LAST_ADAPTIVE_STATE: u8 = 62; + +/// One context variable: probability state and most probable symbol. +#[derive(Debug, Clone, Copy)] +pub struct Context { + state: u8, + mps: bool, +} + +impl Context { + /// The state a slice starts from, given the syntax element's + /// initValue and the slice QP (9.3.2.2). + pub fn new(init_value: u8, slice_qp: i32) -> Self { + let slope = i32::from(init_value >> 4) * 5 - 45; + let offset = (i32::from(init_value & 15) << 3) - 16; + let state = (((slope * slice_qp.clamp(0, 51)) >> 4) + offset).clamp(1, 126); + if state <= 63 { + Self { + state: (63 - state) as u8, + mps: false, + } + } else { + Self { + state: (state - 64) as u8, + mps: true, + } + } + } +} + +/// Arithmetic coder writing into a slice's RBSP after its header. +pub struct CabacWriter { + out: BitWriter, + low: u32, + range: u32, + outstanding: u32, + is_first_bit: bool, +} + +impl CabacWriter { + /// Start coding slice data; `out` holds the byte-aligned slice header. + pub fn new(out: BitWriter) -> Self { + Self { + out, + low: 0, + range: 510, + outstanding: 0, + is_first_bit: true, + } + } + + /// EncodeDecision. + pub fn decision(&mut self, ctx: &mut Context, bin: bool) { + let lps = u32::from(RANGE_TAB_LPS[ctx.state as usize][((self.range >> 6) & 3) as usize]); + self.range -= lps; + if bin == ctx.mps { + ctx.state = (ctx.state + 1).min(LAST_ADAPTIVE_STATE); + } else { + self.low += self.range; + self.range = lps; + if ctx.state == 0 { + ctx.mps = !ctx.mps; + } + ctx.state = TRANS_IDX_LPS[ctx.state as usize]; + } + self.renormalize(); + } + + /// EncodeBypass for each of the low `count` bits of `value`, most + /// significant first. + pub fn bypass_bits(&mut self, value: u32, count: u8) { + for shift in (0..count).rev() { + self.low <<= 1; + if (value >> shift) & 1 == 1 { + self.low += self.range; + } + if self.low >= 1024 { + self.put_bit(true); + self.low -= 1024; + } else if self.low < 512 { + self.put_bit(false); + } else { + self.low -= 512; + self.outstanding += 1; + } + } + } + + /// EncodeTerminate. A 1 ends arithmetic coding (end of slice segment, + /// or a PCM block about to follow), so it flushes the coder; the last + /// bit written is a 1 that doubles as the RBSP stop bit. + pub fn terminate(&mut self, bin: bool) { + self.range -= 2; + if !bin { + self.renormalize(); + return; + } + self.low += self.range; + self.range = 2; + self.renormalize(); + self.put_bit((self.low >> 9) & 1 == 1); + self.out.bits(((self.low >> 7) & 3) | 1, 2); + } + + /// pcm_flag = 1, then the raw samples, then the coder restarts + /// (9.3.2.5). Contexts carry on unchanged. + pub fn pcm(&mut self, samples: &[u8]) { + self.terminate(true); + self.out.align_with_zeros(); + self.out.bytes(samples); + *self = Self::new(std::mem::take(&mut self.out)); + } + + /// end_of_slice_segment_flag = 1 and the byte alignment after it. + pub fn finish(mut self) -> Vec { + self.terminate(true); + self.out.into_aligned_bytes() + } + + /// RenormE. + fn renormalize(&mut self) { + while self.range < 256 { + if self.low < 256 { + self.put_bit(false); + } else if self.low >= 512 { + self.low -= 512; + self.put_bit(true); + } else { + self.low -= 256; + self.outstanding += 1; + } + self.range <<= 1; + self.low <<= 1; + } + } + + /// PutBit: the bit, then any outstanding bits as its complement. + fn put_bit(&mut self, bit: bool) { + if self.is_first_bit { + self.is_first_bit = false; + } else { + self.out.bit(bit); + } + while self.outstanding > 0 { + self.out.bit(!bit); + self.outstanding -= 1; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 9.3.2.2 worked by hand at QP 26. 154: slope 0, offset 64, so state + /// 64 (equiprobable, MPS 1). 139: (-5 * 26) >> 4 = -9 (an arithmetic + /// shift), + 72 = 63 (equiprobable, MPS 0). 197: 390 >> 4 = 24, + 24 = + /// 48, so pStateIdx 15 with MPS 0. + #[test] + fn contexts_start_where_the_spec_says() { + let equal_one = Context::new(154, 26); + assert_eq!((equal_one.state, equal_one.mps), (0, true)); + let equal_zero = Context::new(139, 26); + assert_eq!((equal_zero.state, equal_zero.mps), (0, false)); + let zero_likely = Context::new(197, 26); + assert_eq!((zero_likely.state, zero_likely.mps), (15, false)); + } + + /// A slice that is only end_of_slice_segment_flag = 1. A decoder reads + /// 9 bits as its offset (509 here) and decodes a 1 because that is at + /// least the range left after the terminate (510 - 2). The last of + /// those 9 bits doubles as the stop bit; alignment zeros follow. + #[test] + fn an_empty_slice_flushes_to_its_stop_bit() { + let bytes = CabacWriter::new(BitWriter::new()).finish(); + assert_eq!(bytes, vec![0b1111_1110, 0b1000_0000]); + } +} diff --git a/src/slate/hevc.rs b/src/slate/hevc.rs new file mode 100644 index 0000000..b55d52c --- /dev/null +++ b/src/slate/hevc.rs @@ -0,0 +1,913 @@ +//! Lossless HEVC encoder for flat-colour screens: the H.265 counterpart +//! of `encoder`, for tracks the stream sends as HEVC (Twitch 2K channels). +//! +//! Each 32x32 coding tree unit is one coding unit when it can be, else +//! four 16x16 ones, coded as one of: +//! - intra DC, vertical or horizontal with no residual, when that +//! prediction reproduces the block exactly (flat areas); +//! - PCM, 16x16 only: the raw samples, for anything with detail; +//! - skip: unchanged since the previous picture (zero motion). +//! +//! Every block reconstructs to exactly its source samples, so predictions +//! read the source frame as if it were the decoder's output. Deblocking +//! and SAO are off so nothing touches the reconstruction afterwards. +//! +//! Stream shape: Main profile, 8-bit 4:2:0, CABAC (HEVC has no other +//! entropy coder), one slice per picture, each P picture referencing the +//! one before it, no reordering. + +use super::bitstream::{self, BitWriter}; +use super::cabac::{CabacWriter, Context}; +use super::encoder::StreamShape; +use super::raster::YuvFrame; + +/// Coding tree block side. The coded frame is whole CTBs (the SPS crops +/// the padding), so no coding unit ever crosses the picture edge. +pub const CTB: usize = 32; +/// Side of the coding units a CTB splits into, and of PCM blocks. +const SPLIT_CU: usize = 16; +const LOG2_MAX_POC_LSB: u8 = 8; +/// Parameter-set ids for the reconnect screen. OBS's HEVC uses 0 for its +/// VPS, SPS and PPS; see `encoder::SPS_ID` for why the screen's own must +/// sit beside them instead of replacing them. +const VPS_ID: u32 = 1; +const SPS_ID: u32 = 1; +const PPS_ID: u32 = 1; +const PROFILE_MAIN: u32 = 1; +/// general_profile_compatibility_flag[1] and [2]: Main, and so Main 10. +const PROFILE_COMPATIBILITY: u32 = 0x6000_0000; +const SLICE_QP: i32 = 26; +const SLICE_TYPE_P: u32 = 1; +const SLICE_TYPE_I: u32 = 2; +/// One merge candidate: skipped blocks need no merge_idx, and with every +/// inter block skipped at zero motion the candidate is always zero motion. +const FIVE_MINUS_MAX_MERGE_CANDIDATES: u32 = 4; + +/// Intra prediction mode numbers (8.4.2) the encoder chooses from. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum IntraMode { + Dc = 1, + Horizontal = 10, + Vertical = 26, +} + +const MODE_PLANAR: u8 = 0; +const MODE_DC: u8 = IntraMode::Dc as u8; +const MODE_VERTICAL: u8 = IntraMode::Vertical as u8; +const INTRA_MODES: [IntraMode; 3] = [IntraMode::Dc, IntraMode::Vertical, IntraMode::Horizontal]; + +impl IntraMode { + /// intra_chroma_pred_mode 0..3 naming this mode (Table 8-2). Planar + /// (0) is never chosen. + fn chroma_index(self) -> u32 { + match self { + IntraMode::Vertical => 1, + IntraMode::Horizontal => 2, + IntraMode::Dc => 3, + } + } +} + +/// VPS, SPS and PPS NAL units (headers included, no length prefix). +pub fn parameter_sets(shape: StreamShape) -> [Vec; 3] { + let level = level_idc(shape); + [ + bitstream::hevc_nal_unit(bitstream::HEVC_NAL_VPS, &vps_rbsp(level)), + bitstream::hevc_nal_unit(bitstream::HEVC_NAL_SPS, &sps_rbsp(shape, level)), + bitstream::hevc_nal_unit(bitstream::HEVC_NAL_PPS, &pps_rbsp()), + ] +} + +fn write_profile_tier_level(w: &mut BitWriter, level: u32) { + w.bits(0, 2); // general_profile_space + w.bit(false); // general_tier_flag: Main tier + w.bits(PROFILE_MAIN, 5); + w.bits(PROFILE_COMPATIBILITY, 32); + w.bit(true); // general_progressive_source_flag + w.bit(false); // general_interlaced_source_flag + w.bit(false); // general_non_packed_constraint_flag + w.bit(true); // general_frame_only_constraint_flag + w.bits(0, 32); // general_reserved_zero_43bits, first 32 + w.bits(0, 11); // and the other 11 + w.bit(false); // general_inbld_flag + w.bits(level, 8); +} + +/// max_dec_pic_buffering_minus1 (the picture being decoded plus its one +/// reference), max_num_reorder_pics, max_latency_increase_plus1. +fn write_sub_layer_ordering(w: &mut BitWriter) { + w.bit(true); // sub_layer_ordering_info_present_flag + w.ue(1); + w.ue(0); + w.ue(0); +} + +fn vps_rbsp(level: u32) -> Vec { + let mut w = BitWriter::new(); + w.bits(VPS_ID, 4); + w.bit(true); // vps_base_layer_internal_flag + w.bit(true); // vps_base_layer_available_flag + w.bits(0, 6); // vps_max_layers_minus1 + w.bits(0, 3); // vps_max_sub_layers_minus1 + w.bit(true); // vps_temporal_id_nesting_flag + w.bits(0xFFFF, 16); // vps_reserved_0xffff_16bits + write_profile_tier_level(&mut w, level); + write_sub_layer_ordering(&mut w); + w.bits(0, 6); // vps_max_layer_id + w.ue(0); // vps_num_layer_sets_minus1 + w.bit(false); // vps_timing_info_present_flag: the SPS VUI has it + w.bit(false); // vps_extension_flag + w.finish() +} + +fn sps_rbsp(shape: StreamShape, level: u32) -> Vec { + let (coded_width, coded_height) = coded_size(shape); + let mut w = BitWriter::new(); + w.bits(VPS_ID, 4); + w.bits(0, 3); // sps_max_sub_layers_minus1 + w.bit(true); // sps_temporal_id_nesting_flag + write_profile_tier_level(&mut w, level); + w.ue(SPS_ID); + w.ue(1); // chroma_format_idc: 4:2:0 + w.ue(coded_width as u32); + w.ue(coded_height as u32); + let crop_right = (coded_width - shape.width) / 2; + let crop_bottom = (coded_height - shape.height) / 2; + let needs_crop = crop_right > 0 || crop_bottom > 0; + w.bit(needs_crop); // conformance_window_flag + if needs_crop { + // Offsets are in chroma samples: 2 luma samples in 4:2:0. + w.ue(0); + w.ue(crop_right as u32); + w.ue(0); + w.ue(crop_bottom as u32); + } + w.ue(0); // bit_depth_luma_minus8 + w.ue(0); // bit_depth_chroma_minus8 + w.ue(LOG2_MAX_POC_LSB as u32 - 4); + write_sub_layer_ordering(&mut w); + w.ue(0); // log2_min_luma_coding_block_size_minus3: 8 + w.ue(2); // log2_diff_max_min_luma_coding_block_size: CTB 32 + w.ue(0); // log2_min_luma_transform_block_size_minus2: 4 + w.ue(3); // log2_diff_max_min_luma_transform_block_size: 32 + w.ue(0); // max_transform_hierarchy_depth_inter + w.ue(0); // max_transform_hierarchy_depth_intra + w.bit(false); // scaling_list_enabled_flag + w.bit(false); // amp_enabled_flag + w.bit(false); // sample_adaptive_offset_enabled_flag + w.bit(true); // pcm_enabled_flag + w.bits(7, 4); // pcm_sample_bit_depth_luma_minus1 + w.bits(7, 4); // pcm_sample_bit_depth_chroma_minus1 + w.ue(1); // log2_min_pcm_luma_coding_block_size_minus3: 16 + w.ue(0); // log2_diff_max_min_pcm_luma_coding_block_size: 16 only + w.bit(true); // pcm_loop_filter_disabled_flag + // One reference picture set: the picture just before. + w.ue(1); // num_short_term_ref_pic_sets + w.ue(1); // num_negative_pics + w.ue(0); // num_positive_pics + w.ue(0); // delta_poc_s0_minus1 + w.bit(true); // used_by_curr_pic_s0_flag + w.bit(false); // long_term_ref_pics_present_flag + w.bit(false); // sps_temporal_mvp_enabled_flag + w.bit(false); // strong_intra_smoothing_enabled_flag + w.bit(true); // vui_parameters_present_flag + write_vui(&mut w, shape.fps); + w.bit(false); // sps_extension_present_flag + w.finish() +} + +/// BT.709 limited range (matches `raster`) and a fixed frame rate, like +/// the H.264 screen's VUI. +fn write_vui(w: &mut BitWriter, fps: u32) { + w.bit(false); // aspect_ratio_info_present_flag + w.bit(false); // overscan_info_present_flag + w.bit(true); // video_signal_type_present_flag + w.bits(5, 3); // video_format: unspecified + w.bit(false); // video_full_range_flag + w.bit(true); // colour_description_present_flag + w.bits(1, 8); // colour_primaries: BT.709 + w.bits(1, 8); // transfer_characteristics: BT.709 + w.bits(1, 8); // matrix_coeffs: BT.709 + w.bit(false); // chroma_loc_info_present_flag + w.bit(false); // neutral_chroma_indication_flag + w.bit(false); // field_seq_flag + w.bit(false); // frame_field_info_present_flag + w.bit(false); // default_display_window_flag + w.bit(true); // vui_timing_info_present_flag + w.bits(1, 32); // vui_num_units_in_tick + w.bits(fps, 32); // vui_time_scale: one tick per frame + w.bit(false); // vui_poc_proportional_to_timing_flag + w.bit(false); // vui_hrd_parameters_present_flag + w.bit(true); // bitstream_restriction_flag + w.bit(false); // tiles_fixed_structure_flag + w.bit(true); // motion_vectors_over_pic_boundaries_flag + w.bit(true); // restricted_ref_pic_lists_flag + w.ue(0); // min_spatial_segmentation_idc + w.ue(0); // max_bytes_per_pic_denom: no limit + w.ue(0); // max_bits_per_min_cu_denom: no limit + w.ue(15); // log2_max_mv_length_horizontal + w.ue(15); // log2_max_mv_length_vertical +} + +fn pps_rbsp() -> Vec { + let mut w = BitWriter::new(); + w.ue(PPS_ID); + w.ue(SPS_ID); + w.bit(false); // dependent_slice_segments_enabled_flag + w.bit(false); // output_flag_present_flag + w.bits(0, 3); // num_extra_slice_header_bits + w.bit(false); // sign_data_hiding_enabled_flag + w.bit(false); // cabac_init_present_flag + w.ue(0); // num_ref_idx_l0_default_active_minus1 + w.ue(0); // num_ref_idx_l1_default_active_minus1 + w.se(0); // init_qp_minus26 + w.bit(false); // constrained_intra_pred_flag + w.bit(false); // transform_skip_enabled_flag + w.bit(false); // cu_qp_delta_enabled_flag + w.se(0); // pps_cb_qp_offset + w.se(0); // pps_cr_qp_offset + w.bit(false); // pps_slice_chroma_qp_offsets_present_flag + w.bit(false); // weighted_pred_flag + w.bit(false); // weighted_bipred_flag + w.bit(false); // transquant_bypass_enabled_flag + w.bit(false); // tiles_enabled_flag + w.bit(false); // entropy_coding_sync_enabled_flag + w.bit(false); // pps_loop_filter_across_slices_enabled_flag + w.bit(true); // deblocking_filter_control_present_flag + w.bit(false); // deblocking_filter_override_enabled_flag + w.bit(true); // pps_deblocking_filter_disabled_flag + w.bit(false); // pps_scaling_list_data_present_flag + w.bit(false); // lists_modification_present_flag + w.ue(0); // log2_parallel_merge_level_minus2 + w.bit(false); // slice_segment_header_extension_present_flag + w.bit(false); // pps_extension_present_flag + w.finish() +} + +/// The visible size rounded up to whole CTBs. +pub fn coded_size(shape: StreamShape) -> (usize, usize) { + ( + shape.width.div_ceil(CTB) * CTB, + shape.height.div_ceil(CTB) * CTB, + ) +} + +/// Smallest level (Table A.8) whose picture size, sample rate and +/// longest side fit; 6.2 when none does. +fn level_idc(shape: StreamShape) -> u32 { + const LEVELS: [(u32, u64, u64); 13] = [ + (30, 36_864, 552_960), + (60, 122_880, 3_686_400), + (63, 245_760, 7_372_800), + (90, 552_960, 16_588_800), + (93, 983_040, 33_177_600), + (120, 2_228_224, 66_846_720), + (123, 2_228_224, 133_693_440), + (150, 8_912_896, 267_386_880), + (153, 8_912_896, 534_773_760), + (156, 8_912_896, 1_069_547_520), + (180, 35_651_584, 1_069_547_520), + (183, 35_651_584, 2_139_095_040), + (186, 35_651_584, 4_278_190_080), + ]; + let (width, height) = coded_size(shape); + let picture = (width * height) as u64; + let longest = width.max(height) as u64; + let rate = picture * u64::from(shape.fps); + LEVELS + .iter() + .find(|(_, max_picture, max_rate)| { + picture <= *max_picture && rate <= *max_rate && longest * longest <= max_picture * 8 + }) + .map_or(186, |(level, _, _)| *level) +} + +/// IDR picture: every block intra or PCM. Consecutive HEVC IDRs need no +/// distinguishing id, so a replayed loop can reuse it as is. +pub fn encode_idr(frame: &YuvFrame) -> Vec { + let mut w = BitWriter::new(); + w.bit(true); // first_slice_segment_in_pic_flag + w.bit(false); // no_output_of_prior_pics_flag + w.ue(PPS_ID); + w.ue(SLICE_TYPE_I); + write_slice_tail(&mut w); + let data = PictureCoder::new(frame, None, w).code(); + bitstream::hevc_nal_unit(bitstream::HEVC_NAL_IDR_N_LP, &data) +} + +/// P picture `poc` pictures after the IDR: blocks identical to +/// `previous` are skipped, the rest are intra coded. Every P picture is a +/// reference so the next one can skip against it. +pub fn encode_p(frame: &YuvFrame, previous: &YuvFrame, poc: u32) -> Vec { + let mut w = BitWriter::new(); + w.bit(true); // first_slice_segment_in_pic_flag + w.ue(PPS_ID); + w.ue(SLICE_TYPE_P); + w.bits(poc % (1 << LOG2_MAX_POC_LSB), LOG2_MAX_POC_LSB); + w.bit(true); // short_term_ref_pic_set_sps_flag: the SPS's only set + w.bit(false); // num_ref_idx_active_override_flag + w.ue(FIVE_MINUS_MAX_MERGE_CANDIDATES); + write_slice_tail(&mut w); + let data = PictureCoder::new(frame, Some(previous), w).code(); + bitstream::hevc_nal_unit(bitstream::HEVC_NAL_TRAIL_R, &data) +} + +/// slice_qp_delta, then byte_alignment() so slice data starts on a byte. +fn write_slice_tail(w: &mut BitWriter) { + w.se(0); // slice_qp_delta + w.bit(true); // alignment_bit_equal_to_one + w.align_with_zeros(); +} + +/// A square block of luma samples; its chroma is half the side. +#[derive(Debug, Clone, Copy)] +struct Block { + x: usize, + y: usize, + size: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CuCoding { + Skip, + Predicted { luma: IntraMode, chroma: IntraMode }, + Pcm, +} + +/// What later blocks' syntax reads about each 16x16 area. +#[derive(Debug, Clone, Copy, Default)] +struct AreaState { + /// Covered by a 16x16 coding unit (depth 1) rather than a 32x32 one. + is_split: bool, + is_skipped: bool, + /// Luma mode of an intra, non-PCM coding unit. + intra_mode: Option, +} + +/// Context variables for the syntax elements the encoder writes, by +/// ctxInc. initValues are from H.265 Tables 9-5 to 9-37: initType 0 for +/// I slices, 1 for P slices. +struct Contexts { + split_cu: [Context; 3], + cu_skip: [Context; 3], + pred_mode: Context, + prev_intra_luma_pred: Context, + intra_chroma_pred_mode: Context, + /// ctxInc 1: the transform tree is never split (trafoDepth 0). + cbf_luma: Context, + /// ctxInc 0 (trafoDepth 0), shared by cbf_cb and cbf_cr. + cbf_chroma: Context, +} + +impl Contexts { + fn new(is_p_slice: bool) -> Self { + let pick = + |for_i: u8, for_p: u8| Context::new(if is_p_slice { for_p } else { for_i }, SLICE_QP); + Self { + split_cu: [pick(139, 107), pick(141, 139), pick(157, 126)], + // Only coded in P slices. + cu_skip: [197, 185, 201].map(|value| Context::new(value, SLICE_QP)), + pred_mode: Context::new(149, SLICE_QP), + prev_intra_luma_pred: pick(184, 154), + intra_chroma_pred_mode: pick(63, 152), + cbf_luma: pick(141, 111), + cbf_chroma: pick(94, 149), + } + } +} + +/// Codes one picture's slice data, CTU by CTU in raster order. +struct PictureCoder<'a> { + frame: &'a YuvFrame, + /// The reference picture: P slices only. + previous: Option<&'a YuvFrame>, + cabac: CabacWriter, + contexts: Contexts, + areas_wide: usize, + areas: Vec, +} + +impl<'a> PictureCoder<'a> { + fn new(frame: &'a YuvFrame, previous: Option<&'a YuvFrame>, header: BitWriter) -> Self { + let areas_wide = frame.width / SPLIT_CU; + Self { + frame, + previous, + cabac: CabacWriter::new(header), + contexts: Contexts::new(previous.is_some()), + areas_wide, + areas: vec![AreaState::default(); areas_wide * (frame.height / SPLIT_CU)], + } + } + + fn code(mut self) -> Vec { + let ctbs_wide = self.frame.width / CTB; + let ctb_count = ctbs_wide * (self.frame.height / CTB); + for index in 0..ctb_count { + self.code_ctu(index % ctbs_wide * CTB, index / ctbs_wide * CTB); + if index + 1 < ctb_count { + self.cabac.terminate(false); // end_of_slice_segment_flag + } + } + self.cabac.finish() + } + + fn code_ctu(&mut self, x: usize, y: usize) { + let whole = Block { x, y, size: CTB }; + let coding = choose_coding(self.frame, self.previous, whole); + // split_cu_flag: ctxInc counts neighbours coded deeper (16x16). + let ctx_inc = usize::from(x > 0 && self.area(x - 1, y).is_split) + + usize::from(y > 0 && self.area(x, y - 1).is_split); + self.cabac + .decision(&mut self.contexts.split_cu[ctx_inc], coding.is_none()); + if let Some(coding) = coding { + self.code_cu(whole, coding); + return; + } + for (dx, dy) in [(0, 0), (SPLIT_CU, 0), (0, SPLIT_CU), (SPLIT_CU, SPLIT_CU)] { + let block = Block { + x: x + dx, + y: y + dy, + size: SPLIT_CU, + }; + // Nothing is deeper than 16x16, so its split_cu_flag is + // always 0 with ctxInc 0. + self.cabac.decision(&mut self.contexts.split_cu[0], false); + let coding = choose_coding(self.frame, self.previous, block) + .expect("16x16 blocks can always fall back to PCM"); + self.code_cu(block, coding); + } + } + + fn code_cu(&mut self, block: Block, coding: CuCoding) { + if self.previous.is_some() { + let ctx_inc = usize::from(block.x > 0 && self.area(block.x - 1, block.y).is_skipped) + + usize::from(block.y > 0 && self.area(block.x, block.y - 1).is_skipped); + let is_skip = coding == CuCoding::Skip; + self.cabac + .decision(&mut self.contexts.cu_skip[ctx_inc], is_skip); + if !is_skip { + // pred_mode_flag: MODE_INTRA. + self.cabac.decision(&mut self.contexts.pred_mode, true); + } + } + let is_split = block.size == SPLIT_CU; + let intra_mode = match coding { + CuCoding::Skip => None, + CuCoding::Pcm => { + self.cabac.pcm(&pcm_samples(self.frame, block)); + None + } + CuCoding::Predicted { luma, chroma } => { + // Only 16x16 coding units carry a pcm_flag. + if is_split { + self.cabac.terminate(false); + } + self.code_intra_modes(block, luma, chroma); + // cbf_cb, cbf_cr, cbf_luma: no residual anywhere. + self.cabac.decision(&mut self.contexts.cbf_chroma, false); + self.cabac.decision(&mut self.contexts.cbf_chroma, false); + self.cabac.decision(&mut self.contexts.cbf_luma, false); + Some(luma) + } + }; + let state = AreaState { + is_split, + is_skipped: coding == CuCoding::Skip, + intra_mode, + }; + for y in (block.y..block.y + block.size).step_by(SPLIT_CU) { + for x in (block.x..block.x + block.size).step_by(SPLIT_CU) { + let index = self.area_index(x, y); + self.areas[index] = state; + } + } + } + + /// The luma mode as a most-probable-mode index or a remainder, then + /// the chroma mode, which costs one bin when it repeats the luma one. + fn code_intra_modes(&mut self, block: Block, luma: IntraMode, chroma: IntraMode) { + let candidates = self.most_probable_modes(block); + let luma_number = luma as u8; + match candidates.iter().position(|&mode| mode == luma_number) { + Some(mpm_idx) => { + self.cabac + .decision(&mut self.contexts.prev_intra_luma_pred, true); + // Truncated rice, cMax 2: 0, 10, 11. + match mpm_idx { + 0 => self.cabac.bypass_bits(0, 1), + 1 => self.cabac.bypass_bits(0b10, 2), + _ => self.cabac.bypass_bits(0b11, 2), + } + } + None => { + self.cabac + .decision(&mut self.contexts.prev_intra_luma_pred, false); + let below = candidates + .iter() + .filter(|&&mode| mode < luma_number) + .count(); + self.cabac + .bypass_bits(u32::from(luma_number) - below as u32, 5); + } + } + if chroma == luma { + self.cabac + .decision(&mut self.contexts.intra_chroma_pred_mode, false); + } else { + self.cabac + .decision(&mut self.contexts.intra_chroma_pred_mode, true); + self.cabac.bypass_bits(chroma.chroma_index(), 2); + } + } + + /// candModeList (8.4.2) from the blocks left of and above `block`. + /// The one above only counts inside the same CTB; PCM, skipped and + /// missing neighbours count as DC. + fn most_probable_modes(&self, block: Block) -> [u8; 3] { + let mode_at = |x: usize, y: usize| self.area(x, y).intra_mode.map_or(MODE_DC, |m| m as u8); + let left = if block.x > 0 { + mode_at(block.x - 1, block.y) + } else { + MODE_DC + }; + let above = if !block.y.is_multiple_of(CTB) { + mode_at(block.x, block.y - 1) + } else { + MODE_DC + }; + if left != above { + let third = if left != MODE_PLANAR && above != MODE_PLANAR { + MODE_PLANAR + } else if left != MODE_DC && above != MODE_DC { + MODE_DC + } else { + MODE_VERTICAL + }; + return [left, above, third]; + } + if left < 2 { + return [MODE_PLANAR, MODE_DC, MODE_VERTICAL]; + } + // The two angular modes beside it. + [left, 2 + ((left + 29) % 32), 2 + ((left - 2 + 1) % 32)] + } + + fn area_index(&self, x: usize, y: usize) -> usize { + (y / SPLIT_CU) * self.areas_wide + x / SPLIT_CU + } + + fn area(&self, x: usize, y: usize) -> AreaState { + self.areas[self.area_index(x, y)] + } +} + +/// Skip when unchanged, else an exact intra prediction, else PCM where +/// the size allows it; `None` means the block must split. +fn choose_coding(frame: &YuvFrame, previous: Option<&YuvFrame>, block: Block) -> Option { + if previous.is_some_and(|previous| is_unchanged(frame, previous, block)) { + return Some(CuCoding::Skip); + } + if let Some((luma, chroma)) = exact_prediction(frame, block) { + return Some(CuCoding::Predicted { luma, chroma }); + } + (block.size == SPLIT_CU).then_some(CuCoding::Pcm) +} + +/// Luma and chroma modes that reproduce the block exactly. Chroma +/// prefers the luma mode, which is the cheapest to signal. +fn exact_prediction(frame: &YuvFrame, block: Block) -> Option<(IntraMode, IntraMode)> { + let luma = PlaneBlock::luma(frame, block); + let luma_mode = INTRA_MODES + .into_iter() + .find(|mode| luma.is_reproduced_by(*mode))?; + let chroma = PlaneBlock::chroma(frame, block); + let chroma_reproduced = + |mode: IntraMode| chroma.iter().all(|plane| plane.is_reproduced_by(mode)); + let chroma_mode = std::iter::once(luma_mode) + .chain(INTRA_MODES) + .find(|mode| chroma_reproduced(*mode))?; + Some((luma_mode, chroma_mode)) +} + +/// One plane's square block, and whether intra prediction filters its +/// edges (luma below 32x32, 8.4.4.2.6). +struct PlaneBlock<'a> { + samples: &'a [u8], + stride: usize, + x0: usize, + y0: usize, + size: usize, + filters_edges: bool, +} + +impl<'a> PlaneBlock<'a> { + fn luma(frame: &'a YuvFrame, block: Block) -> Self { + Self { + samples: &frame.y, + stride: frame.width, + x0: block.x, + y0: block.y, + size: block.size, + filters_edges: block.size < 32, + } + } + + fn chroma(frame: &'a YuvFrame, block: Block) -> [Self; 2] { + [&frame.u, &frame.v].map(|samples| Self { + samples, + stride: frame.width / 2, + x0: block.x / 2, + y0: block.y / 2, + size: block.size / 2, + filters_edges: false, + }) + } + + fn at(&self, x: usize, y: usize) -> i32 { + i32::from(self.samples[y * self.stride + x]) + } + + /// The reference samples DC, vertical and horizontal prediction read, + /// after the substitution of 8.4.4.2.2. Left and above are always + /// decoded when inside the picture; a missing side copies the nearest + /// sample of the other, and with neither every sample is 128. + fn references(&self) -> References { + let (x0, y0, size) = (self.x0, self.y0, self.size); + let read_top = || { + (0..size) + .map(|i| self.at(x0 + i, y0 - 1)) + .collect::>() + }; + let read_left = || { + (0..size) + .map(|i| self.at(x0 - 1, y0 + i)) + .collect::>() + }; + match (x0 > 0, y0 > 0) { + (true, true) => References { + top: read_top(), + left: read_left(), + corner: self.at(x0 - 1, y0 - 1), + }, + (true, false) => { + let left = read_left(); + References::filled(left[0], size, |references| references.left = left) + } + (false, true) => { + let top = read_top(); + References::filled(top[0], size, |references| references.top = top) + } + (false, false) => References::filled(128, size, |_| {}), + } + } + + fn is_reproduced_by(&self, mode: IntraMode) -> bool { + let references = self.references(); + let dc = references.dc(); + (0..self.size).all(|y| { + (0..self.size).all(|x| { + let predicted = references.predict(mode, dc, x, y, self.filters_edges); + self.at(self.x0 + x, self.y0 + y) == predicted + }) + }) + } +} + +struct References { + top: Vec, + left: Vec, + /// The sample above-left, p[-1][-1]. + corner: i32, +} + +impl References { + /// Every sample `value`, then `real` puts back the side that exists. + fn filled(value: i32, size: usize, real: impl FnOnce(&mut Self)) -> Self { + let mut references = Self { + top: vec![value; size], + left: vec![value; size], + corner: value, + }; + real(&mut references); + references + } + + /// dcVal (8.4.4.2.5): the mean of both edges, rounded. + fn dc(&self) -> i32 { + let size = self.top.len(); + let sum: i32 = self.top.iter().chain(&self.left).sum(); + (sum + size as i32) >> (size.trailing_zeros() + 1) + } + + /// predSamples[x][y] (8.4.4.2.5 and 8.4.4.2.6), with the edge filters + /// luma blocks below 32x32 get. + fn predict(&self, mode: IntraMode, dc: i32, x: usize, y: usize, filters_edges: bool) -> i32 { + let clip = |value: i32| value.clamp(0, 255); + match (mode, filters_edges) { + (IntraMode::Dc, true) if x == 0 && y == 0 => { + (self.left[0] + 2 * dc + self.top[0] + 2) >> 2 + } + (IntraMode::Dc, true) if y == 0 => (self.top[x] + 3 * dc + 2) >> 2, + (IntraMode::Dc, true) if x == 0 => (self.left[y] + 3 * dc + 2) >> 2, + (IntraMode::Dc, _) => dc, + (IntraMode::Vertical, true) if x == 0 => { + clip(self.top[0] + ((self.left[y] - self.corner) >> 1)) + } + (IntraMode::Vertical, _) => self.top[x], + (IntraMode::Horizontal, true) if y == 0 => { + clip(self.left[0] + ((self.top[x] - self.corner) >> 1)) + } + (IntraMode::Horizontal, _) => self.left[y], + } + } +} + +/// pcm_sample(): the luma block, then Cb, then Cr, row by row. +fn pcm_samples(frame: &YuvFrame, block: Block) -> Vec { + let mut samples = Vec::with_capacity(block.size * block.size * 3 / 2); + for row in 0..block.size { + let start = (block.y + row) * frame.width + block.x; + samples.extend_from_slice(&frame.y[start..start + block.size]); + } + let (chroma_width, chroma_size) = (frame.width / 2, block.size / 2); + for plane in [&frame.u, &frame.v] { + for row in 0..chroma_size { + let start = (block.y / 2 + row) * chroma_width + block.x / 2; + samples.extend_from_slice(&plane[start..start + chroma_size]); + } + } + samples +} + +fn is_unchanged(frame: &YuvFrame, previous: &YuvFrame, block: Block) -> bool { + let rows_match = + |plane: &[u8], before: &[u8], stride: usize, x: usize, y: usize, size: usize| { + (0..size).all(|row| { + let start = (y + row) * stride + x; + plane[start..start + size] == before[start..start + size] + }) + }; + let (chroma_width, half) = (frame.width / 2, block.size / 2); + rows_match( + &frame.y, + &previous.y, + frame.width, + block.x, + block.y, + block.size, + ) && rows_match( + &frame.u, + &previous.u, + chroma_width, + block.x / 2, + block.y / 2, + half, + ) && rows_match( + &frame.v, + &previous.v, + chroma_width, + block.x / 2, + block.y / 2, + half, + ) +} + +#[cfg(test)] +mod tests { + use super::super::raster::{Canvas, Rgb}; + use super::*; + + const BACKGROUND: Rgb = Rgb::new(0x0e, 0x0f, 0x12); + + fn flat_frame(width: usize, height: usize) -> YuvFrame { + Canvas::new(width, height, BACKGROUND).to_yuv420_padded(CTB) + } + + /// Strips the two-byte NAL header and emulation prevention, then reads + /// the first ue(v) fields, as a decoder finds the ids. + fn leading_fields(nal: &[u8], skip_bits: usize, count: usize) -> Vec { + let mut rbsp = Vec::new(); + let mut zeros = 0; + for &byte in &nal[2..] { + if zeros >= 2 && byte == 3 { + zeros = 0; + continue; + } + rbsp.push(byte); + zeros = if byte == 0 { zeros + 1 } else { 0 }; + } + let bits: Vec = rbsp + .iter() + .flat_map(|byte| (0..8).rev().map(move |shift| (byte >> shift) & 1 == 1)) + .collect(); + let mut at = skip_bits; + (0..count) + .map(|_| { + let zeros = bits[at..].iter().take_while(|bit| !**bit).count(); + let value = bits[at + zeros..at + 2 * zeros + 1] + .iter() + .fold(0, |value, bit| (value << 1) | u32::from(*bit)); + at += 2 * zeros + 1; + value - 1 + }) + .collect() + } + + /// Like the H.264 screen, the HEVC one must leave id 0 to the stream, + /// or switching to it overwrites the stream's parameter sets in every + /// decoder. The SPS id sits after the 4+3+1 header bits and the 96-bit + /// profile_tier_level; the PPS names the SPS; slices name the PPS. + #[test] + fn the_screen_never_reuses_the_parameter_set_ids_of_the_stream() { + let [vps, sps, pps] = parameter_sets(StreamShape { + width: 1920, + height: 1080, + fps: 30, + }); + assert_eq!(vps[0] >> 1, bitstream::HEVC_NAL_VPS); + assert_eq!(vps[2] >> 4, VPS_ID as u8, "the VPS id is its first 4 bits"); + assert_eq!(sps[0] >> 1, bitstream::HEVC_NAL_SPS); + assert_eq!( + sps[2] >> 4, + VPS_ID as u8, + "the SPS points at the screen's VPS" + ); + assert_eq!(leading_fields(&sps, 8 + 96, 1), vec![SPS_ID]); + assert_eq!(leading_fields(&pps, 0, 2), vec![PPS_ID, SPS_ID]); + for id in [VPS_ID, SPS_ID, PPS_ID] { + assert_ne!(id, 0, "id 0 belongs to the stream"); + } + + let frame = flat_frame(64, 64); + let idr = encode_idr(&frame); + // first_slice_segment_in_pic_flag, no_output_of_prior_pics_flag. + assert_eq!(leading_fields(&idr, 2, 1), vec![PPS_ID]); + let p = encode_p(&frame, &frame, 1); + assert_eq!(leading_fields(&p, 1, 1), vec![PPS_ID]); + } + + #[test] + fn level_is_the_smallest_in_table_a8_that_fits() { + for (width, height, fps, level) in [ + (320, 180, 30, 60), + (1280, 720, 30, 93), + (1920, 1080, 30, 120), + (1080, 1920, 30, 120), + (1920, 1080, 60, 123), + (2560, 1440, 30, 150), + (3840, 2160, 60, 153), + (8192, 16, 30, 150), + ] { + assert_eq!( + level_idc(StreamShape { width, height, fps }), + level, + "{width}x{height} at {fps} fps" + ); + } + } + + /// The most probable modes the decoder derives, worked from 8.4.2: + /// a DC (or missing) left neighbour gives planar, DC, vertical; an + /// angular one leads its own list; two equal angular modes bring their + /// neighbours. + #[test] + fn most_probable_modes_follow_the_spec() { + let frame = flat_frame(64, 64); + let mut coder = PictureCoder::new(&frame, None, BitWriter::new()); + let block = Block { + x: 16, + y: 16, + size: SPLIT_CU, + }; + assert_eq!(coder.most_probable_modes(block), [0, 1, 26]); + let index = coder.area_index(0, 16); + coder.areas[index].intra_mode = Some(IntraMode::Horizontal); + assert_eq!(coder.most_probable_modes(block), [10, 1, 0]); + let index = coder.area_index(16, 0); + coder.areas[index].intra_mode = Some(IntraMode::Horizontal); + assert_eq!(coder.most_probable_modes(block), [10, 9, 11]); + // Above in another CTB counts as DC. + let top_of_ctb = Block { + x: 32, + y: 32, + size: CTB, + }; + let index = coder.area_index(32, 16); + coder.areas[index].intra_mode = Some(IntraMode::Vertical); + assert_eq!(coder.most_probable_modes(top_of_ctb), [0, 1, 26]); + } + + /// An unchanged picture is one skip per CTU: a few bytes at 1080p, + /// which is what keeps the held screen cheap. + #[test] + fn an_unchanged_p_picture_is_tiny() { + let still = flat_frame(1920, 1080); + let unchanged = encode_p(&still, &still, 1); + assert!( + unchanged.len() < 64, + "{} bytes for an unchanged 1080p picture", + unchanged.len() + ); + } +} diff --git a/src/slate/mod.rs b/src/slate/mod.rs index 0a36930..6fcc27d 100644 --- a/src/slate/mod.rs +++ b/src/slate/mod.rs @@ -1,14 +1,16 @@ //! The crash-protection reconnect screen. //! -//! Draws the streamer's chosen theme and encodes it as a 2 s H.264 loop: -//! one keyframe plus P frames, built once per stream shape and replayed -//! with fresh timestamps for as long as the hold lasts. Keeping the loop +//! Draws the streamer's chosen theme and encodes it as a 2 s H.264 or HEVC +//! loop: one keyframe plus P frames, built once per stream shape and codec +//! and replayed with fresh timestamps for as long as the hold lasts. Keeping the loop //! exactly one keyframe interval long means every replay starts on an //! IDR, so viewers who join mid-hold get a picture within 2 s. mod bitstream; +mod cabac; mod encoder; mod font; +mod hevc; mod pixel_font; mod png; mod raster; @@ -27,15 +29,40 @@ const MIN_SIDE: usize = 16; const MAX_SIDE: usize = 8192; const MAX_FPS: u32 = 240; +/// The codec a loop is encoded in: the one of the track it stands in for. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum SlateCodec { + H264, + Hevc, +} + +impl SlateCodec { + pub fn label(self) -> &'static str { + match self { + SlateCodec::H264 => "H.264", + SlateCodec::Hevc => "HEVC", + } + } + + /// The square the encoder codes in, which frames are padded to. + fn block_size(self) -> usize { + match self { + SlateCodec::H264 => 16, + SlateCodec::Hevc => hevc::CTB, + } + } +} + /// One encoded loop, ready to be wrapped in RTMP video tags. It needs no -/// sequence header of its own: every keyframe carries the loop's SPS and -/// PPS in-band, under ids the stream it interrupts doesn't use. +/// sequence header of its own: every keyframe carries the loop's +/// parameter sets in-band, under ids the stream it interrupts doesn't use. pub struct SlateLoop { - /// The loop's first picture as an AVCC sample (SPS, PPS, IDR), encoded - /// twice with idr_pic_id 0 and 1. Replays alternate them, because - /// back-to-back IDRs must not share an id. + /// The loop's first picture as a length-prefixed sample (H.264: SPS, + /// PPS, IDR; HEVC: VPS, SPS, PPS, IDR), twice. Replays alternate them: + /// back-to-back H.264 IDRs must not share an idr_pic_id, so those two + /// differ; HEVC IDRs carry no id, so those two are the same. pub keyframes: [Vec; 2], - /// The rest of the loop (P pictures) as AVCC samples, in order. + /// The rest of the loop (P pictures) as length-prefixed samples. pub deltas: Vec>, pub shape: StreamShape, } @@ -96,6 +123,7 @@ fn screen_style(settings: &CrashProtection) -> ScreenStyle<'_> { pub fn build_loop( settings: &CrashProtection, shape: StreamShape, + codec: SlateCodec, ) -> Result { let side_ok = |side: usize| (MIN_SIDE..=MAX_SIDE).contains(&side) && side.is_multiple_of(2); if !side_ok(shape.width) || !side_ok(shape.height) || !(1..=MAX_FPS).contains(&shape.fps) { @@ -109,19 +137,32 @@ pub fn build_loop( canvas }; - let parameter_sets = encoder::parameter_sets(shape); let mut previous_canvas = render(0); - let mut previous = previous_canvas.to_yuv420(); - let keyframes = [0, 1].map(|idr_pic_id| { - let idr = encoder::encode_idr(&previous, idr_pic_id); - bitstream::length_prefixed(&[¶meter_sets.sps, ¶meter_sets.pps, &idr]) - }); + let mut previous = previous_canvas.to_yuv420_padded(codec.block_size()); + let keyframes = match codec { + SlateCodec::H264 => { + let sets = encoder::parameter_sets(shape); + [0, 1].map(|idr_pic_id| { + let idr = encoder::encode_idr(&previous, idr_pic_id); + bitstream::length_prefixed(&[&sets.sps, &sets.pps, &idr]) + }) + } + SlateCodec::Hevc => { + let [vps, sps, pps] = hevc::parameter_sets(shape); + let idr = hevc::encode_idr(&previous); + let keyframe = bitstream::length_prefixed(&[&vps, &sps, &pps, &idr]); + [keyframe.clone(), keyframe] + } + }; let mut deltas = Vec::with_capacity(frame_count - 1); for index in 1..frame_count { let canvas = render(index); // Only the animated rows change between frames. let frame = canvas.to_yuv420_after(&previous_canvas, &previous); - let picture = encoder::encode_p(&frame, &previous, index as u32); + let picture = match codec { + SlateCodec::H264 => encoder::encode_p(&frame, &previous, index as u32), + SlateCodec::Hevc => hevc::encode_p(&frame, &previous, index as u32), + }; deltas.push(bitstream::length_prefixed(&[&picture])); (previous_canvas, previous) = (canvas, frame); } @@ -154,6 +195,42 @@ pub fn test_sequence_header(width: usize, height: usize) -> Vec { tag } +/// An HEVCDecoderConfigurationRecord (hvcC) for a `width` x `height` +/// 8-bit 4:2:0 stream with 4-byte NAL lengths, for tests that need HEVC +/// video the reconnect screen can cover. +#[cfg(test)] +pub fn test_hevc_config(width: usize, height: usize) -> Vec { + let sets = hevc::parameter_sets(StreamShape { + width, + height, + fps: 30, + }); + // Version 1, Main profile, Main + Main 10 compatible, progressive and + // frame-only, level 4. + let mut record = vec![1, 0x01, 0x60, 0, 0, 0, 0x90, 0, 0, 0, 0, 0, 120]; + // No segmentation or parallelism, 4:2:0, 8-bit luma and chroma, no + // frame rate, one temporal layer, 4-byte NAL lengths, three arrays. + record.extend_from_slice(&[0xF0, 0x00, 0xFC, 0xFD, 0xF8, 0xF8, 0, 0, 0x0F, 3]); + for nal in &sets { + // array_completeness and the NAL type, then one NAL unit. + record.push(0x80 | (nal[0] >> 1)); + record.extend_from_slice(&1u16.to_be_bytes()); + record.extend_from_slice(&(nal.len() as u16).to_be_bytes()); + record.extend_from_slice(nal); + } + record +} + +/// Tests run ffmpeg one at a time: several at once load the machine +/// enough to fail the wall-clock simulations (`controller::user_sim`) +/// running beside them. Hold the guard for the whole ffmpeg run. +#[cfg(test)] +fn ffmpeg_turn() -> std::sync::MutexGuard<'static, ()> { + static TURN: std::sync::Mutex<()> = std::sync::Mutex::new(()); + // A test that panicked holding the turn leaves nothing to protect. + TURN.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) +} + /// Whether ffmpeg runs here, for the tests that decode with it. Missing /// on a dev machine skips them with a note; missing on CI (`CI` set) /// fails, so a runner without ffmpeg can't turn them into silent passes. @@ -185,6 +262,8 @@ mod tests { use crate::crash_protection::SlateTheme; use std::process::Command; + const CODECS: [SlateCodec; 2] = [SlateCodec::H264, SlateCodec::Hevc]; + fn settings(theme: SlateTheme) -> CrashProtection { CrashProtection { theme, @@ -199,9 +278,11 @@ mod tests { height: 180, fps: 30, }; - let slate = build_loop(&settings(SlateTheme::Arcade), shape).unwrap(); + let slate = build_loop(&settings(SlateTheme::Arcade), shape, SlateCodec::H264).unwrap(); assert_eq!(1 + slate.deltas.len(), 60); assert_ne!(slate.keyframes[0], slate.keyframes[1], "idr_pic_id differs"); + let hevc = build_loop(&settings(SlateTheme::Arcade), shape, SlateCodec::Hevc).unwrap(); + assert_eq!(1 + hevc.deltas.len(), 60); } #[test] @@ -211,18 +292,30 @@ mod tests { height: 180, fps: 10, }; - let slate = build_loop(&settings(SlateTheme::Whisper), shape).unwrap(); - let types = |sample: &[u8]| { - nal_units(sample) - .iter() - .map(|nal| nal[0] & 0x1F) - .collect::>() - }; - for keyframe in &slate.keyframes { - assert_eq!(types(keyframe), vec![7, 8, 5], "SPS, PPS, IDR"); - } - for delta in &slate.deltas { - assert_eq!(types(delta), vec![1], "P frames are just the slice"); + for (codec, keyframe_types) in [ + // SPS, PPS, IDR. + (SlateCodec::H264, vec![7, 8, 5]), + // VPS, SPS, PPS, IDR_N_LP. + (SlateCodec::Hevc, vec![32, 33, 34, 20]), + ] { + let slate = build_loop(&settings(SlateTheme::Whisper), shape, codec).unwrap(); + let types = |sample: &[u8]| { + nal_units(sample) + .iter() + .map(|nal| nal_type(codec, nal)) + .collect::>() + }; + for keyframe in &slate.keyframes { + assert_eq!(types(keyframe), keyframe_types, "{codec:?} keyframe"); + } + for delta in &slate.deltas { + // A non-IDR slice in H.264, TRAIL_R in HEVC: both type 1. + assert_eq!( + types(delta), + vec![1], + "{codec:?} P frames are just the slice" + ); + } } } @@ -255,29 +348,32 @@ mod tests { for (width, height, fps) in [(321, 180, 30), (320, 8, 30), (320, 180, 0), (9000, 180, 30)] { let shape = StreamShape { width, height, fps }; assert_eq!( - build_loop(&base, shape).err(), + build_loop(&base, shape, SlateCodec::H264).err(), Some(UnsupportedShape(shape)) ); } } - /// Every theme's loop at `width` x `height`, built at the rate a hold - /// really uses, stays under `max_kbps`. + /// Every theme's loop at `width` x `height`, in both codecs, built at + /// the rate a hold really uses, stays under `max_kbps`. fn assert_every_theme_fits(width: usize, height: usize, max_kbps: usize) { let shape = StreamShape { width, height, fps: crate::crash_hold::SLATE_FPS, }; - for theme in SlateTheme::ALL { - let slate = build_loop(&settings(theme), shape).unwrap(); + let every_pair = SlateTheme::ALL + .into_iter() + .flat_map(|theme| CODECS.map(|codec| (theme, codec))); + for (theme, codec) in every_pair { + let slate = build_loop(&settings(theme), shape, codec).unwrap(); let loop_bytes: usize = slate.keyframes[0].len() + slate.deltas.iter().map(Vec::len).sum::(); // Bits per millisecond are kbps. let kbps = loop_bytes * 8 / LOOP_MS; assert!( kbps < max_kbps, - "{theme:?} loop at {width}x{height} runs at {kbps} kbps" + "{theme:?} {codec:?} loop at {width}x{height} runs at {kbps} kbps" ); } } @@ -306,34 +402,47 @@ mod tests { } /// Decode two replays of the loop with ffmpeg and require every pixel - /// to match what we rendered. Every theme at 320x180 (180 isn't whole - /// macroblocks), then one portrait 200x360 so a real decoder checks the - /// cropping of the width too; each ffmpeg run costs about half a second. + /// to match what we rendered, in both codecs. Every theme at 320x180 + /// (180 isn't whole macroblocks or CTBs), then one portrait 200x360 so + /// a real decoder checks the cropping of the width too, then one + /// 640x360 so larger text meets every block coding. Sizes stay small: + /// at 720p the raw video written to the temp dir slows the wall-clock + /// simulations (`controller::user_sim`) running beside this test + /// enough to fail them. Each ffmpeg run costs about half a second. #[test] fn ffmpeg_decodes_the_loop_losslessly() { if !ffmpeg_is_available("the reconnect screen decode round trip") { return; } let landscape = SlateTheme::ALL.map(|theme| (theme, 320, 180)); - let portrait = (SlateTheme::Studio, 200, 360); - for (theme, width, height) in landscape.into_iter().chain([portrait]) { - let shape = StreamShape { - width, - height, - fps: 10, - }; - let slate = build_loop(&settings(theme), shape).unwrap(); - let decoded = decode_with_ffmpeg(&annex_b(&slate), shape.fps, theme); - let expected = expected_frames(&settings(theme), shape); - assert_eq!( - decoded.len(), - expected.len(), - "{theme:?} at {width}x{height}: frame count" - ); - assert!( - decoded == expected, - "{theme:?} at {width}x{height}: decoded pixels differ from the render" - ); + let others = [ + (SlateTheme::Studio, 200, 360), + (SlateTheme::Arcade, 640, 360), + ]; + for codec in CODECS { + for (theme, width, height) in landscape.into_iter().chain(others) { + let shape = StreamShape { + width, + height, + fps: 10, + }; + let slate = build_loop(&settings(theme), shape, codec).unwrap(); + let decoded = decode_with_ffmpeg(&annex_b(&slate), shape.fps, codec, theme); + let expected = expected_frames(&settings(theme), shape); + let what = format!("{theme:?} {codec:?} at {width}x{height}"); + assert_eq!(decoded.len(), expected.len(), "{what}: frame count"); + assert!( + decoded == expected, + "{what}: decoded pixels differ from the render" + ); + } + } + } + + fn nal_type(codec: SlateCodec, nal: &[u8]) -> u8 { + match codec { + SlateCodec::H264 => nal[0] & 0x1F, + SlateCodec::Hevc => (nal[0] >> 1) & 0x3F, } } @@ -364,41 +473,259 @@ mod tests { units } - /// Two loop replays as an Annex B elementary stream. - fn annex_b(slate: &SlateLoop) -> Vec { + /// What a destination's decoder sees through a hold: the stream, whose + /// parameter sets (id 0) came once at its start; one loop of the screen + /// with its own in-band; then the stream again from a keyframe that + /// carries none. Every frame must decode without an error, the screen + /// exactly as rendered, and the stream after it exactly as it decodes + /// alone: the screen must not disturb the stream's parameter sets. + /// x264 and x265 stand in for OBS's encoder. + #[test] + fn a_stream_cut_to_the_screen_and_back_decodes_cleanly() { + if !ffmpeg_is_available("the cut to the reconnect screen and back") { + return; + } + let shape = StreamShape { + width: 320, + height: 180, + fps: 10, + }; + let theme = SlateTheme::Beacon; + let frame_bytes = shape.width * shape.height * 3 / 2; + for codec in CODECS { + let Some(stream) = encode_stand_in_stream(codec, shape) else { + continue; + }; + let cut = second_keyframe_offset(&stream, codec); + let before = &stream[..cut]; + // OBS's keyframes reach a destination without parameter sets: + // drop any the encoder repeated. + let after = without_parameter_sets(&stream[cut..], codec); + let slate = build_loop(&settings(theme), shape, codec).unwrap(); + let screen = annex_b_samples(std::iter::once(&slate.keyframes[0]).chain(&slate.deltas)); + let spliced = [before, &screen[..], &after[..]].concat(); + + let continuous = [before, &after[..]].concat(); + let alone = decode_with_ffmpeg(&continuous, shape.fps, codec, theme); + let decoded = decode_with_ffmpeg(&spliced, shape.fps, codec, theme); + let loop_bytes = frames_per_loop(shape.fps) * frame_bytes; + let before_bytes = STAND_IN_GOP * frame_bytes; + assert_eq!( + decoded.len(), + alone.len() + loop_bytes, + "{codec:?}: frame count" + ); + let (stream_before, rest) = decoded.split_at(before_bytes); + let (screen_frames, stream_after) = rest.split_at(loop_bytes); + assert!( + stream_before == &alone[..before_bytes], + "{codec:?}: stream before the cut" + ); + let expected_screen = expected_frames(&settings(theme), shape); + assert!( + screen_frames == &expected_screen[..loop_bytes], + "{codec:?}: the screen" + ); + assert!( + stream_after == &alone[before_bytes..], + "{codec:?}: stream after the screen" + ); + } + } + + /// Keyframe interval of the stand-in stream, in frames. + const STAND_IN_GOP: usize = 10; + + /// Two keyframe intervals (IDR, no B-frames) of a test pattern from + /// x264 or x265, tagged BT.709 limited range as OBS tags its streams by + /// default and the screen tags itself. Matching tags matter to the + /// check, not to viewers: when they differ, ffmpeg converts frames + /// after the screen and they stop matching the stream decoded alone. + /// `None`, with a note, when this ffmpeg can't run the encoder. + fn encode_stand_in_stream(codec: SlateCodec, shape: StreamShape) -> Option> { + let gop = STAND_IN_GOP.to_string(); + let (encoder, format, options): (&str, &str, Vec) = match codec { + SlateCodec::H264 => ( + "libx264", + "h264", + [ + "-g", + &gop, + "-keyint_min", + &gop, + "-sc_threshold", + "0", + "-bf", + "0", + ] + .map(String::from) + .to_vec(), + ), + SlateCodec::Hevc => ( + "libx265", + "hevc", + vec![ + "-x265-params".into(), + format!( + "keyint={gop}:min-keyint={gop}:scenecut=0:bframes=0:\ + open-gop=0:pools=1:frame-threads=1:log-level=error" + ), + ], + ), + }; + let source = format!( + "testsrc2=size={}x{}:rate={}", + shape.width, shape.height, shape.fps + ); + let _turn = ffmpeg_turn(); + let output = Command::new("ffmpeg") + .args(["-v", "error", "-f", "lavfi", "-i", &source, "-frames:v"]) + .arg((2 * STAND_IN_GOP).to_string()) + .args(["-pix_fmt", "yuv420p", "-c:v", encoder]) + .args(&options) + .args(["-color_primaries", "bt709", "-color_trc", "bt709"]) + .args(["-colorspace", "bt709", "-color_range", "tv"]) + .args(["-f", format, "-"]) + .output() + .unwrap(); + if !output.status.success() || output.stdout.is_empty() { + let why = String::from_utf8_lossy(&output.stderr); + eprintln!( + "NOTE: ffmpeg couldn't run {encoder} ({}), so the {codec:?} cut check didn't run", + why.trim() + ); + return None; + } + Some(output.stdout) + } + + /// Byte offset of each Annex B NAL unit's start code, with its type. + fn nal_starts(stream: &[u8], codec: SlateCodec) -> Vec<(usize, u8)> { + let mut starts = Vec::new(); + for at in 0..stream.len().saturating_sub(3) { + if stream[at..at + 3] != [0, 0, 1] { + continue; + } + let start = if at > 0 && stream[at - 1] == 0 { + at - 1 + } else { + at + }; + starts.push((start, nal_type(codec, &stream[at + 3..]))); + } + starts + } + + /// `stream` with its VPS, SPS and PPS NAL units left out. + fn without_parameter_sets(stream: &[u8], codec: SlateCodec) -> Vec { + let starts = nal_starts(stream, codec); + let ends = starts + .iter() + .skip(1) + .map(|(start, _)| *start) + .chain([stream.len()]); + let is_parameter_set = |nal_type: u8| match codec { + SlateCodec::H264 => matches!(nal_type, 7 | 8), + SlateCodec::Hevc => matches!(nal_type, 32..=34), + }; + starts + .iter() + .zip(ends) + .filter(|((_, nal_type), _)| !is_parameter_set(*nal_type)) + .flat_map(|((start, _), end)| &stream[*start..end]) + .copied() + .collect() + } + + /// Where the access unit of the stream's second keyframe begins: its + /// first slice, or the SEI or delimiter NAL units just before it. + fn second_keyframe_offset(stream: &[u8], codec: SlateCodec) -> usize { + type NalTest = fn(u8) -> bool; + let (is_keyframe, is_prefix): (NalTest, NalTest) = match codec { + SlateCodec::H264 => (|t| t == 5, |t| matches!(t, 6 | 9)), + SlateCodec::Hevc => (|t| (16..=21).contains(&t), |t| matches!(t, 35 | 39)), + }; + let nals = nal_starts(stream, codec); + let (mut first, _) = nals + .iter() + .enumerate() + .filter(|(_, (_, nal_type))| is_keyframe(*nal_type)) + .nth(1) + .expect("the stand-in stream has a second keyframe"); + while first > 0 && is_prefix(nals[first - 1].1) { + first -= 1; + } + nals[first].0 + } + + /// Samples as an Annex B elementary stream. + fn annex_b_samples<'a>(samples: impl Iterator>) -> Vec { let mut out = Vec::new(); - for keyframe in &slate.keyframes { - for sample in std::iter::once(keyframe).chain(&slate.deltas) { - for nal in nal_units(sample) { - out.extend_from_slice(&[0, 0, 0, 1]); - out.extend_from_slice(nal); - } + for sample in samples { + for nal in nal_units(sample) { + out.extend_from_slice(&[0, 0, 0, 1]); + out.extend_from_slice(nal); } } out } - /// Raw H.264 carries no timestamps, so the input rate is given - /// explicitly and `-vsync 0` stops ffmpeg dropping identical frames. - fn decode_with_ffmpeg(stream: &[u8], fps: u32, theme: SlateTheme) -> Vec { + /// Two loop replays as an Annex B elementary stream. + fn annex_b(slate: &SlateLoop) -> Vec { + let replay = |keyframe| std::iter::once(keyframe).chain(&slate.deltas); + annex_b_samples(slate.keyframes.iter().flat_map(replay)) + } + + /// Raw H.264 and HEVC carry no timestamps, so the input rate is given + /// explicitly, and `-fps_mode passthrough` stops ffmpeg dropping + /// identical frames. One decoder thread keeps the round trip light on + /// a machine that runs the whole suite in parallel. + fn decode_with_ffmpeg( + stream: &[u8], + fps: u32, + codec: SlateCodec, + theme: SlateTheme, + ) -> Vec { + let format = match codec { + SlateCodec::H264 => "h264", + SlateCodec::Hevc => "hevc", + }; + // Tests decode in parallel, sometimes the same theme: number each + // run so no two share a temp file. + static RUNS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + let run = RUNS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); let dir = std::env::temp_dir(); - let tag = format!("instantclone-slate-{}-{}", std::process::id(), theme.id()); - let input = dir.join(format!("{tag}.h264")); + let tag = format!( + "instantclone-slate-{}-{run}-{}", + std::process::id(), + theme.id() + ); + let input = dir.join(format!("{tag}.{format}")); let output = dir.join(format!("{tag}.yuv")); std::fs::write(&input, stream).unwrap(); + let _turn = ffmpeg_turn(); let result = Command::new("ffmpeg") .args([ "-v", "error", "-y", + "-threads", + "1", "-r", &fps.to_string(), "-f", - "h264", + format, "-i", ]) .arg(&input) - .args(["-vsync", "0", "-f", "rawvideo", "-pix_fmt", "yuv420p"]) + .args([ + "-fps_mode", + "passthrough", + "-f", + "rawvideo", + "-pix_fmt", + "yuv420p", + ]) .arg(&output) .output() .unwrap(); @@ -408,7 +735,8 @@ mod tests { let _ = std::fs::remove_file(&output); assert!( result.status.success() && stderr.trim().is_empty(), - "ffmpeg: {stderr}" + "ffmpeg ({}): {stderr}", + result.status ); decoded } diff --git a/src/slate/png.rs b/src/slate/png.rs index 765ab84..d904fbf 100644 --- a/src/slate/png.rs +++ b/src/slate/png.rs @@ -251,6 +251,7 @@ mod tests { .collect(); let path = std::env::temp_dir().join(format!("ic-png-test-{}.png", std::process::id())); std::fs::write(&path, encode(width, height, &rgb)).unwrap(); + let _turn = super::super::ffmpeg_turn(); let decoded = Command::new("ffmpeg") .args(["-v", "error", "-i"]) .arg(&path) diff --git a/src/slate/raster.rs b/src/slate/raster.rs index 55208af..053788f 100644 --- a/src/slate/raster.rs +++ b/src/slate/raster.rs @@ -128,11 +128,18 @@ impl Canvas { }; } - /// Convert to 4:2:0 planes padded to whole macroblocks. Padding - /// repeats the edge pixels, which the SPS crops away again. + /// Convert to 4:2:0 planes padded to whole H.264 macroblocks. + #[cfg(test)] pub fn to_yuv420(&self) -> YuvFrame { - let coded_width = self.width.div_ceil(16) * 16; - let coded_height = self.height.div_ceil(16) * 16; + self.to_yuv420_padded(16) + } + + /// Convert to 4:2:0 planes padded to whole `block` x `block` squares + /// (16 for H.264 macroblocks, 32 for HEVC CTBs). Padding repeats the + /// edge pixels, which the SPS crops away again. + pub fn to_yuv420_padded(&self, block: usize) -> YuvFrame { + let coded_width = self.width.div_ceil(block) * block; + let coded_height = self.height.div_ceil(block) * block; let mut frame = YuvFrame::blank(coded_width, coded_height); for pair in 0..coded_height / 2 { self.convert_row_pair(&mut frame, pair); @@ -230,7 +237,7 @@ impl Canvas { } } -/// Planar 4:2:0 frame at the coded size (multiples of 16). +/// Planar 4:2:0 frame at the coded size (whole macroblocks or CTBs). #[derive(Clone)] pub struct YuvFrame { pub width: usize,