Add long export/writer to c-api (#8161)

* Add long export/writer c-api

* Add `PyLong_AsNativeBytes`

* Fix PyLongWriter_Create errors

* Set PyLongWriter_Create digits to null on err

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::<PyResult<isize>, _>(|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::<PyInt>(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::<u8>(), 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<u32>,
+}
+
+pub struct PyLongWriter {
+    negative: bool,
+    digits: Vec<u32>,
+}
+
+#[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::<PyResult<()>, _>(|vm| {
+        let py_int = unsafe { &*obj }.try_downcast_ref::<PyInt>(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::<PyResult<*mut c_void>, _>(|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::<u32>();
+        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::<u128>().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::<PyOverflowError>(py));
+        })
+    }
 }