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)]
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)]
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 match pref {
508 SurfaceEncoderPreference::VulkanVideoH264
509 | SurfaceEncoderPreference::VulkanVideoAV1 => {
510 Err("Vulkan Video encoders are managed by the compositor".into())
511 }
512 SurfaceEncoderPreference::NvencH264 => {
513 let (width, height) = ((width + 1) & !1, (height + 1) & !1);
514 let qp = quality.h264_qp() as u32;
515 Ok(Self {
516 width,
517 height,
518 source_width,
519 source_height,
520 kind: SurfaceEncoderKind::NvencH264(Box::new(
521 crate::nvenc_encode::NvencDirectEncoder::try_new(
522 "h264", width, height, qp, verbose, chroma,
523 )?,
524 )),
525 chroma,
526 })
527 }
528 SurfaceEncoderPreference::NvencAV1 => {
529 let (width, height) = ((width + 1) & !1, (height + 1) & !1);
533 let qp = quality.nvenc_av1_qp();
534 Ok(Self {
535 width,
536 height,
537 source_width,
538 source_height,
539 kind: SurfaceEncoderKind::NvencAV1(Box::new(
540 crate::nvenc_encode::NvencDirectEncoder::try_new(
541 "av1", width, height, qp, verbose, chroma,
542 )?,
543 )),
544 chroma,
545 })
546 }
547 #[cfg(target_os = "linux")]
548 SurfaceEncoderPreference::H264Vaapi => {
549 let (width, height) = ((width + 1) & !1, (height + 1) & !1);
550 Ok(Self {
551 width,
552 height,
553 source_width,
554 source_height,
555 kind: SurfaceEncoderKind::H264Vaapi(Box::new(
556 crate::vaapi_encode::VaapiDirectEncoder::try_new(
557 width,
558 height,
559 vaapi_device,
560 quality.h264_qp(),
561 verbose,
562 chroma,
563 )?,
564 )),
565 chroma,
566 })
567 }
568 #[cfg(not(target_os = "linux"))]
569 SurfaceEncoderPreference::H264Vaapi => Err("VA-API is only available on Unix".into()),
570 #[cfg(target_os = "linux")]
571 SurfaceEncoderPreference::AV1Vaapi => {
572 const VAAPI_AV1_MIN: u32 = 512;
589 let enc_w = width.div_ceil(64) * 64;
590 let enc_h = height.div_ceil(64) * 64;
591 let (width, height) = (enc_w.max(VAAPI_AV1_MIN), enc_h.max(VAAPI_AV1_MIN));
592 let source_area = (source_width as u64) * (source_height as u64);
593 let padded_area = (width as u64) * (height as u64);
594 if padded_area > source_area.saturating_mul(2) {
595 return Err(format!(
596 "AV1Vaapi padding {width}x{height} > 2× source \
597 {source_width}x{source_height} — falling back",
598 ));
599 }
600 Ok(Self {
601 width,
602 height,
603 source_width,
604 source_height,
605 kind: SurfaceEncoderKind::AV1Vaapi(Box::new(
606 crate::vaapi_encode::VaapiAv1Encoder::try_new(
607 width,
608 height,
609 source_width,
610 source_height,
611 vaapi_device,
612 quality.av1_quantizer() as u8,
613 verbose,
614 chroma,
615 )?,
616 )),
617 chroma,
618 })
619 }
620 #[cfg(not(target_os = "linux"))]
621 SurfaceEncoderPreference::AV1Vaapi => Err("VA-API is only available on Linux".into()),
622 SurfaceEncoderPreference::AV1Software => Ok(Self {
623 width,
624 height,
625 source_width,
626 source_height,
627 kind: SurfaceEncoderKind::AV1Software(Box::new(SoftwareAV1Encoder::new(
628 width, height, quality, chroma,
629 )?)),
630 chroma,
631 }),
632 SurfaceEncoderPreference::H264Software => {
633 if chroma.is_444() {
634 return Err("openh264 does not support 4:4:4".into());
635 }
636 let (width, height) = ((width + 1) & !1, (height + 1) & !1);
637 Ok(Self {
638 width,
639 height,
640 source_width,
641 source_height,
642 kind: SurfaceEncoderKind::H264Software(Box::new(SoftwareH264Encoder::new(
643 quality,
644 )?)),
645 chroma,
646 })
647 }
648 }
649 }
650
651 pub fn source_dimensions(&self) -> (u32, u32) {
653 (self.source_width, self.source_height)
654 }
655
656 pub fn encoder_dimensions(&self) -> (u32, u32) {
659 (self.width, self.height)
660 }
661
662 pub fn encoder_name(&self) -> &'static str {
665 match (&self.kind, self.chroma) {
666 (SurfaceEncoderKind::H264Software(_), _) => "h264-software",
667 (SurfaceEncoderKind::NvencH264(_), ChromaSubsampling::Cs444) => "h264-nvenc 4:4:4",
668 (SurfaceEncoderKind::NvencH264(_), _) => "h264-nvenc",
669 (SurfaceEncoderKind::NvencAV1(_), ChromaSubsampling::Cs444) => "av1-nvenc 4:4:4",
670 (SurfaceEncoderKind::NvencAV1(_), _) => "av1-nvenc",
671 #[cfg(target_os = "linux")]
672 (SurfaceEncoderKind::H264Vaapi(_), ChromaSubsampling::Cs444) => "h264-vaapi 4:4:4",
673 #[cfg(target_os = "linux")]
674 (SurfaceEncoderKind::H264Vaapi(_), _) => "h264-vaapi",
675 #[cfg(target_os = "linux")]
676 (SurfaceEncoderKind::AV1Vaapi(_), ChromaSubsampling::Cs444) => "av1-vaapi 4:4:4",
677 #[cfg(target_os = "linux")]
678 (SurfaceEncoderKind::AV1Vaapi(_), _) => "av1-vaapi",
679 (SurfaceEncoderKind::AV1Software(_), ChromaSubsampling::Cs444) => "av1-software 4:4:4",
680 (SurfaceEncoderKind::AV1Software(_), _) => "av1-software",
681 }
682 }
683
684 pub fn webcodecs_codec_string(&self) -> String {
687 match &self.kind {
688 SurfaceEncoderKind::H264Software(_) => {
689 if self.chroma.is_444() {
690 "avc1.F4001f".to_string()
691 } else {
692 "avc1.42001f".to_string()
693 }
694 }
695 #[cfg(target_os = "linux")]
696 SurfaceEncoderKind::H264Vaapi(_) => {
697 if self.chroma.is_444() {
698 "avc1.F4001f".to_string()
699 } else {
700 "avc1.640034".to_string()
701 }
702 }
703 SurfaceEncoderKind::NvencH264(_) => "avc1.640034".to_string(),
704 SurfaceEncoderKind::NvencAV1(_) | SurfaceEncoderKind::AV1Software(_) => {
705 let profile = if self.chroma.is_444() { 2 } else { 0 };
706 let level = av1_level_for(self.source_width, self.source_height);
707 format!("av01.{profile}.{level}M.08")
708 }
709 #[cfg(target_os = "linux")]
710 SurfaceEncoderKind::AV1Vaapi(_) => {
711 let profile = if self.chroma.is_444() { 2 } else { 0 };
712 let level = av1_level_for(self.source_width, self.source_height);
713 format!("av01.{profile}.{level}M.08")
714 }
715 }
716 }
717
718 pub fn codec_flag(&self) -> u8 {
719 match &self.kind {
720 SurfaceEncoderKind::H264Software(_) => SURFACE_FRAME_CODEC_H264,
721 #[cfg(target_os = "linux")]
722 SurfaceEncoderKind::H264Vaapi(_) => SURFACE_FRAME_CODEC_H264,
723 SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => {
724 enc.codec_flag()
725 }
726 #[cfg(target_os = "linux")]
727 SurfaceEncoderKind::AV1Vaapi(_) => SURFACE_FRAME_CODEC_AV1,
728 SurfaceEncoderKind::AV1Software(_) => SURFACE_FRAME_CODEC_AV1,
729 }
730 }
731
732 pub fn request_keyframe(&mut self) {
733 match &mut self.kind {
734 SurfaceEncoderKind::H264Software(enc) => enc.request_keyframe(),
735 SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => {
736 enc.request_keyframe()
737 }
738 #[cfg(target_os = "linux")]
739 SurfaceEncoderKind::H264Vaapi(enc) => enc.request_keyframe(),
740 #[cfg(target_os = "linux")]
741 SurfaceEncoderKind::AV1Vaapi(enc) => enc.request_keyframe(),
742 SurfaceEncoderKind::AV1Software(enc) => enc.request_keyframe(),
743 }
744 }
745
746 #[cfg(target_os = "linux")]
748 pub fn gbm_buffers(&self) -> &[crate::vaapi_encode::GbmExportedBuffer] {
749 match &self.kind {
750 SurfaceEncoderKind::H264Vaapi(enc) => enc.gbm_buffers(),
751 SurfaceEncoderKind::AV1Vaapi(enc) => enc.gbm_buffers(),
752 _ => &[],
753 }
754 }
755
756 #[cfg(target_os = "linux")]
757 pub fn gbm_nv12_buffers(&self) -> &[crate::vaapi_encode::GbmNv12Buffer] {
758 match &self.kind {
759 SurfaceEncoderKind::H264Vaapi(enc) => enc.gbm_nv12_buffers(),
760 SurfaceEncoderKind::AV1Vaapi(enc) => enc.gbm_nv12_buffers(),
761 _ => &[],
762 }
763 }
764
765 #[cfg(target_os = "linux")]
766 pub fn allocate_nv12_buffers(&mut self, drm_fd: std::os::fd::RawFd, count: usize) {
767 match &mut self.kind {
768 SurfaceEncoderKind::H264Vaapi(enc) => {
769 if let Some(vpp) = &mut enc.vpp {
770 vpp.allocate_nv12_buffers(drm_fd, count);
771 }
772 }
773 SurfaceEncoderKind::AV1Vaapi(enc) => {
774 if let Some(vpp) = &mut enc.vpp {
775 vpp.allocate_nv12_buffers(drm_fd, count);
776 }
777 }
778 _ => {}
779 }
780 }
781
782 #[cfg(target_os = "linux")]
783 pub fn drm_fd_raw(&self) -> std::os::fd::RawFd {
784 use std::os::fd::AsRawFd;
785 match &self.kind {
786 SurfaceEncoderKind::H264Vaapi(enc) => enc._drm_fd.as_raw_fd(),
787 SurfaceEncoderKind::AV1Vaapi(enc) => enc._drm_fd.as_raw_fd(),
788 _ => -1,
789 }
790 }
791
792 #[cfg(target_os = "linux")]
794 #[allow(dead_code)]
795 pub fn va_display_usize(&self) -> usize {
796 match &self.kind {
797 SurfaceEncoderKind::H264Vaapi(enc) => enc.va_display_usize(),
798 SurfaceEncoderKind::AV1Vaapi(enc) => enc.va_display_usize(),
799 _ => 0,
800 }
801 }
802
803 pub fn encode(&mut self, rgba: &[u8]) -> Option<(Vec<u8>, bool)> {
804 if let SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) =
810 &mut self.kind
811 {
812 let (sw, sh) = (self.source_width as usize, self.source_height as usize);
813 let mut result = enc.encode_rgba_padded(rgba, sw, sh);
814 self.fixup_keyframe(&mut result);
815 return result;
816 }
817
818 let enc_len = expected_rgba_len(self.width, self.height);
819 let enc_len = match enc_len {
820 Some(v) => v,
821 None => {
822 eprintln!(
823 "[surface-encoder] expected_rgba_len overflow {}x{}",
824 self.width, self.height
825 );
826 return None;
827 }
828 };
829 let rgba = if rgba.len() == enc_len {
830 std::borrow::Cow::Borrowed(rgba)
831 } else {
832 let total_px = rgba.len() / 4;
836 if total_px == 0 {
837 return None;
838 }
839 let src_w = [self.width as usize, (self.width - 1) as usize]
841 .into_iter()
842 .find(|&w| w > 0 && total_px.is_multiple_of(w))?;
843 let src_h = total_px / src_w;
844 if src_h == 0 {
845 return None;
846 }
847 let dst_w = self.width as usize;
848 let dst_h = self.height as usize;
849 let mut padded = vec![0u8; enc_len];
850 for row in 0..dst_h {
851 let src_row = row.min(src_h - 1);
852 for col in 0..dst_w {
853 let src_col = col.min(src_w - 1);
854 let si = (src_row * src_w + src_col) * 4;
855 let di = (row * dst_w + col) * 4;
856 padded[di..di + 4].copy_from_slice(&rgba[si..si + 4]);
857 }
858 }
859 std::borrow::Cow::Owned(padded)
860 };
861
862 match &mut self.kind {
863 SurfaceEncoderKind::H264Software(encoder) => {
864 encoder.encode(&rgba, self.width, self.height)
865 }
866 SurfaceEncoderKind::NvencH264(_) | SurfaceEncoderKind::NvencAV1(_) => unreachable!(),
868 #[cfg(target_os = "linux")]
869 SurfaceEncoderKind::H264Vaapi(enc) => {
870 let mut bgra = rgba.into_owned();
871 for px in bgra.chunks_exact_mut(4) {
872 px.swap(0, 2);
873 }
874 let (sw, sh) = (self.source_width as usize, self.source_height as usize);
875 enc.encode_bgra_padded(&bgra, sw, sh)
876 }
877 #[cfg(target_os = "linux")]
878 SurfaceEncoderKind::AV1Vaapi(enc) => {
879 let mut bgra = rgba.into_owned();
880 for px in bgra.chunks_exact_mut(4) {
881 px.swap(0, 2);
882 }
883 let (sw, sh) = (self.source_width as usize, self.source_height as usize);
884 enc.encode_bgra_padded(&bgra, sw, sh)
885 }
886 SurfaceEncoderKind::AV1Software(encoder) => encoder.encode(&rgba),
887 }
888 }
889
890 pub fn encode_pixels(&mut self, pixels: &PixelData) -> Option<(Vec<u8>, bool)> {
893 match pixels {
894 PixelData::Nv12 {
895 data,
896 y_stride,
897 uv_stride,
898 } => self.encode_nv12(data, *y_stride, *uv_stride),
899 PixelData::Bgra(bgra) => self.encode_bgra(bgra),
900 PixelData::Rgba(rgba) => self.encode(rgba),
901 #[cfg(target_os = "linux")]
902 PixelData::DmaBuf {
903 fd,
904 fourcc,
905 modifier,
906 stride,
907 offset,
908 ..
909 } => self
910 .encode_dmabuf(fd, *fourcc, *modifier, *stride, *offset)
911 .or_else(|| {
912 let w = self.width;
915 let h = self.height;
916 let rgba = pixels.to_rgba(w, h);
917 if !rgba.is_empty() {
918 self.encode(&rgba)
919 } else {
920 None
921 }
922 }),
923 #[cfg(not(target_os = "linux"))]
924 PixelData::DmaBuf { .. } => None,
925 #[cfg(target_os = "linux")]
926 PixelData::Nv12DmaBuf {
927 fd,
928 stride,
929 uv_offset,
930 width,
931 height,
932 sync_fd,
933 } => {
934 if let Some(sfd) = sync_fd {
938 use std::os::fd::AsRawFd;
939 let mut pfd = libc::pollfd {
940 fd: sfd.as_raw_fd(),
941 events: libc::POLLIN,
942 revents: 0,
943 };
944 unsafe { libc::poll(&mut pfd, 1, 5000) };
945 }
946 self.encode_nv12_dmabuf(fd, *stride, *uv_offset, *width, *height)
947 }
948 .or_else(|| {
949 use std::os::fd::AsRawFd;
952 let h = *height as usize;
953 let s = *stride as usize;
954 let uv_off = *uv_offset as usize;
955 let raw = fd.as_raw_fd();
956 let map_size = uv_off + s * h.div_ceil(2);
957 let ptr = unsafe {
958 libc::mmap(
959 std::ptr::null_mut(),
960 map_size,
961 libc::PROT_READ,
962 libc::MAP_SHARED,
963 raw,
964 0,
965 )
966 };
967 if ptr == libc::MAP_FAILED || ptr.is_null() {
968 return None;
969 }
970 let data = unsafe { std::slice::from_raw_parts(ptr as *const u8, map_size) };
971 let result = self.encode_nv12(data, s, s);
972 unsafe { libc::munmap(ptr, map_size) };
973 result
974 }),
975 #[cfg(not(target_os = "linux"))]
976 PixelData::Nv12DmaBuf { .. } => None,
977 PixelData::VaSurface { .. } => None,
978 PixelData::Encoded { .. } => None,
981 }
982 }
983
984 #[cfg(target_os = "linux")]
988 fn encode_nv12_dmabuf(
989 &mut self,
990 fd: &std::sync::Arc<std::os::fd::OwnedFd>,
991 _stride: u32,
992 _uv_offset: u32,
993 _width: u32,
994 _height: u32,
995 ) -> Option<(Vec<u8>, bool)> {
996 use std::os::fd::AsRawFd;
997 let raw_fd = fd.as_raw_fd();
998 let find_surface = |nv12s: &[crate::vaapi_encode::GbmNv12Buffer]| -> Option<u32> {
999 let buf = nv12s.iter().find(|n| n.fd.as_raw_fd() == raw_fd)?;
1000 if buf.va_surface == 0 {
1002 return None;
1003 }
1004 Some(buf.va_surface)
1005 };
1006 let mut result = match &mut self.kind {
1007 SurfaceEncoderKind::AV1Vaapi(enc) => {
1008 let surf = find_surface(enc.gbm_nv12_buffers())?;
1009 enc.encode_surface(surf)
1010 }
1011 SurfaceEncoderKind::H264Vaapi(enc) => {
1012 let surf = find_surface(enc.gbm_nv12_buffers())?;
1013 enc.encode_surface(surf)
1014 }
1015 _ => None,
1016 };
1017 self.fixup_keyframe(&mut result);
1018 result
1019 }
1020
1021 #[cfg(target_os = "linux")]
1024 fn encode_dmabuf(
1025 &mut self,
1026 fd: &std::os::fd::OwnedFd,
1027 fourcc: u32,
1028 modifier: u64,
1029 stride: u32,
1030 offset: u32,
1031 ) -> Option<(Vec<u8>, bool)> {
1032 use std::os::fd::AsRawFd;
1033
1034 let src_w = self.source_width;
1037 let src_h = self.source_height;
1038
1039 let mut gpu_result = match &mut self.kind {
1043 SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => enc
1044 .encode_dmabuf_fd(
1045 fd.as_raw_fd(),
1046 fourcc,
1047 modifier,
1048 stride,
1049 offset,
1050 src_w,
1051 src_h,
1052 ),
1053 _ => None,
1054 };
1055 if gpu_result.is_some() {
1056 self.fixup_keyframe(&mut gpu_result);
1057 return gpu_result;
1058 }
1059
1060 self.encode_dmabuf_cpu_fallback(fd, fourcc, stride, offset)
1065 }
1066
1067 #[cfg(target_os = "linux")]
1070 fn encode_dmabuf_cpu_fallback(
1071 &mut self,
1072 fd: &std::os::fd::OwnedFd,
1073 fourcc: u32,
1074 stride: u32,
1075 _offset: u32,
1076 ) -> Option<(Vec<u8>, bool)> {
1077 use std::os::fd::AsRawFd;
1078
1079 let w = self.source_width as usize;
1080 let h = self.source_height as usize;
1081 let stride = stride as usize;
1082 let raw_fd = fd.as_raw_fd();
1083
1084 let file_size = unsafe { libc::lseek(raw_fd, 0, libc::SEEK_END) };
1086 if file_size <= 0 {
1087 return None;
1088 }
1089 let map_len = file_size as usize;
1090
1091 #[repr(C)]
1093 struct DmaBufSync {
1094 flags: u64,
1095 }
1096 const DMA_BUF_SYNC_READ: u64 = 1;
1097 const DMA_BUF_SYNC_START: u64 = 0;
1098 const DMA_BUF_SYNC_END: u64 = 4;
1099 const DMA_BUF_IOCTL_SYNC: libc::c_ulong = 0x40086200;
1103
1104 {
1108 let mut pfd = libc::pollfd {
1109 fd: raw_fd,
1110 events: libc::POLLIN,
1111 revents: 0,
1112 };
1113 let ready = unsafe { libc::poll(&mut pfd, 1, 0) };
1114 if ready <= 0 {
1115 } else {
1117 let sync_start = DmaBufSync {
1118 flags: DMA_BUF_SYNC_START | DMA_BUF_SYNC_READ,
1119 };
1120 unsafe {
1121 libc::ioctl(raw_fd, DMA_BUF_IOCTL_SYNC as _, &sync_start);
1122 }
1123 }
1124 }
1125
1126 let ptr = unsafe {
1128 libc::mmap(
1129 std::ptr::null_mut(),
1130 map_len,
1131 libc::PROT_READ,
1132 libc::MAP_SHARED,
1133 raw_fd,
1134 0,
1135 )
1136 };
1137 if ptr == libc::MAP_FAILED {
1138 let sync_end = DmaBufSync {
1139 flags: DMA_BUF_SYNC_END | DMA_BUF_SYNC_READ,
1140 };
1141 unsafe {
1142 libc::ioctl(raw_fd, DMA_BUF_IOCTL_SYNC as _, &sync_end);
1143 }
1144 return None;
1145 }
1146 let plane_data = unsafe { std::slice::from_raw_parts(ptr as *const u8, map_len) };
1147
1148 let is_gl_fbo = {
1151 let mut link = [0u8; 128];
1152 let path = format!("/proc/self/fd/{raw_fd}\0");
1153 let n = unsafe {
1154 libc::readlink(path.as_ptr() as *const _, link.as_mut_ptr() as *mut _, 127)
1155 };
1156 !(n > 0 && link[..n as usize].starts_with(b"/dev/dri/"))
1157 };
1158
1159 let result = if fourcc == blit_compositor::drm_fourcc::ARGB8888
1160 || fourcc == blit_compositor::drm_fourcc::XRGB8888
1161 {
1162 let mut packed = Vec::with_capacity(w * h * 4);
1164 for i in 0..h {
1165 let row = if is_gl_fbo { h - 1 - i } else { i };
1167 let start = row * stride;
1168 let end = start + w * 4;
1169 if end <= plane_data.len() {
1170 packed.extend_from_slice(&plane_data[start..end]);
1171 }
1172 }
1173 self.encode_bgra(&packed)
1174 } else if fourcc == blit_compositor::drm_fourcc::ABGR8888
1175 || fourcc == blit_compositor::drm_fourcc::XBGR8888
1176 {
1177 let mut packed = Vec::with_capacity(w * h * 4);
1179 for i in 0..h {
1180 let row = if is_gl_fbo { h - 1 - i } else { i };
1181 let start = row * stride;
1182 let end = start + w * 4;
1183 if end <= plane_data.len() {
1184 packed.extend_from_slice(&plane_data[start..end]);
1185 }
1186 }
1187 self.encode(&packed)
1188 } else if fourcc == blit_compositor::drm_fourcc::NV12 {
1189 let uv_stride = stride; let y_size = stride * h;
1194 let uv_h = h.div_ceil(2);
1195 let uv_size = uv_stride * uv_h;
1196 if map_len >= y_size + uv_size {
1197 let out_stride = w;
1199 let mut data = vec![0u8; out_stride * h + out_stride * uv_h];
1200 for row in 0..h {
1201 let src = row * stride;
1202 let dst = row * out_stride;
1203 if src + w <= plane_data.len() {
1204 data[dst..dst + w].copy_from_slice(&plane_data[src..src + w]);
1205 }
1206 }
1207 let uv_dst_base = out_stride * h;
1208 for row in 0..uv_h {
1209 let src = y_size + row * uv_stride;
1210 let dst = uv_dst_base + row * out_stride;
1211 if src + w <= plane_data.len() {
1212 data[dst..dst + w].copy_from_slice(&plane_data[src..src + w]);
1213 }
1214 }
1215 self.encode_nv12(&data, out_stride, out_stride)
1216 } else {
1217 None
1218 }
1219 } else {
1220 None
1221 };
1222
1223 unsafe {
1225 libc::munmap(ptr, map_len);
1226 }
1227 let sync_end = DmaBufSync {
1229 flags: DMA_BUF_SYNC_END | DMA_BUF_SYNC_READ,
1230 };
1231 unsafe {
1232 libc::ioctl(raw_fd, DMA_BUF_IOCTL_SYNC as _, &sync_end);
1233 }
1234
1235 result
1236 }
1237
1238 fn fixup_keyframe(&self, result: &mut Option<(Vec<u8>, bool)>) {
1243 if let Some((data, is_key)) = result.as_mut()
1244 && !*is_key
1245 {
1246 *is_key = match &self.kind {
1247 SurfaceEncoderKind::NvencH264(_) => h264_stream_contains_idr(data),
1248 SurfaceEncoderKind::NvencAV1(_) => av1_stream_contains_keyframe(data),
1249 #[cfg(target_os = "linux")]
1250 SurfaceEncoderKind::H264Vaapi(_) => h264_stream_contains_idr(data),
1251 #[cfg(target_os = "linux")]
1252 SurfaceEncoderKind::AV1Vaapi(_) => av1_stream_contains_keyframe(data),
1253 _ => false,
1254 };
1255 }
1256 }
1257
1258 fn encode_bgra(&mut self, bgra: &[u8]) -> Option<(Vec<u8>, bool)> {
1260 let enc_w = self.width as usize;
1261 let enc_h = self.height as usize;
1262 let src_w = self.source_width as usize;
1263 let src_h = self.source_height as usize;
1264
1265 let mut result = match &mut self.kind {
1266 SurfaceEncoderKind::H264Software(encoder) => {
1267 let yuv = bgra_to_yuv420_padded(bgra, src_w, src_h, enc_w, enc_h);
1268 let yuv_buf = YUVBuffer::from_vec(yuv, enc_w, enc_h);
1269 encoder.encode_yuv(&yuv_buf, self.width, self.height)
1270 }
1271 SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => {
1272 enc.encode_bgra_padded(bgra, src_w, src_h)
1273 }
1274 #[cfg(target_os = "linux")]
1275 SurfaceEncoderKind::H264Vaapi(enc) => enc.encode_bgra_padded(bgra, src_w, src_h),
1276 #[cfg(target_os = "linux")]
1277 SurfaceEncoderKind::AV1Vaapi(enc) => enc.encode_bgra_padded(bgra, src_w, src_h),
1278 SurfaceEncoderKind::AV1Software(encoder) => {
1279 let yuv = if self.chroma.is_444() {
1280 bgra_to_yuv444_padded(bgra, src_w, src_h, enc_w, enc_h)
1281 } else {
1282 bgra_to_yuv420_padded(bgra, src_w, src_h, enc_w, enc_h)
1283 };
1284 encoder.encode_yuv_planes(&yuv)
1285 }
1286 };
1287 self.fixup_keyframe(&mut result);
1288 result
1289 }
1290
1291 fn encode_nv12(
1294 &mut self,
1295 data: &[u8],
1296 y_stride: usize,
1297 uv_stride: usize,
1298 ) -> Option<(Vec<u8>, bool)> {
1299 let src_w = self.source_width as usize;
1301 let src_h = self.source_height as usize;
1302
1303 let mut result = match &mut self.kind {
1304 SurfaceEncoderKind::H264Software(encoder) => {
1305 let enc_w = self.width as usize;
1306 let enc_h = self.height as usize;
1307 if enc_w == src_w && enc_h == src_h {
1308 let yuv = nv12_to_yuv420(data, y_stride, uv_stride, src_w, src_h);
1309 let yuv_buf = YUVBuffer::from_vec(yuv, enc_w, enc_h);
1310 encoder.encode_yuv(&yuv_buf, self.width, self.height)
1311 } else {
1312 let pd = PixelData::Nv12 {
1313 data: std::sync::Arc::new(data.to_vec()),
1314 y_stride,
1315 uv_stride,
1316 };
1317 let rgba = pd.to_rgba(self.source_width, self.source_height);
1318 return self.encode(&rgba);
1319 }
1320 }
1321 SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => {
1322 enc.encode_nv12(data, y_stride, uv_stride, src_h)
1324 }
1325 #[cfg(target_os = "linux")]
1326 SurfaceEncoderKind::H264Vaapi(enc) => {
1327 let uv_offset = y_stride * src_h;
1328 let y_data = &data[..uv_offset];
1329 let uv_data = &data[uv_offset..];
1330 enc.encode_nv12(y_data, uv_data, y_stride, uv_stride)
1331 }
1332 #[cfg(target_os = "linux")]
1333 SurfaceEncoderKind::AV1Vaapi(enc) => {
1334 let uv_offset = y_stride * src_h;
1335 let y_data = &data[..uv_offset];
1336 let uv_data = &data[uv_offset..];
1337 enc.encode_nv12(y_data, uv_data, y_stride, uv_stride)
1338 }
1339 SurfaceEncoderKind::AV1Software(encoder) => {
1340 encoder.encode_nv12(data, y_stride, uv_stride, src_w, src_h)
1341 }
1342 };
1343 self.fixup_keyframe(&mut result);
1344 result
1345 }
1346}
1347
1348fn validate_surface_dimensions(
1349 width: u32,
1350 height: u32,
1351 _preference: SurfaceEncoderPreference,
1352) -> Result<(), String> {
1353 if width == 0 || height == 0 {
1354 return Err("surface encoder requires non-zero dimensions".into());
1355 }
1356 let _ = expected_rgba_len(width, height)
1359 .ok_or_else(|| format!("surface encoder dimensions overflow for {width}x{height}"))?;
1360 Ok(())
1361}
1362
1363fn expected_rgba_len(width: u32, height: u32) -> Option<usize> {
1364 (width as usize)
1365 .checked_mul(height as usize)?
1366 .checked_mul(4)
1367}
1368
1369#[inline(always)]
1375fn rgb_to_y(r: i32, g: i32, b: i32) -> u8 {
1376 ((66 * r + 129 * g + 25 * b + 128) >> 8)
1377 .wrapping_add(16)
1378 .clamp(0, 255) as u8
1379}
1380
1381#[inline(always)]
1382fn rgb_to_u(r: i32, g: i32, b: i32) -> u8 {
1383 ((-38 * r - 74 * g + 112 * b + 128) >> 8)
1384 .wrapping_add(128)
1385 .clamp(0, 255) as u8
1386}
1387
1388#[inline(always)]
1389fn rgb_to_v(r: i32, g: i32, b: i32) -> u8 {
1390 ((112 * r - 94 * g - 18 * b + 128) >> 8)
1391 .wrapping_add(128)
1392 .clamp(0, 255) as u8
1393}
1394
1395#[inline(always)]
1404fn compute_y_plane(
1405 src: &[u8],
1406 width: usize,
1407 height: usize,
1408 y_plane: &mut [u8],
1409 r_off: usize,
1410 g_off: usize,
1411 b_off: usize,
1412) {
1413 let total = width * height;
1414 for (px, y_out) in y_plane[..total].iter_mut().enumerate() {
1415 let i = px * 4;
1416 let r = src[i + r_off] as i32;
1417 let g = src[i + g_off] as i32;
1418 let b = src[i + b_off] as i32;
1419 *y_out = rgb_to_y(r, g, b);
1420 }
1421}
1422
1423#[inline(always)]
1425fn compute_uv_planes(
1426 src: &[u8],
1427 width: usize,
1428 height: usize,
1429 u_plane: &mut [u8],
1430 v_plane: &mut [u8],
1431 r_off: usize,
1432 g_off: usize,
1433 b_off: usize,
1434) {
1435 let chroma_w = width.div_ceil(2);
1436 let chroma_h = height.div_ceil(2);
1437 for cy in 0..chroma_h {
1438 for cx in 0..chroma_w {
1439 let row = cy * 2;
1440 let col = cx * 2;
1441 let mut u_sum = 0i32;
1443 let mut v_sum = 0i32;
1444 for dy in 0..2u32 {
1445 for dx in 0..2u32 {
1446 let sr = (row + dy as usize).min(height - 1);
1447 let sc = (col + dx as usize).min(width - 1);
1448 let i = (sr * width + sc) * 4;
1449 let r = src[i + r_off] as i32;
1450 let g = src[i + g_off] as i32;
1451 let b = src[i + b_off] as i32;
1452 u_sum += rgb_to_u(r, g, b) as i32;
1453 v_sum += rgb_to_v(r, g, b) as i32;
1454 }
1455 }
1456 let idx = cy * chroma_w + cx;
1457 u_plane[idx] = (u_sum / 4) as u8;
1458 v_plane[idx] = (v_sum / 4) as u8;
1459 }
1460 }
1461}
1462
1463#[inline(always)]
1466fn compute_y_plane_padded(
1467 src: &[u8],
1468 src_w: usize,
1469 src_h: usize,
1470 enc_w: usize,
1471 enc_h: usize,
1472 y_plane: &mut [u8],
1473 r_off: usize,
1474 g_off: usize,
1475 b_off: usize,
1476) {
1477 for row in 0..enc_h {
1478 let sr = row.min(src_h - 1);
1479 for col in 0..enc_w {
1480 let sc = col.min(src_w - 1);
1481 let i = (sr * src_w + sc) * 4;
1482 let r = src[i + r_off] as i32;
1483 let g = src[i + g_off] as i32;
1484 let b = src[i + b_off] as i32;
1485 y_plane[row * enc_w + col] = rgb_to_y(r, g, b);
1486 }
1487 }
1488}
1489
1490#[inline(always)]
1493fn compute_uv_planes_padded(
1494 src: &[u8],
1495 src_w: usize,
1496 src_h: usize,
1497 enc_w: usize,
1498 enc_h: usize,
1499 u_plane: &mut [u8],
1500 v_plane: &mut [u8],
1501 r_off: usize,
1502 g_off: usize,
1503 b_off: usize,
1504) {
1505 let chroma_w = enc_w.div_ceil(2);
1506 let chroma_h = enc_h.div_ceil(2);
1507 for cy in 0..chroma_h {
1508 for cx in 0..chroma_w {
1509 let row = cy * 2;
1510 let col = cx * 2;
1511 let mut u_sum = 0i32;
1512 let mut v_sum = 0i32;
1513 for dy in 0..2u32 {
1514 for dx in 0..2u32 {
1515 let sr = (row + dy as usize).min(src_h - 1);
1516 let sc = (col + dx as usize).min(src_w - 1);
1517 let i = (sr * src_w + sc) * 4;
1518 let r = src[i + r_off] as i32;
1519 let g = src[i + g_off] as i32;
1520 let b = src[i + b_off] as i32;
1521 u_sum += rgb_to_u(r, g, b) as i32;
1522 v_sum += rgb_to_v(r, g, b) as i32;
1523 }
1524 }
1525 let idx = cy * chroma_w + cx;
1526 u_plane[idx] = (u_sum / 4) as u8;
1527 v_plane[idx] = (v_sum / 4) as u8;
1528 }
1529 }
1530}
1531
1532#[inline(always)]
1535fn compute_uv_planes_444_padded(
1536 src: &[u8],
1537 src_w: usize,
1538 src_h: usize,
1539 enc_w: usize,
1540 enc_h: usize,
1541 u_plane: &mut [u8],
1542 v_plane: &mut [u8],
1543 r_off: usize,
1544 g_off: usize,
1545 b_off: usize,
1546) {
1547 for row in 0..enc_h {
1548 let sr = row.min(src_h - 1);
1549 for col in 0..enc_w {
1550 let sc = col.min(src_w - 1);
1551 let i = (sr * src_w + sc) * 4;
1552 let r = src[i + r_off] as i32;
1553 let g = src[i + g_off] as i32;
1554 let b = src[i + b_off] as i32;
1555 let idx = row * enc_w + col;
1556 u_plane[idx] = rgb_to_u(r, g, b);
1557 v_plane[idx] = rgb_to_v(r, g, b);
1558 }
1559 }
1560}
1561
1562fn bgra_to_yuv444_padded(
1564 bgra: &[u8],
1565 src_w: usize,
1566 src_h: usize,
1567 enc_w: usize,
1568 enc_h: usize,
1569) -> Vec<u8> {
1570 let plane_size = enc_w * enc_h;
1571 let mut yuv = vec![0u8; plane_size * 3];
1572 let (y_plane, uv) = yuv.split_at_mut(plane_size);
1573 let (u_plane, v_plane) = uv.split_at_mut(plane_size);
1574 compute_y_plane_padded(bgra, src_w, src_h, enc_w, enc_h, y_plane, 2, 1, 0);
1576 compute_uv_planes_444_padded(bgra, src_w, src_h, enc_w, enc_h, u_plane, v_plane, 2, 1, 0);
1577 yuv
1578}
1579
1580fn rgba_to_yuv444(rgba: &[u8], width: usize, height: usize) -> Vec<u8> {
1582 let plane_size = width * height;
1583 let mut yuv = vec![0u8; plane_size * 3];
1584 let (y_plane, uv) = yuv.split_at_mut(plane_size);
1585 let (u_plane, v_plane) = uv.split_at_mut(plane_size);
1586 compute_y_plane(rgba, width, height, y_plane, 0, 1, 2);
1588 compute_uv_planes_444_padded(
1589 rgba, width, height, width, height, u_plane, v_plane, 0, 1, 2,
1590 );
1591 yuv
1592}
1593
1594fn bgra_to_yuv420_padded(
1598 bgra: &[u8],
1599 src_w: usize,
1600 src_h: usize,
1601 enc_w: usize,
1602 enc_h: usize,
1603) -> Vec<u8> {
1604 let y_size = enc_w * enc_h;
1605 let uv_w = enc_w.div_ceil(2);
1610 let uv_size = uv_w * enc_h.div_ceil(2);
1611 let mut yuv = vec![0u8; y_size + uv_size * 2];
1612 let (y_plane, uv) = yuv.split_at_mut(y_size);
1613 let (u_plane, v_plane) = uv.split_at_mut(uv_size);
1614 compute_y_plane_padded(bgra, src_w, src_h, enc_w, enc_h, y_plane, 2, 1, 0);
1616 compute_uv_planes_padded(bgra, src_w, src_h, enc_w, enc_h, u_plane, v_plane, 2, 1, 0);
1617 yuv
1618}
1619
1620fn rgba_to_yuv420(rgba: &[u8], width: usize, height: usize) -> Vec<u8> {
1622 let y_size = width * height;
1623 let uv_w = width.div_ceil(2);
1624 let uv_size = uv_w * height.div_ceil(2);
1625 let mut yuv = vec![0u8; y_size + uv_size * 2];
1626 let (y_plane, uv) = yuv.split_at_mut(y_size);
1627 let (u_plane, v_plane) = uv.split_at_mut(uv_size);
1628 compute_y_plane(rgba, width, height, y_plane, 0, 1, 2);
1630 compute_uv_planes(rgba, width, height, u_plane, v_plane, 0, 1, 2);
1631 yuv
1632}
1633
1634fn nv12_to_yuv420(
1638 data: &[u8],
1639 y_stride: usize,
1640 uv_stride: usize,
1641 width: usize,
1642 height: usize,
1643) -> Vec<u8> {
1644 let y_size = width * height;
1645 let uv_w = width.div_ceil(2);
1646 let uv_h = height.div_ceil(2);
1647 let uv_size = uv_w * uv_h;
1648 let mut yuv = vec![0u8; y_size + uv_size * 2];
1649 let (y_out, uv_out) = yuv.split_at_mut(y_size);
1650 let (u_out, v_out) = uv_out.split_at_mut(uv_size);
1651
1652 let uv_offset = y_stride * height;
1653
1654 for row in 0..height {
1656 let src = row * y_stride;
1657 let dst = row * width;
1658 y_out[dst..dst + width].copy_from_slice(&data[src..src + width]);
1659 }
1660
1661 let src_uv_pairs = width / 2;
1665 for row in 0..uv_h {
1666 let src_start = uv_offset + row.min(height / 2 - 1) * uv_stride;
1667 let dst_start = row * uv_w;
1668 for col in 0..uv_w {
1669 let sc = col.min(src_uv_pairs.saturating_sub(1));
1670 u_out[dst_start + col] = data[src_start + sc * 2];
1671 v_out[dst_start + col] = data[src_start + sc * 2 + 1];
1672 }
1673 }
1674
1675 yuv
1676}
1677
1678fn h264_stream_contains_idr(data: &[u8]) -> bool {
1680 annex_b_contains_nal(data, |byte| (byte & 0x1f) == 5)
1681}
1682
1683fn annex_b_contains_nal(data: &[u8], pred: impl Fn(u8) -> bool) -> bool {
1685 let mut i = 0usize;
1686 while i < data.len() {
1687 let start_code_len = if data[i..].starts_with(&[0, 0, 0, 1]) {
1688 4
1689 } else if data[i..].starts_with(&[0, 0, 1]) {
1690 3
1691 } else {
1692 i += 1;
1693 continue;
1694 };
1695
1696 let nal_header = i + start_code_len;
1697 if let Some(&byte) = data.get(nal_header)
1698 && pred(byte)
1699 {
1700 return true;
1701 }
1702
1703 i = nal_header.saturating_add(1);
1704 }
1705
1706 false
1707}
1708
1709fn av1_stream_contains_keyframe(data: &[u8]) -> bool {
1717 let mut pos = 0;
1721 while pos < data.len() {
1722 let header = data[pos];
1723 let obu_type = (header >> 3) & 0xF;
1724 let has_extension = (header >> 2) & 1;
1725 let has_size = (header >> 1) & 1;
1726 pos += 1;
1727
1728 if has_extension != 0 {
1730 if pos >= data.len() {
1731 break;
1732 }
1733 pos += 1;
1734 }
1735
1736 if obu_type == 1 {
1738 return true;
1739 }
1740
1741 if has_size != 0 {
1744 let mut size: u64 = 0;
1745 let mut shift = 0u32;
1746 while pos < data.len() {
1747 let byte = data[pos];
1748 pos += 1;
1749 size |= ((byte & 0x7F) as u64) << shift;
1750 if byte & 0x80 == 0 {
1751 break;
1752 }
1753 shift += 7;
1754 if shift >= 56 {
1755 return false; }
1757 }
1758 pos = pos.saturating_add(size as usize);
1759 } else {
1760 break;
1763 }
1764 }
1765 false
1766}
1767
1768struct SoftwareH264Encoder {
1769 encoder: OpenH264Encoder,
1770}
1771
1772impl SoftwareH264Encoder {
1773 fn new(quality: SurfaceQuality) -> Result<Self, String> {
1774 use openh264::encoder::{EncoderConfig, RateControlMode};
1775 let config = EncoderConfig::new()
1776 .set_bitrate_bps(quality.openh264_bitrate())
1777 .rate_control_mode(RateControlMode::Bitrate);
1778 let encoder =
1779 OpenH264Encoder::with_api_config(openh264::OpenH264API::from_source(), config)
1780 .map_err(|err| format!("failed to create OpenH264 encoder: {err:?}"))?;
1781 Ok(Self { encoder })
1782 }
1783
1784 fn request_keyframe(&mut self) {
1785 self.encoder.force_intra_frame();
1786 }
1787
1788 fn encode(&mut self, rgba: &[u8], width: u32, height: u32) -> Option<(Vec<u8>, bool)> {
1789 let yuv = rgba_to_yuv420(rgba, width as usize, height as usize);
1790 let yuv_buf = YUVBuffer::from_vec(yuv, width as usize, height as usize);
1791 self.encode_yuv(&yuv_buf, width, height)
1792 }
1793
1794 fn encode_yuv(
1796 &mut self,
1797 yuv_buf: &YUVBuffer,
1798 width: u32,
1799 height: u32,
1800 ) -> Option<(Vec<u8>, bool)> {
1801 let bitstream = match self.encoder.encode(yuv_buf) {
1802 Ok(bs) => bs,
1803 Err(e) => {
1804 eprintln!("[surface-encoder] openh264 encode failed {width}x{height}: {e:?}");
1805 return None;
1806 }
1807 };
1808 let nal_data = bitstream.to_vec();
1809 if nal_data.is_empty() {
1810 eprintln!("[surface-encoder] openh264 produced empty NAL {width}x{height}");
1811 return None;
1812 }
1813 let is_keyframe = h264_stream_contains_idr(&nal_data);
1814 Some((nal_data, is_keyframe))
1815 }
1816}
1817
1818struct SoftwareAV1Encoder {
1823 ctx: rav1e::Context<u8>,
1824 width: usize,
1825 height: usize,
1826 force_keyframe: bool,
1827 chroma: ChromaSubsampling,
1828}
1829
1830impl SoftwareAV1Encoder {
1831 fn new(
1832 width: u32,
1833 height: u32,
1834 quality: SurfaceQuality,
1835 chroma: ChromaSubsampling,
1836 ) -> Result<Self, String> {
1837 use rav1e::prelude::*;
1838
1839 let chroma_sampling = if chroma.is_444() {
1840 ChromaSampling::Cs444
1841 } else {
1842 ChromaSampling::Cs420
1843 };
1844 let mut speed = SpeedSettings::from_preset(quality.av1_speed());
1845 speed.rdo_lookahead_frames = 1;
1846 let enc = EncoderConfig {
1847 width: width as usize,
1848 height: height as usize,
1849 chroma_sampling,
1850 chroma_sample_position: ChromaSamplePosition::Unknown,
1851 speed_settings: speed,
1852 low_latency: true,
1853 min_key_frame_interval: 0,
1854 max_key_frame_interval: 60,
1855 quantizer: quality.av1_quantizer(),
1856 min_quantizer: quality.av1_min_quantizer(),
1857 bitrate: 0,
1858 ..Default::default()
1859 };
1860 let cfg = Config::new().with_encoder_config(enc);
1861 let ctx = cfg
1862 .new_context()
1863 .map_err(|e| format!("rav1e context creation failed: {e}"))?;
1864 Ok(Self {
1865 ctx,
1866 width: width as usize,
1867 height: height as usize,
1868 force_keyframe: false,
1869 chroma,
1870 })
1871 }
1872
1873 fn request_keyframe(&mut self) {
1874 self.force_keyframe = true;
1875 }
1876
1877 fn encode(&mut self, rgba: &[u8]) -> Option<(Vec<u8>, bool)> {
1878 let yuv = if self.chroma.is_444() {
1879 rgba_to_yuv444(rgba, self.width, self.height)
1880 } else {
1881 rgba_to_yuv420(rgba, self.width, self.height)
1882 };
1883 self.encode_yuv_planes(&yuv)
1884 }
1885
1886 fn encode_nv12(
1887 &mut self,
1888 data: &[u8],
1889 y_stride: usize,
1890 uv_stride: usize,
1891 width: usize,
1892 height: usize,
1893 ) -> Option<(Vec<u8>, bool)> {
1894 let yuv = nv12_to_yuv420(data, y_stride, uv_stride, width, height);
1895 self.encode_yuv_planes(&yuv)
1896 }
1897
1898 fn encode_yuv_planes(&mut self, yuv: &[u8]) -> Option<(Vec<u8>, bool)> {
1901 let width = self.width;
1902 let height = self.height;
1903 let y_size = width * height;
1904 let (uv_w, uv_size) = if self.chroma.is_444() {
1905 (width, width * height)
1906 } else {
1907 let uv_w = width.div_ceil(2);
1908 let uv_h = height.div_ceil(2);
1909 (uv_w, uv_w * uv_h)
1910 };
1911
1912 let y_plane = &yuv[..y_size];
1913 let u_plane = &yuv[y_size..y_size + uv_size];
1914 let v_plane = &yuv[y_size + uv_size..];
1915
1916 let mut frame = self.ctx.new_frame();
1917 frame.planes[0].copy_from_raw_u8(y_plane, width, 1);
1918 frame.planes[1].copy_from_raw_u8(u_plane, uv_w, 1);
1919 frame.planes[2].copy_from_raw_u8(v_plane, uv_w, 1);
1920
1921 self.send_and_receive(frame)
1922 }
1923
1924 fn send_and_receive(&mut self, frame: rav1e::Frame<u8>) -> Option<(Vec<u8>, bool)> {
1925 use rav1e::prelude::*;
1926
1927 if self.force_keyframe {
1928 let params = FrameParameters {
1929 frame_type_override: FrameTypeOverride::Key,
1930 ..Default::default()
1931 };
1932 if self.ctx.send_frame((frame, params)).is_ok() {
1933 self.force_keyframe = false;
1934 }
1935 } else {
1936 let _ = self.ctx.send_frame(frame);
1937 }
1938
1939 match self.ctx.receive_packet() {
1940 Ok(packet) => {
1941 let is_key = packet.frame_type == rav1e::prelude::FrameType::KEY;
1942 Some((packet.data, is_key))
1943 }
1944 Err(rav1e::EncoderStatus::Encoded) | Err(rav1e::EncoderStatus::NeedMoreData) => None,
1945 Err(_) => None,
1946 }
1947 }
1948}
1949
1950#[cfg(test)]
1951mod tests {
1952 use super::*;
1953
1954 fn make_obu(obu_type: u8, payload: &[u8]) -> Vec<u8> {
1956 let header = (obu_type & 0xF) << 3 | 0b10; let mut obu = vec![header];
1959 let mut size = payload.len();
1961 loop {
1962 let mut byte = (size & 0x7F) as u8;
1963 size >>= 7;
1964 if size > 0 {
1965 byte |= 0x80;
1966 }
1967 obu.push(byte);
1968 if size == 0 {
1969 break;
1970 }
1971 }
1972 obu.extend_from_slice(payload);
1973 obu
1974 }
1975
1976 #[test]
1977 fn av1_keyframe_with_sequence_header_only() {
1978 let data = make_obu(1, &[0xAA; 10]);
1980 assert!(av1_stream_contains_keyframe(&data));
1981 }
1982
1983 #[test]
1984 fn av1_keyframe_with_temporal_delimiter_prefix() {
1985 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));
1991 }
1992
1993 #[test]
1994 fn av1_non_keyframe_with_temporal_delimiter() {
1995 let mut data = make_obu(2, &[]);
1997 data.extend(make_obu(6, &[0xDD; 15]));
1998 assert!(!av1_stream_contains_keyframe(&data));
1999 }
2000
2001 #[test]
2002 fn av1_non_keyframe_frame_header_only() {
2003 let data = make_obu(3, &[0xEE; 5]);
2005 assert!(!av1_stream_contains_keyframe(&data));
2006 }
2007
2008 #[test]
2009 fn av1_empty_stream() {
2010 assert!(!av1_stream_contains_keyframe(&[]));
2011 }
2012
2013 #[test]
2014 fn av1_keyframe_large_leb128_size() {
2015 let mut data = make_obu(2, &[0x00; 200]);
2018 data.extend(make_obu(1, &[0xFF; 4]));
2019 assert!(av1_stream_contains_keyframe(&data));
2020 }
2021}