From 7e38f0218aae1e962109ad0a18ce724b25d7c163 Mon Sep 17 00:00:00 2001 From: ijl Date: Sat, 3 Feb 2024 17:44:30 +0000 Subject: [PATCH] PyLongObject small value inline --- src/ffi/long.rs | 67 +++++++++++++++++++++++++++++++---- src/ffi/mod.rs | 2 +- src/serialize/per_type/int.rs | 8 ++--- test/test_type.py | 9 +++++ 4 files changed, 75 insertions(+), 11 deletions(-) diff --git a/src/ffi/long.rs b/src/ffi/long.rs index f3b79f69..9052a7ef 100644 --- a/src/ffi/long.rs +++ b/src/ffi/long.rs @@ -6,34 +6,89 @@ const SIGN_MASK: usize = 3; #[cfg(Py_3_12)] const SIGN_ZERO: usize = 1; + +#[cfg(Py_3_12)] +#[allow(non_upper_case_globals)] +const _PyLong_NON_SIZE_BITS: usize = 3; + #[cfg(Py_3_12)] -const SIGN_POSITIVE: usize = 0; +#[repr(C)] +struct _PyLongValue { + pub lv_tag: usize, + pub ob_digit: u32, +} #[cfg(Py_3_12)] -#[allow(dead_code)] +#[repr(C)] struct PyLongObject { pub ob_refcnt: pyo3_ffi::Py_ssize_t, pub ob_type: *mut pyo3_ffi::PyTypeObject, - pub lv_tag: usize, - pub ob_digit: u8, + pub long_value: _PyLongValue, } #[cfg(Py_3_12)] +#[inline(always)] pub fn pylong_is_zero(ptr: *mut pyo3_ffi::PyObject) -> bool { - unsafe { (*(ptr as *mut PyLongObject)).lv_tag & SIGN_MASK == SIGN_ZERO } + unsafe { (*(ptr as *mut PyLongObject)).long_value.lv_tag & SIGN_MASK == SIGN_ZERO } } #[cfg(not(Py_3_12))] +#[inline(always)] pub fn pylong_is_zero(ptr: *mut pyo3_ffi::PyObject) -> bool { unsafe { (*(ptr as *mut pyo3_ffi::PyVarObject)).ob_size == 0 } } #[cfg(Py_3_12)] +#[inline(always)] pub fn pylong_is_unsigned(ptr: *mut pyo3_ffi::PyObject) -> bool { - unsafe { (*(ptr as *mut PyLongObject)).lv_tag & SIGN_MASK == SIGN_POSITIVE } + unsafe { + 1 - (((*(ptr as *mut PyLongObject)).long_value.lv_tag & _PyLong_NON_SIZE_BITS) as isize) > 0 + } } #[cfg(not(Py_3_12))] +#[inline(always)] pub fn pylong_is_unsigned(ptr: *mut pyo3_ffi::PyObject) -> bool { unsafe { (*(ptr as *mut pyo3_ffi::PyVarObject)).ob_size > 0 } } + +#[cfg(Py_3_12)] +#[inline(always)] +fn pylong_is_compact(ptr: *mut pyo3_ffi::PyObject) -> bool { + unsafe { (*(ptr as *mut PyLongObject)).long_value.lv_tag < (2 << _PyLong_NON_SIZE_BITS) } +} + +#[cfg(Py_3_12)] +#[inline(always)] +pub fn pylong_value_unsigned(ptr: *mut pyo3_ffi::PyObject) -> u64 { + if pylong_is_compact(ptr) == true { + unsafe { (*(ptr as *mut PyLongObject)).long_value.ob_digit as u64 } + } else { + ffi!(PyLong_AsUnsignedLongLong(ptr)) + } +} + +#[cfg(not(Py_3_12))] +#[inline(always)] +pub fn pylong_value_unsigned(ptr: *mut pyo3_ffi::PyObject) -> u64 { + ffi!(PyLong_AsUnsignedLongLong(ptr)) +} + +#[cfg(not(Py_3_12))] +#[inline(always)] +pub fn pylong_value_signed(ptr: *mut pyo3_ffi::PyObject) -> i64 { + ffi!(PyLong_AsLongLong(ptr)) +} + +#[cfg(Py_3_12)] +#[inline(always)] +pub fn pylong_value_signed(ptr: *mut pyo3_ffi::PyObject) -> i64 { + if pylong_is_compact(ptr) == true { + unsafe { + let sign = 1 - ((*(ptr as *mut PyLongObject)).long_value.lv_tag & SIGN_MASK) as i64; + sign * (*(ptr as *mut PyLongObject)).long_value.ob_digit as i64 + } + } else { + ffi!(PyLong_AsLongLong(ptr)) + } +} diff --git a/src/ffi/mod.rs b/src/ffi/mod.rs index ba6dedcd..f2b8198f 100644 --- a/src/ffi/mod.rs +++ b/src/ffi/mod.rs @@ -10,4 +10,4 @@ pub mod yyjson; pub use buffer::*; pub use bytes::*; pub use fragment::{orjson_fragmenttype_new, Fragment}; -pub use long::{pylong_is_unsigned, pylong_is_zero}; +pub use long::{pylong_is_unsigned, pylong_is_zero, pylong_value_signed, pylong_value_unsigned}; diff --git a/src/serialize/per_type/int.rs b/src/serialize/per_type/int.rs index 53a90e86..dc5fe36a 100644 --- a/src/serialize/per_type/int.rs +++ b/src/serialize/per_type/int.rs @@ -1,6 +1,6 @@ // SPDX-License-Identifier: (Apache-2.0 OR MIT) -use crate::ffi::{pylong_is_unsigned, pylong_is_zero}; +use crate::ffi::{pylong_is_unsigned, pylong_is_zero, pylong_value_signed, pylong_value_unsigned}; use crate::serialize::error::SerializeError; use serde::ser::{Serialize, Serializer}; @@ -28,14 +28,14 @@ impl Serialize for IntSerializer { if pylong_is_zero(self.ptr) { serializer.serialize_u64(0) } else if pylong_is_unsigned(self.ptr) { - let val = ffi!(PyLong_AsUnsignedLongLong(self.ptr)); + let val = pylong_value_unsigned(self.ptr); if unlikely!(val == u64::MAX) && !ffi!(PyErr_Occurred()).is_null() { err!(SerializeError::Integer64Bits) } else { serializer.serialize_u64(val) } } else { - let val = ffi!(PyLong_AsLongLong(self.ptr)); + let val = pylong_value_signed(self.ptr); if unlikely!(val == -1) && !ffi!(PyErr_Occurred()).is_null() { err!(SerializeError::Integer64Bits) } @@ -62,7 +62,7 @@ impl Serialize for Int53Serializer { where S: Serializer, { - let val = ffi!(PyLong_AsLongLong(self.ptr)); + let val = pylong_value_signed(self.ptr); if unlikely!(val == -1) { if ffi!(PyErr_Occurred()).is_null() { serializer.serialize_i64(val) diff --git a/test/test_type.py b/test/test_type.py index 5eb9a954..a2969ca8 100644 --- a/test/test_type.py +++ b/test/test_type.py @@ -333,6 +333,15 @@ def test_none(self): assert orjson.dumps(obj) == ref.encode("utf-8") assert orjson.loads(ref) == obj + def test_int(self): + """ + int compact and non-compact + """ + obj = [-5000, -1000, -10, -5, -2, -1, 0, 1, 2, 5, 10, 1000, 50000] + ref = b"[-5000,-1000,-10,-5,-2,-1,0,1,2,5,10,1000,50000]" + assert orjson.dumps(obj) == ref + assert orjson.loads(ref) == obj + def test_null_array(self): """ null array