From 5fc8f085b8ad755facf609e2d67008dc6c05343a Mon Sep 17 00:00:00 2001 From: ijl Date: Tue, 10 Oct 2023 14:24:09 +0000 Subject: [PATCH] Support prefetcher on dict iteration --- src/ffi/dict.rs | 67 -------------------- src/ffi/mod.rs | 2 - src/serialize/per_type/dataclass.rs | 60 +++++++++++++++--- src/serialize/per_type/dict.rs | 95 ++++++++++++++++++++++++----- 4 files changed, 131 insertions(+), 93 deletions(-) delete mode 100644 src/ffi/dict.rs diff --git a/src/ffi/dict.rs b/src/ffi/dict.rs deleted file mode 100644 index 21cc3a70..00000000 --- a/src/ffi/dict.rs +++ /dev/null @@ -1,67 +0,0 @@ -// SPDX-License-Identifier: (Apache-2.0 OR MIT) - -use std::ptr::NonNull; - -pub struct PyDictIter { - dict_ptr: *mut pyo3_ffi::PyObject, - pos: isize, -} - -impl PyDictIter { - #[inline] - pub fn from_pyobject(obj: *mut pyo3_ffi::PyObject) -> Self { - unsafe { - PyDictIter { - dict_ptr: obj, - pos: 0, - } - } - } -} - -impl Iterator for PyDictIter { - type Item = (NonNull, NonNull); - - #[cfg(Py_3_13)] - #[inline] - fn next(&mut self) -> Option { - let mut key: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); - let mut value: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); - unsafe { - if pyo3_ffi::PyDict_Next(self.dict_ptr, &mut self.pos, &mut key, &mut value) == 1 { - Some((nonnull!(key), nonnull!(value))) - } else { - None - } - } - } - - #[cfg(not(Py_3_13))] - #[inline] - fn next(&mut self) -> Option { - let mut key: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); - let mut value: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); - unsafe { - if pyo3_ffi::_PyDict_Next( - self.dict_ptr, - &mut self.pos, - &mut key, - &mut value, - std::ptr::null_mut(), - ) == 1 - { - Some((nonnull!(key), nonnull!(value))) - } else { - None - } - } - } - - fn size_hint(&self) -> (usize, Option) { - let len = ffi!(Py_SIZE(self.dict_ptr)) as usize; - (len, Some(len)) - } -} - -#[cfg(feature = "trusted_len")] -unsafe impl std::iter::TrustedLen for PyDictIter {} diff --git a/src/ffi/mod.rs b/src/ffi/mod.rs index ab9c4192..695323e6 100644 --- a/src/ffi/mod.rs +++ b/src/ffi/mod.rs @@ -2,7 +2,6 @@ mod buffer; mod bytes; -mod dict; mod fragment; #[cfg(Py_3_12)] mod immortal; @@ -13,7 +12,6 @@ pub mod yyjson; pub use buffer::*; pub use bytes::*; -pub use dict::*; pub use fragment::{orjson_fragmenttype_new, Fragment}; #[cfg(Py_3_12)] pub use immortal::_Py_IsImmortal; diff --git a/src/serialize/per_type/dataclass.rs b/src/serialize/per_type/dataclass.rs index aa805c4e..6541d398 100644 --- a/src/serialize/per_type/dataclass.rs +++ b/src/serialize/per_type/dataclass.rs @@ -6,7 +6,6 @@ use crate::serialize::serializer::*; use crate::str::*; use crate::typeref::*; -use crate::ffi::PyDictIter; use serde::ser::{Serialize, SerializeMap, Serializer}; use std::ptr::NonNull; @@ -120,11 +119,32 @@ impl Serialize for DataclassFastSerializer { return serializer.serialize_map(Some(0)).unwrap().end(); } let mut map = serializer.serialize_map(None).unwrap(); - for (key, value) in PyDictIter::from_pyobject(self.ptr) { - if unlikely!(unsafe { ob_type!(key.as_ptr()) != STR_TYPE }) { + let mut next_key: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); + let mut next_value: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); + + let mut pos = 0; + + ffi!(PyDict_Next( + self.ptr, + &mut pos, + &mut next_key, + &mut next_value + )); + for _ in 0..=ffi!(Py_SIZE(self.ptr)) as usize - 1 { + let key = next_key; + let value = next_value; + + ffi!(PyDict_Next( + self.ptr, + &mut pos, + &mut next_key, + &mut next_value + )); + + if unlikely!(unsafe { ob_type!(key) != STR_TYPE }) { err!(SerializeError::KeyMustBeStr) } - let data = unicode_to_str(key.as_ptr()); + let data = unicode_to_str(key); if unlikely!(data.is_none()) { err!(SerializeError::InvalidStr) } @@ -133,7 +153,7 @@ impl Serialize for DataclassFastSerializer { continue; } let pyvalue = PyObjectSerializer::new( - value.as_ptr(), + value, self.opts, self.default_calls, self.recursion, @@ -185,14 +205,36 @@ impl Serialize for DataclassFallbackSerializer { return serializer.serialize_map(Some(0)).unwrap().end(); } let mut map = serializer.serialize_map(None).unwrap(); - for (attr, field) in PyDictIter::from_pyobject(fields) { - let field_type = ffi!(PyObject_GetAttr(field.as_ptr(), FIELD_TYPE_STR)); + + let mut next_key: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); + let mut next_value: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); + + let mut pos = 0; + + ffi!(PyDict_Next( + fields, + &mut pos, + &mut next_key, + &mut next_value + )); + for _ in 0..=ffi!(Py_SIZE(fields)) as usize - 1 { + let attr = next_key; + let field = next_value; + + ffi!(PyDict_Next( + fields, + &mut pos, + &mut next_key, + &mut next_value + )); + + let field_type = ffi!(PyObject_GetAttr(field, FIELD_TYPE_STR)); debug_assert!(ffi!(Py_REFCNT(field_type)) >= 2); ffi!(Py_DECREF(field_type)); if unsafe { field_type as *mut pyo3_ffi::PyTypeObject != FIELD_TYPE } { continue; } - let data = unicode_to_str(attr.as_ptr()); + let data = unicode_to_str(attr); if unlikely!(data.is_none()) { err!(SerializeError::InvalidStr); } @@ -201,7 +243,7 @@ impl Serialize for DataclassFallbackSerializer { continue; } - let value = ffi!(PyObject_GetAttr(self.ptr, attr.as_ptr())); + let value = ffi!(PyObject_GetAttr(self.ptr, attr)); debug_assert!(ffi!(Py_REFCNT(value)) >= 2); ffi!(Py_DECREF(value)); let pyvalue = PyObjectSerializer::new( diff --git a/src/serialize/per_type/dict.rs b/src/serialize/per_type/dict.rs index c5d9448b..a20eaf06 100644 --- a/src/serialize/per_type/dict.rs +++ b/src/serialize/per_type/dict.rs @@ -1,6 +1,5 @@ // SPDX-License-Identifier: (Apache-2.0 OR MIT) -use crate::ffi::PyDictIter; use crate::opt::*; use crate::serialize::error::*; use crate::serialize::per_type::datetimelike::{DateTimeBuffer, DateTimeLike}; @@ -114,16 +113,38 @@ impl Serialize for Dict { S: Serializer, { let mut map = serializer.serialize_map(None).unwrap(); - for (key, value) in PyDictIter::from_pyobject(self.ptr) { - if unlikely!(unsafe { ob_type!(key.as_ptr()) != STR_TYPE }) { + + let mut next_key: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); + let mut next_value: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); + + let mut pos = 0; + + ffi!(PyDict_Next( + self.ptr, + &mut pos, + &mut next_key, + &mut next_value + )); + for _ in 0..=ffi!(Py_SIZE(self.ptr)) as usize - 1 { + let key = next_key; + let value = next_value; + + ffi!(PyDict_Next( + self.ptr, + &mut pos, + &mut next_key, + &mut next_value + )); + + if unlikely!(unsafe { ob_type!(key) != STR_TYPE }) { err!(SerializeError::KeyMustBeStr) } - let key_as_str = unicode_to_str(key.as_ptr()); + let key_as_str = unicode_to_str(key); if unlikely!(key_as_str.is_none()) { err!(SerializeError::InvalidStr) } let pyvalue = PyObjectSerializer::new( - value.as_ptr(), + value, self.opts, self.default_calls, self.recursion, @@ -171,15 +192,37 @@ impl Serialize for DictSortedKey { let len = ffi!(Py_SIZE(self.ptr)) as usize; let mut items: SmallVec<[(&str, *mut pyo3_ffi::PyObject); 8]> = SmallVec::with_capacity(len); - for (key, value) in PyDictIter::from_pyobject(self.ptr) { - if unlikely!(unsafe { ob_type!(key.as_ptr()) != STR_TYPE }) { + + let mut next_key: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); + let mut next_value: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); + + let mut pos = 0; + + ffi!(PyDict_Next( + self.ptr, + &mut pos, + &mut next_key, + &mut next_value + )); + for _ in 0..=len as usize - 1 { + let key = next_key; + let value = next_value; + + ffi!(PyDict_Next( + self.ptr, + &mut pos, + &mut next_key, + &mut next_value + )); + + if unlikely!(unsafe { ob_type!(key) != STR_TYPE }) { err!(SerializeError::KeyMustBeStr) } - let data = unicode_to_str(key.as_ptr()); + let data = unicode_to_str(key); if unlikely!(data.is_none()) { err!(SerializeError::InvalidStr) } - items.push((data.unwrap(), value.as_ptr())); + items.push((data.unwrap(), value)); } items.sort_unstable_by(|a, b| a.0.cmp(b.0)); @@ -337,16 +380,38 @@ impl Serialize for DictNonStrKey { let mut items: SmallVec<[(CompactString, *mut pyo3_ffi::PyObject); 8]> = SmallVec::with_capacity(len); let opts = self.opts & NOT_PASSTHROUGH; - for (key, value) in PyDictIter::from_pyobject(self.ptr) { - if is_type!(ob_type!(key.as_ptr()), STR_TYPE) { - let uni = unicode_to_str(key.as_ptr()); + + let mut next_key: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); + let mut next_value: *mut pyo3_ffi::PyObject = std::ptr::null_mut(); + + let mut pos = 0; + + ffi!(PyDict_Next( + self.ptr, + &mut pos, + &mut next_key, + &mut next_value + )); + for _ in 0..=ffi!(Py_SIZE(self.ptr)) as usize - 1 { + let key = next_key; + let value = next_value; + + ffi!(PyDict_Next( + self.ptr, + &mut pos, + &mut next_key, + &mut next_value + )); + + if is_type!(ob_type!(key), STR_TYPE) { + let uni = unicode_to_str(key); if unlikely!(uni.is_none()) { err!(SerializeError::InvalidStr) } - items.push((CompactString::from(uni.unwrap()), value.as_ptr())); + items.push((CompactString::from(uni.unwrap()), value)); } else { - match Self::pyobject_to_string(key.as_ptr(), opts) { - Ok(key_as_str) => items.push((key_as_str, value.as_ptr())), + match Self::pyobject_to_string(key, opts) { + Ok(key_as_str) => items.push((key_as_str, value)), Err(err) => err!(err), } }