2022-04-24 02:50:28 +08:00

129 lines
3.2 KiB
Rust

use crate::quartz;
use std::marker::PhantomData;
use std::sync::{Arc, Mutex, TryLockError};
use std::{io, mem, ops};
pub struct Capturer {
inner: quartz::Capturer,
frame: Arc<Mutex<Option<quartz::Frame>>>,
use_yuv: bool,
i420: Vec<u8>,
saved_raw_data: Vec<u128>, // for faster compare and copy
}
impl Capturer {
pub fn new(display: Display, use_yuv: bool) -> io::Result<Capturer> {
let frame = Arc::new(Mutex::new(None));
let f = frame.clone();
let inner = quartz::Capturer::new(
display.0,
display.width(),
display.height(),
if use_yuv {
quartz::PixelFormat::YCbCr420Video
} else {
quartz::PixelFormat::Argb8888
},
Default::default(),
move |inner| {
if let Ok(mut f) = f.lock() {
*f = Some(inner);
}
},
)
.map_err(|_| io::Error::from(io::ErrorKind::Other))?;
Ok(Capturer {
inner,
frame,
use_yuv,
i420: Vec::new(),
saved_raw_data: Vec::new(),
})
}
pub fn width(&self) -> usize {
self.inner.width()
}
pub fn height(&self) -> usize {
self.inner.height()
}
pub fn frame<'a>(&'a mut self, _timeout_ms: u32) -> io::Result<Frame<'a>> {
match self.frame.try_lock() {
Ok(mut handle) => {
let mut frame = None;
mem::swap(&mut frame, &mut handle);
match frame {
Some(mut frame) => {
crate::would_block_if_equal(&mut self.saved_raw_data, frame.inner())?;
if self.use_yuv {
frame.nv12_to_i420(self.width(), self.height(), &mut self.i420);
}
Ok(Frame(frame, PhantomData))
}
None => Err(io::ErrorKind::WouldBlock.into()),
}
}
Err(TryLockError::WouldBlock) => Err(io::ErrorKind::WouldBlock.into()),
Err(TryLockError::Poisoned(..)) => Err(io::ErrorKind::Other.into()),
}
}
}
pub struct Frame<'a>(quartz::Frame, PhantomData<&'a [u8]>);
impl<'a> ops::Deref for Frame<'a> {
type Target = [u8];
fn deref(&self) -> &[u8] {
&*self.0
}
}
pub struct Display(quartz::Display);
impl Display {
pub fn primary() -> io::Result<Display> {
Ok(Display(quartz::Display::primary()))
}
pub fn all() -> io::Result<Vec<Display>> {
Ok(quartz::Display::online()
.map_err(|_| io::Error::from(io::ErrorKind::Other))?
.into_iter()
.map(Display)
.collect())
}
pub fn width(&self) -> usize {
self.0.width()
}
pub fn height(&self) -> usize {
self.0.height()
}
pub fn name(&self) -> String {
self.0.id().to_string()
}
pub fn is_online(&self) -> bool {
self.0.is_online()
}
pub fn origin(&self) -> (i32, i32) {
let o = self.0.bounds().origin;
(o.x as _, o.y as _)
}
pub fn is_primary(&self) -> bool {
self.0.is_primary()
}
}