1#![allow(clippy::too_many_arguments)]
2
3use blit_compositor::PixelData;
4use blit_remote::{
5 CODEC_SUPPORT_AV1, CODEC_SUPPORT_AV1_444, CODEC_SUPPORT_H264, CODEC_SUPPORT_H264_444,
6 SURFACE_FRAME_CODEC_AV1, SURFACE_FRAME_CODEC_H264,
7};
8use openh264::encoder::Encoder as OpenH264Encoder;
9use openh264::formats::YUVBuffer;
10
11#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
12pub enum SurfaceEncoderPreference {
13 VulkanVideoH264,
14 VulkanVideoAV1,
15 H264Software,
16 H264Vaapi,
17 AV1Vaapi,
18 NvencH264,
19 NvencAV1,
20 AV1Software,
21}
22
23pub type SurfaceH264EncoderPreference = SurfaceEncoderPreference;
25
26const H264_MAX_WIDTH: u16 = 3840;
28const H264_MAX_HEIGHT: u16 = 2160;
29
30impl SurfaceEncoderPreference {
31 pub fn parse(value: &str) -> Option<Self> {
32 match value.trim() {
33 "h264-vulkan" => Some(Self::VulkanVideoH264),
34 "av1-vulkan" => Some(Self::VulkanVideoAV1),
35 "h264-software" | "software" => Some(Self::H264Software),
36 "h264-vaapi" | "vaapi" => Some(Self::H264Vaapi),
37 "av1-vaapi" => Some(Self::AV1Vaapi),
38 "h264-nvenc" => Some(Self::NvencH264),
39 "av1-nvenc" => Some(Self::NvencAV1),
40 "av1-software" => Some(Self::AV1Software),
41 _ => None,
42 }
43 }
44
45 pub fn parse_list(value: &str) -> Result<Vec<Self>, String> {
47 let mut result = Vec::new();
48 for item in value.split(',') {
49 let item = item.trim();
50 if item.is_empty() {
51 continue;
52 }
53 result.push(Self::parse(item).ok_or_else(|| format!("unknown encoder: {item}"))?);
54 }
55 Ok(result)
56 }
57
58 pub fn defaults() -> Vec<Self> {
63 if let Some(list) = std::env::var("BLIT_SURFACE_ENCODERS")
64 .ok()
65 .and_then(|v| Self::parse_list(&v).ok())
66 {
67 return list;
68 }
69 vec![
70 Self::NvencAV1,
75 Self::NvencH264,
76 Self::AV1Vaapi,
77 Self::H264Vaapi,
78 Self::H264Software,
79 Self::AV1Software,
80 ]
81 }
82
83 pub fn supported_by_client(self, codec_support: u8) -> bool {
86 if codec_support == 0 {
87 return true;
88 }
89 match self {
90 Self::H264Software | Self::H264Vaapi | Self::NvencH264 | Self::VulkanVideoH264 => {
91 codec_support & CODEC_SUPPORT_H264 != 0
92 }
93 Self::AV1Vaapi | Self::AV1Software | Self::NvencAV1 | Self::VulkanVideoAV1 => {
94 codec_support & CODEC_SUPPORT_AV1 != 0
95 }
96 }
97 }
98
99 pub fn supports_444_by_client(self, codec_support: u8) -> bool {
104 if codec_support == 0 {
105 return false;
106 }
107 match self {
108 Self::H264Software | Self::H264Vaapi | Self::NvencH264 | Self::VulkanVideoH264 => {
109 codec_support & CODEC_SUPPORT_H264_444 != 0
110 }
111 Self::AV1Vaapi | Self::AV1Software | Self::NvencAV1 | Self::VulkanVideoAV1 => {
112 codec_support & CODEC_SUPPORT_AV1_444 != 0
113 }
114 }
115 }
116
117 pub fn max_dimensions(self) -> Option<(u16, u16)> {
120 match self {
121 Self::H264Software | Self::H264Vaapi | Self::NvencH264 | Self::VulkanVideoH264 => {
122 Some((H264_MAX_WIDTH, H264_MAX_HEIGHT))
123 }
124 Self::AV1Vaapi | Self::NvencAV1 | Self::AV1Software | Self::VulkanVideoAV1 => None,
125 }
126 }
127
128 pub fn max_dimensions_for_list(prefs: &[Self]) -> Option<(u16, u16)> {
130 let mut result: Option<(u16, u16)> = None;
131 for p in prefs {
132 if let Some((w, h)) = p.max_dimensions() {
133 result = Some(match result {
134 Some((rw, rh)) => (rw.min(w), rh.min(h)),
135 None => (w, h),
136 });
137 }
138 }
139 result
140 }
141
142 pub fn is_vulkan_video(self) -> bool {
144 matches!(self, Self::VulkanVideoH264 | Self::VulkanVideoAV1)
145 }
146
147 pub fn vulkan_codec(self) -> u8 {
149 match self {
150 Self::VulkanVideoAV1 => 0x02,
151 _ => 0x01,
152 }
153 }
154
155 pub fn codec_flag(self) -> u8 {
157 match self {
158 Self::H264Software | Self::H264Vaapi | Self::NvencH264 | Self::VulkanVideoH264 => {
159 SURFACE_FRAME_CODEC_H264
160 }
161 Self::AV1Vaapi | Self::AV1Software | Self::NvencAV1 | Self::VulkanVideoAV1 => {
162 SURFACE_FRAME_CODEC_AV1
163 }
164 }
165 }
166}
167
168#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Hash)]
178pub enum ChromaSubsampling {
179 Cs420,
180 #[default]
181 Cs444,
182}
183
184impl ChromaSubsampling {
185 pub fn label(self) -> &'static str {
186 match self {
187 Self::Cs420 => "4:2:0",
188 Self::Cs444 => "4:4:4",
189 }
190 }
191
192 pub fn from_env() -> Self {
193 match std::env::var("BLIT_CHROMA").ok().as_deref() {
194 Some("420") => Self::Cs420,
195 _ => Self::Cs444,
196 }
197 }
198
199 pub fn is_444(self) -> bool {
200 matches!(self, Self::Cs444)
201 }
202}
203
204pub fn av1_level_for(width: u32, height: u32) -> &'static str {
207 let sps = width as u64 * height as u64 * 60;
208 const SPECS: &[(&str, u32, u32, u64)] = &[
210 ("00", 2048, 1152, 5_529_600),
211 ("01", 2816, 1152, 10_454_400),
212 ("04", 4352, 2448, 24_969_600),
213 ("05", 5504, 3096, 39_938_400),
214 ("08", 6144, 3456, 77_856_768),
215 ("09", 6144, 3456, 155_713_536),
216 ("12", 8192, 4352, 273_715_200),
217 ("13", 8192, 4352, 547_430_400),
218 ("16", 16384, 8704, 1_176_502_272),
219 ];
220 for &(level, max_w, max_h, max_rate) in SPECS {
221 if width <= max_w && height <= max_h && sps <= max_rate {
222 return level;
223 }
224 }
225 "16"
226}
227
228#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
238pub enum SurfaceQuality {
239 Low,
240 #[default]
241 Medium,
242 High,
243 Ultra,
244 Custom {
247 quantizer: u8,
248 },
249}
250
251impl SurfaceQuality {
252 pub fn parse(value: &str) -> Option<Self> {
253 match value {
254 "low" => Some(Self::Low),
255 "medium" => Some(Self::Medium),
256 "high" => Some(Self::High),
257 "ultra" | "lossless" => Some(Self::Ultra),
258 _ => None,
259 }
260 }
261
262 pub fn from_wire(value: u8) -> Option<Self> {
269 match value {
270 1 => Some(Self::Low),
271 2 => Some(Self::Medium),
272 3 => Some(Self::High),
273 4 => Some(Self::Ultra),
274 v @ 10..=255 => Some(Self::Custom { quantizer: v }),
275 _ => None,
276 }
277 }
278
279 fn av1_speed(self) -> u8 {
281 match self {
282 Self::Low => 10,
283 Self::Medium => 10,
284 Self::High => 8,
285 Self::Ultra => 6,
286 Self::Custom { quantizer } => {
287 if quantizer <= 40 {
288 6
289 } else if quantizer <= 80 {
290 8
291 } else {
292 10
293 }
294 }
295 }
296 }
297
298 fn av1_quantizer(self) -> usize {
301 match self {
302 Self::Low => 180,
303 Self::Medium => 120,
304 Self::High => 80,
305 Self::Ultra => 1,
306 Self::Custom { quantizer } => quantizer as usize,
307 }
308 }
309
310 fn av1_min_quantizer(self) -> u8 {
312 match self {
313 Self::Low => 120,
314 Self::Medium => 80,
315 Self::High => 40,
316 Self::Ultra => 0,
317 Self::Custom { quantizer } => quantizer.saturating_sub(40),
318 }
319 }
320
321 pub fn h264_qp(self) -> u8 {
324 match self {
325 Self::Low => 35,
326 Self::Medium => 28,
327 Self::High => 20,
328 Self::Ultra => 10,
329 Self::Custom { quantizer } => ((quantizer as u32) * 51 / 255).min(51) as u8,
330 }
331 }
332
333 pub fn nvenc_av1_qp(self) -> u32 {
336 self.av1_quantizer() as u32
337 }
338
339 pub fn av1_qp_for_vulkan(self) -> u8 {
342 self.av1_quantizer().min(255) as u8
343 }
344
345 fn openh264_bitrate(self) -> u32 {
348 match self {
349 Self::Low => 500_000,
350 Self::Medium => 2_000_000,
351 Self::High => 8_000_000,
352 Self::Ultra => 20_000_000,
353 Self::Custom { quantizer } => {
354 let q = quantizer as u32;
356 20_000_000 - q * (20_000_000 - 500_000) / 255
357 }
358 }
359 }
360}
361
362pub struct SurfaceEncoder {
363 width: u32,
365 height: u32,
366 source_width: u32,
368 source_height: u32,
369 kind: SurfaceEncoderKind,
370 chroma: ChromaSubsampling,
373}
374
375enum SurfaceEncoderKind {
376 H264Software(Box<SoftwareH264Encoder>),
377 NvencH264(Box<crate::nvenc_encode::NvencDirectEncoder>),
378 NvencAV1(Box<crate::nvenc_encode::NvencDirectEncoder>),
379 #[cfg(target_os = "linux")]
380 H264Vaapi(Box<crate::vaapi_encode::VaapiDirectEncoder>),
381 #[cfg(target_os = "linux")]
382 AV1Vaapi(Box<crate::vaapi_encode::VaapiAv1Encoder>),
383 AV1Software(Box<SoftwareAV1Encoder>),
384}
385
386impl SurfaceEncoder {
387 pub fn new(
391 preferences: &[SurfaceEncoderPreference],
392 width: u32,
393 height: u32,
394 vaapi_device: &str,
395 quality: SurfaceQuality,
396 verbose: bool,
397 codec_support: u8,
398 chroma: ChromaSubsampling,
399 ) -> Result<Self, String> {
400 let source_width = width;
401 let source_height = height;
402 let mut last_err = String::from("no encoders configured");
403
404 let try_444 = chroma.is_444();
409 if try_444 && verbose {
410 eprintln!(
411 "[surface-encoder] 4:4:4 eligible: codec_support={codec_support:#04x} for {source_width}x{source_height}",
412 );
413 }
414
415 for &pref in preferences {
416 if pref.is_vulkan_video() {
417 continue;
418 }
419 if !pref.supported_by_client(codec_support) {
420 continue;
421 }
422
423 if try_444 && pref.supports_444_by_client(codec_support) {
425 match Self::try_one(
426 pref,
427 width,
428 height,
429 source_width,
430 source_height,
431 vaapi_device,
432 quality,
433 verbose,
434 ChromaSubsampling::Cs444,
435 ) {
436 Ok(enc) => {
437 if verbose {
438 eprintln!(
439 "[surface-encoder] using {:?} 4:4:4 for {source_width}x{source_height}",
440 pref
441 );
442 }
443 return Ok(enc);
444 }
445 Err(err) => {
446 if verbose {
447 eprintln!(
448 "[surface-encoder] {:?} 4:4:4 unavailable for {source_width}x{source_height}: {err}",
449 pref
450 );
451 }
452 }
455 }
456 }
457
458 match Self::try_one(
460 pref,
461 width,
462 height,
463 source_width,
464 source_height,
465 vaapi_device,
466 quality,
467 verbose,
468 ChromaSubsampling::Cs420,
469 ) {
470 Ok(enc) => {
471 if verbose {
472 eprintln!(
473 "[surface-encoder] using {:?} 4:2:0 for {source_width}x{source_height}",
474 pref
475 );
476 }
477 return Ok(enc);
478 }
479 Err(err) => {
480 if verbose {
481 eprintln!(
482 "[surface-encoder] {:?} 4:2:0 unavailable for {source_width}x{source_height}: {err}",
483 pref
484 );
485 }
486 last_err = err;
487 }
488 }
489 }
490 Err(last_err)
491 }
492
493 fn try_one(
494 pref: SurfaceEncoderPreference,
495 width: u32,
496 height: u32,
497 source_width: u32,
498 source_height: u32,
499 vaapi_device: &str,
500 quality: SurfaceQuality,
501 verbose: bool,
502 chroma: ChromaSubsampling,
503 ) -> Result<Self, String> {
504 let _ = vaapi_device;
505 validate_surface_dimensions(width, height, pref)?;
506
507 if let Some(err) = cached_unavailable(pref, chroma) {
515 return Err(err);
516 }
517
518 let result = Self::try_one_inner(
519 pref,
520 width,
521 height,
522 source_width,
523 source_height,
524 vaapi_device,
525 quality,
526 verbose,
527 chroma,
528 );
529 if let Err(err) = &result {
530 record_unavailable(pref, chroma, err);
531 }
532 result
533 }
534
535 fn try_one_inner(
536 pref: SurfaceEncoderPreference,
537 width: u32,
538 height: u32,
539 source_width: u32,
540 source_height: u32,
541 vaapi_device: &str,
542 quality: SurfaceQuality,
543 verbose: bool,
544 chroma: ChromaSubsampling,
545 ) -> Result<Self, String> {
546 let _ = vaapi_device;
547 match pref {
548 SurfaceEncoderPreference::VulkanVideoH264
549 | SurfaceEncoderPreference::VulkanVideoAV1 => {
550 Err("Vulkan Video encoders are managed by the compositor".into())
551 }
552 SurfaceEncoderPreference::NvencH264 => {
553 let (width, height) = ((width + 1) & !1, (height + 1) & !1);
554 let qp = quality.h264_qp() as u32;
555 Ok(Self {
556 width,
557 height,
558 source_width,
559 source_height,
560 kind: SurfaceEncoderKind::NvencH264(Box::new(
561 crate::nvenc_encode::NvencDirectEncoder::try_new(
562 "h264", width, height, qp, verbose, chroma,
563 )?,
564 )),
565 chroma,
566 })
567 }
568 SurfaceEncoderPreference::NvencAV1 => {
569 let (width, height) = ((width + 1) & !1, (height + 1) & !1);
573 let qp = quality.nvenc_av1_qp();
574 Ok(Self {
575 width,
576 height,
577 source_width,
578 source_height,
579 kind: SurfaceEncoderKind::NvencAV1(Box::new(
580 crate::nvenc_encode::NvencDirectEncoder::try_new(
581 "av1", width, height, qp, verbose, chroma,
582 )?,
583 )),
584 chroma,
585 })
586 }
587 #[cfg(target_os = "linux")]
588 SurfaceEncoderPreference::H264Vaapi => {
589 let (width, height) = ((width + 1) & !1, (height + 1) & !1);
590 Ok(Self {
591 width,
592 height,
593 source_width,
594 source_height,
595 kind: SurfaceEncoderKind::H264Vaapi(Box::new(
596 crate::vaapi_encode::VaapiDirectEncoder::try_new(
597 width,
598 height,
599 vaapi_device,
600 quality.h264_qp(),
601 verbose,
602 chroma,
603 )?,
604 )),
605 chroma,
606 })
607 }
608 #[cfg(not(target_os = "linux"))]
609 SurfaceEncoderPreference::H264Vaapi => Err("VA-API is only available on Unix".into()),
610 #[cfg(target_os = "linux")]
611 SurfaceEncoderPreference::AV1Vaapi => {
612 const VAAPI_AV1_MIN: u32 = 512;
629 let enc_w = width.div_ceil(64) * 64;
630 let enc_h = height.div_ceil(64) * 64;
631 let (width, height) = (enc_w.max(VAAPI_AV1_MIN), enc_h.max(VAAPI_AV1_MIN));
632 let source_area = (source_width as u64) * (source_height as u64);
633 let padded_area = (width as u64) * (height as u64);
634 if padded_area > source_area.saturating_mul(2) {
635 return Err(format!(
636 "AV1Vaapi padding {width}x{height} > 2× source \
637 {source_width}x{source_height} — falling back",
638 ));
639 }
640 Ok(Self {
641 width,
642 height,
643 source_width,
644 source_height,
645 kind: SurfaceEncoderKind::AV1Vaapi(Box::new(
646 crate::vaapi_encode::VaapiAv1Encoder::try_new(
647 width,
648 height,
649 source_width,
650 source_height,
651 vaapi_device,
652 quality.av1_quantizer() as u8,
653 verbose,
654 chroma,
655 )?,
656 )),
657 chroma,
658 })
659 }
660 #[cfg(not(target_os = "linux"))]
661 SurfaceEncoderPreference::AV1Vaapi => Err("VA-API is only available on Linux".into()),
662 SurfaceEncoderPreference::AV1Software => Ok(Self {
663 width,
664 height,
665 source_width,
666 source_height,
667 kind: SurfaceEncoderKind::AV1Software(Box::new(SoftwareAV1Encoder::new(
668 width, height, quality, chroma,
669 )?)),
670 chroma,
671 }),
672 SurfaceEncoderPreference::H264Software => {
673 if chroma.is_444() {
674 return Err("openh264 does not support 4:4:4".into());
675 }
676 let (width, height) = ((width + 1) & !1, (height + 1) & !1);
677 Ok(Self {
678 width,
679 height,
680 source_width,
681 source_height,
682 kind: SurfaceEncoderKind::H264Software(Box::new(SoftwareH264Encoder::new(
683 quality,
684 )?)),
685 chroma,
686 })
687 }
688 }
689 }
690
691 pub fn source_dimensions(&self) -> (u32, u32) {
693 (self.source_width, self.source_height)
694 }
695
696 pub fn encoder_dimensions(&self) -> (u32, u32) {
699 (self.width, self.height)
700 }
701
702 pub fn encoder_name(&self) -> &'static str {
705 match (&self.kind, self.chroma) {
706 (SurfaceEncoderKind::H264Software(_), _) => "h264-software",
707 (SurfaceEncoderKind::NvencH264(_), ChromaSubsampling::Cs444) => "h264-nvenc 4:4:4",
708 (SurfaceEncoderKind::NvencH264(_), _) => "h264-nvenc",
709 (SurfaceEncoderKind::NvencAV1(_), ChromaSubsampling::Cs444) => "av1-nvenc 4:4:4",
710 (SurfaceEncoderKind::NvencAV1(_), _) => "av1-nvenc",
711 #[cfg(target_os = "linux")]
712 (SurfaceEncoderKind::H264Vaapi(_), ChromaSubsampling::Cs444) => "h264-vaapi 4:4:4",
713 #[cfg(target_os = "linux")]
714 (SurfaceEncoderKind::H264Vaapi(_), _) => "h264-vaapi",
715 #[cfg(target_os = "linux")]
716 (SurfaceEncoderKind::AV1Vaapi(_), ChromaSubsampling::Cs444) => "av1-vaapi 4:4:4",
717 #[cfg(target_os = "linux")]
718 (SurfaceEncoderKind::AV1Vaapi(_), _) => "av1-vaapi",
719 (SurfaceEncoderKind::AV1Software(_), ChromaSubsampling::Cs444) => "av1-software 4:4:4",
720 (SurfaceEncoderKind::AV1Software(_), _) => "av1-software",
721 }
722 }
723
724 pub fn webcodecs_codec_string(&self) -> String {
727 match &self.kind {
728 SurfaceEncoderKind::H264Software(_) => {
729 if self.chroma.is_444() {
730 "avc1.F4001f".to_string()
731 } else {
732 "avc1.42001f".to_string()
733 }
734 }
735 #[cfg(target_os = "linux")]
736 SurfaceEncoderKind::H264Vaapi(_) => {
737 if self.chroma.is_444() {
738 "avc1.F4001f".to_string()
739 } else {
740 "avc1.640034".to_string()
741 }
742 }
743 SurfaceEncoderKind::NvencH264(_) => "avc1.640034".to_string(),
744 SurfaceEncoderKind::NvencAV1(_) | SurfaceEncoderKind::AV1Software(_) => {
745 let profile = if self.chroma.is_444() { 2 } else { 0 };
746 let level = av1_level_for(self.source_width, self.source_height);
747 format!("av01.{profile}.{level}M.08")
748 }
749 #[cfg(target_os = "linux")]
750 SurfaceEncoderKind::AV1Vaapi(_) => {
751 let profile = if self.chroma.is_444() { 2 } else { 0 };
752 let level = av1_level_for(self.source_width, self.source_height);
753 format!("av01.{profile}.{level}M.08")
754 }
755 }
756 }
757
758 pub fn codec_flag(&self) -> u8 {
759 match &self.kind {
760 SurfaceEncoderKind::H264Software(_) => SURFACE_FRAME_CODEC_H264,
761 #[cfg(target_os = "linux")]
762 SurfaceEncoderKind::H264Vaapi(_) => SURFACE_FRAME_CODEC_H264,
763 SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => {
764 enc.codec_flag()
765 }
766 #[cfg(target_os = "linux")]
767 SurfaceEncoderKind::AV1Vaapi(_) => SURFACE_FRAME_CODEC_AV1,
768 SurfaceEncoderKind::AV1Software(_) => SURFACE_FRAME_CODEC_AV1,
769 }
770 }
771
772 pub fn request_keyframe(&mut self) {
773 match &mut self.kind {
774 SurfaceEncoderKind::H264Software(enc) => enc.request_keyframe(),
775 SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => {
776 enc.request_keyframe()
777 }
778 #[cfg(target_os = "linux")]
779 SurfaceEncoderKind::H264Vaapi(enc) => enc.request_keyframe(),
780 #[cfg(target_os = "linux")]
781 SurfaceEncoderKind::AV1Vaapi(enc) => enc.request_keyframe(),
782 SurfaceEncoderKind::AV1Software(enc) => enc.request_keyframe(),
783 }
784 }
785
786 #[cfg(target_os = "linux")]
788 pub fn gbm_buffers(&self) -> &[crate::vaapi_encode::GbmExportedBuffer] {
789 match &self.kind {
790 SurfaceEncoderKind::H264Vaapi(enc) => enc.gbm_buffers(),
791 SurfaceEncoderKind::AV1Vaapi(enc) => enc.gbm_buffers(),
792 _ => &[],
793 }
794 }
795
796 #[cfg(target_os = "linux")]
797 pub fn gbm_nv12_buffers(&self) -> &[crate::vaapi_encode::GbmNv12Buffer] {
798 match &self.kind {
799 SurfaceEncoderKind::H264Vaapi(enc) => enc.gbm_nv12_buffers(),
800 SurfaceEncoderKind::AV1Vaapi(enc) => enc.gbm_nv12_buffers(),
801 _ => &[],
802 }
803 }
804
805 #[cfg(target_os = "linux")]
806 pub fn allocate_nv12_buffers(&mut self, drm_fd: std::os::fd::RawFd, count: usize) {
807 match &mut self.kind {
808 SurfaceEncoderKind::H264Vaapi(enc) => {
809 if let Some(vpp) = &mut enc.vpp {
810 vpp.allocate_nv12_buffers(drm_fd, count);
811 }
812 }
813 SurfaceEncoderKind::AV1Vaapi(enc) => {
814 if let Some(vpp) = &mut enc.vpp {
815 vpp.allocate_nv12_buffers(drm_fd, count);
816 }
817 }
818 _ => {}
819 }
820 }
821
822 #[cfg(target_os = "linux")]
823 pub fn drm_fd_raw(&self) -> std::os::fd::RawFd {
824 use std::os::fd::AsRawFd;
825 match &self.kind {
826 SurfaceEncoderKind::H264Vaapi(enc) => enc._drm_fd.as_raw_fd(),
827 SurfaceEncoderKind::AV1Vaapi(enc) => enc._drm_fd.as_raw_fd(),
828 _ => -1,
829 }
830 }
831
832 #[cfg(target_os = "linux")]
834 #[allow(dead_code)]
835 pub fn va_display_usize(&self) -> usize {
836 match &self.kind {
837 SurfaceEncoderKind::H264Vaapi(enc) => enc.va_display_usize(),
838 SurfaceEncoderKind::AV1Vaapi(enc) => enc.va_display_usize(),
839 _ => 0,
840 }
841 }
842
843 pub fn encode(&mut self, rgba: &[u8]) -> Option<(Vec<u8>, bool)> {
844 if let SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) =
850 &mut self.kind
851 {
852 let (sw, sh) = (self.source_width as usize, self.source_height as usize);
853 let mut result = enc.encode_rgba_padded(rgba, sw, sh);
854 self.fixup_keyframe(&mut result);
855 return result;
856 }
857
858 let enc_len = expected_rgba_len(self.width, self.height);
859 let enc_len = match enc_len {
860 Some(v) => v,
861 None => {
862 eprintln!(
863 "[surface-encoder] expected_rgba_len overflow {}x{}",
864 self.width, self.height
865 );
866 return None;
867 }
868 };
869 let rgba = if rgba.len() == enc_len {
870 std::borrow::Cow::Borrowed(rgba)
871 } else {
872 let total_px = rgba.len() / 4;
876 if total_px == 0 {
877 return None;
878 }
879 let src_w = [self.width as usize, (self.width - 1) as usize]
881 .into_iter()
882 .find(|&w| w > 0 && total_px.is_multiple_of(w))?;
883 let src_h = total_px / src_w;
884 if src_h == 0 {
885 return None;
886 }
887 let dst_w = self.width as usize;
888 let dst_h = self.height as usize;
889 let mut padded = vec![0u8; enc_len];
890 for row in 0..dst_h {
891 let src_row = row.min(src_h - 1);
892 for col in 0..dst_w {
893 let src_col = col.min(src_w - 1);
894 let si = (src_row * src_w + src_col) * 4;
895 let di = (row * dst_w + col) * 4;
896 padded[di..di + 4].copy_from_slice(&rgba[si..si + 4]);
897 }
898 }
899 std::borrow::Cow::Owned(padded)
900 };
901
902 match &mut self.kind {
903 SurfaceEncoderKind::H264Software(encoder) => {
904 encoder.encode(&rgba, self.width, self.height)
905 }
906 SurfaceEncoderKind::NvencH264(_) | SurfaceEncoderKind::NvencAV1(_) => unreachable!(),
908 #[cfg(target_os = "linux")]
909 SurfaceEncoderKind::H264Vaapi(enc) => {
910 let mut bgra = rgba.into_owned();
911 for px in bgra.chunks_exact_mut(4) {
912 px.swap(0, 2);
913 }
914 let (sw, sh) = (self.source_width as usize, self.source_height as usize);
915 enc.encode_bgra_padded(&bgra, sw, sh)
916 }
917 #[cfg(target_os = "linux")]
918 SurfaceEncoderKind::AV1Vaapi(enc) => {
919 let mut bgra = rgba.into_owned();
920 for px in bgra.chunks_exact_mut(4) {
921 px.swap(0, 2);
922 }
923 let (sw, sh) = (self.source_width as usize, self.source_height as usize);
924 enc.encode_bgra_padded(&bgra, sw, sh)
925 }
926 SurfaceEncoderKind::AV1Software(encoder) => encoder.encode(&rgba),
927 }
928 }
929
930 pub fn encode_pixels(&mut self, pixels: &PixelData) -> Option<(Vec<u8>, bool)> {
933 match pixels {
934 PixelData::Nv12 {
935 data,
936 y_stride,
937 uv_stride,
938 } => self.encode_nv12(data, *y_stride, *uv_stride),
939 PixelData::Bgra(bgra) => self.encode_bgra(bgra),
940 PixelData::Rgba(rgba) => self.encode(rgba),
941 #[cfg(target_os = "linux")]
942 PixelData::DmaBuf {
943 fd,
944 fourcc,
945 modifier,
946 stride,
947 offset,
948 ..
949 } => self
950 .encode_dmabuf(fd, *fourcc, *modifier, *stride, *offset)
951 .or_else(|| {
952 let w = self.width;
955 let h = self.height;
956 let rgba = pixels.to_rgba(w, h);
957 if !rgba.is_empty() {
958 self.encode(&rgba)
959 } else {
960 None
961 }
962 }),
963 #[cfg(not(target_os = "linux"))]
964 PixelData::DmaBuf { .. } => None,
965 #[cfg(target_os = "linux")]
966 PixelData::Nv12DmaBuf {
967 fd,
968 stride,
969 uv_offset,
970 width,
971 height,
972 sync_fd,
973 } => {
974 if let Some(sfd) = sync_fd {
978 use std::os::fd::AsRawFd;
979 let mut pfd = libc::pollfd {
980 fd: sfd.as_raw_fd(),
981 events: libc::POLLIN,
982 revents: 0,
983 };
984 unsafe { libc::poll(&mut pfd, 1, 5000) };
985 }
986 self.encode_nv12_dmabuf(fd, *stride, *uv_offset, *width, *height)
987 }
988 .or_else(|| {
989 use std::os::fd::AsRawFd;
992 let h = *height as usize;
993 let s = *stride as usize;
994 let uv_off = *uv_offset as usize;
995 let raw = fd.as_raw_fd();
996 let map_size = uv_off + s * h.div_ceil(2);
997 let ptr = unsafe {
998 libc::mmap(
999 std::ptr::null_mut(),
1000 map_size,
1001 libc::PROT_READ,
1002 libc::MAP_SHARED,
1003 raw,
1004 0,
1005 )
1006 };
1007 if ptr == libc::MAP_FAILED || ptr.is_null() {
1008 return None;
1009 }
1010 let data = unsafe { std::slice::from_raw_parts(ptr as *const u8, map_size) };
1011 let result = self.encode_nv12(data, s, s);
1012 unsafe { libc::munmap(ptr, map_size) };
1013 result
1014 }),
1015 #[cfg(not(target_os = "linux"))]
1016 PixelData::Nv12DmaBuf { .. } => None,
1017 PixelData::VaSurface { .. } => None,
1018 PixelData::Encoded { .. } => None,
1021 }
1022 }
1023
1024 #[cfg(target_os = "linux")]
1028 fn encode_nv12_dmabuf(
1029 &mut self,
1030 fd: &std::sync::Arc<std::os::fd::OwnedFd>,
1031 _stride: u32,
1032 _uv_offset: u32,
1033 _width: u32,
1034 _height: u32,
1035 ) -> Option<(Vec<u8>, bool)> {
1036 use std::os::fd::AsRawFd;
1037 let raw_fd = fd.as_raw_fd();
1038 let find_surface = |nv12s: &[crate::vaapi_encode::GbmNv12Buffer]| -> Option<u32> {
1039 let buf = nv12s.iter().find(|n| n.fd.as_raw_fd() == raw_fd)?;
1040 if buf.va_surface == 0 {
1042 return None;
1043 }
1044 Some(buf.va_surface)
1045 };
1046 let mut result = match &mut self.kind {
1047 SurfaceEncoderKind::AV1Vaapi(enc) => {
1048 let surf = find_surface(enc.gbm_nv12_buffers())?;
1049 enc.encode_surface(surf)
1050 }
1051 SurfaceEncoderKind::H264Vaapi(enc) => {
1052 let surf = find_surface(enc.gbm_nv12_buffers())?;
1053 enc.encode_surface(surf)
1054 }
1055 _ => None,
1056 };
1057 self.fixup_keyframe(&mut result);
1058 result
1059 }
1060
1061 #[cfg(target_os = "linux")]
1064 fn encode_dmabuf(
1065 &mut self,
1066 fd: &std::os::fd::OwnedFd,
1067 fourcc: u32,
1068 modifier: u64,
1069 stride: u32,
1070 offset: u32,
1071 ) -> Option<(Vec<u8>, bool)> {
1072 use std::os::fd::AsRawFd;
1073
1074 let src_w = self.source_width;
1077 let src_h = self.source_height;
1078
1079 let mut gpu_result = match &mut self.kind {
1083 SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => enc
1084 .encode_dmabuf_fd(
1085 fd.as_raw_fd(),
1086 fourcc,
1087 modifier,
1088 stride,
1089 offset,
1090 src_w,
1091 src_h,
1092 ),
1093 _ => None,
1094 };
1095 if gpu_result.is_some() {
1096 self.fixup_keyframe(&mut gpu_result);
1097 return gpu_result;
1098 }
1099
1100 self.encode_dmabuf_cpu_fallback(fd, fourcc, stride, offset)
1105 }
1106
1107 #[cfg(target_os = "linux")]
1110 fn encode_dmabuf_cpu_fallback(
1111 &mut self,
1112 fd: &std::os::fd::OwnedFd,
1113 fourcc: u32,
1114 stride: u32,
1115 _offset: u32,
1116 ) -> Option<(Vec<u8>, bool)> {
1117 use std::os::fd::AsRawFd;
1118
1119 let w = self.source_width as usize;
1120 let h = self.source_height as usize;
1121 let stride = stride as usize;
1122 let raw_fd = fd.as_raw_fd();
1123
1124 let file_size = unsafe { libc::lseek(raw_fd, 0, libc::SEEK_END) };
1126 if file_size <= 0 {
1127 return None;
1128 }
1129 let map_len = file_size as usize;
1130
1131 #[repr(C)]
1133 struct DmaBufSync {
1134 flags: u64,
1135 }
1136 const DMA_BUF_SYNC_READ: u64 = 1;
1137 const DMA_BUF_SYNC_START: u64 = 0;
1138 const DMA_BUF_SYNC_END: u64 = 4;
1139 const DMA_BUF_IOCTL_SYNC: libc::c_ulong = 0x40086200;
1143
1144 {
1148 let mut pfd = libc::pollfd {
1149 fd: raw_fd,
1150 events: libc::POLLIN,
1151 revents: 0,
1152 };
1153 let ready = unsafe { libc::poll(&mut pfd, 1, 0) };
1154 if ready <= 0 {
1155 } else {
1157 let sync_start = DmaBufSync {
1158 flags: DMA_BUF_SYNC_START | DMA_BUF_SYNC_READ,
1159 };
1160 unsafe {
1161 libc::ioctl(raw_fd, DMA_BUF_IOCTL_SYNC as _, &sync_start);
1162 }
1163 }
1164 }
1165
1166 let ptr = unsafe {
1168 libc::mmap(
1169 std::ptr::null_mut(),
1170 map_len,
1171 libc::PROT_READ,
1172 libc::MAP_SHARED,
1173 raw_fd,
1174 0,
1175 )
1176 };
1177 if ptr == libc::MAP_FAILED {
1178 let sync_end = DmaBufSync {
1179 flags: DMA_BUF_SYNC_END | DMA_BUF_SYNC_READ,
1180 };
1181 unsafe {
1182 libc::ioctl(raw_fd, DMA_BUF_IOCTL_SYNC as _, &sync_end);
1183 }
1184 return None;
1185 }
1186 let plane_data = unsafe { std::slice::from_raw_parts(ptr as *const u8, map_len) };
1187
1188 let is_gl_fbo = {
1191 let mut link = [0u8; 128];
1192 let path = format!("/proc/self/fd/{raw_fd}\0");
1193 let n = unsafe {
1194 libc::readlink(path.as_ptr() as *const _, link.as_mut_ptr() as *mut _, 127)
1195 };
1196 !(n > 0 && link[..n as usize].starts_with(b"/dev/dri/"))
1197 };
1198
1199 let result = if fourcc == blit_compositor::drm_fourcc::ARGB8888
1200 || fourcc == blit_compositor::drm_fourcc::XRGB8888
1201 {
1202 let mut packed = Vec::with_capacity(w * h * 4);
1204 for i in 0..h {
1205 let row = if is_gl_fbo { h - 1 - i } else { i };
1207 let start = row * stride;
1208 let end = start + w * 4;
1209 if end <= plane_data.len() {
1210 packed.extend_from_slice(&plane_data[start..end]);
1211 }
1212 }
1213 self.encode_bgra(&packed)
1214 } else if fourcc == blit_compositor::drm_fourcc::ABGR8888
1215 || fourcc == blit_compositor::drm_fourcc::XBGR8888
1216 {
1217 let mut packed = Vec::with_capacity(w * h * 4);
1219 for i in 0..h {
1220 let row = if is_gl_fbo { h - 1 - i } else { i };
1221 let start = row * stride;
1222 let end = start + w * 4;
1223 if end <= plane_data.len() {
1224 packed.extend_from_slice(&plane_data[start..end]);
1225 }
1226 }
1227 self.encode(&packed)
1228 } else if fourcc == blit_compositor::drm_fourcc::NV12 {
1229 let uv_stride = stride; let y_size = stride * h;
1234 let uv_h = h.div_ceil(2);
1235 let uv_size = uv_stride * uv_h;
1236 if map_len >= y_size + uv_size {
1237 let out_stride = w;
1239 let mut data = vec![0u8; out_stride * h + out_stride * uv_h];
1240 for row in 0..h {
1241 let src = row * stride;
1242 let dst = row * out_stride;
1243 if src + w <= plane_data.len() {
1244 data[dst..dst + w].copy_from_slice(&plane_data[src..src + w]);
1245 }
1246 }
1247 let uv_dst_base = out_stride * h;
1248 for row in 0..uv_h {
1249 let src = y_size + row * uv_stride;
1250 let dst = uv_dst_base + row * out_stride;
1251 if src + w <= plane_data.len() {
1252 data[dst..dst + w].copy_from_slice(&plane_data[src..src + w]);
1253 }
1254 }
1255 self.encode_nv12(&data, out_stride, out_stride)
1256 } else {
1257 None
1258 }
1259 } else {
1260 None
1261 };
1262
1263 unsafe {
1265 libc::munmap(ptr, map_len);
1266 }
1267 let sync_end = DmaBufSync {
1269 flags: DMA_BUF_SYNC_END | DMA_BUF_SYNC_READ,
1270 };
1271 unsafe {
1272 libc::ioctl(raw_fd, DMA_BUF_IOCTL_SYNC as _, &sync_end);
1273 }
1274
1275 result
1276 }
1277
1278 fn fixup_keyframe(&self, result: &mut Option<(Vec<u8>, bool)>) {
1283 if let Some((data, is_key)) = result.as_mut()
1284 && !*is_key
1285 {
1286 *is_key = match &self.kind {
1287 SurfaceEncoderKind::NvencH264(_) => h264_stream_contains_idr(data),
1288 SurfaceEncoderKind::NvencAV1(_) => av1_stream_contains_keyframe(data),
1289 #[cfg(target_os = "linux")]
1290 SurfaceEncoderKind::H264Vaapi(_) => h264_stream_contains_idr(data),
1291 #[cfg(target_os = "linux")]
1292 SurfaceEncoderKind::AV1Vaapi(_) => av1_stream_contains_keyframe(data),
1293 _ => false,
1294 };
1295 }
1296 }
1297
1298 fn encode_bgra(&mut self, bgra: &[u8]) -> Option<(Vec<u8>, bool)> {
1300 let enc_w = self.width as usize;
1301 let enc_h = self.height as usize;
1302 let src_w = self.source_width as usize;
1303 let src_h = self.source_height as usize;
1304
1305 let mut result = match &mut self.kind {
1306 SurfaceEncoderKind::H264Software(encoder) => {
1307 let yuv = bgra_to_yuv420_padded(bgra, src_w, src_h, enc_w, enc_h);
1308 let yuv_buf = YUVBuffer::from_vec(yuv, enc_w, enc_h);
1309 encoder.encode_yuv(&yuv_buf, self.width, self.height)
1310 }
1311 SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => {
1312 enc.encode_bgra_padded(bgra, src_w, src_h)
1313 }
1314 #[cfg(target_os = "linux")]
1315 SurfaceEncoderKind::H264Vaapi(enc) => enc.encode_bgra_padded(bgra, src_w, src_h),
1316 #[cfg(target_os = "linux")]
1317 SurfaceEncoderKind::AV1Vaapi(enc) => enc.encode_bgra_padded(bgra, src_w, src_h),
1318 SurfaceEncoderKind::AV1Software(encoder) => {
1319 let yuv = if self.chroma.is_444() {
1320 bgra_to_yuv444_padded(bgra, src_w, src_h, enc_w, enc_h)
1321 } else {
1322 bgra_to_yuv420_padded(bgra, src_w, src_h, enc_w, enc_h)
1323 };
1324 encoder.encode_yuv_planes(&yuv)
1325 }
1326 };
1327 self.fixup_keyframe(&mut result);
1328 result
1329 }
1330
1331 fn encode_nv12(
1334 &mut self,
1335 data: &[u8],
1336 y_stride: usize,
1337 uv_stride: usize,
1338 ) -> Option<(Vec<u8>, bool)> {
1339 let src_w = self.source_width as usize;
1341 let src_h = self.source_height as usize;
1342
1343 let mut result = match &mut self.kind {
1344 SurfaceEncoderKind::H264Software(encoder) => {
1345 let enc_w = self.width as usize;
1346 let enc_h = self.height as usize;
1347 if enc_w == src_w && enc_h == src_h {
1348 let yuv = nv12_to_yuv420(data, y_stride, uv_stride, src_w, src_h);
1349 let yuv_buf = YUVBuffer::from_vec(yuv, enc_w, enc_h);
1350 encoder.encode_yuv(&yuv_buf, self.width, self.height)
1351 } else {
1352 let pd = PixelData::Nv12 {
1353 data: std::sync::Arc::new(data.to_vec()),
1354 y_stride,
1355 uv_stride,
1356 };
1357 let rgba = pd.to_rgba(self.source_width, self.source_height);
1358 return self.encode(&rgba);
1359 }
1360 }
1361 SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => {
1362 enc.encode_nv12(data, y_stride, uv_stride, src_h)
1364 }
1365 #[cfg(target_os = "linux")]
1366 SurfaceEncoderKind::H264Vaapi(enc) => {
1367 let uv_offset = y_stride * src_h;
1368 let y_data = &data[..uv_offset];
1369 let uv_data = &data[uv_offset..];
1370 enc.encode_nv12(y_data, uv_data, y_stride, uv_stride)
1371 }
1372 #[cfg(target_os = "linux")]
1373 SurfaceEncoderKind::AV1Vaapi(enc) => {
1374 let uv_offset = y_stride * src_h;
1375 let y_data = &data[..uv_offset];
1376 let uv_data = &data[uv_offset..];
1377 enc.encode_nv12(y_data, uv_data, y_stride, uv_stride)
1378 }
1379 SurfaceEncoderKind::AV1Software(encoder) => {
1380 encoder.encode_nv12(data, y_stride, uv_stride, src_w, src_h)
1381 }
1382 };
1383 self.fixup_keyframe(&mut result);
1384 result
1385 }
1386}
1387
1388fn validate_surface_dimensions(
1389 width: u32,
1390 height: u32,
1391 _preference: SurfaceEncoderPreference,
1392) -> Result<(), String> {
1393 if width == 0 || height == 0 {
1394 return Err("surface encoder requires non-zero dimensions".into());
1395 }
1396 let _ = expected_rgba_len(width, height)
1399 .ok_or_else(|| format!("surface encoder dimensions overflow for {width}x{height}"))?;
1400 Ok(())
1401}
1402
1403static ENCODER_UNAVAILABLE: std::sync::OnceLock<
1414 std::sync::Mutex<
1415 std::collections::HashMap<(SurfaceEncoderPreference, ChromaSubsampling), String>,
1416 >,
1417> = std::sync::OnceLock::new();
1418
1419fn encoder_unavailable_map() -> &'static std::sync::Mutex<
1420 std::collections::HashMap<(SurfaceEncoderPreference, ChromaSubsampling), String>,
1421> {
1422 ENCODER_UNAVAILABLE.get_or_init(Default::default)
1423}
1424
1425fn cached_unavailable(pref: SurfaceEncoderPreference, chroma: ChromaSubsampling) -> Option<String> {
1426 encoder_unavailable_map()
1427 .lock()
1428 .ok()?
1429 .get(&(pref, chroma))
1430 .cloned()
1431}
1432
1433fn record_unavailable(pref: SurfaceEncoderPreference, chroma: ChromaSubsampling, err: &str) {
1434 if let Ok(mut map) = encoder_unavailable_map().lock() {
1435 map.entry((pref, chroma)).or_insert_with(|| err.to_string());
1436 }
1437}
1438
1439fn expected_rgba_len(width: u32, height: u32) -> Option<usize> {
1440 (width as usize)
1441 .checked_mul(height as usize)?
1442 .checked_mul(4)
1443}
1444
1445#[inline(always)]
1451fn rgb_to_y(r: i32, g: i32, b: i32) -> u8 {
1452 ((66 * r + 129 * g + 25 * b + 128) >> 8)
1453 .wrapping_add(16)
1454 .clamp(0, 255) as u8
1455}
1456
1457#[inline(always)]
1458fn rgb_to_u(r: i32, g: i32, b: i32) -> u8 {
1459 ((-38 * r - 74 * g + 112 * b + 128) >> 8)
1460 .wrapping_add(128)
1461 .clamp(0, 255) as u8
1462}
1463
1464#[inline(always)]
1465fn rgb_to_v(r: i32, g: i32, b: i32) -> u8 {
1466 ((112 * r - 94 * g - 18 * b + 128) >> 8)
1467 .wrapping_add(128)
1468 .clamp(0, 255) as u8
1469}
1470
1471#[inline(always)]
1480fn compute_y_plane(
1481 src: &[u8],
1482 width: usize,
1483 height: usize,
1484 y_plane: &mut [u8],
1485 r_off: usize,
1486 g_off: usize,
1487 b_off: usize,
1488) {
1489 let total = width * height;
1490 for (px, y_out) in y_plane[..total].iter_mut().enumerate() {
1491 let i = px * 4;
1492 let r = src[i + r_off] as i32;
1493 let g = src[i + g_off] as i32;
1494 let b = src[i + b_off] as i32;
1495 *y_out = rgb_to_y(r, g, b);
1496 }
1497}
1498
1499#[inline(always)]
1501fn compute_uv_planes(
1502 src: &[u8],
1503 width: usize,
1504 height: usize,
1505 u_plane: &mut [u8],
1506 v_plane: &mut [u8],
1507 r_off: usize,
1508 g_off: usize,
1509 b_off: usize,
1510) {
1511 let chroma_w = width.div_ceil(2);
1512 let chroma_h = height.div_ceil(2);
1513 for cy in 0..chroma_h {
1514 for cx in 0..chroma_w {
1515 let row = cy * 2;
1516 let col = cx * 2;
1517 let mut u_sum = 0i32;
1519 let mut v_sum = 0i32;
1520 for dy in 0..2u32 {
1521 for dx in 0..2u32 {
1522 let sr = (row + dy as usize).min(height - 1);
1523 let sc = (col + dx as usize).min(width - 1);
1524 let i = (sr * width + sc) * 4;
1525 let r = src[i + r_off] as i32;
1526 let g = src[i + g_off] as i32;
1527 let b = src[i + b_off] as i32;
1528 u_sum += rgb_to_u(r, g, b) as i32;
1529 v_sum += rgb_to_v(r, g, b) as i32;
1530 }
1531 }
1532 let idx = cy * chroma_w + cx;
1533 u_plane[idx] = (u_sum / 4) as u8;
1534 v_plane[idx] = (v_sum / 4) as u8;
1535 }
1536 }
1537}
1538
1539#[inline(always)]
1542fn compute_y_plane_padded(
1543 src: &[u8],
1544 src_w: usize,
1545 src_h: usize,
1546 enc_w: usize,
1547 enc_h: usize,
1548 y_plane: &mut [u8],
1549 r_off: usize,
1550 g_off: usize,
1551 b_off: usize,
1552) {
1553 for row in 0..enc_h {
1554 let sr = row.min(src_h - 1);
1555 for col in 0..enc_w {
1556 let sc = col.min(src_w - 1);
1557 let i = (sr * src_w + sc) * 4;
1558 let r = src[i + r_off] as i32;
1559 let g = src[i + g_off] as i32;
1560 let b = src[i + b_off] as i32;
1561 y_plane[row * enc_w + col] = rgb_to_y(r, g, b);
1562 }
1563 }
1564}
1565
1566#[inline(always)]
1569fn compute_uv_planes_padded(
1570 src: &[u8],
1571 src_w: usize,
1572 src_h: usize,
1573 enc_w: usize,
1574 enc_h: usize,
1575 u_plane: &mut [u8],
1576 v_plane: &mut [u8],
1577 r_off: usize,
1578 g_off: usize,
1579 b_off: usize,
1580) {
1581 let chroma_w = enc_w.div_ceil(2);
1582 let chroma_h = enc_h.div_ceil(2);
1583 for cy in 0..chroma_h {
1584 for cx in 0..chroma_w {
1585 let row = cy * 2;
1586 let col = cx * 2;
1587 let mut u_sum = 0i32;
1588 let mut v_sum = 0i32;
1589 for dy in 0..2u32 {
1590 for dx in 0..2u32 {
1591 let sr = (row + dy as usize).min(src_h - 1);
1592 let sc = (col + dx as usize).min(src_w - 1);
1593 let i = (sr * src_w + sc) * 4;
1594 let r = src[i + r_off] as i32;
1595 let g = src[i + g_off] as i32;
1596 let b = src[i + b_off] as i32;
1597 u_sum += rgb_to_u(r, g, b) as i32;
1598 v_sum += rgb_to_v(r, g, b) as i32;
1599 }
1600 }
1601 let idx = cy * chroma_w + cx;
1602 u_plane[idx] = (u_sum / 4) as u8;
1603 v_plane[idx] = (v_sum / 4) as u8;
1604 }
1605 }
1606}
1607
1608#[inline(always)]
1611fn compute_uv_planes_444_padded(
1612 src: &[u8],
1613 src_w: usize,
1614 src_h: usize,
1615 enc_w: usize,
1616 enc_h: usize,
1617 u_plane: &mut [u8],
1618 v_plane: &mut [u8],
1619 r_off: usize,
1620 g_off: usize,
1621 b_off: usize,
1622) {
1623 for row in 0..enc_h {
1624 let sr = row.min(src_h - 1);
1625 for col in 0..enc_w {
1626 let sc = col.min(src_w - 1);
1627 let i = (sr * src_w + sc) * 4;
1628 let r = src[i + r_off] as i32;
1629 let g = src[i + g_off] as i32;
1630 let b = src[i + b_off] as i32;
1631 let idx = row * enc_w + col;
1632 u_plane[idx] = rgb_to_u(r, g, b);
1633 v_plane[idx] = rgb_to_v(r, g, b);
1634 }
1635 }
1636}
1637
1638fn bgra_to_yuv444_padded(
1640 bgra: &[u8],
1641 src_w: usize,
1642 src_h: usize,
1643 enc_w: usize,
1644 enc_h: usize,
1645) -> Vec<u8> {
1646 let plane_size = enc_w * enc_h;
1647 let mut yuv = vec![0u8; plane_size * 3];
1648 let (y_plane, uv) = yuv.split_at_mut(plane_size);
1649 let (u_plane, v_plane) = uv.split_at_mut(plane_size);
1650 compute_y_plane_padded(bgra, src_w, src_h, enc_w, enc_h, y_plane, 2, 1, 0);
1652 compute_uv_planes_444_padded(bgra, src_w, src_h, enc_w, enc_h, u_plane, v_plane, 2, 1, 0);
1653 yuv
1654}
1655
1656fn rgba_to_yuv444(rgba: &[u8], width: usize, height: usize) -> Vec<u8> {
1658 let plane_size = width * height;
1659 let mut yuv = vec![0u8; plane_size * 3];
1660 let (y_plane, uv) = yuv.split_at_mut(plane_size);
1661 let (u_plane, v_plane) = uv.split_at_mut(plane_size);
1662 compute_y_plane(rgba, width, height, y_plane, 0, 1, 2);
1664 compute_uv_planes_444_padded(
1665 rgba, width, height, width, height, u_plane, v_plane, 0, 1, 2,
1666 );
1667 yuv
1668}
1669
1670fn bgra_to_yuv420_padded(
1674 bgra: &[u8],
1675 src_w: usize,
1676 src_h: usize,
1677 enc_w: usize,
1678 enc_h: usize,
1679) -> Vec<u8> {
1680 let y_size = enc_w * enc_h;
1681 let uv_w = enc_w.div_ceil(2);
1686 let uv_size = uv_w * enc_h.div_ceil(2);
1687 let mut yuv = vec![0u8; y_size + uv_size * 2];
1688 let (y_plane, uv) = yuv.split_at_mut(y_size);
1689 let (u_plane, v_plane) = uv.split_at_mut(uv_size);
1690 compute_y_plane_padded(bgra, src_w, src_h, enc_w, enc_h, y_plane, 2, 1, 0);
1692 compute_uv_planes_padded(bgra, src_w, src_h, enc_w, enc_h, u_plane, v_plane, 2, 1, 0);
1693 yuv
1694}
1695
1696fn rgba_to_yuv420(rgba: &[u8], width: usize, height: usize) -> Vec<u8> {
1698 let y_size = width * height;
1699 let uv_w = width.div_ceil(2);
1700 let uv_size = uv_w * height.div_ceil(2);
1701 let mut yuv = vec![0u8; y_size + uv_size * 2];
1702 let (y_plane, uv) = yuv.split_at_mut(y_size);
1703 let (u_plane, v_plane) = uv.split_at_mut(uv_size);
1704 compute_y_plane(rgba, width, height, y_plane, 0, 1, 2);
1706 compute_uv_planes(rgba, width, height, u_plane, v_plane, 0, 1, 2);
1707 yuv
1708}
1709
1710fn nv12_to_yuv420(
1714 data: &[u8],
1715 y_stride: usize,
1716 uv_stride: usize,
1717 width: usize,
1718 height: usize,
1719) -> Vec<u8> {
1720 let y_size = width * height;
1721 let uv_w = width.div_ceil(2);
1722 let uv_h = height.div_ceil(2);
1723 let uv_size = uv_w * uv_h;
1724 let mut yuv = vec![0u8; y_size + uv_size * 2];
1725 let (y_out, uv_out) = yuv.split_at_mut(y_size);
1726 let (u_out, v_out) = uv_out.split_at_mut(uv_size);
1727
1728 let uv_offset = y_stride * height;
1729
1730 for row in 0..height {
1732 let src = row * y_stride;
1733 let dst = row * width;
1734 y_out[dst..dst + width].copy_from_slice(&data[src..src + width]);
1735 }
1736
1737 let src_uv_pairs = width / 2;
1741 for row in 0..uv_h {
1742 let src_start = uv_offset + row.min(height / 2 - 1) * uv_stride;
1743 let dst_start = row * uv_w;
1744 for col in 0..uv_w {
1745 let sc = col.min(src_uv_pairs.saturating_sub(1));
1746 u_out[dst_start + col] = data[src_start + sc * 2];
1747 v_out[dst_start + col] = data[src_start + sc * 2 + 1];
1748 }
1749 }
1750
1751 yuv
1752}
1753
1754fn h264_stream_contains_idr(data: &[u8]) -> bool {
1756 annex_b_contains_nal(data, |byte| (byte & 0x1f) == 5)
1757}
1758
1759fn annex_b_contains_nal(data: &[u8], pred: impl Fn(u8) -> bool) -> bool {
1761 let mut i = 0usize;
1762 while i < data.len() {
1763 let start_code_len = if data[i..].starts_with(&[0, 0, 0, 1]) {
1764 4
1765 } else if data[i..].starts_with(&[0, 0, 1]) {
1766 3
1767 } else {
1768 i += 1;
1769 continue;
1770 };
1771
1772 let nal_header = i + start_code_len;
1773 if let Some(&byte) = data.get(nal_header)
1774 && pred(byte)
1775 {
1776 return true;
1777 }
1778
1779 i = nal_header.saturating_add(1);
1780 }
1781
1782 false
1783}
1784
1785fn av1_stream_contains_keyframe(data: &[u8]) -> bool {
1793 let mut pos = 0;
1797 while pos < data.len() {
1798 let header = data[pos];
1799 let obu_type = (header >> 3) & 0xF;
1800 let has_extension = (header >> 2) & 1;
1801 let has_size = (header >> 1) & 1;
1802 pos += 1;
1803
1804 if has_extension != 0 {
1806 if pos >= data.len() {
1807 break;
1808 }
1809 pos += 1;
1810 }
1811
1812 if obu_type == 1 {
1814 return true;
1815 }
1816
1817 if has_size != 0 {
1820 let mut size: u64 = 0;
1821 let mut shift = 0u32;
1822 while pos < data.len() {
1823 let byte = data[pos];
1824 pos += 1;
1825 size |= ((byte & 0x7F) as u64) << shift;
1826 if byte & 0x80 == 0 {
1827 break;
1828 }
1829 shift += 7;
1830 if shift >= 56 {
1831 return false; }
1833 }
1834 pos = pos.saturating_add(size as usize);
1835 } else {
1836 break;
1839 }
1840 }
1841 false
1842}
1843
1844struct SoftwareH264Encoder {
1845 encoder: OpenH264Encoder,
1846}
1847
1848impl SoftwareH264Encoder {
1849 fn new(quality: SurfaceQuality) -> Result<Self, String> {
1850 use openh264::encoder::{EncoderConfig, RateControlMode};
1851 let config = EncoderConfig::new()
1852 .set_bitrate_bps(quality.openh264_bitrate())
1853 .rate_control_mode(RateControlMode::Bitrate);
1854 let encoder =
1855 OpenH264Encoder::with_api_config(openh264::OpenH264API::from_source(), config)
1856 .map_err(|err| format!("failed to create OpenH264 encoder: {err:?}"))?;
1857 Ok(Self { encoder })
1858 }
1859
1860 fn request_keyframe(&mut self) {
1861 self.encoder.force_intra_frame();
1862 }
1863
1864 fn encode(&mut self, rgba: &[u8], width: u32, height: u32) -> Option<(Vec<u8>, bool)> {
1865 let yuv = rgba_to_yuv420(rgba, width as usize, height as usize);
1866 let yuv_buf = YUVBuffer::from_vec(yuv, width as usize, height as usize);
1867 self.encode_yuv(&yuv_buf, width, height)
1868 }
1869
1870 fn encode_yuv(
1872 &mut self,
1873 yuv_buf: &YUVBuffer,
1874 width: u32,
1875 height: u32,
1876 ) -> Option<(Vec<u8>, bool)> {
1877 let bitstream = match self.encoder.encode(yuv_buf) {
1878 Ok(bs) => bs,
1879 Err(e) => {
1880 eprintln!("[surface-encoder] openh264 encode failed {width}x{height}: {e:?}");
1881 return None;
1882 }
1883 };
1884 let nal_data = bitstream.to_vec();
1885 if nal_data.is_empty() {
1886 eprintln!("[surface-encoder] openh264 produced empty NAL {width}x{height}");
1887 return None;
1888 }
1889 let is_keyframe = h264_stream_contains_idr(&nal_data);
1890 Some((nal_data, is_keyframe))
1891 }
1892}
1893
1894struct SoftwareAV1Encoder {
1899 ctx: rav1e::Context<u8>,
1900 width: usize,
1901 height: usize,
1902 force_keyframe: bool,
1903 chroma: ChromaSubsampling,
1904}
1905
1906impl SoftwareAV1Encoder {
1907 fn new(
1908 width: u32,
1909 height: u32,
1910 quality: SurfaceQuality,
1911 chroma: ChromaSubsampling,
1912 ) -> Result<Self, String> {
1913 use rav1e::prelude::*;
1914
1915 let chroma_sampling = if chroma.is_444() {
1916 ChromaSampling::Cs444
1917 } else {
1918 ChromaSampling::Cs420
1919 };
1920 let mut speed = SpeedSettings::from_preset(quality.av1_speed());
1921 speed.rdo_lookahead_frames = 1;
1922 let enc = EncoderConfig {
1923 width: width as usize,
1924 height: height as usize,
1925 chroma_sampling,
1926 chroma_sample_position: ChromaSamplePosition::Unknown,
1927 speed_settings: speed,
1928 low_latency: true,
1929 min_key_frame_interval: 0,
1930 max_key_frame_interval: 60,
1931 quantizer: quality.av1_quantizer(),
1932 min_quantizer: quality.av1_min_quantizer(),
1933 bitrate: 0,
1934 ..Default::default()
1935 };
1936 let cfg = Config::new().with_encoder_config(enc);
1937 let ctx = cfg
1938 .new_context()
1939 .map_err(|e| format!("rav1e context creation failed: {e}"))?;
1940 Ok(Self {
1941 ctx,
1942 width: width as usize,
1943 height: height as usize,
1944 force_keyframe: false,
1945 chroma,
1946 })
1947 }
1948
1949 fn request_keyframe(&mut self) {
1950 self.force_keyframe = true;
1951 }
1952
1953 fn encode(&mut self, rgba: &[u8]) -> Option<(Vec<u8>, bool)> {
1954 let yuv = if self.chroma.is_444() {
1955 rgba_to_yuv444(rgba, self.width, self.height)
1956 } else {
1957 rgba_to_yuv420(rgba, self.width, self.height)
1958 };
1959 self.encode_yuv_planes(&yuv)
1960 }
1961
1962 fn encode_nv12(
1963 &mut self,
1964 data: &[u8],
1965 y_stride: usize,
1966 uv_stride: usize,
1967 width: usize,
1968 height: usize,
1969 ) -> Option<(Vec<u8>, bool)> {
1970 let yuv = nv12_to_yuv420(data, y_stride, uv_stride, width, height);
1971 self.encode_yuv_planes(&yuv)
1972 }
1973
1974 fn encode_yuv_planes(&mut self, yuv: &[u8]) -> Option<(Vec<u8>, bool)> {
1977 let width = self.width;
1978 let height = self.height;
1979 let y_size = width * height;
1980 let (uv_w, uv_size) = if self.chroma.is_444() {
1981 (width, width * height)
1982 } else {
1983 let uv_w = width.div_ceil(2);
1984 let uv_h = height.div_ceil(2);
1985 (uv_w, uv_w * uv_h)
1986 };
1987
1988 let y_plane = &yuv[..y_size];
1989 let u_plane = &yuv[y_size..y_size + uv_size];
1990 let v_plane = &yuv[y_size + uv_size..];
1991
1992 let mut frame = self.ctx.new_frame();
1993 frame.planes[0].copy_from_raw_u8(y_plane, width, 1);
1994 frame.planes[1].copy_from_raw_u8(u_plane, uv_w, 1);
1995 frame.planes[2].copy_from_raw_u8(v_plane, uv_w, 1);
1996
1997 self.send_and_receive(frame)
1998 }
1999
2000 fn send_and_receive(&mut self, frame: rav1e::Frame<u8>) -> Option<(Vec<u8>, bool)> {
2001 use rav1e::prelude::*;
2002
2003 if self.force_keyframe {
2004 let params = FrameParameters {
2005 frame_type_override: FrameTypeOverride::Key,
2006 ..Default::default()
2007 };
2008 if self.ctx.send_frame((frame, params)).is_ok() {
2009 self.force_keyframe = false;
2010 }
2011 } else {
2012 let _ = self.ctx.send_frame(frame);
2013 }
2014
2015 match self.ctx.receive_packet() {
2016 Ok(packet) => {
2017 let is_key = packet.frame_type == rav1e::prelude::FrameType::KEY;
2018 Some((packet.data, is_key))
2019 }
2020 Err(rav1e::EncoderStatus::Encoded) | Err(rav1e::EncoderStatus::NeedMoreData) => None,
2021 Err(_) => None,
2022 }
2023 }
2024}
2025
2026#[cfg(test)]
2027mod tests {
2028 use super::*;
2029
2030 fn make_obu(obu_type: u8, payload: &[u8]) -> Vec<u8> {
2032 let header = (obu_type & 0xF) << 3 | 0b10; let mut obu = vec![header];
2035 let mut size = payload.len();
2037 loop {
2038 let mut byte = (size & 0x7F) as u8;
2039 size >>= 7;
2040 if size > 0 {
2041 byte |= 0x80;
2042 }
2043 obu.push(byte);
2044 if size == 0 {
2045 break;
2046 }
2047 }
2048 obu.extend_from_slice(payload);
2049 obu
2050 }
2051
2052 #[test]
2053 fn av1_keyframe_with_sequence_header_only() {
2054 let data = make_obu(1, &[0xAA; 10]);
2056 assert!(av1_stream_contains_keyframe(&data));
2057 }
2058
2059 #[test]
2060 fn av1_keyframe_with_temporal_delimiter_prefix() {
2061 let mut data = make_obu(2, &[]); data.extend(make_obu(1, &[0xBB; 8])); data.extend(make_obu(6, &[0xCC; 20])); assert!(av1_stream_contains_keyframe(&data));
2067 }
2068
2069 #[test]
2070 fn av1_non_keyframe_with_temporal_delimiter() {
2071 let mut data = make_obu(2, &[]);
2073 data.extend(make_obu(6, &[0xDD; 15]));
2074 assert!(!av1_stream_contains_keyframe(&data));
2075 }
2076
2077 #[test]
2078 fn av1_non_keyframe_frame_header_only() {
2079 let data = make_obu(3, &[0xEE; 5]);
2081 assert!(!av1_stream_contains_keyframe(&data));
2082 }
2083
2084 #[test]
2085 fn av1_empty_stream() {
2086 assert!(!av1_stream_contains_keyframe(&[]));
2087 }
2088
2089 #[test]
2090 fn av1_keyframe_large_leb128_size() {
2091 let mut data = make_obu(2, &[0x00; 200]);
2094 data.extend(make_obu(1, &[0xFF; 4]));
2095 assert!(av1_stream_contains_keyframe(&data));
2096 }
2097}