1use std::{
2 fmt::Display,
3 ptr::{null_mut, NonNull},
4};
5
6use core_foundation::kCFAllocatorDefault;
7use core_media::{CMAudioFormatDescription, CMAudioFormatDescriptionGetStreamBasicDescription};
8use core_metal::{MTLDevice as Objc2MTLDevice, MTLPixelFormat as Objc2MTLPixelFormat};
9use core_video::{
10 kCVPixelFormatType_32BGRA, kCVPixelFormatType_32RGBA,
11 kCVPixelFormatType_420YpCbCr8BiPlanarFullRange,
12 kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange, kCVPixelFormatType_420YpCbCr8Planar,
13 kCVReturnSuccess, CVMetalTexture, CVMetalTextureCache, CVMetalTextureCacheCreate,
14 CVMetalTextureCacheCreateTextureFromImage, CVMetalTextureCacheFlush, CVMetalTextureGetTexture,
15 CVPixelBuffer, CVPixelBufferGetBaseAddressOfPlane, CVPixelBufferGetBytesPerRowOfPlane,
16 CVPixelBufferGetHeight, CVPixelBufferGetPixelFormatType, CVPixelBufferGetWidth,
17 CVPixelBufferLockBaseAddress, CVPixelBufferLockFlags, CVPixelBufferUnlockBaseAddress,
18};
19
20use objc2::{rc::Retained, runtime::ProtocolObject};
21
22use crate::{frame::VideoFormat, Size};
23
24pub use core_audo_types::AudioStreamBasicDescription;
25pub use metal::{
26 foreign_types::ForeignType, Device, MTLPixelFormat, MTLTexture, MTLTextureType, Texture,
27 TextureRef,
28};
29
30pub type CVPixelBufferRef = *mut CVPixelBuffer;
31
32#[derive(Debug)]
33pub struct Error(i32);
34
35impl std::error::Error for Error {}
36
37impl Display for Error {
38 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39 write!(f, "error code={}", self.0)
40 }
41}
42
43impl From<i32> for Error {
44 fn from(value: i32) -> Self {
45 Self(value)
46 }
47}
48
49#[allow(non_upper_case_globals)]
50pub fn get_pixel_buffer_format(buffer: CVPixelBufferRef) -> VideoFormat {
51 match unsafe { CVPixelBufferGetPixelFormatType(&*buffer) } {
52 kCVPixelFormatType_32RGBA => VideoFormat::RGBA,
53 kCVPixelFormatType_32BGRA => VideoFormat::BGRA,
54 kCVPixelFormatType_420YpCbCr8Planar => VideoFormat::I420,
55 kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange
56 | kCVPixelFormatType_420YpCbCr8BiPlanarFullRange => VideoFormat::NV12,
57 format => unimplemented!("unsupports format = {:?}", format),
58 }
59}
60
61pub fn get_pixel_buffer_size(buffer: CVPixelBufferRef) -> Size {
62 Size {
63 width: unsafe { CVPixelBufferGetWidth(&*buffer) } as u32,
64 height: unsafe { CVPixelBufferGetHeight(&*buffer) } as u32,
65 }
66}
67
68pub fn get_format_description_info<'a>(
69 descr: *const CMAudioFormatDescription,
70) -> Option<&'a AudioStreamBasicDescription> {
71 let ptr = unsafe { CMAudioFormatDescriptionGetStreamBasicDescription(&*descr) };
72
73 if ptr.is_null() {
74 None
75 } else {
76 Some(unsafe { &*ptr })
77 }
78}
79
80pub struct PixelMomeryBuffer<'a> {
81 pub size: Size,
82 pub format: VideoFormat,
83 pub data: [&'a [u8]; 3],
84 pub linesize: [usize; 3],
85 buffer: CVPixelBufferRef,
86}
87
88impl<'a> PixelMomeryBuffer<'a> {
89 pub fn as_ref(&self) -> CVPixelBufferRef {
90 self.buffer
91 }
92}
93
94impl<'a> From<(CVPixelBufferRef, VideoFormat, Size)> for PixelMomeryBuffer<'a> {
95 fn from((buffer, format, size): (CVPixelBufferRef, VideoFormat, Size)) -> Self {
96 unsafe {
97 CVPixelBufferLockBaseAddress(&*buffer, CVPixelBufferLockFlags::ReadOnly);
98 }
99
100 let mut this = Self {
101 size,
102 format,
103 buffer,
104 data: [&[]; 3],
105 linesize: [0; 3],
106 };
107
108 for i in 0..3 {
109 this.linesize[i] = unsafe { CVPixelBufferGetBytesPerRowOfPlane(&*buffer, i) };
110 this.data[i] = unsafe {
111 std::slice::from_raw_parts(
112 CVPixelBufferGetBaseAddressOfPlane(&*buffer, i) as *const _,
113 this.linesize[i]
114 * if format == VideoFormat::I420 {
115 size.height / 2
116 } else {
117 size.height
118 } as usize,
119 )
120 };
121 }
122
123 this
124 }
125}
126
127impl<'a> From<CVPixelBufferRef> for PixelMomeryBuffer<'a> {
128 fn from(buffer: CVPixelBufferRef) -> Self {
129 Self::from((
130 buffer,
131 get_pixel_buffer_format(buffer),
132 get_pixel_buffer_size(buffer),
133 ))
134 }
135}
136
137impl<'a> Drop for PixelMomeryBuffer<'a> {
138 fn drop(&mut self) {
139 unsafe {
140 CVPixelBufferUnlockBaseAddress(&*self.buffer, CVPixelBufferLockFlags::ReadOnly);
141 }
142 }
143}
144
145#[derive(Clone, Copy)]
146pub struct PixelBuffer {
147 buffer: CVPixelBufferRef,
148 pub format: VideoFormat,
149 pub size: Size,
150}
151
152impl PixelBuffer {
153 pub fn as_ref(&self) -> &CVPixelBuffer {
154 unsafe { &*self.buffer }
155 }
156
157 pub fn as_raw(&self) -> CVPixelBufferRef {
158 self.buffer
159 }
160}
161
162impl From<CVPixelBufferRef> for PixelBuffer {
163 fn from(buffer: CVPixelBufferRef) -> Self {
164 Self::from((
165 buffer,
166 get_pixel_buffer_format(buffer),
167 get_pixel_buffer_size(buffer),
168 ))
169 }
170}
171
172impl From<(CVPixelBufferRef, VideoFormat, Size)> for PixelBuffer {
173 fn from((buffer, format, size): (CVPixelBufferRef, VideoFormat, Size)) -> Self {
174 Self {
175 buffer,
176 format,
177 size,
178 }
179 }
180}
181
182pub struct MetalTextureCache(Retained<CVMetalTextureCache>);
183
184impl MetalTextureCache {
185 pub fn new(device: Device) -> Result<Self, Error> {
186 let device: Retained<ProtocolObject<dyn Objc2MTLDevice>> =
187 unsafe { Retained::from_raw(device.into_ptr().cast()).unwrap() };
188
189 let mut cache = null_mut();
190 let code = unsafe {
191 CVMetalTextureCacheCreate(
192 kCFAllocatorDefault,
193 None,
194 device.as_ref(),
195 None,
196 NonNull::new(&mut cache).unwrap(),
197 )
198 };
199
200 if code != kCVReturnSuccess || cache.is_null() {
201 return Err(Error(code));
202 }
203
204 Ok(Self(unsafe { Retained::from_raw(cache).unwrap() }))
205 }
206
207 pub fn map(&self, buffer: PixelBuffer) -> Result<MetalTexture, Error> {
208 let Size { width, height } = buffer.size;
209
210 let mut texture = null_mut();
211 let code = unsafe {
212 CVMetalTextureCacheCreateTextureFromImage(
213 kCFAllocatorDefault,
214 &self.0,
215 buffer.as_ref(),
216 None,
217 match buffer.format {
218 VideoFormat::BGRA => Objc2MTLPixelFormat::BGRA8Unorm,
219 VideoFormat::RGBA => Objc2MTLPixelFormat::RGBA8Unorm,
220 _ => unimplemented!("unsupports format = {:?}", buffer.format),
221 },
222 width as usize,
223 height as usize,
224 0,
225 NonNull::new(&mut texture).unwrap(),
226 )
227 };
228
229 if code != kCVReturnSuccess || texture.is_null() {
230 return Err(Error(code));
231 }
232
233 Ok(MetalTexture(unsafe {
234 Retained::from_raw(texture).unwrap()
235 }))
236 }
237
238 pub fn flush(&self) {
239 unsafe {
240 CVMetalTextureCacheFlush(&self.0, 0);
241 }
242 }
243}
244
245pub struct MetalTexture(Retained<CVMetalTexture>);
246
247impl MetalTexture {
248 pub fn get_texture(&mut self) -> Result<Texture, Error> {
249 if let Some(texture) = unsafe { CVMetalTextureGetTexture(&self.0) } {
250 Ok(unsafe { Texture::from_ptr(Retained::into_raw(texture).cast()).to_owned() })
251 } else {
252 Err(Error(-1))
253 }
254 }
255}