From 48885f7902716c8c65bf34f3bfff22bfe44f11ef Mon Sep 17 00:00:00 2001
From: Volker Mische <volker.mische@gmail.com>
Date: Thu, 5 Dec 2019 14:27:05 +0100
Subject: [PATCH] feat: get GPU devices via FFI

A new FFI call (`util::get_gpu_devices()`) is introduced. It returns
the available GPUs as an array of strings.
---
 rust/Cargo.toml        |  5 +++++
 rust/src/lib.rs        |  1 +
 rust/src/util/api.rs   | 47 ++++++++++++++++++++++++++++++++++++++++++
 rust/src/util/mod.rs   |  2 ++
 rust/src/util/types.rs | 32 ++++++++++++++++++++++++++++
 5 files changed, 87 insertions(+)
 create mode 100644 rust/src/util/api.rs
 create mode 100644 rust/src/util/mod.rs
 create mode 100644 rust/src/util/types.rs

diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index 16c783fe..5fcc1778 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -27,9 +27,14 @@ fil_logger = "0.1.0"
 rand = "0.7"
 rayon = "1"
 anyhow = "1.0.23"
+bellperson = "0.4.4"
 
 [build-dependencies]
 cbindgen = "= 0.10.0"
 
 [dev-dependencies]
 tempfile = "3.0.8"
+
+[features]
+default = ["gpu"]
+gpu = ["bellperson/gpu"]
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 56be4ed4..501235cb 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -8,3 +8,4 @@ extern crate log;
 
 pub mod bls;
 pub mod proofs;
+pub mod util;
diff --git a/rust/src/util/api.rs b/rust/src/util/api.rs
new file mode 100644
index 00000000..d4307eb7
--- /dev/null
+++ b/rust/src/util/api.rs
@@ -0,0 +1,47 @@
+use std::ffi::CString;
+
+use bellperson::GPU_NVIDIA_DEVICES;
+use ffi_toolkit::{catch_panic_response, raw_ptr};
+
+use super::types::GpuDeviceResponse;
+
+/// Returns an array of strings containing the device names that can be used.
+#[no_mangle]
+pub unsafe extern "C" fn get_gpu_devices() -> *mut GpuDeviceResponse {
+    catch_panic_response(|| {
+        let devices: Vec<*const i8> = GPU_NVIDIA_DEVICES
+            .iter()
+            .map(|device| {
+                let name = device.name().unwrap_or("Unknown".to_string());
+                CString::new(&name[..]).unwrap().as_ptr()
+            })
+            .collect();
+        let mut response = GpuDeviceResponse::default();
+        response.devices_len = devices.len();
+        response.devices_ptr = devices.as_ptr();
+
+        raw_ptr(response)
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use std::ffi::CStr;
+    use std::slice::from_raw_parts;
+
+    use crate::util::api::get_gpu_devices;
+    use crate::util::types::destroy_gpu_device_response;
+
+    #[test]
+    fn test_get_gpu_devices() {
+        unsafe {
+            let resp = get_gpu_devices();
+            let devices: Vec<&str> = from_raw_parts((*resp).devices_ptr, (*resp).devices_len)
+                .iter()
+                .map(|name_ptr| CStr::from_ptr(*name_ptr).to_str().unwrap())
+                .collect();
+            assert_eq!(devices.len(), (*resp).devices_len);
+            destroy_gpu_device_response(resp);
+        }
+    }
+}
diff --git a/rust/src/util/mod.rs b/rust/src/util/mod.rs
new file mode 100644
index 00000000..8389f117
--- /dev/null
+++ b/rust/src/util/mod.rs
@@ -0,0 +1,2 @@
+pub mod api;
+pub mod types;
diff --git a/rust/src/util/types.rs b/rust/src/util/types.rs
new file mode 100644
index 00000000..c9a5f2ca
--- /dev/null
+++ b/rust/src/util/types.rs
@@ -0,0 +1,32 @@
+use std::ptr;
+
+use drop_struct_macro_derive::DropStructMacro;
+// `CodeAndMessage` is the trait implemented by `code_and_message_impl
+use ffi_toolkit::{code_and_message_impl, free_c_str, CodeAndMessage, FCPResponseStatus};
+
+#[repr(C)]
+#[derive(DropStructMacro)]
+pub struct GpuDeviceResponse {
+    pub status_code: FCPResponseStatus,
+    pub error_msg: *const libc::c_char,
+    pub devices_len: libc::size_t,
+    pub devices_ptr: *const *const i8,
+}
+
+impl Default for GpuDeviceResponse {
+    fn default() -> Self {
+        Self {
+            error_msg: ptr::null(),
+            status_code: FCPResponseStatus::FCPNoError,
+            devices_len: 0,
+            devices_ptr: ptr::null(),
+        }
+    }
+}
+
+code_and_message_impl!(GpuDeviceResponse);
+
+#[no_mangle]
+pub unsafe extern "C" fn destroy_gpu_device_response(ptr: *mut GpuDeviceResponse) {
+    let _ = Box::from_raw(ptr);
+}