use crate::helpers::symbol_to_atomic_number;
use crate::iterators::{self, ConFrameIterator};
use crate::types::{ConFrame, ConFrameBuilder, meta};
use crate::writer::ConFrameWriter;
use std::ffi::{CStr, CString, c_char};
use std::fs::File;
use std::path::Path;
use std::ptr;
pub const RKR_CON_SPEC_VERSION: u32 = 3;
#[unsafe(no_mangle)]
pub extern "C" fn rkr_con_spec_version() -> u32 {
crate::CON_SPEC_VERSION
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_library_version() -> *const c_char {
const VERSION_NUL: &[u8] = concat!(env!("CARGO_PKG_VERSION"), "\0").as_bytes();
VERSION_NUL.as_ptr() as *const c_char
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_atom_index_by_id(
frame_handle: *const RKRConFrame,
atom_id: u64,
) -> u64 {
let frame = match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f,
None => return u64::MAX,
};
match frame.atom_index_by_id(atom_id) {
Some(idx) => idx as u64,
None => u64::MAX,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_symbol_to_z(symbol: *const c_char) -> u64 {
if symbol.is_null() {
return 0;
}
match unsafe { CStr::from_ptr(symbol) }.to_str() {
Ok(s) => symbol_to_atomic_number(s),
Err(_) => 0,
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_z_to_symbol(z: u64) -> *const c_char {
macro_rules! cstrs {
($($lit:literal),* $(,)?) => {
[$(concat!($lit, "\0").as_bytes()),*]
};
}
const TABLE: [&[u8]; 93] = cstrs![
"X", "H", "He", "Li", "Be", "B", "C", "N", "O", "F", "Ne", "Na", "Mg", "Al", "Si", "P",
"S", "Cl", "Ar", "K", "Ca", "Sc", "Ti", "V", "Cr", "Mn", "Fe", "Co", "Ni", "Cu", "Zn",
"Ga", "Ge", "As", "Se", "Br", "Kr", "Rb", "Sr", "Y", "Zr", "Nb", "Mo", "Tc", "Ru", "Rh",
"Pd", "Ag", "Cd", "In", "Sn", "Sb", "Te", "I", "Xe", "Cs", "Ba", "La", "Ce", "Pr", "Nd",
"Pm", "Sm", "Eu", "Gd", "Tb", "Dy", "Ho", "Er", "Tm", "Yb", "Lu", "Hf", "Ta", "W", "Re",
"Os", "Ir", "Pt", "Au", "Hg", "Tl", "Pb", "Bi", "Po", "At", "Rn", "Fr", "Ra", "Ac", "Th",
"Pa", "U",
];
let idx = if (1..=92).contains(&z) { z as usize } else { 0 };
TABLE[idx].as_ptr() as *const c_char
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_spec_version(frame_handle: *const RKRConFrame) -> u32 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f.header.spec_version,
None => 0,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_metadata_json(frame_handle: *const RKRConFrame) -> *mut c_char {
let frame = match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f,
None => return ptr::null_mut(),
};
let mut obj = serde_json::Map::new();
obj.insert(
meta::CON_SPEC_VERSION.into(),
serde_json::Value::from(frame.header.spec_version),
);
for (k, v) in &frame.header.metadata {
obj.insert(k.clone(), v.clone());
}
let json_str = serde_json::Value::Object(obj).to_string();
match CString::new(json_str) {
Ok(cs) => cs.into_raw(),
Err(_) => ptr::null_mut(),
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_energy(frame_handle: *const RKRConFrame) -> f64 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f.header.energy().unwrap_or(f64::NAN),
None => f64::NAN,
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_index_energy(frame_handle: *const RKRConFrame) -> f64 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => crate::index_proj::finite_energy(f).unwrap_or(f64::NAN),
None => f64::NAN,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_composition_formula(
frame_handle: *const RKRConFrame,
) -> *mut c_char {
let frame = match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f,
None => return ptr::null_mut(),
};
let s = crate::index_proj::frame_composition_formula(frame);
match CString::new(s) {
Ok(cs) => cs.into_raw(),
Err(_) => ptr::null_mut(),
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_total_mass(frame_handle: *const RKRConFrame) -> f64 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => crate::index_proj::frame_total_mass(f).unwrap_or(f64::NAN),
None => f64::NAN,
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_cell_volume(frame_handle: *const RKRConFrame) -> f64 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => crate::index_proj::frame_cell_volume(f).unwrap_or(f64::NAN),
None => f64::NAN,
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_fmax(frame_handle: *const RKRConFrame) -> f64 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => crate::index_proj::frame_fmax(f).unwrap_or(f64::NAN),
None => f64::NAN,
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_sections_mask(frame_handle: *const RKRConFrame) -> u8 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => crate::index_proj::sections_present_mask(f),
None => 0,
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_index_natoms(frame_handle: *const RKRConFrame) -> u32 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f.atom_data.len() as u32,
None => 0,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_index_projection_json(
frame_handle: *const RKRConFrame,
) -> *mut c_char {
let frame = match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f,
None => return ptr::null_mut(),
};
let p = crate::index_proj::FrameIndexProjection::from_frame(frame);
let v = serde_json::json!({
"n_atoms": p.n_atoms,
"formula": p.formula,
"energy": p.energy,
"fmax": p.fmax,
"total_mass": p.total_mass,
"cell_volume": p.cell_volume,
"sections_mask": p.sections_mask,
"has_forces": p.has_forces,
"has_velocities": p.has_velocities,
"has_energy": p.has_energy,
"symbols": p.symbols,
"species_counts": p.species_counts.iter().map(|(s,c)| serde_json::json!([s, c])).collect::<Vec<_>>(),
"time": p.time,
"timestep": p.timestep,
"frame_index": p.frame_index,
"neb_bead": p.neb_bead,
"neb_band": p.neb_band,
"charge": p.charge,
"magmom": p.magmom,
});
match CString::new(v.to_string()) {
Ok(cs) => cs.into_raw(),
Err(_) => ptr::null_mut(),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_potential_type(frame_handle: *const RKRConFrame) -> *mut c_char {
let frame = match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f,
None => return ptr::null_mut(),
};
match frame.header.potential_type() {
Some(pot_type) => match CString::new(pot_type) {
Ok(cs) => cs.into_raw(),
Err(_) => ptr::null_mut(),
},
None => ptr::null_mut(),
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_frame_index(frame_handle: *const RKRConFrame) -> u64 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f.header.frame_index().unwrap_or(u64::MAX),
None => u64::MAX,
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_time(frame_handle: *const RKRConFrame) -> f64 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f.header.time().unwrap_or(f64::NAN),
None => f64::NAN,
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_timestep(frame_handle: *const RKRConFrame) -> f64 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f.header.timestep().unwrap_or(f64::NAN),
None => f64::NAN,
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_neb_bead(frame_handle: *const RKRConFrame) -> u64 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f.header.neb_bead().unwrap_or(u64::MAX),
None => u64::MAX,
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_frame_neb_band(frame_handle: *const RKRConFrame) -> u64 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f.header.neb_band().unwrap_or(u64::MAX),
None => u64::MAX,
}
}
#[repr(C)]
#[allow(non_camel_case_types)]
#[derive(Debug, PartialEq, Eq)]
pub enum RKRStatus {
RKR_STATUS_SUCCESS = 0,
RKR_STATUS_NULL_POINTER = -1,
RKR_STATUS_INVALID_UTF8 = -2,
RKR_STATUS_INVALID_JSON = -3,
RKR_STATUS_IO_ERROR = -4,
RKR_STATUS_INDEX_OUT_OF_BOUNDS = -5,
RKR_STATUS_BUFFER_TOO_SMALL = -6,
RKR_STATUS_INTERNAL_ERROR = -7,
RKR_STATUS_SECTION_ABSENT = -8,
RKR_STATUS_VALIDATION_ERROR = -9,
RKR_STATUS_SELECTION_ERROR = -10,
RKR_STATUS_FEATURE_DISABLED = -11,
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_bond_count(frame_handle: *const RKRConFrame) -> u64 {
match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f.bonds().len() as u64,
None => 0,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_bond_at(
frame_handle: *const RKRConFrame,
index: u64,
out_i: *mut u32,
out_j: *mut u32,
out_has_order: *mut u8,
out_order: *mut i32,
) -> RKRStatus {
if frame_handle.is_null()
|| out_i.is_null()
|| out_j.is_null()
|| out_has_order.is_null()
|| out_order.is_null()
{
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
let bonds = frame.bonds();
let Some(bond) = bonds.get(index as usize) else {
return RKRStatus::RKR_STATUS_INDEX_OUT_OF_BOUNDS;
};
unsafe {
*out_i = bond.i;
*out_j = bond.j;
if let Some(order) = bond.order {
*out_has_order = 1;
*out_order = order;
} else {
*out_has_order = 0;
*out_order = 0;
}
}
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_status_message(status: RKRStatus) -> *const c_char {
match status {
RKRStatus::RKR_STATUS_SUCCESS => c"success".as_ptr(),
RKRStatus::RKR_STATUS_NULL_POINTER => c"null pointer".as_ptr(),
RKRStatus::RKR_STATUS_INVALID_UTF8 => c"invalid UTF-8".as_ptr(),
RKRStatus::RKR_STATUS_INVALID_JSON => c"invalid JSON".as_ptr(),
RKRStatus::RKR_STATUS_IO_ERROR => c"I/O error".as_ptr(),
RKRStatus::RKR_STATUS_INDEX_OUT_OF_BOUNDS => c"index out of bounds".as_ptr(),
RKRStatus::RKR_STATUS_BUFFER_TOO_SMALL => c"buffer too small".as_ptr(),
RKRStatus::RKR_STATUS_INTERNAL_ERROR => c"internal error".as_ptr(),
RKRStatus::RKR_STATUS_SECTION_ABSENT => c"section absent".as_ptr(),
RKRStatus::RKR_STATUS_VALIDATION_ERROR => c"validation error".as_ptr(),
RKRStatus::RKR_STATUS_SELECTION_ERROR => c"selection error".as_ptr(),
RKRStatus::RKR_STATUS_FEATURE_DISABLED => c"feature disabled in this build".as_ptr(),
}
}
#[repr(C)]
pub struct RKRConFrame {
_private: [u8; 0],
}
#[repr(C)]
pub struct RKRConFrameWriter {
_private: [u8; 0],
}
#[repr(C)]
pub struct CFrame {
pub atoms: *mut CAtom,
pub num_atoms: usize,
pub cell: [f64; 3],
pub angles: [f64; 3],
pub has_velocities: bool,
pub has_forces: bool,
pub has_energies: bool,
}
#[repr(C)]
pub struct CAtom {
pub atomic_number: u64,
pub x: f64,
pub y: f64,
pub z: f64,
pub atom_id: u64,
pub mass: f64,
pub is_fixed: bool,
pub fixed_x: bool,
pub fixed_y: bool,
pub fixed_z: bool,
pub vx: f64,
pub vy: f64,
pub vz: f64,
pub has_velocity: bool,
pub fx: f64,
pub fy: f64,
pub fz: f64,
pub has_forces: bool,
pub energy: f64,
pub has_energy: bool,
}
#[repr(C)]
pub struct CConFrameIterator {
iterator: *mut ConFrameIterator<'static>,
file_contents: *mut String,
}
fn c_iterator_from_owned_string(contents: String) -> *mut CConFrameIterator {
let file_contents_box = Box::new(contents);
let file_contents_ptr = Box::into_raw(file_contents_box);
let static_file_contents: &'static str = unsafe { &*file_contents_ptr };
let iterator = Box::new(ConFrameIterator::new(static_file_contents));
let c_iterator = Box::new(CConFrameIterator {
iterator: Box::into_raw(iterator),
file_contents: file_contents_ptr,
});
Box::into_raw(c_iterator)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn read_con_file_iterator(
filename_c: *const c_char,
) -> *mut CConFrameIterator {
if filename_c.is_null() {
return ptr::null_mut();
}
let filename = match unsafe { CStr::from_ptr(filename_c).to_str() } {
Ok(s) => s,
Err(_) => return ptr::null_mut(),
};
let owned = match crate::compression::read_file_contents(Path::new(filename)) {
Ok(fc) => match fc.as_str() {
Ok(s) => s.to_owned(),
Err(_) => return ptr::null_mut(),
},
Err(_) => return ptr::null_mut(),
};
c_iterator_from_owned_string(owned)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn read_con_string_iterator(
contents_c: *const c_char,
) -> *mut CConFrameIterator {
if contents_c.is_null() {
return ptr::null_mut();
}
let contents = match unsafe { CStr::from_ptr(contents_c).to_str() } {
Ok(s) => s.to_owned(),
Err(_) => return ptr::null_mut(),
};
c_iterator_from_owned_string(contents)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn read_con_buffer_iterator(
data: *const u8,
len: usize,
) -> *mut CConFrameIterator {
if data.is_null() && len > 0 {
return ptr::null_mut();
}
if len == 0 {
return c_iterator_from_owned_string(String::new());
}
let slice = unsafe { std::slice::from_raw_parts(data, len) };
let contents = match std::str::from_utf8(slice) {
Ok(s) => s.to_owned(),
Err(_) => return ptr::null_mut(),
};
c_iterator_from_owned_string(contents)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn con_frame_iterator_next(
iterator: *mut CConFrameIterator,
) -> *mut RKRConFrame {
if iterator.is_null() {
return ptr::null_mut();
}
let iter = unsafe { &mut *(*iterator).iterator };
match iter.next() {
Some(Ok(frame)) => Box::into_raw(Box::new(frame)) as *mut RKRConFrame,
_ => ptr::null_mut(),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn free_rkr_frame(frame_handle: *mut RKRConFrame) {
if !frame_handle.is_null() {
let _ = unsafe { Box::from_raw(frame_handle as *mut ConFrame) };
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn free_con_frame_iterator(iterator: *mut CConFrameIterator) {
if iterator.is_null() {
return;
}
unsafe {
let c_iterator_box = Box::from_raw(iterator);
let _ = Box::from_raw(c_iterator_box.iterator);
let _ = Box::from_raw(c_iterator_box.file_contents);
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_to_c_frame(frame_handle: *const RKRConFrame) -> *mut CFrame {
let frame = match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f,
None => return ptr::null_mut(),
};
let masses_iter = frame
.header
.natms_per_type
.iter()
.zip(frame.header.masses_per_type.iter())
.flat_map(|(num_atoms, mass)| std::iter::repeat_n(*mass, *num_atoms));
let has_velocities = frame.has_velocities();
let mut c_atoms: Vec<CAtom> = frame
.atom_data
.iter()
.zip(masses_iter)
.map(|(atom_datum, mass)| {
let [vx, vy, vz] = atom_datum.velocity.unwrap_or([0.0; 3]);
let [fx, fy, fz] = atom_datum.force.unwrap_or([0.0; 3]);
CAtom {
atomic_number: symbol_to_atomic_number(&atom_datum.symbol),
x: atom_datum.x,
y: atom_datum.y,
z: atom_datum.z,
is_fixed: atom_datum.is_fixed(),
fixed_x: atom_datum.fixed[0],
fixed_y: atom_datum.fixed[1],
fixed_z: atom_datum.fixed[2],
atom_id: atom_datum.atom_id,
mass,
vx,
vy,
vz,
has_velocity: atom_datum.has_velocity(),
fx,
fy,
fz,
has_forces: atom_datum.has_forces(),
energy: atom_datum.energy.unwrap_or(0.0),
has_energy: atom_datum.has_energy(),
}
})
.collect();
let atoms_ptr = c_atoms.as_mut_ptr();
let num_atoms = c_atoms.len();
std::mem::forget(c_atoms);
let has_forces = frame.has_forces();
let has_energies = frame.has_energies();
let c_frame = Box::new(CFrame {
atoms: atoms_ptr,
num_atoms,
cell: frame.header.boxl,
angles: frame.header.angles,
has_velocities,
has_forces,
has_energies,
});
Box::into_raw(c_frame)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn free_c_frame(frame: *mut CFrame) {
if frame.is_null() {
return;
}
unsafe {
let frame_box = Box::from_raw(frame);
let _ = Vec::from_raw_parts(frame_box.atoms, frame_box.num_atoms, frame_box.num_atoms);
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_get_header_line(
frame_handle: *const RKRConFrame,
is_prebox: bool,
line_index: usize,
buffer: *mut c_char,
buffer_len: usize,
) -> RKRStatus {
let frame = match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f,
None => return RKRStatus::RKR_STATUS_NULL_POINTER,
};
if buffer.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
if buffer_len == 0 {
return RKRStatus::RKR_STATUS_BUFFER_TOO_SMALL;
}
let line_to_copy: Option<&str> = if is_prebox {
match line_index {
0 => Some(frame.header.prebox_header.user.as_str()),
1 => Some(frame.header.prebox_header.metadata_line()),
_ => None,
}
} else {
frame.header.postbox_header.get(line_index).map(String::as_str)
};
if let Some(line) = line_to_copy {
let bytes = line.as_bytes();
let len_to_copy = std::cmp::min(bytes.len(), buffer_len - 1);
unsafe {
ptr::copy_nonoverlapping(bytes.as_ptr(), buffer as *mut u8, len_to_copy);
*buffer.add(len_to_copy) = 0;
}
RKRStatus::RKR_STATUS_SUCCESS
} else {
RKRStatus::RKR_STATUS_INDEX_OUT_OF_BOUNDS
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_get_header_line_cpp(
frame_handle: *const RKRConFrame,
is_prebox: bool,
line_index: usize,
) -> *mut c_char {
let frame = match unsafe { (frame_handle as *const ConFrame).as_ref() } {
Some(f) => f,
None => return ptr::null_mut(),
};
let line_to_copy: Option<&str> = if is_prebox {
match line_index {
0 => Some(frame.header.prebox_header.user.as_str()),
1 => Some(frame.header.prebox_header.metadata_line()),
_ => None,
}
} else {
frame.header.postbox_header.get(line_index).map(String::as_str)
};
if let Some(line) = line_to_copy {
match CString::new(line) {
Ok(c_string) => c_string.into_raw(), Err(_) => ptr::null_mut(), }
} else {
ptr::null_mut() }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_free_string(s: *mut c_char) {
if !s.is_null() {
let _ = unsafe { CString::from_raw(s) };
}
}
type RkrWriter = ConFrameWriter<Box<dyn std::io::Write>>;
#[inline]
fn into_rkr_writer(
sink: Box<dyn std::io::Write>,
precision: Option<u8>,
) -> *mut RKRConFrameWriter {
let writer: RkrWriter = match precision {
Some(p) => ConFrameWriter::with_precision(sink, p as usize),
None => ConFrameWriter::new(sink),
};
Box::into_raw(Box::new(writer)) as *mut RKRConFrameWriter
}
#[inline]
unsafe fn cstr_path<'a>(filename_c: *const c_char) -> Option<&'a str> {
if filename_c.is_null() {
return None;
}
unsafe { CStr::from_ptr(filename_c).to_str().ok() }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn create_writer_from_path_c(
filename_c: *const c_char,
) -> *mut RKRConFrameWriter {
let filename = match unsafe { cstr_path(filename_c) } {
Some(s) => s,
None => return ptr::null_mut(),
};
match File::create(filename) {
Ok(file) => into_rkr_writer(Box::new(file), None),
Err(_) => ptr::null_mut(),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn free_rkr_writer(writer_handle: *mut RKRConFrameWriter) {
if !writer_handle.is_null() {
let _ = unsafe { Box::from_raw(writer_handle as *mut RkrWriter) };
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_writer_extend(
writer_handle: *mut RKRConFrameWriter,
frame_handles: *const *const RKRConFrame,
num_frames: usize,
) -> RKRStatus {
let writer = match unsafe { (writer_handle as *mut RkrWriter).as_mut() } {
Some(w) => w,
None => return RKRStatus::RKR_STATUS_NULL_POINTER,
};
if frame_handles.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let handles_slice = unsafe { std::slice::from_raw_parts(frame_handles, num_frames) };
let mut rust_frames: Vec<&ConFrame> = Vec::with_capacity(num_frames);
if handles_slice.iter().any(|&handle| handle.is_null()) {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
for &handle in handles_slice.iter() {
match unsafe { (handle as *const ConFrame).as_ref() } {
Some(frame) => rust_frames.push(frame),
None => return RKRStatus::RKR_STATUS_NULL_POINTER,
}
}
match writer.extend(rust_frames.into_iter()) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(_) => RKRStatus::RKR_STATUS_IO_ERROR,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_writer_set_canonical(
writer_handle: *mut RKRConFrameWriter,
canonical: u8,
) -> RKRStatus {
let writer = match unsafe { (writer_handle as *mut RkrWriter).as_mut() } {
Some(w) => w,
None => return RKRStatus::RKR_STATUS_NULL_POINTER,
};
writer.set_canonical(canonical != 0);
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_writer_is_canonical(writer_handle: *const RKRConFrameWriter) -> u8 {
match unsafe { (writer_handle as *const RkrWriter).as_ref() } {
Some(w) => u8::from(w.is_canonical()),
None => 0,
}
}
#[cfg(test)]
mod index_proj_ffi_tests {
use super::*;
use std::ffi::CStr;
use std::fs;
use std::path::PathBuf;
fn fixture_path() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("resources/test/tiny_cuh2.con")
}
#[test]
fn ffi_projection_matches_index_proj() {
let frames = crate::iterators::read_all_frames(&fixture_path()).unwrap();
let fr = &frames[0];
let handle = fr as *const ConFrame as *const RKRConFrame;
let proj = crate::index_proj::FrameIndexProjection::from_frame(fr);
assert_eq!(rkr_frame_index_natoms(handle), proj.n_atoms);
assert_eq!(rkr_frame_sections_mask(handle), proj.sections_mask);
let formula_c = unsafe { rkr_frame_composition_formula(handle) };
assert!(!formula_c.is_null());
let formula = unsafe { CStr::from_ptr(formula_c) }.to_str().unwrap();
assert_eq!(formula, proj.formula);
unsafe { rkr_free_string(formula_c) };
let ie = rkr_frame_index_energy(handle);
match proj.energy {
Some(e) => assert!((ie - e).abs() < 1e-12 || (ie.is_nan() && e.is_nan())),
None => assert!(ie.is_nan()),
}
let tm = rkr_frame_total_mass(handle);
match proj.total_mass {
Some(m) => assert!((tm - m).abs() < 1e-9, "total_mass C={tm} proj={m}"),
None => assert!(tm.is_nan()),
}
let cv = rkr_frame_cell_volume(handle);
match proj.cell_volume {
Some(v) => assert!((cv - v).abs() < 1e-6 * v.max(1.0), "cell_volume C={cv} proj={v}"),
None => assert!(cv.is_nan()),
}
let fm = rkr_frame_fmax(handle);
match proj.fmax {
Some(f) => assert!((fm - f).abs() < 1e-12 || (fm.is_nan() && f.is_nan())),
None => assert!(fm.is_nan()),
}
let json_c = unsafe { rkr_frame_index_projection_json(handle) };
assert!(!json_c.is_null());
let json = unsafe { CStr::from_ptr(json_c) }.to_str().unwrap();
assert!(json.contains("\"formula\""));
assert!(json.contains(&proj.formula) || proj.formula.is_empty());
assert!(json.contains("\"n_atoms\""));
assert!(json.contains("\"total_mass\"") || json.contains("\"cell_volume\""));
unsafe { rkr_free_string(json_c) };
}
#[test]
fn ffi_canonical_writer_byte_identical() {
let frames = crate::iterators::read_all_frames(&fixture_path()).unwrap();
let fr = &frames[0];
let dir = tempfile::tempdir().unwrap();
let p1 = dir.path().join("a.con");
let p2 = dir.path().join("b.con");
for p in [&p1, &p2] {
let path_c = std::ffi::CString::new(p.to_str().unwrap()).unwrap();
let w = unsafe { create_writer_from_path_c(path_c.as_ptr()) };
assert!(!w.is_null());
assert_eq!(
unsafe { rkr_writer_set_canonical(w, 1) },
RKRStatus::RKR_STATUS_SUCCESS
);
assert_eq!(unsafe { rkr_writer_is_canonical(w) }, 1);
let handles = [fr as *const ConFrame as *const RKRConFrame];
assert_eq!(
unsafe { rkr_writer_extend(w, handles.as_ptr(), 1) },
RKRStatus::RKR_STATUS_SUCCESS
);
unsafe { free_rkr_writer(w) };
}
let b1 = fs::read(&p1).unwrap();
let b2 = fs::read(&p2).unwrap();
assert_eq!(b1, b2);
assert!(!b1.is_empty());
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn create_writer_from_path_with_precision_c(
filename_c: *const c_char,
precision: u8,
) -> *mut RKRConFrameWriter {
let filename = match unsafe { cstr_path(filename_c) } {
Some(s) => s,
None => return ptr::null_mut(),
};
match File::create(filename) {
Ok(file) => into_rkr_writer(Box::new(file), Some(precision)),
Err(_) => ptr::null_mut(),
}
}
#[repr(C)]
pub struct RKRConFrameBuilder {
_private: [u8; 0],
}
#[allow(clippy::too_many_arguments)]
unsafe fn add_builder_atom(
builder_handle: *mut RKRConFrameBuilder,
symbol: *const c_char,
x: f64,
y: f64,
z: f64,
fixed: [bool; 3],
atom_id: u64,
mass: f64,
velocity: Option<[f64; 3]>,
forces: Option<[f64; 3]>,
) -> RKRStatus {
if builder_handle.is_null() || symbol.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
let sym = match unsafe { CStr::from_ptr(symbol).to_str() } {
Ok(s) => s,
Err(_) => return RKRStatus::RKR_STATUS_INVALID_UTF8,
};
builder.add_atom(sym, x, y, z, fixed, atom_id, mass);
if let Some(v) = velocity {
builder.with_velocity(v);
}
if let Some(f) = forces {
builder.with_force(f);
}
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_last_velocity(
builder_handle: *mut RKRConFrameBuilder,
velocity: *const f64,
) -> RKRStatus {
if builder_handle.is_null() || velocity.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
let v = unsafe { [*velocity, *velocity.add(1), *velocity.add(2)] };
builder.with_velocity(v);
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_last_force(
builder_handle: *mut RKRConFrameBuilder,
force: *const f64,
) -> RKRStatus {
if builder_handle.is_null() || force.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
let f = unsafe { [*force, *force.add(1), *force.add(2)] };
builder.with_force(f);
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_last_energy(
builder_handle: *mut RKRConFrameBuilder,
energy: f64,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
builder.with_energy(energy);
RKRStatus::RKR_STATUS_SUCCESS
}
fn map_builder_err(e: crate::error::ParseError) -> RKRStatus {
use crate::error::ParseError;
match e {
ParseError::IndexOutOfBounds { .. } | ParseError::InvalidVectorLength { .. } => {
RKRStatus::RKR_STATUS_INDEX_OUT_OF_BOUNDS
}
_ => RKRStatus::RKR_STATUS_INTERNAL_ERROR,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_atom_count(
builder_handle: *const RKRConFrameBuilder,
) -> usize {
if builder_handle.is_null() {
return 0;
}
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
builder.atom_count()
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_atom_position(
builder_handle: *mut RKRConFrameBuilder,
index: usize,
x: f64,
y: f64,
z: f64,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
match builder.set_atom_position(index, x, y, z) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_atom_velocity(
builder_handle: *mut RKRConFrameBuilder,
index: usize,
velocity: *const f64,
) -> RKRStatus {
if builder_handle.is_null() || velocity.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
let v = unsafe { [*velocity, *velocity.add(1), *velocity.add(2)] };
match builder.set_atom_velocity(index, v) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_atom_force(
builder_handle: *mut RKRConFrameBuilder,
index: usize,
force: *const f64,
) -> RKRStatus {
if builder_handle.is_null() || force.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
let f = unsafe { [*force, *force.add(1), *force.add(2)] };
match builder.set_atom_force(index, f) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_atom_energy(
builder_handle: *mut RKRConFrameBuilder,
index: usize,
energy: f64,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
match builder.set_atom_energy(index, energy) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_atom_fixed(
builder_handle: *mut RKRConFrameBuilder,
index: usize,
fixed_x: bool,
fixed_y: bool,
fixed_z: bool,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
match builder.set_atom_fixed(index, [fixed_x, fixed_y, fixed_z]) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_atom_mass(
builder_handle: *mut RKRConFrameBuilder,
index: usize,
mass: f64,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
match builder.set_atom_mass(index, mass) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_atom_id(
builder_handle: *mut RKRConFrameBuilder,
index: usize,
atom_id: u64,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
match builder.set_atom_id(index, atom_id) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_clear_atom_velocity(
builder_handle: *mut RKRConFrameBuilder,
index: usize,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
match builder.clear_atom_velocity(index) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_clear_atom_force(
builder_handle: *mut RKRConFrameBuilder,
index: usize,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
match builder.clear_atom_force(index) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_clear_atom_energy(
builder_handle: *mut RKRConFrameBuilder,
index: usize,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
match builder.clear_atom_energy(index) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_positions_from_flat(
builder_handle: *mut RKRConFrameBuilder,
positions: *const f64,
len: usize,
) -> RKRStatus {
if builder_handle.is_null() || positions.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
let slice = unsafe { std::slice::from_raw_parts(positions, len) };
match builder.set_positions_from_flat(slice) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_forces_from_flat(
builder_handle: *mut RKRConFrameBuilder,
forces: *const f64,
len: usize,
) -> RKRStatus {
if builder_handle.is_null() || forces.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
let slice = unsafe { std::slice::from_raw_parts(forces, len) };
match builder.set_forces_from_flat(slice) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_atom_energies_from_flat(
builder_handle: *mut RKRConFrameBuilder,
energies: *const f64,
len: usize,
) -> RKRStatus {
if builder_handle.is_null() || energies.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
let slice = unsafe { std::slice::from_raw_parts(energies, len) };
match builder.set_atom_energies_from_flat(slice) {
Ok(_) => RKRStatus::RKR_STATUS_SUCCESS,
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_get_atom_position(
builder_handle: *const RKRConFrameBuilder,
index: usize,
out_xyz: *mut f64,
) -> RKRStatus {
if builder_handle.is_null() || out_xyz.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
match builder.get_atom_position(index) {
Ok((x, y, z)) => unsafe {
*out_xyz = x;
*out_xyz.add(1) = y;
*out_xyz.add(2) = z;
RKRStatus::RKR_STATUS_SUCCESS
},
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_get_atom_velocity(
builder_handle: *const RKRConFrameBuilder,
index: usize,
out_xyz: *mut f64,
has_value: *mut bool,
) -> RKRStatus {
if builder_handle.is_null() || out_xyz.is_null() || has_value.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
match builder.get_atom_velocity(index) {
Ok(Some(v)) => unsafe {
*out_xyz = v[0];
*out_xyz.add(1) = v[1];
*out_xyz.add(2) = v[2];
*has_value = true;
RKRStatus::RKR_STATUS_SUCCESS
},
Ok(None) => unsafe {
*has_value = false;
RKRStatus::RKR_STATUS_SUCCESS
},
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_get_atom_force(
builder_handle: *const RKRConFrameBuilder,
index: usize,
out_xyz: *mut f64,
has_value: *mut bool,
) -> RKRStatus {
if builder_handle.is_null() || out_xyz.is_null() || has_value.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
match builder.get_atom_force(index) {
Ok(Some(f)) => unsafe {
*out_xyz = f[0];
*out_xyz.add(1) = f[1];
*out_xyz.add(2) = f[2];
*has_value = true;
RKRStatus::RKR_STATUS_SUCCESS
},
Ok(None) => unsafe {
*has_value = false;
RKRStatus::RKR_STATUS_SUCCESS
},
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_get_atom_energy(
builder_handle: *const RKRConFrameBuilder,
index: usize,
out_value: *mut f64,
has_value: *mut bool,
) -> RKRStatus {
if builder_handle.is_null() || out_value.is_null() || has_value.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
match builder.get_atom_energy(index) {
Ok(Some(e)) => unsafe {
*out_value = e;
*has_value = true;
RKRStatus::RKR_STATUS_SUCCESS
},
Ok(None) => unsafe {
*has_value = false;
RKRStatus::RKR_STATUS_SUCCESS
},
Err(e) => map_builder_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_get_atom_mass(
builder_handle: *const RKRConFrameBuilder,
index: usize,
out_mass: *mut f64,
) -> RKRStatus {
if builder_handle.is_null() || out_mass.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
match builder.get_atom_mass(index) {
Ok(m) => unsafe {
*out_mass = m;
RKRStatus::RKR_STATUS_SUCCESS
},
Err(e) => map_builder_err(e),
}
}
pub use dlpk::sys::DLManagedTensorVersioned as RKRDLManagedTensorVersioned;
fn map_dlpack_err(e: crate::error::ParseError) -> RKRStatus {
use crate::error::ParseError;
match e {
ParseError::ValidationError(_) => RKRStatus::RKR_STATUS_VALIDATION_ERROR,
_ => RKRStatus::RKR_STATUS_INTERNAL_ERROR,
}
}
pub mod rkr_dl_type_code {
pub const RKR_DL_INT: u8 = 0;
pub const RKR_DL_UINT: u8 = 1;
pub const RKR_DL_FLOAT: u8 = 2;
pub const RKR_DL_OPAQUE_HANDLE: u8 = 3;
pub const RKR_DL_BFLOAT: u8 = 4;
pub const RKR_DL_COMPLEX: u8 = 5;
pub const RKR_DL_BOOL: u8 = 6;
}
pub mod rkr_dl_device_type {
pub const RKR_DL_CPU: i32 = 1;
pub const RKR_DL_CUDA: i32 = 2;
pub const RKR_DL_CUDA_HOST: i32 = 3;
}
#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct RKRDLDataType {
pub code: u8,
pub bits: u8,
pub lanes: u16,
}
#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct RKRDLDevice {
pub device_type: i32,
pub device_id: i32,
}
#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct RKRDlpackExportOptions {
pub dtype: RKRDLDataType,
pub device: RKRDLDevice,
}
impl Default for RKRDlpackExportOptions {
fn default() -> Self {
Self {
dtype: RKRDLDataType {
code: rkr_dl_type_code::RKR_DL_FLOAT,
bits: 64,
lanes: 1,
},
device: RKRDLDevice {
device_type: rkr_dl_device_type::RKR_DL_CPU,
device_id: 0,
},
}
}
}
fn resolve_dlpack_opts(opts: *const RKRDlpackExportOptions) -> Result<RKRDlpackExportOptions, RKRStatus> {
let o = if opts.is_null() {
RKRDlpackExportOptions::default()
} else {
unsafe { *opts }
};
if o.device.device_type != rkr_dl_device_type::RKR_DL_CPU {
return Err(RKRStatus::RKR_STATUS_FEATURE_DISABLED);
}
if o.dtype.lanes != 1 {
return Err(RKRStatus::RKR_STATUS_VALIDATION_ERROR);
}
Ok(o)
}
fn finish_dlpack_tensor<E: std::fmt::Display>(
result: Result<dlpk::DLPackTensor, E>,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
match result {
Ok(tensor) => {
let raw = tensor.into_raw();
unsafe {
*out_tensor = raw.as_ptr();
}
RKRStatus::RKR_STATUS_SUCCESS
}
Err(e) => map_dlpack_err(crate::error::ParseError::ValidationError(format!(
"DLPack export failed: {e}"
))),
}
}
fn export_f64_slice_as_dlpack(
data: &[f64],
shape: &[usize],
dtype: RKRDLDataType,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
use rkr_dl_type_code::*;
if dtype.lanes != 1 {
return RKRStatus::RKR_STATUS_VALIDATION_ERROR;
}
macro_rules! arc_export {
($ty:ty, $map:expr) => {{
let v: Vec<$ty> = data.iter().map($map).collect();
match *shape {
[r, c] => match ndarray::ArcArray2::from_shape_vec((r, c), v) {
Ok(a) => dlpk::DLPackTensor::try_from(a),
Err(_) => return RKRStatus::RKR_STATUS_VALIDATION_ERROR,
},
[_] => dlpk::DLPackTensor::try_from(ndarray::ArcArray1::from_vec(v)),
_ => return RKRStatus::RKR_STATUS_VALIDATION_ERROR,
}
}};
}
let tensor = match (dtype.code, dtype.bits) {
(RKR_DL_FLOAT, 64) => arc_export!(f64, |&x| x),
(RKR_DL_FLOAT, 32) => arc_export!(f32, |&x| x as f32),
(RKR_DL_INT, 8) => arc_export!(i8, |&x| x as i8),
(RKR_DL_INT, 16) => arc_export!(i16, |&x| x as i16),
(RKR_DL_INT, 32) => arc_export!(i32, |&x| x as i32),
(RKR_DL_INT, 64) => arc_export!(i64, |&x| x as i64),
(RKR_DL_UINT, 8) => arc_export!(u8, |&x| x as u8),
(RKR_DL_UINT, 16) => arc_export!(u16, |&x| x as u16),
(RKR_DL_UINT, 32) => arc_export!(u32, |&x| x as u32),
(RKR_DL_UINT, 64) => arc_export!(u64, |&x| x as u64),
(RKR_DL_BOOL, 8) => {
let v: Vec<bool> = data.iter().map(|&x| x != 0.0).collect();
return finish_dlpack_tensor(dlpk::DLPackTensor::try_from(v), out_tensor);
}
_ => return RKRStatus::RKR_STATUS_VALIDATION_ERROR,
};
finish_dlpack_tensor(tensor, out_tensor)
}
fn export_owned_array2_dlpack_opts(
arr: &ndarray::ArcArray2<f64>,
opts: &RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
let (r, c) = arr.dim();
let flat: Vec<f64> = arr.iter().copied().collect();
export_f64_slice_as_dlpack(&flat, &[r, c], opts.dtype, out_tensor)
}
fn export_owned_array1_f64_dlpack_opts(
arr: &ndarray::ArcArray1<f64>,
opts: &RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
let flat = arr.to_vec();
let n = flat.len();
export_f64_slice_as_dlpack(&flat, &[n], opts.dtype, out_tensor)
}
fn export_owned_array1_u64_dlpack(
arr: &ndarray::ArcArray1<u64>,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
finish_dlpack_tensor(
dlpk::DLPackTensor::try_from(arr.clone()),
out_tensor,
)
}
fn export_owned_array1_u64_dlpack_opts(
arr: &ndarray::ArcArray1<u64>,
opts: &RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
use rkr_dl_type_code::*;
if opts.dtype.code == RKR_DL_UINT && opts.dtype.bits == 64 && opts.dtype.lanes == 1 {
return export_owned_array1_u64_dlpack(arr, out_tensor);
}
let as_f64: Vec<f64> = arr.iter().map(|&x| x as f64).collect();
let n = as_f64.len();
export_f64_slice_as_dlpack(&as_f64, &[n], opts.dtype, out_tensor)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_positions_dlpack(
builder_handle: *const RKRConFrameBuilder,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
unsafe { rkr_frame_builder_positions_dlpack_ex(builder_handle, std::ptr::null(), out_tensor) }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_positions_dlpack_ex(
builder_handle: *const RKRConFrameBuilder,
opts: *const RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
if builder_handle.is_null() || out_tensor.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let o = match resolve_dlpack_opts(opts) {
Ok(o) => o,
Err(st) => return st,
};
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
export_owned_array2_dlpack_opts(builder.positions_2d_ref(), &o, out_tensor)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_velocities_dlpack(
builder_handle: *const RKRConFrameBuilder,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
unsafe { rkr_frame_builder_velocities_dlpack_ex(builder_handle, std::ptr::null(), out_tensor) }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_velocities_dlpack_ex(
builder_handle: *const RKRConFrameBuilder,
opts: *const RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
if builder_handle.is_null() || out_tensor.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let o = match resolve_dlpack_opts(opts) {
Ok(o) => o,
Err(st) => return st,
};
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
if !builder.has_velocities_section() {
return RKRStatus::RKR_STATUS_SECTION_ABSENT;
}
export_owned_array2_dlpack_opts(builder.velocities_2d_ref(), &o, out_tensor)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_forces_dlpack(
builder_handle: *const RKRConFrameBuilder,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
unsafe { rkr_frame_builder_forces_dlpack_ex(builder_handle, std::ptr::null(), out_tensor) }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_forces_dlpack_ex(
builder_handle: *const RKRConFrameBuilder,
opts: *const RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
if builder_handle.is_null() || out_tensor.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let o = match resolve_dlpack_opts(opts) {
Ok(o) => o,
Err(st) => return st,
};
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
if !builder.has_forces_section() {
return RKRStatus::RKR_STATUS_SECTION_ABSENT;
}
export_owned_array2_dlpack_opts(builder.forces_2d_ref(), &o, out_tensor)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_atom_energies_dlpack(
builder_handle: *const RKRConFrameBuilder,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
unsafe {
rkr_frame_builder_atom_energies_dlpack_ex(builder_handle, std::ptr::null(), out_tensor)
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_atom_energies_dlpack_ex(
builder_handle: *const RKRConFrameBuilder,
opts: *const RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
if builder_handle.is_null() || out_tensor.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let o = match resolve_dlpack_opts(opts) {
Ok(o) => o,
Err(st) => return st,
};
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
if !builder.has_energies_section() {
return RKRStatus::RKR_STATUS_SECTION_ABSENT;
}
export_owned_array1_f64_dlpack_opts(builder.atom_energies_1d_ref(), &o, out_tensor)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_masses_dlpack(
builder_handle: *const RKRConFrameBuilder,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
unsafe { rkr_frame_builder_masses_dlpack_ex(builder_handle, std::ptr::null(), out_tensor) }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_masses_dlpack_ex(
builder_handle: *const RKRConFrameBuilder,
opts: *const RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
if builder_handle.is_null() || out_tensor.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let o = match resolve_dlpack_opts(opts) {
Ok(o) => o,
Err(st) => return st,
};
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
export_owned_array1_f64_dlpack_opts(builder.masses_1d_ref(), &o, out_tensor)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_atom_ids_dlpack(
builder_handle: *const RKRConFrameBuilder,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
unsafe { rkr_frame_builder_atom_ids_dlpack_ex(builder_handle, std::ptr::null(), out_tensor) }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_atom_ids_dlpack_ex(
builder_handle: *const RKRConFrameBuilder,
opts: *const RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
if builder_handle.is_null() || out_tensor.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let o = match resolve_dlpack_opts(opts) {
Ok(o) => o,
Err(st) => return st,
};
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
export_owned_array1_u64_dlpack_opts(builder.atom_ids_1d_ref(), &o, out_tensor)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_positions_data(
builder_handle: *mut RKRConFrameBuilder,
) -> *mut f64 {
if builder_handle.is_null() {
return std::ptr::null_mut();
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
builder
.positions_view_mut()
.as_slice_memory_order_mut()
.map(|s| s.as_mut_ptr())
.unwrap_or(std::ptr::null_mut())
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_velocities_data(
builder_handle: *mut RKRConFrameBuilder,
) -> *mut f64 {
if builder_handle.is_null() {
return std::ptr::null_mut();
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
if !builder.has_velocities_section() {
return std::ptr::null_mut();
}
let slice = builder.velocities_mut();
if slice.is_empty() {
std::ptr::null_mut()
} else {
slice.as_mut_ptr()
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_forces_data(
builder_handle: *mut RKRConFrameBuilder,
) -> *mut f64 {
if builder_handle.is_null() {
return std::ptr::null_mut();
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
if !builder.has_forces_section() {
return std::ptr::null_mut();
}
let slice = builder.forces_mut();
if slice.is_empty() {
std::ptr::null_mut()
} else {
slice.as_mut_ptr()
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_atom_energies_data(
builder_handle: *mut RKRConFrameBuilder,
) -> *mut f64 {
if builder_handle.is_null() {
return std::ptr::null_mut();
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
if !builder.has_energies_section() {
return std::ptr::null_mut();
}
let slice = builder.atom_energies_mut();
if slice.is_empty() {
std::ptr::null_mut()
} else {
slice.as_mut_ptr()
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_masses_data(
builder_handle: *mut RKRConFrameBuilder,
) -> *mut f64 {
if builder_handle.is_null() {
return std::ptr::null_mut();
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
let slice = builder.masses_mut();
if slice.is_empty() {
std::ptr::null_mut()
} else {
slice.as_mut_ptr()
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_atom_ids_data(
builder_handle: *const RKRConFrameBuilder,
) -> *const u64 {
if builder_handle.is_null() {
return std::ptr::null();
}
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
let slice = builder.atom_ids();
if slice.is_empty() {
std::ptr::null()
} else {
slice.as_ptr()
}
}
#[unsafe(no_mangle)]
#[allow(clippy::too_many_arguments)]
pub unsafe extern "C" fn rkr_frame_add_atom_full(
builder_handle: *mut RKRConFrameBuilder,
symbol: *const c_char,
x: f64,
y: f64,
z: f64,
fixed_x: bool,
fixed_y: bool,
fixed_z: bool,
atom_id: u64,
mass: f64,
velocity: *const f64,
force: *const f64,
) -> RKRStatus {
let velocity = if velocity.is_null() {
None
} else {
Some(unsafe { [*velocity, *velocity.add(1), *velocity.add(2)] })
};
let force = if force.is_null() {
None
} else {
Some(unsafe { [*force, *force.add(1), *force.add(2)] })
};
unsafe {
add_builder_atom(
builder_handle,
symbol,
x,
y,
z,
[fixed_x, fixed_y, fixed_z],
atom_id,
mass,
velocity,
force,
)
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_new(
cell: *const f64,
angles: *const f64,
prebox0: *const c_char,
prebox1: *const c_char,
postbox0: *const c_char,
postbox1: *const c_char,
) -> *mut RKRConFrameBuilder {
if cell.is_null() || angles.is_null() {
return ptr::null_mut();
}
let cell_arr = unsafe { [*cell, *cell.add(1), *cell.add(2)] };
let angles_arr = unsafe { [*angles, *angles.add(1), *angles.add(2)] };
let get_str = |p: *const c_char| -> String {
if p.is_null() {
String::new()
} else {
unsafe { CStr::from_ptr(p) }
.to_str()
.unwrap_or("")
.to_string()
}
};
let _ = get_str(prebox1);
let mut builder = ConFrameBuilder::new(cell_arr, angles_arr);
builder
.prebox_header(get_str(prebox0))
.postbox_header([get_str(postbox0), get_str(postbox1)]);
Box::into_raw(Box::new(builder)) as *mut RKRConFrameBuilder
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_metadata_json(
builder_handle: *mut RKRConFrameBuilder,
metadata_json: *const c_char,
) -> RKRStatus {
if builder_handle.is_null() || metadata_json.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
let metadata_json = match unsafe { CStr::from_ptr(metadata_json).to_str() } {
Ok(s) => s,
Err(_) => return RKRStatus::RKR_STATUS_INVALID_UTF8,
};
match builder.set_metadata_json(metadata_json) {
Ok(()) => RKRStatus::RKR_STATUS_SUCCESS,
Err(_) => RKRStatus::RKR_STATUS_INVALID_JSON,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_scalar_metadata(
builder_handle: *mut RKRConFrameBuilder,
key: *const c_char,
value: f64,
) -> RKRStatus {
if builder_handle.is_null() || key.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
let key = match unsafe { CStr::from_ptr(key).to_str() } {
Ok(s) => s,
Err(_) => return RKRStatus::RKR_STATUS_INVALID_UTF8,
};
builder.set_scalar_metadata(key, value);
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_string_metadata(
builder_handle: *mut RKRConFrameBuilder,
key: *const c_char,
value: *const c_char,
) -> RKRStatus {
if builder_handle.is_null() || key.is_null() || value.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
let key = match unsafe { CStr::from_ptr(key).to_str() } {
Ok(s) => s,
Err(_) => return RKRStatus::RKR_STATUS_INVALID_UTF8,
};
let value = match unsafe { CStr::from_ptr(value).to_str() } {
Ok(s) => s,
Err(_) => return RKRStatus::RKR_STATUS_INVALID_UTF8,
};
builder.set_string_metadata(key, value);
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_energy(
builder_handle: *mut RKRConFrameBuilder,
energy: f64,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
builder.set_energy(energy);
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_frame_index(
builder_handle: *mut RKRConFrameBuilder,
idx: u64,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
builder.set_frame_index(idx);
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_time(
builder_handle: *mut RKRConFrameBuilder,
time: f64,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
builder.set_time(time);
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_timestep(
builder_handle: *mut RKRConFrameBuilder,
dt: f64,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
builder.set_timestep(dt);
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_neb_bead(
builder_handle: *mut RKRConFrameBuilder,
bead: u64,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
builder.set_neb_bead(bead);
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_set_neb_band(
builder_handle: *mut RKRConFrameBuilder,
band: u64,
) -> RKRStatus {
if builder_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let builder = unsafe { &mut *(builder_handle as *mut ConFrameBuilder) };
builder.set_neb_band(band);
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_add_atom(
builder_handle: *mut RKRConFrameBuilder,
symbol: *const c_char,
x: f64,
y: f64,
z: f64,
is_fixed: bool,
atom_id: u64,
mass: f64,
) -> RKRStatus {
unsafe {
add_builder_atom(
builder_handle,
symbol,
x,
y,
z,
[is_fixed; 3],
atom_id,
mass,
None,
None,
)
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_add_atom_with_fixed_mask(
builder_handle: *mut RKRConFrameBuilder,
symbol: *const c_char,
x: f64,
y: f64,
z: f64,
fixed_x: bool,
fixed_y: bool,
fixed_z: bool,
atom_id: u64,
mass: f64,
) -> RKRStatus {
unsafe {
add_builder_atom(
builder_handle,
symbol,
x,
y,
z,
[fixed_x, fixed_y, fixed_z],
atom_id,
mass,
None,
None,
)
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_add_atom_with_velocity(
builder_handle: *mut RKRConFrameBuilder,
symbol: *const c_char,
x: f64,
y: f64,
z: f64,
is_fixed: bool,
atom_id: u64,
mass: f64,
vx: f64,
vy: f64,
vz: f64,
) -> RKRStatus {
unsafe {
add_builder_atom(
builder_handle,
symbol,
x,
y,
z,
[is_fixed; 3],
atom_id,
mass,
Some([vx, vy, vz]),
None,
)
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_add_atom_with_velocity_fixed_mask(
builder_handle: *mut RKRConFrameBuilder,
symbol: *const c_char,
x: f64,
y: f64,
z: f64,
fixed_x: bool,
fixed_y: bool,
fixed_z: bool,
atom_id: u64,
mass: f64,
vx: f64,
vy: f64,
vz: f64,
) -> RKRStatus {
unsafe {
add_builder_atom(
builder_handle,
symbol,
x,
y,
z,
[fixed_x, fixed_y, fixed_z],
atom_id,
mass,
Some([vx, vy, vz]),
None,
)
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_add_atom_with_forces(
builder_handle: *mut RKRConFrameBuilder,
symbol: *const c_char,
x: f64,
y: f64,
z: f64,
is_fixed: bool,
atom_id: u64,
mass: f64,
fx: f64,
fy: f64,
fz: f64,
) -> RKRStatus {
unsafe {
add_builder_atom(
builder_handle,
symbol,
x,
y,
z,
[is_fixed; 3],
atom_id,
mass,
None,
Some([fx, fy, fz]),
)
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_add_atom_with_forces_fixed_mask(
builder_handle: *mut RKRConFrameBuilder,
symbol: *const c_char,
x: f64,
y: f64,
z: f64,
fixed_x: bool,
fixed_y: bool,
fixed_z: bool,
atom_id: u64,
mass: f64,
fx: f64,
fy: f64,
fz: f64,
) -> RKRStatus {
unsafe {
add_builder_atom(
builder_handle,
symbol,
x,
y,
z,
[fixed_x, fixed_y, fixed_z],
atom_id,
mass,
None,
Some([fx, fy, fz]),
)
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_add_atom_with_velocity_and_forces(
builder_handle: *mut RKRConFrameBuilder,
symbol: *const c_char,
x: f64,
y: f64,
z: f64,
is_fixed: bool,
atom_id: u64,
mass: f64,
vx: f64,
vy: f64,
vz: f64,
fx: f64,
fy: f64,
fz: f64,
) -> RKRStatus {
unsafe {
add_builder_atom(
builder_handle,
symbol,
x,
y,
z,
[is_fixed; 3],
atom_id,
mass,
Some([vx, vy, vz]),
Some([fx, fy, fz]),
)
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_add_atom_with_velocity_and_forces_fixed_mask(
builder_handle: *mut RKRConFrameBuilder,
symbol: *const c_char,
x: f64,
y: f64,
z: f64,
fixed_x: bool,
fixed_y: bool,
fixed_z: bool,
atom_id: u64,
mass: f64,
vx: f64,
vy: f64,
vz: f64,
fx: f64,
fy: f64,
fz: f64,
) -> RKRStatus {
unsafe {
add_builder_atom(
builder_handle,
symbol,
x,
y,
z,
[fixed_x, fixed_y, fixed_z],
atom_id,
mass,
Some([vx, vy, vz]),
Some([fx, fy, fz]),
)
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_build(
builder_handle: *mut RKRConFrameBuilder,
) -> *mut RKRConFrame {
if builder_handle.is_null() {
return ptr::null_mut();
}
let builder = unsafe { *Box::from_raw(builder_handle as *mut ConFrameBuilder) };
let frame = builder.build();
Box::into_raw(Box::new(frame)) as *mut RKRConFrame
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn free_rkr_frame_builder(builder_handle: *mut RKRConFrameBuilder) {
if !builder_handle.is_null() {
let _ = unsafe { Box::from_raw(builder_handle as *mut ConFrameBuilder) };
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_builder_clone(
builder_handle: *const RKRConFrameBuilder,
) -> *mut RKRConFrameBuilder {
if builder_handle.is_null() {
return std::ptr::null_mut();
}
let builder = unsafe { &*(builder_handle as *const ConFrameBuilder) };
let cloned = builder.clone();
Box::into_raw(Box::new(cloned)) as *mut RKRConFrameBuilder
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn create_writer_gzip_c(filename_c: *const c_char) -> *mut RKRConFrameWriter {
let filename = match unsafe { cstr_path(filename_c) } {
Some(s) => s,
None => return ptr::null_mut(),
};
match crate::compression::gzip_writer(Path::new(filename)) {
Ok(encoder) => into_rkr_writer(Box::new(encoder), None),
Err(_) => ptr::null_mut(),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn create_writer_gzip_with_precision_c(
filename_c: *const c_char,
precision: u8,
) -> *mut RKRConFrameWriter {
let filename = match unsafe { cstr_path(filename_c) } {
Some(s) => s,
None => return ptr::null_mut(),
};
match crate::compression::gzip_writer(Path::new(filename)) {
Ok(encoder) => into_rkr_writer(Box::new(encoder), Some(precision)),
Err(_) => ptr::null_mut(),
}
}
#[cfg(feature = "zstd")]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn create_writer_zstd_c(filename_c: *const c_char) -> *mut RKRConFrameWriter {
let filename = match unsafe { cstr_path(filename_c) } {
Some(s) => s,
None => return ptr::null_mut(),
};
match crate::compression::zstd_writer(Path::new(filename)) {
Ok(encoder) => into_rkr_writer(Box::new(encoder), None),
Err(_) => ptr::null_mut(),
}
}
#[cfg(feature = "zstd")]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn create_writer_zstd_with_precision_c(
filename_c: *const c_char,
precision: u8,
) -> *mut RKRConFrameWriter {
let filename = match unsafe { cstr_path(filename_c) } {
Some(s) => s,
None => return ptr::null_mut(),
};
match crate::compression::zstd_writer(Path::new(filename)) {
Ok(encoder) => into_rkr_writer(Box::new(encoder), Some(precision)),
Err(_) => ptr::null_mut(),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_read_first_frame(filename_c: *const c_char) -> *mut RKRConFrame {
if filename_c.is_null() {
return ptr::null_mut();
}
let filename = match unsafe { CStr::from_ptr(filename_c).to_str() } {
Ok(s) => s,
Err(_) => return ptr::null_mut(),
};
match iterators::read_first_frame(Path::new(filename)) {
Ok(frame) => Box::into_raw(Box::new(frame)) as *mut RKRConFrame,
Err(_) => ptr::null_mut(),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_read_all_frames(
filename_c: *const c_char,
num_frames: *mut usize,
) -> *mut *mut RKRConFrame {
if filename_c.is_null() || num_frames.is_null() {
return ptr::null_mut();
}
let filename = match unsafe { CStr::from_ptr(filename_c).to_str() } {
Ok(s) => s,
Err(_) => return ptr::null_mut(),
};
match iterators::read_all_frames(Path::new(filename)) {
Ok(frames) => {
let count = frames.len();
let mut handles: Vec<*mut RKRConFrame> = frames
.into_iter()
.map(|f| Box::into_raw(Box::new(f)) as *mut RKRConFrame)
.collect();
handles.shrink_to_fit();
debug_assert_eq!(handles.len(), handles.capacity());
let ptr = handles.as_mut_ptr();
std::mem::forget(handles);
unsafe { *num_frames = count };
ptr
}
Err(_) => ptr::null_mut(),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn free_rkr_frame_array(frames: *mut *mut RKRConFrame, num_frames: usize) {
if frames.is_null() {
return;
}
unsafe {
let handles = Vec::from_raw_parts(frames, num_frames, num_frames);
for handle in handles {
if !handle.is_null() {
let _ = Box::from_raw(handle as *mut ConFrame);
}
}
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn free_rkr_frame_ptr_array(frames: *mut *mut RKRConFrame, num_frames: usize) {
if frames.is_null() {
return;
}
unsafe {
let _ptrs = Vec::from_raw_parts(frames, num_frames, num_frames);
}
}
#[cfg(feature = "metatensor")]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_mts_block_free(block: *mut metatensor::c_api::mts_block_t) {
unsafe { crate::metatensor_export::mts_block_free_sys(block) };
}
#[cfg(feature = "metatensor")]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_metatensor_positions_block(
frame_handle: *const RKRConFrame,
out_block: *mut *mut metatensor::c_api::mts_block_t,
) -> RKRStatus {
if frame_handle.is_null() || out_block.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
unsafe { *out_block = std::ptr::null_mut() };
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
crate::metatensor_export::frame_positions_block(frame)
})) {
Ok(Ok(b)) => {
unsafe { *out_block = crate::metatensor_export::tensor_block_into_raw_mts(b) };
RKRStatus::RKR_STATUS_SUCCESS
}
Ok(Err(_)) | Err(_) => RKRStatus::RKR_STATUS_INTERNAL_ERROR,
}
}
#[cfg(feature = "metatensor")]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_metatensor_velocities_block(
frame_handle: *const RKRConFrame,
out_block: *mut *mut metatensor::c_api::mts_block_t,
) -> RKRStatus {
if frame_handle.is_null() || out_block.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
unsafe { *out_block = std::ptr::null_mut() };
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
crate::metatensor_export::frame_velocities_block(frame)
})) {
Ok(Ok(Some(b))) => {
unsafe { *out_block = crate::metatensor_export::tensor_block_into_raw_mts(b) };
RKRStatus::RKR_STATUS_SUCCESS
}
Ok(Ok(None)) => RKRStatus::RKR_STATUS_SECTION_ABSENT,
Ok(Err(_)) | Err(_) => RKRStatus::RKR_STATUS_INTERNAL_ERROR,
}
}
#[cfg(feature = "metatensor")]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_metatensor_forces_block(
frame_handle: *const RKRConFrame,
out_block: *mut *mut metatensor::c_api::mts_block_t,
) -> RKRStatus {
if frame_handle.is_null() || out_block.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
unsafe { *out_block = std::ptr::null_mut() };
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
crate::metatensor_export::frame_forces_block(frame)
})) {
Ok(Ok(Some(b))) => {
unsafe { *out_block = crate::metatensor_export::tensor_block_into_raw_mts(b) };
RKRStatus::RKR_STATUS_SUCCESS
}
Ok(Ok(None)) => RKRStatus::RKR_STATUS_SECTION_ABSENT,
Ok(Err(_)) | Err(_) => RKRStatus::RKR_STATUS_INTERNAL_ERROR,
}
}
#[cfg(feature = "metatensor")]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_metatensor_atom_energies_block(
frame_handle: *const RKRConFrame,
out_block: *mut *mut metatensor::c_api::mts_block_t,
) -> RKRStatus {
if frame_handle.is_null() || out_block.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
unsafe { *out_block = std::ptr::null_mut() };
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
crate::metatensor_export::frame_energies_block(frame)
})) {
Ok(Ok(Some(b))) => {
unsafe { *out_block = crate::metatensor_export::tensor_block_into_raw_mts(b) };
RKRStatus::RKR_STATUS_SUCCESS
}
Ok(Ok(None)) => RKRStatus::RKR_STATUS_SECTION_ABSENT,
Ok(Err(_)) | Err(_) => RKRStatus::RKR_STATUS_INTERNAL_ERROR,
}
}
#[cfg(not(feature = "zstd"))]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn create_writer_zstd_c(_filename_c: *const c_char) -> *mut RKRConFrameWriter {
ptr::null_mut()
}
#[cfg(not(feature = "zstd"))]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn create_writer_zstd_with_precision_c(
_filename_c: *const c_char,
_precision: u8,
) -> *mut RKRConFrameWriter {
ptr::null_mut()
}
#[cfg(not(feature = "metatensor"))]
#[repr(C)]
pub struct mts_block_t {
_private: [u8; 0],
}
#[cfg(not(feature = "metatensor"))]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_mts_block_free(_block: *mut mts_block_t) {}
#[cfg(not(feature = "metatensor"))]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_metatensor_positions_block(
_frame_handle: *const RKRConFrame,
out_block: *mut *mut mts_block_t,
) -> RKRStatus {
if !out_block.is_null() {
unsafe { *out_block = std::ptr::null_mut() };
}
RKRStatus::RKR_STATUS_FEATURE_DISABLED
}
#[cfg(not(feature = "metatensor"))]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_metatensor_velocities_block(
_frame_handle: *const RKRConFrame,
out_block: *mut *mut mts_block_t,
) -> RKRStatus {
if !out_block.is_null() {
unsafe { *out_block = std::ptr::null_mut() };
}
RKRStatus::RKR_STATUS_FEATURE_DISABLED
}
#[cfg(not(feature = "metatensor"))]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_metatensor_forces_block(
_frame_handle: *const RKRConFrame,
out_block: *mut *mut mts_block_t,
) -> RKRStatus {
if !out_block.is_null() {
unsafe { *out_block = std::ptr::null_mut() };
}
RKRStatus::RKR_STATUS_FEATURE_DISABLED
}
#[cfg(not(feature = "metatensor"))]
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_metatensor_atom_energies_block(
_frame_handle: *const RKRConFrame,
out_block: *mut *mut mts_block_t,
) -> RKRStatus {
if !out_block.is_null() {
unsafe { *out_block = std::ptr::null_mut() };
}
RKRStatus::RKR_STATUS_FEATURE_DISABLED
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_atom_count(frame_handle: *const RKRConFrame) -> usize {
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return 0;
};
frame.atom_data.len()
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_copy_positions(
frame_handle: *const RKRConFrame,
out: *mut f64,
out_len: usize,
) -> RKRStatus {
if frame_handle.is_null() || out.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
let n = frame.atom_data.len();
let need = n.saturating_mul(3);
if out_len < need {
return RKRStatus::RKR_STATUS_BUFFER_TOO_SMALL;
}
let slice = unsafe { std::slice::from_raw_parts_mut(out, need) };
for (i, a) in frame.atom_data.iter().enumerate() {
slice[i * 3] = a.x;
slice[i * 3 + 1] = a.y;
slice[i * 3 + 2] = a.z;
}
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_copy_velocities(
frame_handle: *const RKRConFrame,
out: *mut f64,
out_len: usize,
) -> RKRStatus {
if frame_handle.is_null() || out.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
if !frame.has_velocities() {
return RKRStatus::RKR_STATUS_SECTION_ABSENT;
}
let n = frame.atom_data.len();
let need = n.saturating_mul(3);
if out_len < need {
return RKRStatus::RKR_STATUS_BUFFER_TOO_SMALL;
}
let slice = unsafe { std::slice::from_raw_parts_mut(out, need) };
for (i, a) in frame.atom_data.iter().enumerate() {
let [vx, vy, vz] = a.velocity.unwrap_or([0.0; 3]);
slice[i * 3] = vx;
slice[i * 3 + 1] = vy;
slice[i * 3 + 2] = vz;
}
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_copy_forces(
frame_handle: *const RKRConFrame,
out: *mut f64,
out_len: usize,
) -> RKRStatus {
if frame_handle.is_null() || out.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
if !frame.has_forces() {
return RKRStatus::RKR_STATUS_SECTION_ABSENT;
}
let n = frame.atom_data.len();
let need = n.saturating_mul(3);
if out_len < need {
return RKRStatus::RKR_STATUS_BUFFER_TOO_SMALL;
}
let slice = unsafe { std::slice::from_raw_parts_mut(out, need) };
for (i, a) in frame.atom_data.iter().enumerate() {
let [fx, fy, fz] = a.force.unwrap_or([0.0; 3]);
slice[i * 3] = fx;
slice[i * 3 + 1] = fy;
slice[i * 3 + 2] = fz;
}
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_copy_atom_energies(
frame_handle: *const RKRConFrame,
out: *mut f64,
out_len: usize,
) -> RKRStatus {
if frame_handle.is_null() || out.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
if !frame.has_energies() {
return RKRStatus::RKR_STATUS_SECTION_ABSENT;
}
let n = frame.atom_data.len();
if out_len < n {
return RKRStatus::RKR_STATUS_BUFFER_TOO_SMALL;
}
let slice = unsafe { std::slice::from_raw_parts_mut(out, n) };
for (i, a) in frame.atom_data.iter().enumerate() {
slice[i] = a.energy.unwrap_or(0.0);
}
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_copy_masses(
frame_handle: *const RKRConFrame,
out: *mut f64,
out_len: usize,
) -> RKRStatus {
if frame_handle.is_null() || out.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
let n = frame.atom_data.len();
if out_len < n {
return RKRStatus::RKR_STATUS_BUFFER_TOO_SMALL;
}
let slice = unsafe { std::slice::from_raw_parts_mut(out, n) };
let mut expanded: Vec<f64> = Vec::with_capacity(n);
for (ti, &count) in frame.header.natms_per_type.iter().enumerate() {
let m = frame.header.masses_per_type.get(ti).copied().unwrap_or(0.0);
expanded.extend(std::iter::repeat_n(m, count));
}
if expanded.len() != n {
expanded.resize(n, 0.0);
}
slice.copy_from_slice(&expanded[..n]);
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_copy_atom_ids(
frame_handle: *const RKRConFrame,
out: *mut u64,
out_len: usize,
) -> RKRStatus {
if frame_handle.is_null() || out.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
let n = frame.atom_data.len();
if out_len < n {
return RKRStatus::RKR_STATUS_BUFFER_TOO_SMALL;
}
let slice = unsafe { std::slice::from_raw_parts_mut(out, n) };
for (i, a) in frame.atom_data.iter().enumerate() {
slice[i] = a.atom_id;
}
RKRStatus::RKR_STATUS_SUCCESS
}
fn frame_positions_arc(frame: &ConFrame) -> ndarray::ArcArray2<f64> {
let n = frame.atom_data.len();
let mut data = Vec::with_capacity(n * 3);
for a in &frame.atom_data {
data.extend_from_slice(&[a.x, a.y, a.z]);
}
ndarray::ArcArray2::from_shape_vec((n, 3), data)
.unwrap_or_else(|_| ndarray::ArcArray2::zeros((0, 3)))
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_positions_as_dlpack(
frame_handle: *const RKRConFrame,
device_type: i32,
device_id: i32,
_stream: i64,
_max_version_major: u32,
_max_version_minor: u32,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
if frame_handle.is_null() || out_tensor.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
unsafe { *out_tensor = std::ptr::null_mut() };
if device_type != rkr_dl_device_type::RKR_DL_CPU {
return RKRStatus::RKR_STATUS_FEATURE_DISABLED;
}
let _ = device_id;
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
match frame.positions_as_dlpack(dlpk::sys::DLDevice::cpu()) {
Ok(tensor) => {
let raw = tensor.into_raw();
unsafe {
*out_tensor = raw.as_ptr();
}
RKRStatus::RKR_STATUS_SUCCESS
}
Err(e) => map_dlpack_err(e),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_positions_from_dlpack(
frame_handle: *mut RKRConFrame,
tensor: *const RKRDLManagedTensorVersioned,
) -> RKRStatus {
if frame_handle.is_null() || tensor.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let frame = unsafe { &mut *(frame_handle as *mut ConFrame) };
let dl = unsafe { &(*tensor).dl_tensor };
if dl.device.device_type != dlpk::sys::DLDeviceType::kDLCPU {
return RKRStatus::RKR_STATUS_FEATURE_DISABLED;
}
let n = frame.atom_data.len();
let need = n.saturating_mul(3);
let ndim = dl.ndim as usize;
let shape = if dl.shape.is_null() {
return RKRStatus::RKR_STATUS_VALIDATION_ERROR;
} else {
unsafe { std::slice::from_raw_parts(dl.shape, ndim) }
};
let nelem = if ndim == 2 && shape[0] == n as i64 && shape[1] == 3 {
need
} else if ndim == 1 && shape[0] == need as i64 {
need
} else {
return RKRStatus::RKR_STATUS_VALIDATION_ERROR;
};
let code = dl.dtype.code as u8;
let bits = dl.dtype.bits;
if dl.data.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let vals: Vec<f64> = if code == rkr_dl_type_code::RKR_DL_FLOAT && bits == 64 {
let s = unsafe { std::slice::from_raw_parts(dl.data as *const f64, nelem) };
s.to_vec()
} else if code == rkr_dl_type_code::RKR_DL_FLOAT && bits == 32 {
let s = unsafe { std::slice::from_raw_parts(dl.data as *const f32, nelem) };
s.iter().map(|&x| x as f64).collect()
} else {
return RKRStatus::RKR_STATUS_VALIDATION_ERROR;
};
if frame.positions.nrows() != n {
return RKRStatus::RKR_STATUS_VALIDATION_ERROR;
}
for i in 0..n {
frame.positions.set_f64_row(
i,
[vals[i * 3], vals[i * 3 + 1], vals[i * 3 + 2]],
);
}
frame.sync_atom_data_from_arrays();
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_positions_dlpack(
frame_handle: *const RKRConFrame,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
unsafe {
rkr_frame_positions_as_dlpack(frame_handle, rkr_dl_device_type::RKR_DL_CPU, 0, 0, 1, 0, out_tensor)
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_positions_dlpack_ex(
frame_handle: *const RKRConFrame,
opts: *const RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
if frame_handle.is_null() || out_tensor.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
unsafe { *out_tensor = std::ptr::null_mut() };
let o = match resolve_dlpack_opts(opts) {
Ok(o) => o,
Err(st) => return st,
};
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
let arr = frame_positions_arc(frame);
export_owned_array2_dlpack_opts(&arr, &o, out_tensor)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_velocities_dlpack(
frame_handle: *const RKRConFrame,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
unsafe { rkr_frame_velocities_dlpack_ex(frame_handle, std::ptr::null(), out_tensor) }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_velocities_dlpack_ex(
frame_handle: *const RKRConFrame,
opts: *const RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
if frame_handle.is_null() || out_tensor.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
unsafe { *out_tensor = std::ptr::null_mut() };
let o = match resolve_dlpack_opts(opts) {
Ok(o) => o,
Err(st) => return st,
};
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
if !frame.has_velocities() {
return RKRStatus::RKR_STATUS_SECTION_ABSENT;
}
let n = frame.atom_data.len();
let mut data = Vec::with_capacity(n * 3);
for a in &frame.atom_data {
let v = a.velocity.unwrap_or([0.0; 3]);
data.extend_from_slice(&v);
}
let arr = ndarray::ArcArray2::from_shape_vec((n, 3), data)
.unwrap_or_else(|_| ndarray::ArcArray2::zeros((0, 3)));
export_owned_array2_dlpack_opts(&arr, &o, out_tensor)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_forces_dlpack(
frame_handle: *const RKRConFrame,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
unsafe { rkr_frame_forces_dlpack_ex(frame_handle, std::ptr::null(), out_tensor) }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_forces_dlpack_ex(
frame_handle: *const RKRConFrame,
opts: *const RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
if frame_handle.is_null() || out_tensor.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
unsafe { *out_tensor = std::ptr::null_mut() };
let o = match resolve_dlpack_opts(opts) {
Ok(o) => o,
Err(st) => return st,
};
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
if !frame.has_forces() {
return RKRStatus::RKR_STATUS_SECTION_ABSENT;
}
let n = frame.atom_data.len();
let mut data = Vec::with_capacity(n * 3);
for a in &frame.atom_data {
let f = a.force.unwrap_or([0.0; 3]);
data.extend_from_slice(&f);
}
let arr = ndarray::ArcArray2::from_shape_vec((n, 3), data)
.unwrap_or_else(|_| ndarray::ArcArray2::zeros((0, 3)));
export_owned_array2_dlpack_opts(&arr, &o, out_tensor)
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_atom_energies_dlpack(
frame_handle: *const RKRConFrame,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
unsafe { rkr_frame_atom_energies_dlpack_ex(frame_handle, std::ptr::null(), out_tensor) }
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_atom_energies_dlpack_ex(
frame_handle: *const RKRConFrame,
opts: *const RKRDlpackExportOptions,
out_tensor: *mut *mut RKRDLManagedTensorVersioned,
) -> RKRStatus {
if frame_handle.is_null() || out_tensor.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
unsafe { *out_tensor = std::ptr::null_mut() };
let o = match resolve_dlpack_opts(opts) {
Ok(o) => o,
Err(st) => return st,
};
let Some(frame) = (unsafe { (frame_handle as *const ConFrame).as_ref() }) else {
return RKRStatus::RKR_STATUS_NULL_POINTER;
};
if !frame.has_energies() {
return RKRStatus::RKR_STATUS_SECTION_ABSENT;
}
let data: Vec<f64> = frame
.atom_data
.iter()
.map(|a| a.energy.unwrap_or(0.0))
.collect();
let arr = ndarray::ArcArray1::from_vec(data);
export_owned_array1_f64_dlpack_opts(&arr, &o, out_tensor)
}
pub struct RKRSelectionResult;
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_frame_select(
frame_handle: *const RKRConFrame,
selection: *const c_char,
out_result: *mut *mut RKRSelectionResult,
) -> RKRStatus {
if frame_handle.is_null() || selection.is_null() || out_result.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let frame = unsafe { &*(frame_handle as *const ConFrame) };
let sel_str = match unsafe { CStr::from_ptr(selection) }.to_str() {
Ok(s) => s,
Err(_) => return RKRStatus::RKR_STATUS_INVALID_UTF8,
};
match crate::chemfiles_selection::evaluate_selection_on_con_frame(sel_str, frame) {
Ok(result) => {
let boxed = Box::new(result);
unsafe {
*out_result = Box::into_raw(boxed) as *mut RKRSelectionResult;
}
RKRStatus::RKR_STATUS_SUCCESS
}
Err(_) => RKRStatus::RKR_STATUS_SELECTION_ERROR,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_selection_result_match_count(
result_handle: *const RKRSelectionResult,
) -> u64 {
if result_handle.is_null() {
return 0;
}
let result =
unsafe { &*(result_handle as *const crate::chemfiles_selection::SelectionResult) };
result.matches.len() as u64
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_selection_result_context_size(
result_handle: *const RKRSelectionResult,
) -> u32 {
if result_handle.is_null() {
return 0;
}
let result =
unsafe { &*(result_handle as *const crate::chemfiles_selection::SelectionResult) };
result.context_size as u32
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_selection_result_match_at(
result_handle: *const RKRSelectionResult,
match_index: u64,
out_atoms: *mut u64,
out_size: *mut u32,
) -> RKRStatus {
if result_handle.is_null() || out_atoms.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let result =
unsafe { &*(result_handle as *const crate::chemfiles_selection::SelectionResult) };
let idx = match_index as usize;
if idx >= result.matches.len() {
return RKRStatus::RKR_STATUS_INDEX_OUT_OF_BOUNDS;
}
let m = &result.matches[idx];
unsafe {
for i in 0..4 {
*out_atoms.add(i) = if i < m.size {
m.atoms[i] as u64
} else {
u64::MAX
};
}
if !out_size.is_null() {
*out_size = m.size as u32;
}
}
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_selection_result_primary_indices(
result_handle: *const RKRSelectionResult,
out_indices: *mut u64,
capacity: u64,
out_written: *mut u64,
) -> RKRStatus {
if result_handle.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
let result =
unsafe { &*(result_handle as *const crate::chemfiles_selection::SelectionResult) };
let n = result.matches.len() as u64;
if !out_written.is_null() {
unsafe {
*out_written = n;
}
}
if n == 0 {
return RKRStatus::RKR_STATUS_SUCCESS;
}
if out_indices.is_null() {
return RKRStatus::RKR_STATUS_NULL_POINTER;
}
if capacity < n {
return RKRStatus::RKR_STATUS_BUFFER_TOO_SMALL;
}
unsafe {
for (i, m) in result.matches.iter().enumerate() {
*out_indices.add(i) = m.atoms[0] as u64;
}
}
RKRStatus::RKR_STATUS_SUCCESS
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_selection_result_free(result_handle: *mut RKRSelectionResult) {
if result_handle.is_null() {
return;
}
unsafe {
drop(Box::from_raw(
result_handle as *mut crate::chemfiles_selection::SelectionResult,
));
}
}
#[unsafe(no_mangle)]
pub extern "C" fn rkr_has_chemfiles_support() -> u8 {
#[cfg(feature = "chemfiles")]
{
1
}
#[cfg(not(feature = "chemfiles"))]
{
0
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_read_chemfiles_first(path_c: *const c_char) -> *mut RKRConFrame {
if path_c.is_null() {
return std::ptr::null_mut();
}
let Ok(path_str) = unsafe { CStr::from_ptr(path_c) }.to_str() else {
return std::ptr::null_mut();
};
match crate::chemfiles_import::con_frame_from_trajectory_path(path_str) {
Ok(frame) => Box::into_raw(Box::new(frame)) as *mut RKRConFrame,
Err(_) => std::ptr::null_mut(),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_read_chemfiles_memory(
data_c: *const c_char,
format_c: *const c_char,
num_frames: *mut usize,
) -> *mut *mut RKRConFrame {
if data_c.is_null() || format_c.is_null() || num_frames.is_null() {
return std::ptr::null_mut();
}
let Ok(data) = unsafe { CStr::from_ptr(data_c) }.to_str() else {
return std::ptr::null_mut();
};
let Ok(format) = unsafe { CStr::from_ptr(format_c) }.to_str() else {
return std::ptr::null_mut();
};
match crate::chemfiles_import::con_frames_from_memory(data, format) {
Ok(frames) => {
let n = frames.len();
unsafe { *num_frames = n };
let mut ptrs: Vec<*mut RKRConFrame> = frames
.into_iter()
.map(|f| Box::into_raw(Box::new(f)) as *mut RKRConFrame)
.collect();
let p = ptrs.as_mut_ptr();
std::mem::forget(ptrs);
p
}
Err(_) => std::ptr::null_mut(),
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn rkr_dlpack_delete(tensor: *mut RKRDLManagedTensorVersioned) {
if tensor.is_null() {
return;
}
unsafe {
let t = &mut *tensor;
if let Some(del) = t.deleter {
del(tensor);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::ffi::{CStr, CString};
#[test]
fn frame_copy_positions_without_cframe() {
let handle = test_frame_handle();
let n = unsafe { rkr_frame_atom_count(handle) };
assert_eq!(n, 1);
let mut buf = vec![0.0f64; n * 3];
assert_eq!(
unsafe { rkr_frame_copy_positions(handle, buf.as_mut_ptr(), buf.len()) },
RKRStatus::RKR_STATUS_SUCCESS
);
let mut tensor: *mut RKRDLManagedTensorVersioned = std::ptr::null_mut();
assert_eq!(
unsafe { rkr_frame_positions_dlpack(handle, &mut tensor) },
RKRStatus::RKR_STATUS_SUCCESS
);
assert!(!tensor.is_null());
unsafe { rkr_dlpack_delete(tensor) };
assert_eq!(
unsafe { rkr_frame_copy_velocities(handle, buf.as_mut_ptr(), buf.len()) },
RKRStatus::RKR_STATUS_SECTION_ABSENT
);
unsafe { free_rkr_frame(handle) };
}
#[test]
fn read_all_frames_c_abi_tiny() {
let path = std::ffi::CString::new("resources/test/tiny_cuh2.con").unwrap();
let mut n: usize = 0;
let arr = unsafe { rkr_read_all_frames(path.as_ptr(), &mut n) };
assert!(!arr.is_null() && n >= 1);
let first = unsafe { *arr };
let nat = unsafe { rkr_frame_atom_count(first) };
let mut buf = vec![0.0f64; nat * 3];
assert_eq!(
unsafe { rkr_frame_copy_positions(first, buf.as_mut_ptr(), buf.len()) },
RKRStatus::RKR_STATUS_SUCCESS
);
unsafe { free_rkr_frame_array(arr, n) };
}
#[test]
fn free_rkr_frame_ptr_array_keeps_frames() {
let path = std::ffi::CString::new("resources/test/tiny_cuh2.con").unwrap();
let mut n: usize = 0;
let arr = unsafe { rkr_read_all_frames(path.as_ptr(), &mut n) };
assert!(!arr.is_null() && n >= 1);
let first = unsafe { *arr };
let rest: Vec<*mut RKRConFrame> = (1..n).map(|i| unsafe { *arr.add(i) }).collect();
unsafe { free_rkr_frame_ptr_array(arr, n) };
assert!(unsafe { rkr_frame_atom_count(first) } >= 1);
unsafe { free_rkr_frame(first) };
for h in rest {
if !h.is_null() {
unsafe { free_rkr_frame(h) };
}
}
}
fn test_frame_handle() -> *mut RKRConFrame {
let mut builder = ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
builder
.prebox_header("Generated by test")
.postbox_header(["0 0".to_string(), "0 0 0".to_string()]);
builder.add_atom("Cu", 0.0, 0.0, 0.0, [false, false, false], 0, 63.546);
Box::into_raw(Box::new(builder.build())) as *mut RKRConFrame
}
#[test]
fn header_line_rejects_null_buffer() {
let frame = test_frame_handle();
let status = unsafe { rkr_frame_get_header_line(frame, true, 0, std::ptr::null_mut(), 16) };
unsafe { free_rkr_frame(frame) };
assert_eq!(status, RKRStatus::RKR_STATUS_NULL_POINTER);
}
#[test]
fn header_line_rejects_empty_buffer() {
let frame = test_frame_handle();
let mut buffer = [0 as c_char; 1];
let status = unsafe { rkr_frame_get_header_line(frame, true, 0, buffer.as_mut_ptr(), 0) };
unsafe { free_rkr_frame(frame) };
assert_eq!(status, RKRStatus::RKR_STATUS_BUFFER_TOO_SMALL);
}
#[test]
fn header_line_truncates_and_terminates_buffer() {
let frame = test_frame_handle();
let mut buffer = [0 as c_char; 10];
let status =
unsafe { rkr_frame_get_header_line(frame, true, 0, buffer.as_mut_ptr(), buffer.len()) };
unsafe { free_rkr_frame(frame) };
assert_eq!(status, RKRStatus::RKR_STATUS_SUCCESS);
let copied = unsafe { CStr::from_ptr(buffer.as_ptr()) };
assert_eq!(copied.to_str().unwrap(), "Generated");
}
fn test_builder_handle() -> *mut RKRConFrameBuilder {
let cell = [10.0, 11.0, 12.0];
let angles = [90.0, 91.0, 92.0];
unsafe {
rkr_frame_new(
cell.as_ptr(),
angles.as_ptr(),
ptr::null(),
ptr::null(),
ptr::null(),
ptr::null(),
)
}
}
fn c_string(s: &str) -> CString {
CString::new(s).unwrap()
}
unsafe fn assert_single_atom(
frame: *mut RKRConFrame,
fixed: [bool; 3],
velocity: Option<[f64; 3]>,
forces: Option<[f64; 3]>,
) {
let c_frame = unsafe { rkr_frame_to_c_frame(frame) };
assert!(!c_frame.is_null());
let c_frame_ref = unsafe { &*c_frame };
assert_eq!(c_frame_ref.num_atoms, 1);
assert_eq!(c_frame_ref.has_velocities, velocity.is_some());
assert_eq!(c_frame_ref.has_forces, forces.is_some());
let atom = unsafe { &*c_frame_ref.atoms };
assert_eq!(atom.fixed_x, fixed[0]);
assert_eq!(atom.fixed_y, fixed[1]);
assert_eq!(atom.fixed_z, fixed[2]);
assert_eq!(atom.is_fixed, fixed.iter().any(|&value| value));
assert_eq!(atom.has_velocity, velocity.is_some());
assert_eq!(atom.has_forces, forces.is_some());
if let Some([vx, vy, vz]) = velocity {
assert_eq!([atom.vx, atom.vy, atom.vz], [vx, vy, vz]);
}
if let Some([fx, fy, fz]) = forces {
assert_eq!([atom.fx, atom.fy, atom.fz], [fx, fy, fz]);
}
unsafe { free_c_frame(c_frame) };
unsafe { free_rkr_frame(frame) };
}
#[test]
fn builder_preserves_fixed_mask_for_atom_without_velocity_or_forces() {
let builder = test_builder_handle();
let symbol = c_string("Cu");
let status = unsafe {
rkr_frame_add_atom_with_fixed_mask(
builder,
symbol.as_ptr(),
1.0,
2.0,
3.0,
true,
false,
true,
7,
63.546,
)
};
assert_eq!(status, RKRStatus::RKR_STATUS_SUCCESS);
let frame = unsafe { rkr_frame_builder_build(builder) };
unsafe { assert_single_atom(frame, [true, false, true], None, None) };
}
#[test]
fn builder_preserves_fixed_mask_for_atom_with_velocity() {
let builder = test_builder_handle();
let symbol = c_string("H");
let status = unsafe {
rkr_frame_add_atom_with_velocity_fixed_mask(
builder,
symbol.as_ptr(),
1.0,
2.0,
3.0,
false,
true,
false,
9,
1.008,
0.1,
0.2,
0.3,
)
};
assert_eq!(status, RKRStatus::RKR_STATUS_SUCCESS);
let frame = unsafe { rkr_frame_builder_build(builder) };
unsafe { assert_single_atom(frame, [false, true, false], Some([0.1, 0.2, 0.3]), None) };
}
#[test]
fn builder_preserves_fixed_mask_for_atom_with_forces() {
let builder = test_builder_handle();
let symbol = c_string("O");
let status = unsafe {
rkr_frame_add_atom_with_forces_fixed_mask(
builder,
symbol.as_ptr(),
1.0,
2.0,
3.0,
true,
true,
false,
11,
15.999,
-0.1,
-0.2,
-0.3,
)
};
assert_eq!(status, RKRStatus::RKR_STATUS_SUCCESS);
let frame = unsafe { rkr_frame_builder_build(builder) };
unsafe { assert_single_atom(frame, [true, true, false], None, Some([-0.1, -0.2, -0.3])) };
}
#[test]
fn builder_preserves_fixed_mask_for_atom_with_velocity_and_forces() {
let builder = test_builder_handle();
let symbol = c_string("N");
let status = unsafe {
rkr_frame_add_atom_with_velocity_and_forces_fixed_mask(
builder,
symbol.as_ptr(),
1.0,
2.0,
3.0,
false,
true,
true,
13,
14.007,
0.4,
0.5,
0.6,
-0.4,
-0.5,
-0.6,
)
};
assert_eq!(status, RKRStatus::RKR_STATUS_SUCCESS);
let frame = unsafe { rkr_frame_builder_build(builder) };
unsafe {
assert_single_atom(
frame,
[false, true, true],
Some([0.4, 0.5, 0.6]),
Some([-0.4, -0.5, -0.6]),
)
};
}
#[test]
fn builder_bool_fixed_functions_set_all_axes_together() {
let builder = test_builder_handle();
let cu = c_string("Cu");
let h = c_string("H");
let atom_status =
unsafe { rkr_frame_add_atom(builder, cu.as_ptr(), 1.0, 2.0, 3.0, true, 1, 63.546) };
assert_eq!(atom_status, RKRStatus::RKR_STATUS_SUCCESS);
let velocity_status = unsafe {
rkr_frame_add_atom_with_velocity(
builder,
h.as_ptr(),
4.0,
5.0,
6.0,
false,
2,
1.008,
0.7,
0.8,
0.9,
)
};
assert_eq!(velocity_status, RKRStatus::RKR_STATUS_SUCCESS);
let frame = unsafe { rkr_frame_builder_build(builder) };
let c_frame = unsafe { rkr_frame_to_c_frame(frame) };
assert!(!c_frame.is_null());
let c_frame_ref = unsafe { &*c_frame };
assert_eq!(c_frame_ref.num_atoms, 2);
let atoms = unsafe { std::slice::from_raw_parts(c_frame_ref.atoms, c_frame_ref.num_atoms) };
assert_eq!(
[atoms[0].fixed_x, atoms[0].fixed_y, atoms[0].fixed_z],
[true, true, true]
);
assert_eq!(
[atoms[1].fixed_x, atoms[1].fixed_y, atoms[1].fixed_z],
[false, false, false]
);
unsafe { free_c_frame(c_frame) };
unsafe { free_rkr_frame(frame) };
}
#[test]
fn status_message_returns_static_strings_for_all_status_values() {
let cases = [
(RKRStatus::RKR_STATUS_SUCCESS, "success"),
(RKRStatus::RKR_STATUS_NULL_POINTER, "null pointer"),
(RKRStatus::RKR_STATUS_INVALID_UTF8, "invalid UTF-8"),
(RKRStatus::RKR_STATUS_INVALID_JSON, "invalid JSON"),
(RKRStatus::RKR_STATUS_IO_ERROR, "I/O error"),
(
RKRStatus::RKR_STATUS_INDEX_OUT_OF_BOUNDS,
"index out of bounds",
),
(RKRStatus::RKR_STATUS_BUFFER_TOO_SMALL, "buffer too small"),
(RKRStatus::RKR_STATUS_INTERNAL_ERROR, "internal error"),
(RKRStatus::RKR_STATUS_SECTION_ABSENT, "section absent"),
(RKRStatus::RKR_STATUS_VALIDATION_ERROR, "validation error"),
(RKRStatus::RKR_STATUS_SELECTION_ERROR, "selection error"),
(
RKRStatus::RKR_STATUS_FEATURE_DISABLED,
"feature disabled in this build",
),
];
for (status, expected) in cases {
let message = unsafe { CStr::from_ptr(rkr_status_message(status)) };
assert_eq!(message.to_str().unwrap(), expected);
}
}
unsafe fn assert_dlpack_cpu_float(
t: *mut RKRDLManagedTensorVersioned,
expect_ndim: i32,
expect_shape: &[i64],
expect_bits: u8,
) {
assert!(!t.is_null());
let dl = unsafe { &(*t).dl_tensor };
assert_eq!(dl.ndim, expect_ndim);
let shape = unsafe { std::slice::from_raw_parts(dl.shape, expect_ndim as usize) };
assert_eq!(shape, expect_shape);
assert_eq!(dl.dtype.code, dlpk::sys::DLDataTypeCode::kDLFloat);
assert_eq!(dl.dtype.bits, expect_bits);
assert_eq!(dl.dtype.lanes, 1);
assert_eq!(dl.device, dlpk::sys::DLDevice::cpu());
assert!(!dl.data.is_null());
}
#[test]
fn frame_optional_section_dlpack_present_and_absent() {
let handle = test_frame_handle();
let mut t: *mut RKRDLManagedTensorVersioned = std::ptr::null_mut();
assert_eq!(
unsafe { rkr_frame_velocities_dlpack(handle, &mut t) },
RKRStatus::RKR_STATUS_SECTION_ABSENT
);
assert!(t.is_null());
assert_eq!(
unsafe { rkr_frame_forces_dlpack(handle, &mut t) },
RKRStatus::RKR_STATUS_SECTION_ABSENT
);
assert_eq!(
unsafe { rkr_frame_atom_energies_dlpack(handle, &mut t) },
RKRStatus::RKR_STATUS_SECTION_ABSENT
);
unsafe { free_rkr_frame(handle) };
let path = CString::new("resources/test/tiny_cuh2.convel").unwrap();
let fr = unsafe { rkr_read_first_frame(path.as_ptr()) };
assert!(!fr.is_null());
let n = unsafe { rkr_frame_atom_count(fr) } as i64;
assert!(n > 0);
let mut vel: *mut RKRDLManagedTensorVersioned = std::ptr::null_mut();
let st = unsafe { rkr_frame_velocities_dlpack(fr, &mut vel) };
assert_eq!(st, RKRStatus::RKR_STATUS_SUCCESS);
unsafe {
assert_dlpack_cpu_float(vel, 2, &[n, 3], 64);
rkr_dlpack_delete(vel);
free_rkr_frame(fr);
}
let cell = [10.0f64; 3];
let ang = [90.0f64; 3];
let b = unsafe {
rkr_frame_new(
cell.as_ptr(),
ang.as_ptr(),
std::ptr::null(),
std::ptr::null(),
std::ptr::null(),
std::ptr::null(),
)
};
assert!(!b.is_null());
let sym = CString::new("H").unwrap();
unsafe {
rkr_frame_add_atom_with_velocity_and_forces_fixed_mask(
b,
sym.as_ptr(),
0.0,
0.0,
0.0,
false,
false,
false,
0,
1.0,
0.1,
0.0,
0.0,
0.0,
0.0,
-1.0,
);
rkr_frame_builder_set_last_energy(b, -0.5);
}
let built = unsafe { rkr_frame_builder_build(b) };
assert!(!built.is_null());
let n_built = unsafe { rkr_frame_atom_count(built) } as i64;
let mut frc: *mut RKRDLManagedTensorVersioned = std::ptr::null_mut();
let mut eng: *mut RKRDLManagedTensorVersioned = std::ptr::null_mut();
assert_eq!(
unsafe { rkr_frame_forces_dlpack(built, &mut frc) },
RKRStatus::RKR_STATUS_SUCCESS
);
unsafe { assert_dlpack_cpu_float(frc, 2, &[n_built, 3], 64) };
assert_eq!(
unsafe { rkr_frame_atom_energies_dlpack(built, &mut eng) },
RKRStatus::RKR_STATUS_SUCCESS
);
unsafe { assert_dlpack_cpu_float(eng, 1, &[n_built], 64) };
let opts32 = RKRDlpackExportOptions {
dtype: RKRDLDataType {
code: rkr_dl_type_code::RKR_DL_FLOAT,
bits: 32,
lanes: 1,
},
device: RKRDLDevice {
device_type: rkr_dl_device_type::RKR_DL_CPU,
device_id: 0,
},
};
let mut pos32: *mut RKRDLManagedTensorVersioned = std::ptr::null_mut();
assert_eq!(
unsafe { rkr_frame_positions_dlpack_ex(built, &opts32, &mut pos32) },
RKRStatus::RKR_STATUS_SUCCESS
);
unsafe {
assert_dlpack_cpu_float(pos32, 2, &[n_built, 3], 32);
rkr_dlpack_delete(pos32);
}
let bad_dev = RKRDlpackExportOptions {
dtype: RKRDLDataType {
code: rkr_dl_type_code::RKR_DL_FLOAT,
bits: 64,
lanes: 1,
},
device: RKRDLDevice {
device_type: rkr_dl_device_type::RKR_DL_CUDA,
device_id: 0,
},
};
let mut junk: *mut RKRDLManagedTensorVersioned = std::ptr::null_mut();
assert_eq!(
unsafe { rkr_frame_positions_dlpack_ex(built, &bad_dev, &mut junk) },
RKRStatus::RKR_STATUS_FEATURE_DISABLED
);
let bad_bits = RKRDlpackExportOptions {
dtype: RKRDLDataType {
code: rkr_dl_type_code::RKR_DL_FLOAT,
bits: 16,
lanes: 1,
},
device: RKRDLDevice {
device_type: rkr_dl_device_type::RKR_DL_CPU,
device_id: 0,
},
};
assert_eq!(
unsafe { rkr_frame_positions_dlpack_ex(built, &bad_bits, &mut junk) },
RKRStatus::RKR_STATUS_VALIDATION_ERROR
);
let mut as_t: *mut RKRDLManagedTensorVersioned = std::ptr::null_mut();
assert_eq!(
unsafe {
rkr_frame_positions_as_dlpack(
built,
rkr_dl_device_type::RKR_DL_CPU,
0,
0,
1,
0,
&mut as_t,
)
},
RKRStatus::RKR_STATUS_SUCCESS
);
unsafe {
assert_dlpack_cpu_float(as_t, 2, &[n_built, 3], 64);
}
{
let fr = unsafe { &mut *(built as *mut ConFrame) };
for a in fr.atom_data.iter_mut() {
a.x = -99.0;
a.y = -99.0;
a.z = -99.0;
}
}
assert_eq!(
unsafe { rkr_frame_positions_from_dlpack(built, as_t) },
RKRStatus::RKR_STATUS_SUCCESS
);
{
let fr = unsafe { &*(built as *const ConFrame) };
assert!((fr.atom_data[0].x - 0.0).abs() < 1e-12);
assert!((fr.atom_data[0].y - 0.0).abs() < 1e-12);
assert!((fr.atom_data[0].z - 0.0).abs() < 1e-12);
}
let mut again: *mut RKRDLManagedTensorVersioned = std::ptr::null_mut();
assert_eq!(
unsafe {
rkr_frame_positions_as_dlpack(
built,
rkr_dl_device_type::RKR_DL_CPU,
0,
0,
1,
0,
&mut again,
)
},
RKRStatus::RKR_STATUS_SUCCESS
);
unsafe {
let dl = &(*again).dl_tensor;
let data = std::slice::from_raw_parts(dl.data as *const f64, 3);
assert!((data[0] - 0.0).abs() < 1e-12);
assert!((data[1] - 0.0).abs() < 1e-12);
assert!((data[2] - 0.0).abs() < 1e-12);
rkr_dlpack_delete(again);
rkr_dlpack_delete(as_t);
rkr_dlpack_delete(frc);
rkr_dlpack_delete(eng);
free_rkr_frame(built);
}
}
#[test]
fn ffi_positions_dlpack_round_trip() {
let handle = test_builder_handle();
let sym = c_string("Cu");
unsafe {
rkr_frame_add_atom_full(
handle,
sym.as_ptr(),
1.0,
2.0,
3.0,
false,
false,
false,
7,
63.5,
ptr::null(),
ptr::null(),
)
};
let mut t: *mut RKRDLManagedTensorVersioned = ptr::null_mut();
let status = unsafe { rkr_frame_builder_positions_dlpack(handle, &mut t) };
assert_eq!(status, RKRStatus::RKR_STATUS_SUCCESS);
assert!(!t.is_null());
let dl = unsafe { &(*t).dl_tensor };
assert_eq!(dl.ndim, 2);
let shape = unsafe { std::slice::from_raw_parts(dl.shape, 2) };
assert_eq!(shape, &[1, 3]);
assert_eq!(dl.dtype.code, dlpk::sys::DLDataTypeCode::kDLFloat);
assert_eq!(dl.dtype.bits, 64);
assert_eq!(dl.dtype.lanes, 1);
assert_eq!(dl.device, dlpk::sys::DLDevice::cpu());
let data = unsafe { std::slice::from_raw_parts(dl.data as *const f64, 3) };
assert_eq!(data, &[1.0, 2.0, 3.0]);
let deleter = unsafe { (*t).deleter };
if let Some(del) = deleter {
unsafe { del(t) };
}
unsafe { free_rkr_frame_builder(handle) };
}
#[test]
fn ffi_velocities_dlpack_section_absent() {
let handle = test_builder_handle();
let sym = c_string("Cu");
unsafe {
rkr_frame_add_atom_full(
handle,
sym.as_ptr(),
0.0,
0.0,
0.0,
false,
false,
false,
0,
63.5,
ptr::null(),
ptr::null(),
)
};
let mut t: *mut RKRDLManagedTensorVersioned = ptr::null_mut();
let status = unsafe { rkr_frame_builder_velocities_dlpack(handle, &mut t) };
assert_eq!(status, RKRStatus::RKR_STATUS_SECTION_ABSENT);
assert!(t.is_null());
unsafe { free_rkr_frame_builder(handle) };
}
#[test]
fn ffi_dlpack_null_handle_rejects() {
let mut t: *mut RKRDLManagedTensorVersioned = ptr::null_mut();
let status = unsafe { rkr_frame_builder_positions_dlpack(ptr::null(), &mut t) };
assert_eq!(status, RKRStatus::RKR_STATUS_NULL_POINTER);
assert!(t.is_null());
}
#[cfg(feature = "chemfiles")]
#[test]
fn rkr_frame_select_finds_oxygen() {
use crate::types::ConFrameBuilder;
let mut b = ConFrameBuilder::new([10.0; 3], [90.0; 3]);
b.add_atom("O", 0.0, 0.0, 0.0, [false; 3], 0, 16.0);
b.add_atom("H", 1.0, 0.0, 0.0, [false; 3], 1, 1.0);
let frame = b.build();
let frame_ptr = Box::into_raw(Box::new(frame)) as *mut RKRConFrame;
let sel = CString::new("name O").unwrap();
let mut out: *mut RKRSelectionResult = ptr::null_mut();
let st = unsafe { rkr_frame_select(frame_ptr, sel.as_ptr(), &mut out) };
assert_eq!(st, RKRStatus::RKR_STATUS_SUCCESS);
assert!(!out.is_null());
let n = unsafe { rkr_selection_result_match_count(out) };
assert_eq!(n, 1);
let mut atoms = [u64::MAX; 4];
let mut size = 0u32;
let st2 = unsafe { rkr_selection_result_match_at(out, 0, atoms.as_mut_ptr(), &mut size) };
assert_eq!(st2, RKRStatus::RKR_STATUS_SUCCESS);
assert_eq!(size, 1);
assert_eq!(atoms[0], 0);
unsafe {
rkr_selection_result_free(out);
free_rkr_frame(frame_ptr);
}
}
#[cfg(feature = "chemfiles")]
#[test]
fn rkr_frame_select_cpp_topology_bonds_angles_dihedrals() {
use crate::types::{Bond, ConFrameBuilder};
let mut b = ConFrameBuilder::new([10.0; 3], [90.0; 3]);
b.add_atom("H", 0.0, 1.0, 2.0, [false; 3], 0, 1.0);
b.add_atom("O", 1.0, 2.0, 3.0, [false; 3], 1, 16.0);
b.add_atom("O", 2.0, 3.0, 4.0, [false; 3], 2, 16.0);
b.add_atom("H", 3.0, 4.0, 5.0, [false; 3], 3, 1.0);
let mut frame = b.build();
let id_to = |id: u64| {
frame
.atom_data
.iter()
.position(|a| a.atom_id == id)
.unwrap() as u32
};
frame.header.set_bonds(&[
Bond::new(id_to(0), id_to(1)),
Bond::new(id_to(1), id_to(2)),
Bond::new(id_to(2), id_to(3)),
]);
let frame_ptr = Box::into_raw(Box::new(frame)) as *mut RKRConFrame;
let run = |sel: &str| -> (u64, u32) {
let csel = CString::new(sel).unwrap();
let mut out: *mut RKRSelectionResult = ptr::null_mut();
let st = unsafe { rkr_frame_select(frame_ptr, csel.as_ptr(), &mut out) };
assert_eq!(st, RKRStatus::RKR_STATUS_SUCCESS, "select failed: {sel}");
let n = unsafe { rkr_selection_result_match_count(out) };
let ctx = unsafe { rkr_selection_result_context_size(out) };
unsafe { rkr_selection_result_free(out) };
(n, ctx)
};
assert_eq!(run("bonds: all"), (3, 2));
assert_eq!(run("angles: all"), (2, 3));
assert_eq!(run("dihedrals: all"), (1, 4));
assert_eq!(run("bonds: name(#1) O and type(#2) H").0, 2);
assert_eq!(
run("two: type(#1) H and name(#2) O and is_bonded(#1, #2)").0,
run("bonds: type(#1) H and name(#2) O").0
);
unsafe { free_rkr_frame(frame_ptr) };
}
#[cfg(feature = "metatensor")]
fn assert_mts_block_shape(block: *mut metatensor::c_api::mts_block_t, n: usize, props: usize) {
assert!(!block.is_null());
let mut array = unsafe { std::mem::zeroed::<metatensor::c_api::mts_array_t>() };
let status = unsafe { metatensor::c_api::mts_block_data(block, &mut array) };
assert_eq!(status, metatensor::c_api::MTS_SUCCESS);
let shape_fn = array.shape.expect("mts_array_t.shape from metatensor C API");
let mut shape_ptr: *const usize = std::ptr::null();
let mut shape_count: usize = 0;
let st_shape = unsafe { shape_fn(array.ptr, &mut shape_ptr, &mut shape_count) };
assert_eq!(st_shape, metatensor::c_api::MTS_SUCCESS);
assert_eq!(shape_count, 2);
let shape = unsafe { std::slice::from_raw_parts(shape_ptr, shape_count) };
assert_eq!(shape[0], n);
assert_eq!(shape[1], props);
let samples = unsafe { metatensor::c_api::mts_block_labels(block, 0) };
let prop_lab = unsafe { metatensor::c_api::mts_block_labels(block, 1) };
assert!(!samples.is_null() && !prop_lab.is_null());
}
#[cfg(feature = "metatensor")]
#[test]
fn metatensor_positions_via_c_abi() {
let handle = test_frame_handle();
let mut out: *mut metatensor::c_api::mts_block_t = std::ptr::null_mut();
let st = unsafe { rkr_frame_metatensor_positions_block(handle, &mut out) };
assert_eq!(st, RKRStatus::RKR_STATUS_SUCCESS);
assert_mts_block_shape(out, 1, 3);
unsafe { rkr_mts_block_free(out) };
for (name, export) in [
(
"velocities",
rkr_frame_metatensor_velocities_block
as unsafe extern "C" fn(
*const RKRConFrame,
*mut *mut metatensor::c_api::mts_block_t,
) -> RKRStatus,
),
(
"forces",
rkr_frame_metatensor_forces_block
as unsafe extern "C" fn(
*const RKRConFrame,
*mut *mut metatensor::c_api::mts_block_t,
) -> RKRStatus,
),
(
"atom_energies",
rkr_frame_metatensor_atom_energies_block
as unsafe extern "C" fn(
*const RKRConFrame,
*mut *mut metatensor::c_api::mts_block_t,
) -> RKRStatus,
),
] {
let mut o: *mut metatensor::c_api::mts_block_t = std::ptr::null_mut();
let st_abs = unsafe { export(handle, &mut o) };
assert_eq!(
st_abs,
RKRStatus::RKR_STATUS_SECTION_ABSENT,
"{name} must be SECTION_ABSENT on minimal test frame"
);
assert!(o.is_null());
}
unsafe { free_rkr_frame(handle) };
}
#[cfg(feature = "metatensor")]
#[test]
fn metatensor_optional_sections_via_c_abi() {
let mut builder = ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
builder.add_atom("H", 0.0, 0.0, 0.0, [false; 3], 1, 1.0);
builder.add_atom("O", 1.0, 0.0, 0.0, [false; 3], 2, 16.0);
builder.set_atom_velocity(0, [0.1, 0.2, 0.3]).unwrap();
builder.set_atom_velocity(1, [0.0, 0.1, 0.0]).unwrap();
builder.set_atom_force(0, [1.0, 0.0, 0.0]).unwrap();
builder.set_atom_force(1, [0.0, 1.0, 0.0]).unwrap();
builder.set_atom_energy(0, -0.5).unwrap();
builder.set_atom_energy(1, -1.0).unwrap();
let frame = builder.build();
let handle = Box::into_raw(Box::new(frame)) as *mut RKRConFrame;
let mut pos: *mut metatensor::c_api::mts_block_t = std::ptr::null_mut();
assert_eq!(
unsafe { rkr_frame_metatensor_positions_block(handle, &mut pos) },
RKRStatus::RKR_STATUS_SUCCESS
);
assert_mts_block_shape(pos, 2, 3);
unsafe { rkr_mts_block_free(pos) };
let mut vel: *mut metatensor::c_api::mts_block_t = std::ptr::null_mut();
assert_eq!(
unsafe { rkr_frame_metatensor_velocities_block(handle, &mut vel) },
RKRStatus::RKR_STATUS_SUCCESS
);
assert_mts_block_shape(vel, 2, 3);
unsafe { rkr_mts_block_free(vel) };
let mut frc: *mut metatensor::c_api::mts_block_t = std::ptr::null_mut();
assert_eq!(
unsafe { rkr_frame_metatensor_forces_block(handle, &mut frc) },
RKRStatus::RKR_STATUS_SUCCESS
);
assert_mts_block_shape(frc, 2, 3);
unsafe { rkr_mts_block_free(frc) };
let mut eng: *mut metatensor::c_api::mts_block_t = std::ptr::null_mut();
assert_eq!(
unsafe { rkr_frame_metatensor_atom_energies_block(handle, &mut eng) },
RKRStatus::RKR_STATUS_SUCCESS
);
assert_mts_block_shape(eng, 2, 1);
unsafe { rkr_mts_block_free(eng) };
unsafe { free_rkr_frame(handle) };
}
#[test]
fn string_iterator_yields_frames_from_buffer() {
let text = std::fs::read_to_string("resources/test/tiny_cuh2.con")
.expect("fixture tiny_cuh2.con");
let c_text = CString::new(text.as_str()).unwrap();
let it = unsafe { read_con_string_iterator(c_text.as_ptr()) };
assert!(!it.is_null());
let mut n = 0usize;
loop {
let fr = unsafe { con_frame_iterator_next(it) };
if fr.is_null() {
break;
}
n += 1;
unsafe { free_rkr_frame(fr) };
}
unsafe { free_con_frame_iterator(it) };
assert!(n >= 1, "string iterator should yield >=1 frame");
let bytes = text.as_bytes();
let it2 = unsafe { read_con_buffer_iterator(bytes.as_ptr(), bytes.len()) };
assert!(!it2.is_null());
let fr2 = unsafe { con_frame_iterator_next(it2) };
assert!(!fr2.is_null());
unsafe {
free_rkr_frame(fr2);
free_con_frame_iterator(it2);
}
}
#[test]
fn file_iterator_reads_gzip_when_present() {
use flate2::write::GzEncoder;
use flate2::Compression;
use std::io::Write;
let plain = std::fs::read("resources/test/tiny_cuh2.con").expect("fixture");
let dir = tempfile::tempdir().expect("tempdir");
let gz_path = dir.path().join("tiny_cuh2.con.gz");
{
let f = std::fs::File::create(&gz_path).unwrap();
let mut enc = GzEncoder::new(f, Compression::default());
enc.write_all(&plain).unwrap();
enc.finish().unwrap();
}
let c_path = CString::new(gz_path.to_str().unwrap()).unwrap();
let it = unsafe { read_con_file_iterator(c_path.as_ptr()) };
assert!(
!it.is_null(),
"path iterator must decompress .con.gz transparently"
);
let fr = unsafe { con_frame_iterator_next(it) };
assert!(!fr.is_null());
let n = unsafe { rkr_frame_atom_count(fr) };
assert!(n > 0);
unsafe {
free_rkr_frame(fr);
free_con_frame_iterator(it);
}
}
}