ustreamer_capture/
lib.rs

1//! GPU frame capture from wgpu render targets.
2//!
3//! Provides a trait [`FrameCapture`] with platform-specific implementations:
4//! - **Metal/IOSurface** (macOS): zero-copy via `wgpu-hal` Metal interop
5//! - **Vulkan/CUDA** (NVIDIA): zero-copy via external memory export
6//! - **Staging buffer** (fallback): triple-buffered `copy_texture_to_buffer`
7
8#[cfg(target_os = "macos")]
9pub mod metal;
10pub mod staging;
11#[cfg(all(
12    feature = "vulkan-external",
13    any(target_os = "linux", target_os = "windows")
14))]
15pub mod vulkan_external;
16
17#[cfg(target_os = "macos")]
18use objc2_core_foundation::CFRetained;
19#[cfg(target_os = "macos")]
20use objc2_core_video::CVPixelBuffer;
21#[cfg(target_os = "macos")]
22use objc2_io_surface::IOSurfaceRef;
23#[cfg(all(
24    feature = "vulkan-external",
25    any(target_os = "linux", target_os = "windows")
26))]
27pub use vulkan_external::{
28    VulkanCaptureSyncMode, VulkanExternalCapture, VulkanExternalImage, VulkanExternalMemoryHandle,
29    VulkanExternalSync, VulkanExternalSyncHandle,
30};
31
32/// Diagnostic checksum over canonical RGBA8 pixel bytes.
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub struct FrameChecksum {
35    pub width: u32,
36    pub height: u32,
37    pub rgba8_fnv1a64: u64,
38}
39
40impl FrameChecksum {
41    pub fn hex_string(&self) -> String {
42        format!("{:016x}", self.rgba8_fnv1a64)
43    }
44}
45
46/// Handle to a captured frame, ready for the encoder.
47pub enum CapturedFrame {
48    /// Raw pixel data in CPU memory (from staging buffer path).
49    CpuBuffer {
50        data: Vec<u8>,
51        width: u32,
52        height: u32,
53        stride: u32,
54        format: wgpu::TextureFormat,
55    },
56    /// macOS zero-copy capture via IOSurface → CVPixelBuffer.
57    #[cfg(target_os = "macos")]
58    MetalPixelBuffer {
59        /// Retained IOSurface backing the rendered Metal texture.
60        surface: CFRetained<IOSurfaceRef>,
61        /// CoreVideo wrapper around the IOSurface, ready for VideoToolbox input.
62        pixel_buffer: CFRetained<CVPixelBuffer>,
63        /// Frame width in pixels.
64        width: u32,
65        /// Frame height in pixels.
66        height: u32,
67        /// Row stride in bytes.
68        stride: u32,
69        /// CoreVideo / IOSurface pixel format fourcc.
70        pixel_format: u32,
71    },
72    /// Vulkan image + external memory handles for future CUDA/NVENC import.
73    #[cfg(all(
74        feature = "vulkan-external",
75        any(target_os = "linux", target_os = "windows")
76    ))]
77    VulkanExternalImage(vulkan_external::VulkanExternalImage),
78    // CudaMappedResource { ptr: ... },
79}
80
81impl CapturedFrame {
82    /// Compute a diagnostic checksum over canonical RGBA8 bytes when CPU pixel data is available.
83    pub fn diagnostic_checksum(&self) -> Result<Option<FrameChecksum>, CaptureError> {
84        match self {
85            Self::CpuBuffer {
86                data,
87                width,
88                height,
89                stride,
90                format,
91            } => checksum_cpu_buffer(data, *width, *height, *stride, *format).map(Some),
92            #[cfg(target_os = "macos")]
93            Self::MetalPixelBuffer { .. } => Ok(None),
94            #[cfg(all(
95                feature = "vulkan-external",
96                any(target_os = "linux", target_os = "windows")
97            ))]
98            Self::VulkanExternalImage(..) => Ok(None),
99        }
100    }
101}
102
103/// Trait for frame capture implementations.
104pub trait FrameCapture: Send + Sync {
105    /// Capture the current render target contents.
106    fn capture(
107        &mut self,
108        instance: &wgpu::Instance,
109        device: &wgpu::Device,
110        queue: &wgpu::Queue,
111        texture: &wgpu::Texture,
112    ) -> Result<CapturedFrame, CaptureError>;
113}
114
115#[derive(Debug, thiserror::Error)]
116pub enum CaptureError {
117    #[error("buffer mapping failed: {0}")]
118    MapFailed(String),
119    #[error("texture format unsupported: {0:?}")]
120    UnsupportedFormat(wgpu::TextureFormat),
121    #[error("capture backend unsupported: {0}")]
122    UnsupportedBackend(&'static str),
123    #[error("texture is not backed by IOSurface")]
124    NotIosurfaceBacked,
125    #[error("failed to wrap IOSurface in CVPixelBuffer (status {0})")]
126    PixelBufferCreateFailed(i32),
127    #[error("invalid IOSurface metadata: {0}")]
128    InvalidSurface(String),
129    #[error("invalid CPU buffer metadata: {0}")]
130    InvalidCpuBuffer(String),
131    #[error("invalid source texture for capture: {0}")]
132    InvalidTexture(String),
133    #[error("Vulkan texture is not backed by exportable external memory: {0}")]
134    ExternalMemoryUnavailable(String),
135    #[error("Vulkan interop failed: {0}")]
136    VulkanInteropFailed(String),
137}
138
139fn checksum_cpu_buffer(
140    data: &[u8],
141    width: u32,
142    height: u32,
143    stride: u32,
144    format: wgpu::TextureFormat,
145) -> Result<FrameChecksum, CaptureError> {
146    let row_bytes = match format {
147        wgpu::TextureFormat::Bgra8Unorm | wgpu::TextureFormat::Bgra8UnormSrgb => width
148            .checked_mul(4)
149            .ok_or_else(|| CaptureError::InvalidCpuBuffer("row byte size overflow".into()))?
150            as usize,
151        other => return Err(CaptureError::UnsupportedFormat(other)),
152    };
153
154    let stride = stride as usize;
155    if stride < row_bytes {
156        return Err(CaptureError::InvalidCpuBuffer(format!(
157            "stride {stride} is smaller than required row width {row_bytes}"
158        )));
159    }
160
161    let required_len = stride
162        .checked_mul(height as usize)
163        .ok_or_else(|| CaptureError::InvalidCpuBuffer("buffer size overflow".into()))?;
164    if data.len() < required_len {
165        return Err(CaptureError::InvalidCpuBuffer(format!(
166            "buffer length {} is smaller than expected {required_len}",
167            data.len()
168        )));
169    }
170
171    let mut hash = 0xcbf2_9ce4_8422_2325u64;
172    for row in 0..height as usize {
173        let row_start = row * stride;
174        let row_data = &data[row_start..row_start + row_bytes];
175        for pixel in row_data.chunks_exact(4) {
176            hash = fnv1a64_byte(hash, pixel[2]);
177            hash = fnv1a64_byte(hash, pixel[1]);
178            hash = fnv1a64_byte(hash, pixel[0]);
179            hash = fnv1a64_byte(hash, pixel[3]);
180        }
181    }
182
183    Ok(FrameChecksum {
184        width,
185        height,
186        rgba8_fnv1a64: hash,
187    })
188}
189
190#[inline]
191fn fnv1a64_byte(hash: u64, byte: u8) -> u64 {
192    (hash ^ byte as u64).wrapping_mul(0x0000_0100_0000_01b3)
193}
194
195#[cfg(test)]
196mod tests {
197    use super::{CapturedFrame, FrameChecksum};
198
199    #[test]
200    fn checksum_normalizes_bgra_to_rgba() {
201        let frame = CapturedFrame::CpuBuffer {
202            data: vec![1, 2, 3, 4],
203            width: 1,
204            height: 1,
205            stride: 4,
206            format: wgpu::TextureFormat::Bgra8Unorm,
207        };
208
209        let checksum = frame.diagnostic_checksum().unwrap().unwrap();
210        assert_eq!(
211            checksum,
212            FrameChecksum {
213                width: 1,
214                height: 1,
215                rgba8_fnv1a64: 0xdbd0_9687_ea36_bd25,
216            }
217        );
218        assert_eq!(checksum.hex_string(), "dbd09687ea36bd25");
219    }
220
221    #[test]
222    fn checksum_ignores_padding_bytes() {
223        let frame = CapturedFrame::CpuBuffer {
224            data: vec![1, 2, 3, 4, 7, 6, 5, 8, 255, 255, 255, 255],
225            width: 2,
226            height: 1,
227            stride: 12,
228            format: wgpu::TextureFormat::Bgra8Unorm,
229        };
230
231        let checksum = frame.diagnostic_checksum().unwrap().unwrap();
232        assert_eq!(checksum.hex_string(), "608d68d9e0fd3305");
233    }
234}
ustreamer_capture/lib.rs

ustreamer_capture/
lib.rs
