diff --git a/crates/capi/src/longobject.rs b/crates/capi/src/longobject.rs index d88523a7d..c853c5351 100644 --- a/crates/capi/src/longobject.rs +++ b/crates/capi/src/longobject.rs @@ -2,7 +2,7 @@ use crate::object::define_py_check; use crate::{PyObject, pystate::with_vm}; use bitflags::bitflags; use core::ffi::{CStr, c_char, c_double, c_int, c_long, c_longlong, c_ulong, c_ulonglong, c_void}; -use malachite_bigint::{BigInt, Sign}; +use malachite_bigint::{BigInt, BigUint, Sign}; use rustpython_vm::builtins::{PyInt, try_bigint_to_f64, try_f64_to_bigint}; use rustpython_vm::common::int::bytes_to_int; use rustpython_vm::protocol::handle_bytes_to_int_err; @@ -79,12 +79,19 @@ bitflags! { } impl AsNativeBytesFlags { + const ENDIAN_MASK: c_int = 0b11; + + #[inline] + fn is_native_endian(self) -> bool { + (self.bits() & Self::ENDIAN_MASK) == Self::NATIVE_ENDIAN.bits() + } + #[inline] fn is_little_endian(self) -> bool { - if self.contains(Self::NATIVE_ENDIAN) { + if self.is_native_endian() { cfg!(target_endian = "little") } else { - self.contains(Self::LITTLE_ENDIAN) + (self.bits() & Self::ENDIAN_MASK) == Self::LITTLE_ENDIAN.bits() } } @@ -96,8 +103,8 @@ impl AsNativeBytesFlags { let flags = Self::from_bits(raw_flags) .ok_or_else(|| vm.new_value_error("Invalid NativeBytes flags"))?; - if flags.contains(Self::LITTLE_ENDIAN) & flags.contains(Self::NATIVE_ENDIAN) { - Err(vm.new_value_error("Cannot specify both LITTLE_ENDIAN and NATIVE_ENDIAN")) + if (flags.bits() & Self::ENDIAN_MASK) == 0b10 { + Err(vm.new_value_error("Passing 2 is reserved")) } else { Ok(flags) } @@ -157,6 +164,71 @@ pub unsafe extern "C" fn PyLong_FromUnsignedNativeBytes( }) } +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyLong_AsNativeBytes( + obj: *mut PyObject, + buffer: *mut c_void, + n_bytes: isize, + flags: c_int, +) -> isize { + with_vm::, _>(|vm| { + if n_bytes < 0 { + return Err(vm.new_value_error("n_bytes must not be negative")); + } + + let flags = AsNativeBytesFlags::from_bits_or_default(vm, flags)?; + let value = if flags.contains(AsNativeBytesFlags::ALLOW_INDEX) { + unsafe { &*obj }.to_owned().try_index(vm)? + } else { + unsafe { &*obj }.try_downcast_ref::(vm)?.to_owned() + }; + let bigint = value.as_bigint(); + + if flags.contains(AsNativeBytesFlags::REJECT_NEGATIVE) && bigint.sign() == Sign::Minus { + return Err(vm.new_value_error("cannot convert negative int to unsigned")); + } + + let little_endian = flags.is_little_endian(); + let unsigned_buffer = flags.contains(AsNativeBytesFlags::UNSIGNED_BUFFER); + let is_negative = bigint.sign() == Sign::Minus; + + let bytes = if is_negative || !unsigned_buffer { + if little_endian { + bigint.to_signed_bytes_le() + } else { + bigint.to_signed_bytes_be() + } + } else if little_endian { + bigint.to_bytes_le().1 + } else { + bigint.to_bytes_be().1 + }; + + let required_size = bytes.len().max(1); + let out_len = + usize::try_from(n_bytes).map_err(|_| vm.new_overflow_error("n_bytes out of range"))?; + + if out_len > 0 { + if buffer.is_null() { + return Err(vm.new_system_error("buffer must not be NULL when n_bytes > 0")); + } + let out = unsafe { core::slice::from_raw_parts_mut(buffer.cast::(), out_len) }; + let fill = if is_negative { 0xff } else { 0x00 }; + out.fill(fill); + + let copy_len = out_len.min(bytes.len()); + if little_endian { + out[..copy_len].copy_from_slice(&bytes[..copy_len]); + } else { + out[out_len - copy_len..].copy_from_slice(&bytes[bytes.len() - copy_len..]); + } + } + + isize::try_from(required_size) + .map_err(|_| vm.new_overflow_error("required size too large for Py_ssize_t")) + }) +} + #[unsafe(no_mangle)] pub extern "C" fn PyLong_FromVoidPtr(ptr: *mut c_void) -> *mut PyObject { with_vm(|vm| vm.ctx.new_int(ptr as usize)) @@ -372,9 +444,193 @@ pub unsafe extern "C" fn PyLong_AsUnsignedLongLong(obj: *mut PyObject) -> c_ulon }) } +#[repr(C)] +pub struct PyLongLayout { + pub bits_per_digit: u8, + pub digit_size: u8, + pub digits_order: i8, + pub digit_endianness: i8, +} + +#[repr(C)] +#[derive(Default)] +pub struct PyLongExport { + pub value: i64, + pub negative: u8, + pub ndigits: isize, + pub digits: *const c_void, + _reserved: *mut Vec, +} + +pub struct PyLongWriter { + negative: bool, + digits: Vec, +} + +#[unsafe(no_mangle)] +pub extern "C" fn PyLong_GetNativeLayout() -> *const PyLongLayout { + const NATIVE_LONG_LAYOUT: PyLongLayout = PyLongLayout { + bits_per_digit: 32, + digit_size: 4, + digits_order: -1, + digit_endianness: if cfg!(target_endian = "little") { + -1 + } else { + 1 + }, + }; + &NATIVE_LONG_LAYOUT +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyLong_Export( + obj: *mut PyObject, + export_long: *mut PyLongExport, +) -> c_int { + with_vm::, _>(|vm| { + let py_int = unsafe { &*obj }.try_downcast_ref::(vm)?; + let bigint = py_int.as_bigint(); + + if let Ok(value) = i64::try_from(bigint) { + unsafe { + *export_long = PyLongExport { + value, + ..Default::default() + }; + } + return Ok(()); + } + + let (sign, digits) = bigint.to_u32_digits(); + let boxed_digits = Box::new(digits); + let ndigits = boxed_digits.len().try_into().map_err(|_| { + vm.new_overflow_error("PyLong_Export: too many digits to fit into Py_ssize_t") + })?; + + unsafe { + *export_long = PyLongExport { + value: 0, + negative: u8::from(matches!(sign, Sign::Minus)), + ndigits, + digits: boxed_digits.as_ptr().cast(), + _reserved: Box::into_raw(boxed_digits), + }; + } + Ok(()) + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyLong_FreeExport(export_long: *mut PyLongExport) { + if export_long.is_null() { + return; + } + + let export_long = unsafe { &mut *export_long }; + if !export_long._reserved.is_null() { + unsafe { + drop(Box::from_raw(export_long._reserved)); + } + } + core::mem::take(export_long); +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyLongWriter_Create( + negative: c_int, + ndigits: isize, + digits: *mut *mut c_void, +) -> *mut PyLongWriter { + with_vm::, _>(|vm| { + unsafe { + *digits = core::ptr::null_mut(); + } + + if ndigits <= 0 { + return Err(vm.new_value_error("ndigits must be greater than 0")); + } + if digits.is_null() { + return Err(vm.new_system_error("digits must not be null")); + } + if negative != 0 && negative != 1 { + return Err(vm.new_value_error("negative must be 0 or 1")); + } + + let ndigits = ndigits + .try_into() + .map_err(|_| vm.new_overflow_error("ndigits out of range"))?; + + let max_ndigits = (isize::MAX as usize) / core::mem::size_of::(); + if ndigits > max_ndigits { + return Err(vm.new_overflow_error("too many digits in integer")); + } + + let mut digits_vec = Vec::new(); + digits_vec + .try_reserve_exact(ndigits) + .map_err(|_| vm.new_memory_error("could not allocate digit storage"))?; + digits_vec.resize(ndigits, 0); + + let mut writer = Box::new(PyLongWriter { + negative: negative == 1, + digits: digits_vec, + }); + + unsafe { + *digits = writer.digits.as_mut_ptr().cast(); + } + + Ok(Box::into_raw(writer).cast()) + }) + .cast() +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyLongWriter_Finish(writer: *mut PyLongWriter) -> *mut PyObject { + with_vm(|vm| { + if writer.is_null() { + return Err(vm.new_system_error("PyLongWriter_Finish: writer must not be null")); + } + + let writer = unsafe { Box::from_raw(writer) }; + let mut digits = writer.digits; + while matches!(digits.last(), Some(0)) { + digits.pop(); + } + + if digits.is_empty() { + return Ok(vm.ctx.new_int(0)); + } + + let magnitude = BigUint::new(digits); + let sign = if writer.negative { + Sign::Minus + } else { + Sign::Plus + }; + + let value = BigInt::from_biguint(sign, magnitude); + Ok(vm.ctx.new_int(value)) + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyLongWriter_Discard(writer: *mut PyLongWriter) { + if writer.is_null() { + return; + } + + unsafe { + drop(Box::from_raw(writer)); + } +} + #[cfg(test)] mod tests { + use super::*; + use pyo3::exceptions::PyOverflowError; use pyo3::prelude::*; + use pyo3::types::PyBool; use pyo3::types::PyInt; #[test] @@ -403,4 +659,64 @@ mod tests { assert_eq!(number.extract::().unwrap(), value); }) } + + #[test] + fn as_native_bytes_defaults_unsigned_size_and_copy() { + Python::attach(|py| unsafe { + let obj: Bound<'_, PyInt> = 128.into_pyobject(py).unwrap(); + + let required = PyLong_AsNativeBytes(obj.as_ptr().cast(), core::ptr::null_mut(), 0, -1); + assert_eq!(required, 1); + + let mut out = [0xAAu8; 1]; + let written = PyLong_AsNativeBytes(obj.as_ptr().cast(), out.as_mut_ptr().cast(), 1, -1); + assert_eq!(written, 1); + assert_eq!(out, [0x80]); + }) + } + + #[test] + fn as_native_bytes_sign_extends_when_signed() { + Python::attach(|py| unsafe { + let obj: Bound<'_, PyInt> = (-1).into_pyobject(py).unwrap(); + + let mut out = [0u8; 4]; + let written = PyLong_AsNativeBytes(obj.as_ptr().cast(), out.as_mut_ptr().cast(), 4, 0); + assert_eq!(written, 1); + assert_eq!(out, [0xFF, 0xFF, 0xFF, 0xFF]); + }) + } + + #[test] + fn as_native_bytes_reject_negative() { + Python::attach(|py| unsafe { + let obj: Bound<'_, PyInt> = (-7).into_pyobject(py).unwrap(); + + let rc = PyLong_AsNativeBytes(obj.as_ptr().cast(), core::ptr::null_mut(), 0, 8); + assert_eq!(rc, -1); + assert!(PyErr::take(py).is_some()); + }) + } + + #[test] + fn as_native_bytes_allow_index() { + Python::attach(|py| unsafe { + let instance = PyBool::new(py, true); + let mut out = [0u8; 2]; + let rc = PyLong_AsNativeBytes(instance.as_ptr().cast(), out.as_mut_ptr().cast(), 2, 17); + assert_eq!(rc, 1); + assert_eq!(out, [1, 0]); + }) + } + + #[test] + fn longwriter_create_ndigits_capacity_overflow() { + Python::attach(|py| unsafe { + let mut digits = core::ptr::null_mut(); + let writer = PyLongWriter_Create(0, isize::MAX, &mut digits); + assert!(writer.is_null()); + let err = PyErr::take(py).expect("expected an error"); + assert!(err.is_instance_of::(py)); + }) + } }