Skip to content

Commit

Permalink
Support prefetcher on dict iteration
Browse files Browse the repository at this point in the history
  • Loading branch information
@ijl
ijl committed Oct 10, 2023
1 parent ce9bae8 commit 5fc8f08
Show file tree
Hide file tree
Showing 4 changed files with 131 additions and 93 deletions.
67 changes: 0 additions & 67 deletions src/ffi/dict.rs
View file

This file was deleted.

2 changes: 0 additions & 2 deletions src/ffi/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@

mod buffer;
mod bytes;
mod dict;
mod fragment;
#[cfg(Py_3_12)]
mod immortal;
Expand All @@ -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;
Expand Down
60 changes: 51 additions & 9 deletions src/serialize/per_type/dataclass.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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)
}
Expand All @@ -133,7 +153,7 @@ impl Serialize for DataclassFastSerializer {
continue;
}
let pyvalue = PyObjectSerializer::new(
value.as_ptr(),
value,
self.opts,
self.default_calls,
self.recursion,
Expand Down Expand Up @@ -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);
}
Expand All @@ -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(
Expand Down
95 changes: 80 additions & 15 deletions src/serialize/per_type/dict.rs
View file
Original file line number Diff line number Diff line change
@@ -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};
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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),
}
}
Expand Down

0 comments on commit 5fc8f08

Please sign in to comment.