|
13 | 13 | // See the License for the specific language governing permissions and |
14 | 14 | // limitations under the License. |
15 | 15 |
|
16 | | -mod native_conversion; |
17 | | -mod swap_endian; |
| 16 | +use pyo3::exceptions::PyValueError; |
| 17 | +use pyo3::prelude::*; |
| 18 | +use pyo3::types::{PyBytes, PyInt}; |
| 19 | +use pyo3::{pyfunction, Bound, PyErr, PyResult}; |
18 | 20 |
|
19 | 21 | use crate::register_package; |
20 | | -use pyo3::prelude::*; |
| 22 | + |
| 23 | +#[pyfunction] |
| 24 | +fn swap_endian<'py>( |
| 25 | + type_size: Bound<'py, PyInt>, |
| 26 | + data: Bound<'py, PyBytes>, |
| 27 | +) -> PyResult<Bound<'py, PyBytes>> { |
| 28 | + let py = type_size.py(); |
| 29 | + |
| 30 | + let type_size: usize = match type_size.extract::<usize>() { |
| 31 | + Ok(type_size @ 2) | Ok(type_size @ 4) | Ok(type_size @ 8) => type_size, |
| 32 | + _ => { |
| 33 | + return Err(PyErr::new::<PyValueError, _>(format!( |
| 34 | + "Unsupported type size {type_size}", |
| 35 | + ))) |
| 36 | + } |
| 37 | + }; |
| 38 | + let bytes = &data.as_bytes(); |
| 39 | + let len = bytes.len(); |
| 40 | + if len % type_size != 0 { |
| 41 | + return Err(PyErr::new::<PyValueError, _>( |
| 42 | + "Data length not a multiple of type_size", |
| 43 | + )); |
| 44 | + } |
| 45 | + |
| 46 | + PyBytes::new_with(py, bytes.len(), |out| { |
| 47 | + match type_size { |
| 48 | + 2 => swap_n::<2>(bytes, out), |
| 49 | + 4 => swap_n::<4>(bytes, out), |
| 50 | + 8 => swap_n::<8>(bytes, out), |
| 51 | + _ => unreachable!(), |
| 52 | + } |
| 53 | + Ok(()) |
| 54 | + }) |
| 55 | +} |
| 56 | + |
| 57 | +#[inline(always)] |
| 58 | +fn swap_n<const N: usize>(src: &[u8], dst: &mut [u8]) { |
| 59 | + // Doesn't technically need to be a function with a const generic, but this |
| 60 | + // allows the compiler to optimize the code better. |
| 61 | + assert_eq!(src.len(), dst.len()); |
| 62 | + assert_eq!(src.len() % N, 0); |
| 63 | + for i in (0..src.len()).step_by(N) { |
| 64 | + for j in 0..N { |
| 65 | + dst[i + j] = src[i + N - j - 1]; |
| 66 | + } |
| 67 | + } |
| 68 | +} |
21 | 69 |
|
22 | 70 | pub(super) fn init_module(m: &Bound<PyModule>, name: &str) -> PyResult<()> { |
23 | 71 | m.gil_used(false)?; |
24 | 72 | register_package(m, name)?; |
25 | 73 |
|
26 | | - m.add_function(wrap_pyfunction!(swap_endian::swap_endian, m)?)?; |
27 | | - m.add_function(wrap_pyfunction!(native_conversion::vec_f64_from_native, m)?)?; |
28 | | - m.add_function(wrap_pyfunction!(native_conversion::vec_f64_to_native, m)?)?; |
29 | | - m.add_function(wrap_pyfunction!(native_conversion::vec_f32_from_native, m)?)?; |
30 | | - m.add_function(wrap_pyfunction!(native_conversion::vec_f32_to_native, m)?)?; |
31 | | - m.add_function(wrap_pyfunction!(native_conversion::vec_i64_from_native, m)?)?; |
32 | | - m.add_function(wrap_pyfunction!(native_conversion::vec_i64_to_native, m)?)?; |
33 | | - m.add_function(wrap_pyfunction!(native_conversion::vec_i32_from_native, m)?)?; |
34 | | - m.add_function(wrap_pyfunction!(native_conversion::vec_i32_to_native, m)?)?; |
35 | | - m.add_function(wrap_pyfunction!(native_conversion::vec_i16_from_native, m)?)?; |
36 | | - m.add_function(wrap_pyfunction!(native_conversion::vec_i16_to_native, m)?)?; |
37 | | - m.add_function(wrap_pyfunction!(native_conversion::vec_i8_from_native, m)?)?; |
38 | | - m.add_function(wrap_pyfunction!(native_conversion::vec_i8_to_native, m)?)?; |
| 74 | + m.add_function(wrap_pyfunction!(swap_endian, m)?)?; |
39 | 75 |
|
40 | 76 | Ok(()) |
41 | 77 | } |
0 commit comments