Skip to main content

pixelflux/encoders/
session.rs

1/*
2 * This Source Code Form is subject to the terms of the Mozilla Public
3 * License, v. 2.0. If a copy of the MPL was not distributed with this
4 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
5 */
6
7//! What the full-frame software sessions share: the planar picture a packed host frame is
8//! converted into, at 8 or 10 bits, the constant quantizer that follows the session quality index, the
9//! rate-control settings a live change is compared against, and the frames an encoder has
10//! been handed but not yet answered.
11
12use std::collections::VecDeque;
13
14use super::frame_rate::FrameRate;
15use super::software::{convert_to_yuv_mt, convert_to_yuv10_mt};
16use super::{QP_HYSTERESIS_LIMIT, vbv_bits};
17use crate::CaptureSettings;
18
19/// A planar picture, 4:2:0 or 4:4:4, with tightly packed rows: 8-bit samples in `y`, `u`, and
20/// `v`, or 10-bit ones in `y16`, `u16`, and `v16`, the planes of the other depth staying empty.
21pub struct Planes {
22    pub width: usize,
23    pub height: usize,
24    pub i444: bool,
25    pub bit_depth: u32,
26    pub y: Vec<u8>,
27    pub u: Vec<u8>,
28    pub v: Vec<u8>,
29    pub y16: Vec<u16>,
30    pub u16: Vec<u16>,
31    pub v16: Vec<u16>,
32}
33
34impl Planes {
35    pub fn new(width: usize, height: usize, i444: bool, bit_depth: u32) -> Self {
36        let (cw, ch) = if i444 {
37            (width, height)
38        } else {
39            (width.div_ceil(2), height.div_ceil(2))
40        };
41        let (narrow, wide) = if bit_depth > 8 { (0, 1) } else { (1, 0) };
42        Self {
43            width,
44            height,
45            i444,
46            bit_depth,
47            y: vec![0; width * height * narrow],
48            u: vec![0; cw * ch * narrow],
49            v: vec![0; cw * ch * narrow],
50            y16: vec![0; width * height * wide],
51            u16: vec![0; cw * ch * wide],
52            v16: vec![0; cw * ch * wide],
53        }
54    }
55
56    pub fn chroma_width(&self) -> usize {
57        if self.i444 {
58            self.width
59        } else {
60            self.width.div_ceil(2)
61        }
62    }
63
64    /// The planes as a library reads them, whichever depth they hold.
65    pub fn pointers(&mut self) -> [*mut u8; 3] {
66        if self.bit_depth > 8 {
67            [
68                self.y16.as_mut_ptr().cast(),
69                self.u16.as_mut_ptr().cast(),
70                self.v16.as_mut_ptr().cast(),
71            ]
72        } else {
73            [
74                self.y.as_mut_ptr(),
75                self.u.as_mut_ptr(),
76                self.v.as_mut_ptr(),
77            ]
78        }
79    }
80
81    /// The bytes a sample takes.
82    pub fn sample_bytes(&self) -> usize {
83        if self.bit_depth > 8 { 2 } else { 1 }
84    }
85
86    /// The bytes all three planes hold.
87    pub fn byte_len(&self) -> usize {
88        self.y.len()
89            + self.u.len()
90            + self.v.len()
91            + 2 * (self.y16.len() + self.u16.len() + self.v16.len())
92    }
93
94    /// Convert a packed host frame (`stride` bytes per row, R,G,B,A when `rgba`, else B,G,R,A)
95    /// into the planes across `threads` bands, with the BT.709 matrix at `full_range`, or BT.601
96    /// where `bt601` names a codec whose bitstream can carry no other, which has no 10 bits.
97    pub fn convert(
98        &mut self,
99        pixels: &[u8],
100        stride: usize,
101        rgba: bool,
102        full_range: bool,
103        bt601: bool,
104        threads: usize,
105    ) -> Result<(), String> {
106        let cw = self.chroma_width();
107        if self.bit_depth > 8 {
108            check_host_frame(pixels, stride, self.width, self.height)?;
109            convert_to_yuv10_mt(
110                pixels,
111                stride,
112                self.width,
113                self.height,
114                rgba,
115                self.i444,
116                full_range,
117                &mut self.y16,
118                &mut self.u16,
119                &mut self.v16,
120                (self.width, cw),
121                threads,
122            );
123            return Ok(());
124        }
125        convert_into(
126            pixels,
127            stride,
128            self.width,
129            self.height,
130            rgba,
131            self.i444,
132            full_range,
133            bt601,
134            threads,
135            &mut self.y,
136            &mut self.u,
137            &mut self.v,
138            (self.width, cw),
139        )
140    }
141}
142
143/// Convert a packed host frame into planes a library owns, at `strides` (luma, chroma) bytes
144/// per row, as `Planes::convert` does into its own.
145#[allow(clippy::too_many_arguments)]
146pub fn convert_into(
147    pixels: &[u8],
148    stride: usize,
149    width: usize,
150    height: usize,
151    rgba: bool,
152    i444: bool,
153    full_range: bool,
154    bt601: bool,
155    threads: usize,
156    y: &mut [u8],
157    u: &mut [u8],
158    v: &mut [u8],
159    strides: (usize, usize),
160) -> Result<(), String> {
161    check_host_frame(pixels, stride, width, height)?;
162    convert_to_yuv_mt(
163        pixels,
164        stride as u32,
165        width,
166        height,
167        rgba,
168        i444,
169        full_range,
170        bt601,
171        y,
172        u,
173        v,
174        strides,
175        threads,
176    )
177    .map_err(|e| format!("rgb-to-yuv conversion failed: {e:?}"))
178}
179
180/// Whether `pixels` holds a `width` x `height` packed picture at `stride` bytes per row.
181pub fn check_host_frame(
182    pixels: &[u8],
183    stride: usize,
184    width: usize,
185    height: usize,
186) -> Result<(), String> {
187    let row_bytes = width * 4;
188    let needed = if height == 0 {
189        0
190    } else {
191        stride.checked_mul(height - 1).ok_or("stride overflow")? + row_bytes
192    };
193    if stride < row_bytes || pixels.len() < needed {
194        return Err("Input buffer too small".into());
195    }
196    Ok(())
197}
198
199/// The constant quantizer of a session, moved toward the one the quality index selects: a
200/// decrease sharpens the picture and applies at once, an increase waits out
201/// `QP_HYSTERESIS_LIMIT` consecutive requests so transient motion does not make quality blink.
202/// The domain is the codec's own (`Codec::quantizer`).
203pub struct Quality {
204    pub current: u32,
205    counter: u32,
206}
207
208impl Quality {
209    pub fn new(quantizer: u32) -> Self {
210        Self {
211            current: quantizer,
212            counter: 0,
213        }
214    }
215
216    /// The quantizer to program now for `target`, once the hysteresis admits the change.
217    pub fn update(&mut self, target: u32) -> Option<u32> {
218        if target == self.current {
219            self.counter = 0;
220            return None;
221        }
222        if target > self.current {
223            self.counter += 1;
224            if self.counter <= QP_HYSTERESIS_LIMIT {
225                return None;
226            }
227        }
228        self.counter = 0;
229        self.current = target;
230        Some(target)
231    }
232}
233
234/// The rate-control and frame-rate settings a session was configured with, so a live change
235/// re-programs the encoder only when a value it reads actually moved.
236#[derive(Clone, Copy, Debug, PartialEq)]
237pub struct RateSettings {
238    pub cbr: bool,
239    /// Whether a constant-quality session is capped at a peak of `bitrate_kbps` with a VBV of
240    /// `vbv_multiplier` frames (`CaptureSettings::video_vbv_multiplier`), fixed at its open.
241    pub capped: bool,
242    pub bitrate_kbps: i32,
243    pub vbv_multiplier: f64,
244    pub keyframe_interval_s: f64,
245    pub fps: FrameRate,
246    pub min_qp: i32,
247    pub max_qp: i32,
248}
249
250impl RateSettings {
251    pub fn new(settings: &CaptureSettings) -> Self {
252        Self {
253            cbr: settings.video_cbr_mode,
254            capped: !settings.video_cbr_mode && settings.video_vbv_multiplier > 0.0,
255            bitrate_kbps: settings.video_bitrate_kbps,
256            vbv_multiplier: settings.video_vbv_multiplier,
257            keyframe_interval_s: settings.keyframe_interval_s,
258            fps: FrameRate::of(settings.target_fps),
259            min_qp: settings.video_min_qp,
260            max_qp: settings.video_max_qp,
261        }
262    }
263
264    /// The settings of a session that runs its constant quality uncapped whatever a caller asks,
265    /// whose bitrate then moves nothing.
266    pub fn uncapped(self) -> Self {
267        Self {
268            capped: false,
269            ..self
270        }
271    }
272
273    /// The settings as they stand now, if a rate or frame-rate value the encoder was programmed
274    /// with changed: a bitrate or VBV multiplier in CBR and in a capped CRF, the frame rate in
275    /// any mode.
276    pub fn changed(&self, settings: &CaptureSettings) -> Option<Self> {
277        let next = Self {
278            cbr: self.cbr,
279            capped: self.capped,
280            min_qp: self.min_qp,
281            max_qp: self.max_qp,
282            ..Self::new(settings)
283        };
284        let rate_moved = (self.cbr || self.capped)
285            && (next.bitrate_kbps != self.bitrate_kbps
286                || next.vbv_multiplier != self.vbv_multiplier);
287        (rate_moved || next.fps != self.fps).then_some(next)
288    }
289
290    /// The constant-rate target in bits per second.
291    pub fn bps(&self) -> u64 {
292        (self.bitrate_kbps.max(0) as u64) * 1000
293    }
294
295    /// The VBV buffer, in bits, the constant-rate target is held to.
296    pub fn vbv(&self) -> u32 {
297        vbv_bits(
298            self.bps().min(u64::from(u32::MAX)) as u32,
299            self.fps.fps(),
300            self.keyframe_interval_s,
301            self.vbv_multiplier,
302        )
303    }
304}
305
306/// A leaky bucket of a capped constant quality's peak (`RateSettings::capped`), its buffer the
307/// session's VBV, for an encoder whose constant quality takes no cap of its own: the session
308/// codes as an uncapped one while the bucket holds what it codes, binds to the peak through its
309/// constant rate, never finer than its quality, from the frame that overflows the bucket, and
310/// lets go once `RELEASE_FRAMES` frames in a row come out at that quality with the bucket back
311/// within its buffer, which says the content fits under the peak again: a rate control that
312/// takes the peak with its buffer full codes at the quality for a while whatever the content.
313/// The frame that overflows has already gone out: the cap holds the frames after it.
314#[derive(Clone, Copy, Debug)]
315pub struct PeakBucket {
316    pub bps: u32,
317    pub vbv: u32,
318    fps: f64,
319    /// The bits the bucket holds.
320    fill: f64,
321    pub binding: bool,
322    /// Frames in a row the constant rate coded at the quality while binding.
323    at_floor: u32,
324}
325
326impl PeakBucket {
327    pub const RELEASE_FRAMES: u32 = 30;
328
329    pub fn new(bps: u32, vbv: u32, fps: f64) -> Self {
330        Self {
331            bps,
332            vbv,
333            fps: fps.max(1.0),
334            fill: 0.0,
335            binding: false,
336            at_floor: 0,
337        }
338    }
339
340    /// Frames in a row the constant rate has coded at the quality while binding.
341    #[cfg(test)]
342    pub fn floor_run(&self) -> u32 {
343        self.at_floor
344    }
345
346    /// Take a live peak, buffer, and frame rate.
347    pub fn follow(&mut self, bps: u32, vbv: u32, fps: f64) {
348        self.bps = bps;
349        self.vbv = vbv;
350        self.fps = fps.max(1.0);
351    }
352
353    /// Fold in a frame of `bits`, `at_floor` saying whether the rate control coded it at the
354    /// session's quality (`None` for a held frame, whose quantizer is not the rate control's):
355    /// `Some(true)` to bind the session to the peak, `Some(false)` to release it.
356    pub fn note(&mut self, bits: f64, at_floor: Option<bool>) -> Option<bool> {
357        self.fill = (self.fill - f64::from(self.bps) / self.fps).max(0.0) + bits;
358        if !self.binding {
359            self.binding = self.fill > f64::from(self.vbv);
360            self.at_floor = 0;
361            return self.binding.then_some(true);
362        }
363        match at_floor {
364            Some(true) => self.at_floor += 1,
365            Some(false) => self.at_floor = 0,
366            None => {}
367        }
368        if self.at_floor < Self::RELEASE_FRAMES || self.fill > f64::from(self.vbv) {
369            return None;
370        }
371        self.binding = false;
372        Some(false)
373    }
374}
375
376/// The peak of a capped constant quality held beside an encoder that codes its constant quality
377/// per frame and has no rate control to bind to it (x264, SVT-AV1), where `PeakBucket` switches
378/// one that does. A leaky bucket of the peak, its buffer the session's VBV, takes each frame: from
379/// a frame that overflows it the quality index stands `offset` steps above the session's, aimed
380/// after each frame at bringing the bucket back to half full in `DRAIN_FRAMES` frames (a size
381/// halving every six steps) and lowered at most `RELEASE_STEP` a frame, so a session that never
382/// overflows it codes as an uncapped one.
383#[derive(Clone, Copy, Debug)]
384pub struct PeakOffset {
385    /// Bits a frame drains from the bucket, and the bucket's buffer.
386    per_frame: f64,
387    buffer: f64,
388    pub(crate) fill: f64,
389    /// Quality index steps above the session's own.
390    pub(crate) offset: f64,
391}
392
393impl PeakOffset {
394    const DRAIN_FRAMES: f64 = 8.0;
395    const RELEASE_STEP: f64 = 2.0;
396
397    pub fn new(bitrate_kbps: i32, vbv_kbit: i32, fps: f64) -> Self {
398        let mut peak = PeakOffset {
399            per_frame: 0.0,
400            buffer: 0.0,
401            fill: 0.0,
402            offset: 0.0,
403        };
404        peak.follow(bitrate_kbps, vbv_kbit, fps);
405        peak
406    }
407
408    /// Take a live peak and buffer.
409    pub fn follow(&mut self, bitrate_kbps: i32, vbv_kbit: i32, fps: f64) {
410        self.per_frame = f64::from(bitrate_kbps.max(1)) * 1000.0 / fps.max(1.0);
411        self.buffer = f64::from(vbv_kbit.max(1)) * 1000.0;
412    }
413
414    /// Fold in a frame of `bits` and return the offset the next frame takes. A key frame, or a
415    /// frame held at a quantizer of its own (`steers` false), fills the bucket but says nothing
416    /// of what the session's own frames come to.
417    pub fn note(&mut self, bits: f64, steers: bool, max_offset: f64) -> f64 {
418        self.fill = (self.fill - self.per_frame).max(0.0) + bits;
419        if !steers || (self.offset <= 0.0 && self.fill <= self.buffer) {
420            return self.offset;
421        }
422        let target = (self.per_frame + (self.buffer / 2.0 - self.fill) / Self::DRAIN_FRAMES)
423            .max(self.per_frame / 4.0);
424        let step = 6.0 * (bits.max(1.0) / target).log2();
425        self.offset = (self.offset + step.max(-Self::RELEASE_STEP)).clamp(0.0, max_offset.max(0.0));
426        self.offset
427    }
428
429    /// The bits the next frame has before the bucket overflows.
430    pub fn room(&self) -> f64 {
431        self.buffer - (self.fill - self.per_frame).max(0.0)
432    }
433}
434
435/// The frames handed to an encoder and not yet answered, by the timestamp each was submitted
436/// with, so a packet names the frame it encodes rather than the one submitted alongside it.
437/// An encoder that pipelines answers a frame several later, and one that drops a frame never
438/// answers it at all, so taking a timestamp discards every submission it passed over.
439#[derive(Default)]
440pub struct Pending(VecDeque<(u64, u16)>);
441
442impl Pending {
443    /// Record that the frame with wire id `frame_id` went in at `pts`.
444    pub fn push(&mut self, pts: u64, frame_id: u16) {
445        self.0.push_back((pts, frame_id));
446    }
447
448    /// Drop the last submission, for a frame the encoder refused outright.
449    pub fn undo(&mut self) {
450        self.0.pop_back();
451    }
452
453    pub fn is_empty(&self) -> bool {
454        self.0.is_empty()
455    }
456
457    /// The wire id submitted at `pts`, dropping it and everything older; `None` where no
458    /// submission still held carries that timestamp.
459    pub fn take(&mut self, pts: u64) -> Option<u16> {
460        let at = self.0.iter().position(|&(p, _)| p == pts)?;
461        let id = self.0[at].1;
462        self.0.drain(..=at);
463        Some(id)
464    }
465}
466
467/// The encode threads a software session spreads across, between one and `most`: one less than
468/// the host's cores, or one alone where `cpu_encode_threading` is off. The conversion ahead of
469/// the encoder runs on as many.
470fn threads_up_to(settings: &CaptureSettings, most: usize) -> i32 {
471    if !settings.cpu_encode_threading {
472        return 1;
473    }
474    std::thread::available_parallelism()
475        .map(|n| n.get())
476        .unwrap_or(1)
477        .saturating_sub(1)
478        .clamp(1, most) as i32
479}
480
481/// The encode threads a full-frame H.265, VP8, VP9, or AV1 software session spreads across
482/// (`threads_up_to` eight): x265's and kvazaar's wavefront, libvpx's token partitions and
483/// rows, and SVT-AV1's logical processors.
484pub fn encode_threads(settings: &CaptureSettings) -> i32 {
485    threads_up_to(settings, 8)
486}
487
488/// The encode threads a full-frame H.264 software session spreads across (`threads_up_to`
489/// four): x264's sliced threads, and OpenH264's threads over its four slices, more than which
490/// would buy nothing.
491pub fn h264_fullframe_threads(settings: &CaptureSettings) -> i32 {
492    threads_up_to(settings, 4)
493}
494
495#[cfg(test)]
496mod tests {
497    use super::*;
498
499    /// Every software session's thread count follows `cpu_encode_threading`: one thread with it
500    /// off, and with it on one less than the cores within the session's ceiling, one at least.
501    #[test]
502    fn cpu_encode_threading_off_leaves_every_session_one_thread() {
503        let off = CaptureSettings {
504            cpu_encode_threading: false,
505            ..Default::default()
506        };
507        assert_eq!(encode_threads(&off), 1);
508        assert_eq!(h264_fullframe_threads(&off), 1);
509        let on = CaptureSettings::default();
510        let cores = std::thread::available_parallelism()
511            .map(|n| n.get())
512            .unwrap_or(1);
513        assert_eq!(
514            encode_threads(&on),
515            cores.saturating_sub(1).clamp(1, 8) as i32
516        );
517        assert_eq!(
518            h264_fullframe_threads(&on),
519            cores.saturating_sub(1).clamp(1, 4) as i32
520        );
521    }
522
523    #[test]
524    fn a_quality_increase_applies_at_once_and_a_drop_waits() {
525        let mut q = Quality::new(30);
526        assert_eq!(q.update(30), None);
527        assert_eq!(q.update(20), Some(20), "a lower quantizer applies at once");
528        for _ in 0..QP_HYSTERESIS_LIMIT {
529            assert_eq!(q.update(40), None);
530        }
531        assert_eq!(
532            q.update(40),
533            Some(40),
534            "a higher one waits out the hysteresis"
535        );
536        assert_eq!(q.update(45), None);
537        assert_eq!(
538            q.update(40),
539            None,
540            "a request back at the current value resets the count"
541        );
542        assert_eq!(q.current, 40);
543    }
544
545    #[test]
546    fn a_rate_change_is_noticed_only_where_the_encoder_reads_it() {
547        let mut settings = CaptureSettings {
548            target_fps: 30.0,
549            video_bitrate_kbps: 4000,
550            ..Default::default()
551        };
552        let rate = RateSettings::new(&settings);
553        assert_eq!(rate.changed(&settings), None);
554        settings.video_bitrate_kbps = 8000;
555        assert_eq!(
556            rate.changed(&settings),
557            None,
558            "a bitrate moves nothing in constant-quality mode"
559        );
560        settings.target_fps = 60.0;
561        assert_eq!(
562            rate.changed(&settings).map(|r| r.fps),
563            Some(FrameRate { num: 60, den: 1 })
564        );
565        let ntsc = RateSettings::new(&CaptureSettings {
566            target_fps: 60000.0 / 1001.0,
567            ..settings.clone()
568        });
569        assert_eq!(
570            ntsc.fps,
571            FrameRate {
572                num: 60000,
573                den: 1001
574            }
575        );
576        assert!(
577            ntsc.changed(&settings).is_some(),
578            "60 fps moves a session at 59.94"
579        );
580        settings.video_cbr_mode = true;
581        let cbr = RateSettings::new(&settings);
582        settings.video_bitrate_kbps = 9000;
583        assert_eq!(cbr.changed(&settings).map(|r| r.bitrate_kbps), Some(9000));
584        assert_eq!(cbr.bps(), 8_000_000);
585        assert!(!cbr.capped);
586    }
587
588    /// A capped constant-quality session reads its peak as a constant rate reads its target; a
589    /// session that cannot cap reads neither.
590    #[test]
591    fn a_capped_crf_session_notices_its_peak_move() {
592        let mut settings = CaptureSettings {
593            target_fps: 60.0,
594            video_bitrate_kbps: 4000,
595            video_vbv_multiplier: crate::encoders::CAPPED_VBV_FRAMES,
596            ..Default::default()
597        };
598        let capped = RateSettings::new(&settings);
599        assert!(capped.capped && !capped.cbr);
600        let uncapped = capped.uncapped();
601        settings.video_bitrate_kbps = 2000;
602        assert_eq!(
603            capped
604                .changed(&settings)
605                .map(|r| (r.bitrate_kbps, r.capped)),
606            Some((2000, true))
607        );
608        assert_eq!(uncapped.changed(&settings), None);
609        assert!(
610            !RateSettings::new(&CaptureSettings {
611                video_vbv_multiplier: 0.0,
612                ..settings
613            })
614            .capped,
615            "no buffer leaves the rate factor uncapped"
616        );
617    }
618
619    /// A packet names the frame it encodes however deep the encoder pipelines, the frames it
620    /// passed over go with it, and a timestamp never submitted names nothing.
621    #[test]
622    fn a_packet_names_the_frame_it_encodes() {
623        let mut pending = Pending::default();
624        assert!(pending.is_empty());
625        for (pts, id) in [(0u64, 7u16), (1, 8), (2, 9)] {
626            pending.push(pts, id);
627        }
628        assert_eq!(pending.take(99), None, "a timestamp never submitted");
629        assert_eq!(
630            pending.take(1),
631            Some(8),
632            "the encoder answered two frames in"
633        );
634        assert_eq!(
635            pending.take(0),
636            None,
637            "the frame it passed over went with it"
638        );
639        assert_eq!(pending.take(2), Some(9));
640        assert!(pending.is_empty());
641        pending.push(3, 10);
642        pending.undo();
643        assert_eq!(pending.take(3), None, "a refused frame was never submitted");
644    }
645
646    #[test]
647    fn a_frame_shorter_than_its_stride_says_so() {
648        assert!(check_host_frame(&[0; 16 * 4 * 4], 64, 16, 4).is_ok());
649        assert!(check_host_frame(&[0; 16 * 4 * 4 - 1], 64, 16, 4).is_err());
650        assert!(
651            check_host_frame(&[0; 16 * 4 * 4], 60, 16, 4).is_err(),
652            "a stride shorter than a row"
653        );
654        let mut planes = Planes::new(4, 2, false, 8);
655        assert_eq!((planes.u.len(), planes.chroma_width()), (2, 2));
656        assert!(
657            planes
658                .convert(&[0; 4 * 4 * 2], 16, false, false, false, 1)
659                .is_ok()
660        );
661        assert!(
662            Planes::new(4, 2, true, 8)
663                .convert(&[0; 4 * 4], 16, false, true, false, 1)
664                .is_err()
665        );
666    }
667}