Skip to content

Commit ecd6fe1

Browse files
bschoenmaeckersbschoenmaeckers
committed
Flatten Py_slot_subslots and Py_tp_slots in PySlot iterator
1 parent 398dd70 commit ecd6fe1

2 files changed

Lines changed: 114 additions & 98 deletions

File tree

‎crates/capi/src/object/pytype.rs‎

Lines changed: 34 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,10 @@ use core::ffi::{c_char, c_int, c_uint, c_ulong, c_void};
99
use core::ptr;
1010
use rustpython_vm::builtins::{PyStr, PyType};
1111
use rustpython_vm::function::{ItemDoc, PyMethodFlags};
12-
use rustpython_vm::types::{AtomicPyTypeFlags, PyTypeFlags, PyTypeSlots, SlotAccessor};
12+
use rustpython_vm::types::{
13+
AtomicPyTypeFlags, PyTypeFlags, PyTypeSlots,
14+
SlotAccessor::{TpBase, TpDealloc, TpDoc, TpGetset, TpMembers, TpMethods, TpNew, TpTraverse},
15+
};
1316
use rustpython_vm::{AsObject, Py, PyObject};
1417

1518
pub type PyTypeObject = Py<PyType>;
@@ -32,25 +35,6 @@ pub struct PyType_Spec {
3235
pub slots: *mut PyType_Slot,
3336
}
3437

35-
impl PyType_Slot {
36-
pub(crate) fn iter<'a>(mut slots: *const Self) -> impl Iterator<Item = &'a Self> {
37-
core::iter::from_fn(move || {
38-
let slot = unsafe { &*slots };
39-
if slot.slot == 0 {
40-
None
41-
} else {
42-
slots = unsafe { slots.add(1) };
43-
Some(slot)
44-
}
45-
})
46-
}
47-
48-
fn as_slot_accessor(&self) -> Option<SlotAccessor> {
49-
let slot_id: u8 = self.slot.try_into().ok()?;
50-
slot_id.try_into().ok()
51-
}
52-
}
53-
5438
#[unsafe(no_mangle)]
5539
pub unsafe extern "C" fn Py_TYPE(op: *mut PyObject) -> *const PyTypeObject {
5640
unsafe { op.assume_borrowed() }.class()
@@ -148,64 +132,38 @@ pub unsafe extern "C" fn PyType_FromSlots(slots: *const PySlot) -> *mut PyObject
148132
));
149133
}
150134
}
151-
PySlotType::Slots { value, .. } => {
152-
for slot in PyType_Slot::iter(value) {
153-
let accessor = slot.as_slot_accessor().ok_or_else(|| {
154-
vm.new_value_error(format!(
155-
"Invalid slot id: {} for PyType_FromSlots",
156-
slot.slot
157-
))
158-
})?;
159-
match accessor {
160-
SlotAccessor::TpDoc => {
161-
let doc = unsafe {
162-
slot.pfunc.cast::<c_char>().try_as_str_opt(vm)?
163-
}
164-
.map_or(
165-
ItemDoc::NONE,
166-
|doc| {
167-
let text: &'static str =
168-
Box::leak(doc.to_owned().into_boxed_str());
169-
ItemDoc::static_text(text)
170-
},
171-
);
172-
type_slots.doc = doc;
173-
}
174-
SlotAccessor::TpNew => {
175-
type_slots.new.store(Some(|ty, _args, vm| {
176-
Err(vm.new_not_implemented_error(format!("tp_new is not yet implemented in PyType_FromSlots for {ty:?}")))
177-
}));
178-
}
179-
SlotAccessor::TpBase => {
180-
base = unsafe { Some(&*slot.pfunc.cast::<PyTypeObject>()) }
181-
}
182-
SlotAccessor::TpDealloc => {
183-
//TODO
184-
}
185-
SlotAccessor::TpMethods => {
186-
for def in PyMethodDef::iter(slot.pfunc.cast()) {
187-
let name = unsafe { def.ml_name.try_as_str(vm)? };
188-
let is_static =
189-
PyMethodFlags::from_bits_retain(def.ml_flags as _)
190-
.contains(PyMethodFlags::STATIC);
191-
let method = build_method_def(vm, def, !is_static)?;
192-
methods.push((name, method));
193-
}
194-
}
195-
SlotAccessor::TpGetset => {
196-
getsets.extend(PyGetSetDef::iter(slot.pfunc.cast()));
197-
}
198-
SlotAccessor::TpMembers => {
199-
members.extend(PyMemberDef::iter(slot.pfunc.cast()));
200-
}
201-
slot => {
202-
return Err(vm.new_not_implemented_error(format!(
203-
"PyType_FromSlots with PyType_Slot {slot:?} not implemented yet"
204-
)));
205-
}
206-
}
135+
PySlotType::TypeSlot(TpDoc, pfunc) => {
136+
let doc = unsafe { pfunc.cast::<c_char>().try_as_str_opt(vm)? }.map_or(
137+
ItemDoc::NONE,
138+
|doc| {
139+
let text: &'static str =
140+
Box::leak(doc.to_owned().into_boxed_str());
141+
ItemDoc::static_text(text)
142+
},
143+
);
144+
type_slots.doc = doc;
145+
}
146+
PySlotType::TypeSlot(TpBase, pfunc) => {
147+
base = unsafe { Some(&*pfunc.cast::<PyTypeObject>()) }
148+
}
149+
PySlotType::TypeSlot(TpMembers, pfunc) => {
150+
members.extend(PyMemberDef::iter(pfunc.cast()));
151+
}
152+
PySlotType::TypeSlot(TpGetset, pfunc) => {
153+
getsets.extend(PyGetSetDef::iter(pfunc.cast()));
154+
}
155+
PySlotType::TypeSlot(TpMethods, pfunc) => {
156+
for def in PyMethodDef::iter(pfunc.cast()) {
157+
let name = unsafe { def.ml_name.try_as_str(vm)? };
158+
let is_static = PyMethodFlags::from_bits_retain(def.ml_flags as _)
159+
.contains(PyMethodFlags::STATIC);
160+
let method = build_method_def(vm, def, !is_static)?;
161+
methods.push((name, method));
207162
}
208163
}
164+
PySlotType::TypeSlot(TpNew | TpDealloc | TpTraverse, _) => {
165+
// TODO
166+
}
209167
_ => {
210168
return Err(vm.new_not_implemented_error(format!(
211169
"PyType_FromSlots with slot {kind:?} not implemented yet"

‎crates/capi/src/slots.rs‎

Lines changed: 80 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -3,19 +3,36 @@ use crate::methodobject::PyMethodDef;
33
use crate::modsupport::PyABIInfo;
44
use crate::object::PyType_Slot;
55
use crate::util::CStrExt;
6+
use alloc::borrow::Cow;
67
use core::ffi::{c_char, c_int, c_void};
7-
use core::slice::from_ref;
8+
use itertools::Either;
89
use rustpython_vm::builtins::{PyModule, PyType};
10+
use rustpython_vm::types::SlotAccessor;
911
use rustpython_vm::{Py, PyResult, VirtualMachine};
12+
use std::num::TryFromIntError;
1013

1114
#[repr(C)]
15+
#[derive(Clone)]
1216
pub struct PySlot {
1317
pub sl_id: u16,
1418
pub sl_flags: u16,
1519
_reserved: u32,
1620
pub value: PySlotValue,
1721
}
1822

23+
impl TryFrom<&PyType_Slot> for PySlot {
24+
type Error = TryFromIntError;
25+
26+
fn try_from(slot: &PyType_Slot) -> Result<Self, Self::Error> {
27+
Ok(Self {
28+
sl_id: slot.slot.try_into()?,
29+
sl_flags: 0,
30+
_reserved: 0,
31+
value: slot.pfunc.into(),
32+
})
33+
}
34+
}
35+
1936
impl PySlot {
2037
const SLOT_OPTIONAL: u16 = 0x0001;
2138
const SLOT_STATIC: u16 = 0x0002;
@@ -31,14 +48,51 @@ impl PySlot {
3148
}
3249
}
3350

34-
pub(crate) fn iter<'a>(mut slots: *const Self) -> impl Iterator<Item = &'a Self> {
51+
pub(crate) fn iter<'a>(slots: *const Self) -> impl Iterator<Item = Cow<'a, Self>> {
52+
let mut stack: Vec<Either<*const Self, *const PyType_Slot>> = vec![Either::Left(slots)];
3553
core::iter::from_fn(move || {
36-
let slot = unsafe { &*slots };
37-
if slot.sl_id == 0 {
38-
None
39-
} else {
40-
slots = unsafe { slots.add(1) };
41-
Some(slot)
54+
loop {
55+
return match stack.last_mut()? {
56+
Either::Left(slot_ptr) => {
57+
let slot = unsafe { &**slot_ptr };
58+
if slot.sl_id == 0 {
59+
stack.pop();
60+
continue;
61+
}
62+
63+
*slot_ptr = unsafe { slot_ptr.add(1) };
64+
65+
if slot.sl_id == 92 {
66+
let subslots = unsafe { slot.value.sl_ptr }.cast_const().cast::<Self>();
67+
if !subslots.is_null() {
68+
stack.push(Either::Left(subslots));
69+
}
70+
continue;
71+
}
72+
73+
if slot.sl_id == 93 {
74+
let type_slots = unsafe { slot.value.sl_ptr }
75+
.cast_const()
76+
.cast::<PyType_Slot>();
77+
if !type_slots.is_null() {
78+
stack.push(Either::Right(type_slots));
79+
}
80+
continue;
81+
}
82+
83+
Some(Cow::Borrowed(slot))
84+
}
85+
Either::Right(slot_ptr) => {
86+
let type_slot = unsafe { &**slot_ptr };
87+
if type_slot.slot == 0 {
88+
stack.pop();
89+
continue;
90+
}
91+
92+
*slot_ptr = unsafe { slot_ptr.add(1) };
93+
Some(Cow::Owned(type_slot.try_into().expect("invalid slot")))
94+
}
95+
};
4296
}
4397
})
4498
}
@@ -58,12 +112,15 @@ impl PySlot {
58112
self.sl_flags & Self::SLOT_INTPTR != 0
59113
}
60114

61-
pub(crate) fn as_kind(&self, vm: &VirtualMachine) -> PyResult<PySlotKind<'_>> {
115+
pub(crate) fn as_kind<'a>(&self, vm: &VirtualMachine) -> PyResult<PySlotKind<'a>> {
62116
let value_ptr = unsafe { self.value.sl_ptr };
63117
let is_static = self.is_static();
64118
let kind = match self.sl_id {
65-
48 => PySlotKind::Type(PySlotType::Base(unsafe { &*value_ptr.cast() })),
66-
49 => PySlotKind::Type(PySlotType::Bases(unsafe { from_ref(&*value_ptr.cast()) })),
119+
1..82 => {
120+
let slot_id: u8 = self.sl_id.try_into().unwrap();
121+
let accessor = slot_id.try_into().unwrap();
122+
PySlotKind::Type(PySlotType::TypeSlot(accessor, value_ptr))
123+
}
67124
84 => {
68125
let create = unsafe {
69126
core::mem::transmute::<
@@ -89,11 +146,9 @@ impl PySlot {
89146
87 => PySlotKind::Module(PySlotModule::Gil {
90147
gil_used: !value_ptr.is_null(),
91148
}),
92-
// 92 => Py_slot_subslots
93-
93 => PySlotKind::Type(PySlotType::Slots {
94-
value: value_ptr.cast(),
95-
is_static,
96-
}),
149+
92 | 93 => {
150+
unreachable!("Py_slot_subslots and Py_tp_slots should be handled in the iterator")
151+
}
97152
95 => PySlotKind::Type(PySlotType::Name(unsafe {
98153
value_ptr.cast::<c_char>().try_as_str(vm)?
99154
})),
@@ -140,6 +195,14 @@ pub union PySlotValue {
140195
pub sl_uint64: u64,
141196
}
142197

198+
impl Clone for PySlotValue {
199+
fn clone(&self) -> Self {
200+
Self {
201+
sl_ptr: unsafe { self.sl_ptr },
202+
}
203+
}
204+
}
205+
143206
impl<T> From<*mut T> for PySlotValue {
144207
fn from(ptr: *mut T) -> Self {
145208
Self { sl_ptr: ptr.cast() }
@@ -174,12 +237,7 @@ pub(crate) enum PySlotModule<'a> {
174237
#[allow(dead_code)]
175238
#[derive(Debug, Copy, Clone)]
176239
pub(crate) enum PySlotType<'a> {
177-
Base(&'a Py<PyType>),
178-
Bases(&'a [Py<PyType>]),
179-
Slots {
180-
value: *mut PyType_Slot,
181-
is_static: bool,
182-
},
240+
TypeSlot(SlotAccessor, *mut c_void),
183241
Name(&'a str),
184242
Flags(u64),
185243
BasicSize(isize),

0 commit comments

Comments
 (0)