Refactor code to add a SystemCapturer.

nvfbc-sys/src/lib.rs

@@ -14,6 +14,8 @@ pub const NVFBC_TOGL_SETUP_PARAMS_VER: u32 = nvfbc_struct_version::<NVFBC_TOGL_S
 pub const NVFBC_TOGL_GRAB_FRAME_PARAMS_VER: u32 = nvfbc_struct_version::<NVFBC_TOGL_GRAB_FRAME_PARAMS>(2);
 pub const NVFBC_TOCUDA_SETUP_PARAMS_VER: u32 = nvfbc_struct_version::<NVFBC_TOCUDA_SETUP_PARAMS>(1);
 pub const NVFBC_TOCUDA_GRAB_FRAME_PARAMS_VER: u32 = nvfbc_struct_version::<NVFBC_TOCUDA_GRAB_FRAME_PARAMS>(2);
+pub const NVFBC_TOSYS_SETUP_PARAMS_VER: u32 = nvfbc_struct_version::<NVFBC_TOSYS_SETUP_PARAMS>(3);
+pub const NVFBC_TOSYS_GRAB_FRAME_PARAMS_VER: u32 = nvfbc_struct_version::<NVFBC_TOSYS_GRAB_FRAME_PARAMS>(2);
 
 pub const fn nvfbc_struct_version<T>(version: u32) -> u32 {
 	std::mem::size_of::<T>() as u32 | ((version) << 16 | NVFBC_VERSION << 24)

nvfbc/Cargo.toml

@@ -11,3 +11,6 @@ repository = "https://github.com/hgaiser/nvfbc-rs"
 
 [dependencies]
 nvfbc-sys = { path = "../nvfbc-sys" }
+
+[dev-dependencies]
+image = "0.24.2"

nvfbc/src/common.rs

@@ -0,0 +1,68 @@
+use std::os::raw::c_uint;
+use std::{mem::MaybeUninit, ffi::CStr};
+
+use nvfbc_sys::_NVFBCSTATUS_NVFBC_SUCCESS as SUCCESS;
+use nvfbc_sys::NVFBC_SESSION_HANDLE;
+
+use crate::CaptureType;
+use crate::Error;
+use crate::Status;
+
+pub type Handle = NVFBC_SESSION_HANDLE;
+
+pub(crate) fn check_ret(handle: Handle, ret: nvfbc_sys::_NVFBCSTATUS) -> Result<(), Error> {
+	if ret != SUCCESS {
+		return Err(Error::new(ret, get_last_error(handle)));
+	}
+	Ok(())
+}
+
+pub(crate) fn create_handle() -> Result<nvfbc_sys::NVFBC_SESSION_HANDLE, Error> {
+	let mut params: nvfbc_sys::_NVFBC_CREATE_HANDLE_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
+	params.dwVersion = nvfbc_sys::NVFBC_CREATE_HANDLE_PARAMS_VER;
+	let mut handle = 0;
+	let ret = unsafe { nvfbc_sys::NvFBCCreateHandle(
+		&mut handle,
+		&mut params
+	)};
+	if ret != SUCCESS {
+		return Err(Error::new(ret, None));
+	}
+
+	Ok(handle)
+}
+
+pub(crate) fn destroy_handle(handle: Handle) -> Result<(), Error> {
+	let mut params: nvfbc_sys::_NVFBC_DESTROY_HANDLE_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
+	params.dwVersion = nvfbc_sys::NVFBC_DESTROY_HANDLE_PARAMS_VER;
+	check_ret(handle, unsafe { nvfbc_sys::NvFBCDestroyHandle(handle, &mut params) })
+}
+
+pub(crate) fn get_last_error(handle: Handle) -> Option<String> {
+	let error = unsafe { nvfbc_sys::NvFBCGetLastErrorStr(handle) };
+	let error = unsafe { CStr::from_ptr(error) };
+	error.to_str().ok().map(|e| e.to_string())
+}
+
+pub(crate) fn status(handle: Handle) -> Result<Status, Error> {
+	let mut params: nvfbc_sys::_NVFBC_GET_STATUS_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
+	params.dwVersion = nvfbc_sys::NVFBC_GET_STATUS_PARAMS_VER;
+	check_ret(handle, unsafe { nvfbc_sys::NvFBCGetStatus(handle, &mut params) })?;
+	Ok(params.into())
+}
+
+pub(crate) fn create_capture_session(handle: Handle, capture_type: CaptureType) -> Result<(), Error> {
+	let mut params: nvfbc_sys::_NVFBC_CREATE_CAPTURE_SESSION_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
+	params.dwVersion = nvfbc_sys::NVFBC_CREATE_CAPTURE_SESSION_PARAMS_VER;
+	params.eCaptureType = capture_type as c_uint;
+	params.bWithCursor = nvfbc_sys::_NVFBC_BOOL_NVFBC_TRUE;
+	params.frameSize = nvfbc_sys::NVFBC_SIZE { w: 0, h: 0 };
+	params.eTrackingType = nvfbc_sys::NVFBC_TRACKING_TYPE_NVFBC_TRACKING_DEFAULT;
+	check_ret(handle, unsafe { nvfbc_sys::NvFBCCreateCaptureSession(handle, &mut params) })
+}
+
+pub(crate) fn destroy_capture_session(handle: Handle) -> Result<(), Error> {
+	let mut params: nvfbc_sys::_NVFBC_DESTROY_CAPTURE_SESSION_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
+	params.dwVersion = nvfbc_sys::NVFBC_DESTROY_CAPTURE_SESSION_PARAMS_VER;
+	check_ret(handle, unsafe { nvfbc_sys::NvFBCDestroyCaptureSession(handle, &mut params) })
+}

nvfbc/src/cuda.rs

@@ -1,116 +1,62 @@
-use std::{ffi::{CStr, c_void}, mem::MaybeUninit, ptr::null_mut, os::raw::c_uint};
-
-use nvfbc_sys::_NVFBCSTATUS_NVFBC_SUCCESS as SUCCESS;
-
-use crate::{BufferFormat, Error, Status, CaptureType, CudaFrameInfo};
+use std::ffi::c_void;
+use std::mem::MaybeUninit;
+use std::ptr::null_mut;
+
+use crate::{
+	BufferFormat,
+	CaptureType,
+	CudaFrameInfo,
+	Error,
+	Status,
+};
+
+use crate::common::{
+	Handle,
+	check_ret,
+	create_capture_session,
+	create_handle,
+	destroy_capture_session,
+	destroy_handle,
+	status,
+};
 
 pub struct CudaCapturer {
-	handle: nvfbc_sys::NVFBC_SESSION_HANDLE,
+	handle: Handle,
 }
 
 impl CudaCapturer {
 	pub fn new() -> Result<Self, Error> {
-		let handle = Self::create_handle()?;
-
+		let handle = create_handle()?;
 		let self_ = Self { handle };
-		self_.setup()?;
-
 		Ok(self_)
 	}
 
-	fn create_handle() -> Result<nvfbc_sys::NVFBC_SESSION_HANDLE, Error> {
-		let mut params: nvfbc_sys::_NVFBC_CREATE_HANDLE_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
-		params.dwVersion = nvfbc_sys::NVFBC_CREATE_HANDLE_PARAMS_VER;
-		let mut handle = 0;
-		let ret = unsafe { nvfbc_sys::NvFBCCreateHandle(
-			&mut handle,
-			&mut params
-		)};
-		if ret != SUCCESS {
-			return Err(Error::new(ret, None));
-		}
-
-		Ok(handle)
-	}
-
-	fn destroy_handle(&self) -> Result<(), Error> {
-		let mut params: nvfbc_sys::_NVFBC_DESTROY_HANDLE_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
-		params.dwVersion = nvfbc_sys::NVFBC_DESTROY_HANDLE_PARAMS_VER;
-		let ret = unsafe { nvfbc_sys::NvFBCDestroyHandle(self.handle, &mut params) };
-		if ret != SUCCESS {
-			return Err(Error::new(ret, self.get_last_error()));
-		}
-
-		Ok(())
-	}
-
-	fn get_last_error(&self) -> Option<String> {
-		let error = unsafe { nvfbc_sys::NvFBCGetLastErrorStr(self.handle) };
-		let error = unsafe { CStr::from_ptr(error) };
-		error.to_str().ok().map(|e| e.to_string())
-	}
-
-	fn setup(&self) -> Result<(), Error> {
-		let mut params: nvfbc_sys::_NVFBC_CREATE_CAPTURE_SESSION_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
-		params.dwVersion = nvfbc_sys::NVFBC_CREATE_CAPTURE_SESSION_PARAMS_VER;
-		params.eCaptureType = CaptureType::SharedCuda as c_uint;
-		params.bWithCursor = nvfbc_sys::_NVFBC_BOOL_NVFBC_TRUE;
-		params.frameSize = nvfbc_sys::NVFBC_SIZE { w: 0, h: 0 };
-		params.eTrackingType = nvfbc_sys::NVFBC_TRACKING_TYPE_NVFBC_TRACKING_DEFAULT;
-		let ret = unsafe { nvfbc_sys::NvFBCCreateCaptureSession(self.handle, &mut params) };
-		if ret != SUCCESS {
-			return Err(Error::new(ret, self.get_last_error()));
-		}
-
-		Ok(())
+	pub fn status(&self) -> Result<Status, Error> {
+		status(self.handle)
 	}
 
 	pub fn start(&self, buffer_format: BufferFormat) -> Result<(), Error> {
+		create_capture_session(self.handle, CaptureType::SharedCuda)?;
+
 		let mut params: nvfbc_sys::NVFBC_TOCUDA_SETUP_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
 		params.dwVersion = nvfbc_sys::NVFBC_TOCUDA_SETUP_PARAMS_VER;
 		params.eBufferFormat = buffer_format as u32;
-		let ret = unsafe { nvfbc_sys::NvFBCToCudaSetUp(self.handle, &mut params) };
-		if ret != SUCCESS {
-			return Err(Error::new(ret, self.get_last_error()));
-		}
-
-		Ok(())
-	}
-
-	pub fn status(&self) -> Result<Status, Error> {
-		let mut params: nvfbc_sys::_NVFBC_GET_STATUS_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
-		params.dwVersion = nvfbc_sys::NVFBC_GET_STATUS_PARAMS_VER;
-		let ret = unsafe { nvfbc_sys::NvFBCGetStatus(self.handle, &mut params) };
-		if ret != SUCCESS {
-			return Err(Error::new(ret, self.get_last_error()));
-		}
-
-		Ok(params.into())
+		check_ret(self.handle, unsafe { nvfbc_sys::NvFBCToCudaSetUp(self.handle, &mut params) })
 	}
 
 	pub fn stop(&self) -> Result<(), Error> {
-		let mut params: nvfbc_sys::_NVFBC_DESTROY_CAPTURE_SESSION_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
-		params.dwVersion = nvfbc_sys::NVFBC_DESTROY_CAPTURE_SESSION_PARAMS_VER;
-		let ret = unsafe { nvfbc_sys::NvFBCDestroyCaptureSession(self.handle, &mut params) };
-		if ret != SUCCESS {
-			return Err(Error::new(ret, self.get_last_error()));
-		}
-
-		Ok(())
+		destroy_capture_session(self.handle)
 	}
 
-	pub fn grab(&self) -> Result<CudaFrameInfo, Error> {
-		let mut device_buffer: *mut c_void = null_mut();
+	pub fn next(&self) -> Result<CudaFrameInfo, Error> {
+		let device_buffer: *mut c_void = null_mut();
 		let mut frame_info: nvfbc_sys::NVFBC_FRAME_GRAB_INFO = unsafe { MaybeUninit::zeroed().assume_init() };
 		let mut params: nvfbc_sys::NVFBC_TOCUDA_GRAB_FRAME_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
 		params.dwVersion = nvfbc_sys::NVFBC_TOCUDA_GRAB_FRAME_PARAMS_VER;
-		params.dwFlags = nvfbc_sys::NVFBC_TOCUDA_FLAGS_NVFBC_TOCUDA_GRAB_FLAGS_NOWAIT;
+		params.dwFlags = nvfbc_sys::NVFBC_TOCUDA_FLAGS_NVFBC_TOCUDA_GRAB_FLAGS_NOFLAGS;
 		params.pFrameGrabInfo = &mut frame_info;
-		params.pCUDADeviceBuffer = &mut device_buffer as *mut _ as *mut c_void;
-		let ret = unsafe { nvfbc_sys::NvFBCToCudaGrabFrame(self.handle, &mut params) };
-		if ret != SUCCESS {
-			return Err(Error::new(ret, self.get_last_error()));
-		}
+		params.pCUDADeviceBuffer = device_buffer;
+		check_ret(self.handle, unsafe { nvfbc_sys::NvFBCToCudaGrabFrame(self.handle, &mut params) })?;
 
 		Ok(CudaFrameInfo {
 			device_buffer,
@@ -124,8 +70,7 @@ impl CudaCapturer {
 
 impl Drop for CudaCapturer {
 	fn drop(&mut self) {
-		self.stop().ok();
 		// TODO: Figure out why this crashes (nvfbc examples also fail here..)
-		self.destroy_handle().ok();
+		destroy_handle(self.handle).ok();
 	}
 }

nvfbc/src/lib.rs

@@ -1,7 +1,10 @@
-mod types;
-mod error;
+mod common;
 mod cuda;
+mod error;
+mod system;
+mod types;
 
 pub use types::*;
 pub use error::Error;
 pub use cuda::CudaCapturer;
+pub use system::SystemCapturer;

nvfbc/src/system.rs

@@ -0,0 +1,75 @@
+use std::ffi::c_void;
+use std::mem::MaybeUninit;
+use std::ptr::null_mut;
+
+use crate::common::{
+	Handle,
+	check_ret,
+	create_capture_session,
+	create_handle,
+	destroy_capture_session,
+	destroy_handle,
+	status,
+};
+use crate::{
+	BufferFormat,
+	Error,
+	Status,
+	CaptureType,
+	SystemFrameInfo
+};
+
+pub struct SystemCapturer {
+	handle: Handle,
+	buffer: *mut c_void,
+}
+
+impl SystemCapturer {
+	pub fn new() -> Result<Self, Error> {
+		let handle = create_handle()?;
+		let self_ = Self { handle, buffer: null_mut() };
+		Ok(self_)
+	}
+
+	pub fn status(&self) -> Result<Status, Error> {
+		status(self.handle)
+	}
+
+	pub fn start(&mut self, buffer_format: BufferFormat) -> Result<(), Error> {
+		create_capture_session(self.handle, CaptureType::ToSystem)?;
+
+		let mut params: nvfbc_sys::NVFBC_TOSYS_SETUP_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
+		params.dwVersion = nvfbc_sys::NVFBC_TOSYS_SETUP_PARAMS_VER;
+		params.eBufferFormat = buffer_format as u32;
+		params.ppBuffer = &mut self.buffer as *mut _;
+		check_ret(self.handle, unsafe { nvfbc_sys::NvFBCToSysSetUp(self.handle, &mut params) })
+	}
+
+	pub fn stop(&self) -> Result<(), Error> {
+		destroy_capture_session(self.handle)
+	}
+
+	pub fn next(&self) -> Result<SystemFrameInfo, Error> {
+		let mut frame_info: nvfbc_sys::NVFBC_FRAME_GRAB_INFO = unsafe { MaybeUninit::zeroed().assume_init() };
+		let mut params: nvfbc_sys::NVFBC_TOSYS_GRAB_FRAME_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
+		params.dwVersion = nvfbc_sys::NVFBC_TOSYS_GRAB_FRAME_PARAMS_VER;
+		params.dwFlags = nvfbc_sys::NVFBC_TOSYS_GRAB_FLAGS_NVFBC_TOSYS_GRAB_FLAGS_NOFLAGS;
+		params.pFrameGrabInfo = &mut frame_info;
+		check_ret(self.handle, unsafe { nvfbc_sys::NvFBCToSysGrabFrame(self.handle, &mut params) })?;
+
+		Ok(SystemFrameInfo {
+			buffer: self.buffer,
+			width: frame_info.dwWidth,
+			height: frame_info.dwHeight,
+			byte_size: frame_info.dwByteSize as usize,
+			current_frame: frame_info.dwCurrentFrame,
+		})
+	}
+}
+
+impl Drop for SystemCapturer {
+	fn drop(&mut self) {
+		// TODO: Figure out why this crashes (nvfbc examples also fail here..)
+		destroy_handle(self.handle).ok();
+	}
+}

nvfbc/src/types.rs

@@ -16,6 +16,22 @@ pub struct CudaFrameInfo {
 	pub current_frame: u32,
 }
 
+#[derive(Debug, Copy, Clone)]
+pub struct SystemFrameInfo {
+	/// Pointer to the frame that is grabbed.
+	pub buffer: *mut c_void,
+	/// Width of the captured frame.
+	pub width: u32,
+	/// Height of the captured frame.
+	pub height: u32,
+	/// Size of the frame in bytes.
+	pub byte_size: usize,
+	/// Incremental ID of the current frame.
+	///
+	/// This can be used to identify a frame.
+	pub current_frame: u32,
+}
+
 #[derive(Debug, Copy, Clone)]
 pub enum CaptureType {
 	/// Capture frames to a buffer in system memory.