Skip to content

Commit 2895c81

Browse files
qingshi163youknowone
qingshi163
authored and
youknowone
committed
use heaptypeext for number protocol
1 parent 44daeef commit 2895c81

File tree

2 files changed

+83
-44
lines changed

2 files changed

+83
-44
lines changed

‎vm/src/protocol/number.rs‎

Lines changed: 7 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,9 @@ pub struct PyNumberMethods {
6363
}
6464

6565
impl PyNumberMethods {
66+
/// this is NOT a global variable
67+
// TODO: weak order read for performance
68+
#[allow(clippy::declare_interior_mutable_const)]
6669
pub const NOT_IMPLEMENTED: PyNumberMethods = PyNumberMethods {
6770
add: AtomicCell::new(None),
6871
subtract: AtomicCell::new(None),
@@ -101,28 +104,6 @@ impl PyNumberMethods {
101104
matrix_multiply: AtomicCell::new(None),
102105
inplace_matrix_multiply: AtomicCell::new(None),
103106
};
104-
105-
fn int(num: &PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyInt>> {
106-
let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __int__), ())?;
107-
ret.downcast::<PyInt>().map_err(|obj| {
108-
vm.new_type_error(format!("__int__ returned non-int (type {})", obj.class()))
109-
})
110-
}
111-
fn float(num: &PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyFloat>> {
112-
let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __float__), ())?;
113-
ret.downcast::<PyFloat>().map_err(|obj| {
114-
vm.new_type_error(format!(
115-
"__float__ returned non-float (type {})",
116-
obj.class()
117-
))
118-
})
119-
}
120-
fn index(num: &PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyInt>> {
121-
let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __index__), ())?;
122-
ret.downcast::<PyInt>().map_err(|obj| {
123-
vm.new_type_error(format!("__index__ returned non-int (type {})", obj.class()))
124-
})
125-
}
126107
}
127108

128109
pub struct PyNumber<'a> {
@@ -142,18 +123,19 @@ impl<'a> From<&'a PyObject> for PyNumber<'a> {
142123

143124
impl PyNumber<'_> {
144125
pub fn methods(&self) -> &PyNumberMethods {
126+
static GLOBAL_NOT_IMPLEMENTED: PyNumberMethods = PyNumberMethods::NOT_IMPLEMENTED;
145127
let as_number = self.methods.get_or_init(|| {
146-
Self::find_methods(self.obj).unwrap_or(NonNull::from(&PyNumberMethods::NOT_IMPLEMENTED))
128+
Self::find_methods(self.obj).unwrap_or_else(|| NonNull::from(&GLOBAL_NOT_IMPLEMENTED))
147129
});
148130
unsafe { as_number.as_ref() }
149131
}
150132

151-
fn find_methods<'a>(obj: &'a PyObject) -> Option<NonNull<PyNumberMethods>> {
133+
fn find_methods(obj: &PyObject) -> Option<NonNull<PyNumberMethods>> {
152134
obj.class().mro_find_map(|x| x.slots.as_number.load())
153135
}
154136

155137
// PyNumber_Check
156-
pub fn check<'a>(obj: &'a PyObject, vm: &VirtualMachine) -> bool {
138+
pub fn check(obj: &PyObject) -> bool {
157139
let num = PyNumber::from(obj);
158140
let methods = num.methods();
159141
methods.int.load().is_some()

‎vm/src/types/slot.rs‎

Lines changed: 76 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::common::{hash::PyHash, lock::PyRwLock};
22
use crate::{
3-
builtins::{PyInt, PyStrInterned, PyStrRef, PyType, PyTypeRef},
3+
builtins::{PyFloat, PyInt, PyStrInterned, PyStrRef, PyType, PyTypeRef},
44
bytecode::ComparisonOperator,
55
convert::ToPyResult,
66
function::Either,
@@ -205,6 +205,30 @@ fn slot_as_sequence(zelf: &PyObject, vm: &VirtualMachine) -> &'static PySequence
205205
PySequenceMethods::generic(has_length, has_ass_item)
206206
}
207207

208+
fn int_wrapper(num: &PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyInt>> {
209+
let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __int__), ())?;
210+
ret.downcast::<PyInt>().map_err(|obj| {
211+
vm.new_type_error(format!("__int__ returned non-int (type {})", obj.class()))
212+
})
213+
}
214+
215+
fn index_wrapper(num: &PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyInt>> {
216+
let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __index__), ())?;
217+
ret.downcast::<PyInt>().map_err(|obj| {
218+
vm.new_type_error(format!("__index__ returned non-int (type {})", obj.class()))
219+
})
220+
}
221+
222+
fn float_wrapper(num: &PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyFloat>> {
223+
let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __float__), ())?;
224+
ret.downcast::<PyFloat>().map_err(|obj| {
225+
vm.new_type_error(format!(
226+
"__float__ returned non-float (type {})",
227+
obj.class()
228+
))
229+
})
230+
}
231+
208232
fn hash_wrapper(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<PyHash> {
209233
let hash_obj = vm.call_special_method(zelf.to_owned(), identifier!(vm, __hash__), ())?;
210234
match hash_obj.payload_if_subclass::<PyInt>(vm) {
@@ -318,21 +342,37 @@ impl PyType {
318342
debug_assert!(name.as_str().starts_with("__"));
319343
debug_assert!(name.as_str().ends_with("__"));
320344

321-
macro_rules! update_slot {
345+
macro_rules! toggle_slot {
322346
($name:ident, $func:expr) => {{
323347
self.slots.$name.store(if add { Some($func) } else { None });
324348
}};
325349
}
350+
351+
macro_rules! update_slot {
352+
($name:ident, $func:expr) => {{
353+
self.slots.$name.store(Some($func));
354+
}};
355+
}
356+
357+
macro_rules! update_pointer_slot {
358+
($name:ident, $pointed:ident) => {{
359+
self.slots.$name.store(
360+
self.heaptype_ext
361+
.as_ref()
362+
.map(|ext| NonNull::from(&ext.$pointed)),
363+
);
364+
}};
365+
}
326366
match name.as_str() {
327367
"__len__" | "__getitem__" | "__setitem__" | "__delitem__" => {
328368
update_slot!(as_mapping, slot_as_mapping);
329369
update_slot!(as_sequence, slot_as_sequence);
330370
}
331371
"__hash__" => {
332-
update_slot!(hash, hash_wrapper);
372+
toggle_slot!(hash, hash_wrapper);
333373
}
334374
"__call__" => {
335-
update_slot!(call, call_wrapper);
375+
toggle_slot!(call, call_wrapper);
336376
}
337377
"__getattr__" | "__getattribute__" => {
338378
update_slot!(getattro, getattro_wrapper);
@@ -344,28 +384,52 @@ impl PyType {
344384
update_slot!(richcompare, richcompare_wrapper);
345385
}
346386
"__iter__" => {
347-
update_slot!(iter, iter_wrapper);
387+
toggle_slot!(iter, iter_wrapper);
348388
}
349389
"__next__" => {
350-
update_slot!(iternext, iternext_wrapper);
390+
toggle_slot!(iternext, iternext_wrapper);
351391
}
352392
"__get__" => {
353-
update_slot!(descr_get, descr_get_wrapper);
393+
toggle_slot!(descr_get, descr_get_wrapper);
354394
}
355395
"__set__" | "__delete__" => {
356396
update_slot!(descr_set, descr_set_wrapper);
357397
}
358398
"__init__" => {
359-
update_slot!(init, init_wrapper);
399+
toggle_slot!(init, init_wrapper);
360400
}
361401
"__new__" => {
362-
update_slot!(new, new_wrapper);
402+
toggle_slot!(new, new_wrapper);
363403
}
364404
"__del__" => {
365-
update_slot!(del, del_wrapper);
405+
toggle_slot!(del, del_wrapper);
406+
}
407+
"__int__" => {
408+
self.heaptype_ext
409+
.as_ref()
410+
.unwrap()
411+
.number_methods
412+
.int
413+
.store(Some(int_wrapper));
414+
update_pointer_slot!(as_number, number_methods);
366415
}
367-
"__int__" | "__index__" | "__float__" => {
368-
// update_slot!(as_number, slot_as_number);
416+
"__index__" => {
417+
self.heaptype_ext
418+
.as_ref()
419+
.unwrap()
420+
.number_methods
421+
.index
422+
.store(Some(index_wrapper));
423+
update_pointer_slot!(as_number, number_methods);
424+
}
425+
"__float__" => {
426+
self.heaptype_ext
427+
.as_ref()
428+
.unwrap()
429+
.number_methods
430+
.float
431+
.store(Some(float_wrapper));
432+
update_pointer_slot!(as_number, number_methods);
369433
}
370434
_ => {}
371435
}
@@ -871,15 +935,8 @@ pub trait AsSequence: PyPayload {
871935

872936
#[pyimpl]
873937
pub trait AsNumber: PyPayload {
874-
// const AS_NUMBER: PyNumberMethods;
875-
876938
#[pyslot]
877939
fn as_number() -> &'static PyNumberMethods;
878-
// #[inline]
879-
// #[pyslot]
880-
// fn as_number() -> &'static PyNumberMethods {
881-
// &Self::AS_NUMBER
882-
// }
883940

884941
fn number_downcast<'a>(number: &'a PyNumber) -> &'a Py<Self> {
885942
unsafe { number.obj.downcast_unchecked_ref() }

0 commit comments

Comments
 (0)