PyLongObject small value inline

ijl · Feb 3, 2024 · 41018c8 · 41018c8
1 parent 323a930
commit 41018c8
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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
@@ -329,6 +329,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