1use 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
19pub 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 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 pub fn sample_bytes(&self) -> usize {
83 if self.bit_depth > 8 { 2 } else { 1 }
84 }
85
86 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 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#[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
180pub 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
199pub 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 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#[derive(Clone, Copy, Debug, PartialEq)]
237pub struct RateSettings {
238 pub cbr: bool,
239 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 pub fn uncapped(self) -> Self {
267 Self {
268 capped: false,
269 ..self
270 }
271 }
272
273 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 pub fn bps(&self) -> u64 {
292 (self.bitrate_kbps.max(0) as u64) * 1000
293 }
294
295 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#[derive(Clone, Copy, Debug)]
315pub struct PeakBucket {
316 pub bps: u32,
317 pub vbv: u32,
318 fps: f64,
319 fill: f64,
321 pub binding: bool,
322 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 #[cfg(test)]
342 pub fn floor_run(&self) -> u32 {
343 self.at_floor
344 }
345
346 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 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#[derive(Clone, Copy, Debug)]
384pub struct PeakOffset {
385 per_frame: f64,
387 buffer: f64,
388 pub(crate) fill: f64,
389 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 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 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 pub fn room(&self) -> f64 {
431 self.buffer - (self.fill - self.per_frame).max(0.0)
432 }
433}
434
435#[derive(Default)]
440pub struct Pending(VecDeque<(u64, u16)>);
441
442impl Pending {
443 pub fn push(&mut self, pts: u64, frame_id: u16) {
445 self.0.push_back((pts, frame_id));
446 }
447
448 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 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
467fn 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
481pub fn encode_threads(settings: &CaptureSettings) -> i32 {
485 threads_up_to(settings, 8)
486}
487
488pub 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 #[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 #[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 #[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}