metatensor/io/
mod.rs

1//! Input/Output facilities for storing [`crate::TensorMap`] and
2//! [`crate::Labels`] on disk
3
4use std::os::raw::c_void;
5
6use crate::c_api::{MTS_SUCCESS, mts_array_t, mts_status_t};
7use crate::Array;
8
9mod tensor;
10pub use self::tensor::{load, save, load_buffer, save_buffer};
11
12mod block;
13pub use self::block::{load_block, load_block_buffer, save_block, save_block_buffer};
14
15mod labels;
16pub use self::labels::{load_labels, load_labels_buffer, save_labels, save_labels_buffer};
17
18
19/// Implementation of realloc for `Vec<u8>`, used in `save_buffer`
20unsafe extern "C" fn realloc_vec(user_data: *mut c_void, _ptr: *mut u8, new_size: usize) -> *mut u8 {
21    let mut result = std::ptr::null_mut();
22    let unwind_wrapper = std::panic::AssertUnwindSafe(&mut result);
23
24    let status = crate::errors::catch_unwind(move || {
25        let vector = &mut *user_data.cast::<Vec<u8>>();
26        vector.resize(new_size, 0);
27
28        // force the closure to capture the full unwind_wrapper, not just
29        // unwind_wrapper.0
30        let _ = &unwind_wrapper;
31        *(unwind_wrapper.0) = vector.as_mut_ptr();
32    });
33
34    if status != MTS_SUCCESS {
35        return std::ptr::null_mut();
36    }
37
38    return result;
39}
40
41/// callback used to create `ndarray::ArrayD` when loading a `TensorMap`
42unsafe extern "C" fn create_ndarray(
43    shape_ptr: *const usize,
44    shape_count: usize,
45    c_array: *mut mts_array_t,
46) -> mts_status_t {
47    crate::errors::catch_unwind(|| {
48        assert!(shape_count != 0);
49        let shape = std::slice::from_raw_parts(shape_ptr, shape_count);
50        let array = ndarray::ArrayD::from_elem(shape, 0.0);
51        *c_array = (Box::new(array) as Box<dyn Array>).into();
52    })
53}