Skip to content

Commit 0c9d9a5

Browse files
youknowoneyouknowone
committed
Use static ref
1 parent 2b43c92 commit 0c9d9a5

File tree

6 files changed

+95
-77
lines changed

6 files changed

+95
-77
lines changed

‎vm/src/builtins/float.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,11 +59,11 @@ impl From<f64> for PyFloat {
5959

6060
impl PyObject {
6161
pub fn try_float_opt(&self, vm: &VirtualMachine) -> PyResult<Option<PyRef<PyFloat>>> {
62-
PyNumber::from(self).float_opt(vm)
62+
PyNumber::new(self, vm).float_opt(vm)
6363
}
6464

6565
pub fn try_float(&self, vm: &VirtualMachine) -> PyResult<PyRef<PyFloat>> {
66-
PyNumber::from(self).float(vm)
66+
PyNumber::new(self, vm).float(vm)
6767
}
6868
}
6969

‎vm/src/builtins/int.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -265,7 +265,7 @@ impl Constructor for PyInt {
265265
val
266266
};
267267

268-
PyNumber::from(val.as_ref())
268+
PyNumber::new(val.as_ref(), vm)
269269
.int(vm)
270270
.map(|x| x.as_bigint().clone())
271271
}

‎vm/src/function/number.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ impl Deref for ArgIntoFloat {
8282
impl TryFromObject for ArgIntoFloat {
8383
// Equivalent to PyFloat_AsDouble.
8484
fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
85-
let value = PyNumber::from(obj.as_ref()).float(vm)?.to_f64();
85+
let value = PyNumber::new(obj.as_ref(), vm).float(vm)?.to_f64();
8686
Ok(ArgIntoFloat { value })
8787
}
8888
}

‎vm/src/protocol/number.rs

Lines changed: 25 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,5 @@
1-
use std::borrow::Cow;
2-
31
use crate::{
42
builtins::{int, PyByteArray, PyBytes, PyComplex, PyFloat, PyInt, PyIntRef, PyStr},
5-
common::lock::OnceCell,
63
function::ArgBytesLike,
74
stdlib::warnings,
85
AsObject, PyObject, PyPayload, PyRef, PyResult, TryFromBorrowedObject, VirtualMachine,
@@ -100,45 +97,42 @@ impl PyNumberMethods {
10097
pub struct PyNumber<'a> {
10198
pub obj: &'a PyObject,
10299
// some fast path do not need methods, so we do lazy initialize
103-
methods: OnceCell<Cow<'static, PyNumberMethods>>,
100+
pub methods: Option<&'static PyNumberMethods>,
104101
}
105102

106-
impl<'a> From<&'a PyObject> for PyNumber<'a> {
107-
fn from(obj: &'a PyObject) -> Self {
103+
impl<'a> PyNumber<'a> {
104+
pub fn new(obj: &'a PyObject, vm: &VirtualMachine) -> Self {
108105
Self {
109106
obj,
110-
methods: OnceCell::new(),
107+
methods: Self::find_methods(obj, vm),
111108
}
112109
}
113110
}
114111

115112
impl PyNumber<'_> {
116-
pub fn methods(&self, vm: &VirtualMachine) -> &PyNumberMethods {
117-
&*self.methods_cow(vm)
113+
pub fn find_methods(obj: &PyObject, vm: &VirtualMachine) -> Option<&'static PyNumberMethods> {
114+
obj.class()
115+
.mro_find_map(|x| x.slots.as_number.load())
116+
.map(|f| f(obj, vm))
118117
}
119118

120-
pub fn methods_cow(&self, vm: &VirtualMachine) -> &Cow<'static, PyNumberMethods> {
121-
self.methods.get_or_init(|| {
122-
self.obj
123-
.class()
124-
.mro_find_map(|x| x.slots.as_number.load())
125-
.map(|f| f(self.obj, vm))
126-
.unwrap_or_else(|| Cow::Borrowed(&PyNumberMethods::NOT_IMPLEMENTED))
127-
})
119+
pub fn methods(&self) -> &'static PyNumberMethods {
120+
self.methods.unwrap_or(&PyNumberMethods::NOT_IMPLEMENTED)
128121
}
129122

130123
// PyNumber_Check
131-
pub fn check(&self, vm: &VirtualMachine) -> bool {
132-
let methods = self.methods(vm);
133-
methods.int.is_some()
134-
|| methods.index.is_some()
135-
|| methods.float.is_some()
136-
|| self.obj.payload_is::<PyComplex>()
124+
pub fn check(obj: &PyObject, vm: &VirtualMachine) -> bool {
125+
Self::find_methods(obj, vm).map_or(false, |methods| {
126+
methods.int.is_some()
127+
|| methods.index.is_some()
128+
|| methods.float.is_some()
129+
|| obj.payload_is::<PyComplex>()
130+
})
137131
}
138132

139133
// PyIndex_Check
140-
pub fn is_index(&self, vm: &VirtualMachine) -> bool {
141-
self.methods(vm).index.is_some()
134+
pub fn is_index(&self) -> bool {
135+
self.methods().index.is_some()
142136
}
143137

144138
pub fn int(&self, vm: &VirtualMachine) -> PyResult<PyIntRef> {
@@ -156,7 +150,7 @@ impl PyNumber<'_> {
156150

157151
if let Some(i) = self.obj.downcast_ref_if_exact::<PyInt>(vm) {
158152
Ok(i.to_owned())
159-
} else if let Some(f) = self.methods(vm).int {
153+
} else if let Some(f) = self.methods().int {
160154
let ret = f(self, vm)?;
161155
if !ret.class().is(PyInt::class(vm)) {
162156
warnings::warn(
@@ -174,7 +168,7 @@ impl PyNumber<'_> {
174168
} else {
175169
Ok(ret)
176170
}
177-
} else if self.methods(vm).index.is_some() {
171+
} else if self.methods().index.is_some() {
178172
self.index(vm)
179173
} else if let Ok(Ok(f)) =
180174
vm.get_special_method(self.obj.to_owned(), identifier!(vm, __trunc__))
@@ -187,7 +181,7 @@ impl PyNumber<'_> {
187181
// vm,
188182
// )?;
189183
let ret = f.invoke((), vm)?;
190-
PyNumber::from(ret.as_ref()).index(vm).map_err(|_| {
184+
PyNumber::new(ret.as_ref(), vm).index(vm).map_err(|_| {
191185
vm.new_type_error(format!(
192186
"__trunc__ returned non-Integral (type {})",
193187
ret.class()
@@ -215,7 +209,7 @@ impl PyNumber<'_> {
215209
Ok(Some(i.to_owned()))
216210
} else if let Some(i) = self.obj.payload::<PyInt>() {
217211
Ok(Some(vm.ctx.new_bigint(i.as_bigint())))
218-
} else if let Some(f) = self.methods(vm).index {
212+
} else if let Some(f) = self.methods().index {
219213
let ret = f(self, vm)?;
220214
if !ret.class().is(PyInt::class(vm)) {
221215
warnings::warn(
@@ -250,7 +244,7 @@ impl PyNumber<'_> {
250244
pub fn float_opt(&self, vm: &VirtualMachine) -> PyResult<Option<PyRef<PyFloat>>> {
251245
if let Some(float) = self.obj.downcast_ref_if_exact::<PyFloat>(vm) {
252246
Ok(Some(float.to_owned()))
253-
} else if let Some(f) = self.methods(vm).float {
247+
} else if let Some(f) = self.methods().float {
254248
let ret = f(self, vm)?;
255249
if !ret.class().is(PyFloat::class(vm)) {
256250
warnings::warn(
@@ -268,7 +262,7 @@ impl PyNumber<'_> {
268262
} else {
269263
Ok(Some(ret))
270264
}
271-
} else if self.methods(vm).index.is_some() {
265+
} else if self.methods().index.is_some() {
272266
let i = self.index(vm)?;
273267
let value = int::try_to_float(i.as_bigint(), vm)?;
274268
Ok(Some(vm.ctx.new_float(value)))

‎vm/src/types/slot.rs

Lines changed: 62 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ impl Default for PyTypeFlags {
139139

140140
pub(crate) type GenericMethod = fn(&PyObject, FuncArgs, &VirtualMachine) -> PyResult;
141141
pub(crate) type AsMappingFunc = fn(&PyObject, &VirtualMachine) -> &'static PyMappingMethods;
142-
pub(crate) type AsNumberFunc = fn(&PyObject, &VirtualMachine) -> Cow<'static, PyNumberMethods>;
142+
pub(crate) type AsNumberFunc = fn(&PyObject, &VirtualMachine) -> &'static PyNumberMethods;
143143
pub(crate) type HashFunc = fn(&PyObject, &VirtualMachine) -> PyResult<PyHash>;
144144
// CallFunc = GenericMethod
145145
pub(crate) type GetattroFunc = fn(&PyObject, PyStrRef, &VirtualMachine) -> PyResult;
@@ -340,43 +340,65 @@ fn as_sequence_generic(zelf: &PyObject, vm: &VirtualMachine) -> &'static PySeque
340340
static_as_sequence_generic(has_length, has_ass_item)
341341
}
342342

343-
fn as_number_wrapper(zelf: &PyObject, vm: &VirtualMachine) -> Cow<'static, PyNumberMethods> {
344-
Cow::Owned(PyNumberMethods {
345-
int: then_some_closure!(
346-
zelf.class().has_attr(identifier!(vm, __int__)),
347-
|num, vm| {
348-
let ret =
349-
vm.call_special_method(num.obj.to_owned(), identifier!(vm, __int__), ())?;
350-
ret.downcast::<PyInt>().map_err(|obj| {
351-
vm.new_type_error(format!("__int__ returned non-int (type {})", obj.class()))
352-
})
353-
}
354-
),
355-
float: then_some_closure!(
356-
zelf.class().has_attr(identifier!(vm, __float__)),
357-
|num, vm| {
358-
let ret =
359-
vm.call_special_method(num.obj.to_owned(), identifier!(vm, __float__), ())?;
360-
ret.downcast::<PyFloat>().map_err(|obj| {
361-
vm.new_type_error(format!(
362-
"__float__ returned non-float (type {})",
363-
obj.class()
364-
))
365-
})
366-
}
367-
),
368-
index: then_some_closure!(
369-
zelf.class().has_attr(identifier!(vm, __index__)),
370-
|num, vm| {
371-
let ret =
372-
vm.call_special_method(num.obj.to_owned(), identifier!(vm, __index__), ())?;
373-
ret.downcast::<PyInt>().map_err(|obj| {
374-
vm.new_type_error(format!("__index__ returned non-int (type {})", obj.class()))
375-
})
376-
}
377-
),
378-
..PyNumberMethods::NOT_IMPLEMENTED
379-
})
343+
pub(crate) fn static_as_number_generic(
344+
has_int: bool,
345+
has_float: bool,
346+
has_index: bool,
347+
) -> &'static PyNumberMethods {
348+
static METHODS: &[PyNumberMethods] = &[
349+
new_generic(false, false, false),
350+
new_generic(true, false, false),
351+
new_generic(false, true, false),
352+
new_generic(true, true, false),
353+
new_generic(false, false, true),
354+
new_generic(true, false, true),
355+
new_generic(false, true, true),
356+
new_generic(true, true, true),
357+
];
358+
359+
fn int(num: &PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyInt>> {
360+
let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __int__), ())?;
361+
ret.downcast::<PyInt>().map_err(|obj| {
362+
vm.new_type_error(format!("__int__ returned non-int (type {})", obj.class()))
363+
})
364+
}
365+
fn float(num: &PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyFloat>> {
366+
let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __float__), ())?;
367+
ret.downcast::<PyFloat>().map_err(|obj| {
368+
vm.new_type_error(format!(
369+
"__float__ returned non-float (type {})",
370+
obj.class()
371+
))
372+
})
373+
}
374+
fn index(num: &PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyInt>> {
375+
let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __index__), ())?;
376+
ret.downcast::<PyInt>().map_err(|obj| {
377+
vm.new_type_error(format!("__index__ returned non-int (type {})", obj.class()))
378+
})
379+
}
380+
381+
const fn new_generic(has_int: bool, has_float: bool, has_index: bool) -> PyNumberMethods {
382+
PyNumberMethods {
383+
int: if has_int { Some(int) } else { None },
384+
float: if has_float { Some(float) } else { None },
385+
index: if has_index { Some(index) } else { None },
386+
..PyNumberMethods::NOT_IMPLEMENTED
387+
}
388+
}
389+
390+
let key = bool_int(has_int) | (bool_int(has_float) << 1) | (bool_int(has_index) << 2);
391+
392+
&METHODS[key]
393+
}
394+
395+
fn as_number_wrapper(zelf: &PyObject, vm: &VirtualMachine) -> &'static PyNumberMethods {
396+
let (has_int, has_float, has_index) = (
397+
zelf.class().has_attr(identifier!(vm, __int__)),
398+
zelf.class().has_attr(identifier!(vm, __float__)),
399+
zelf.class().has_attr(identifier!(vm, __index__)),
400+
);
401+
static_as_number_generic(has_int, has_float, has_index)
380402
}
381403

382404
fn hash_wrapper(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<PyHash> {
@@ -1041,8 +1063,8 @@ pub trait AsNumber: PyPayload {
10411063

10421064
#[inline]
10431065
#[pyslot]
1044-
fn as_number(_zelf: &PyObject, _vm: &VirtualMachine) -> Cow<'static, PyNumberMethods> {
1045-
Cow::Borrowed(&Self::AS_NUMBER)
1066+
fn as_number(_zelf: &PyObject, _vm: &VirtualMachine) -> &'static PyNumberMethods {
1067+
&Self::AS_NUMBER
10461068
}
10471069

10481070
fn number_downcast<'a>(number: &'a PyNumber) -> &'a Py<Self> {

‎vm/src/vm/vm_ops.rs

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,13 @@ use crate::{
1010
/// Collection of operators
1111
impl VirtualMachine {
1212
pub fn to_index_opt(&self, obj: PyObjectRef) -> Option<PyResult<PyIntRef>> {
13-
PyNumber::from(obj.as_ref()).index_opt(self).transpose()
13+
PyNumber::new(obj.as_ref(), self)
14+
.index_opt(self)
15+
.transpose()
1416
}
1517

1618
pub fn to_index(&self, obj: &PyObject) -> PyResult<PyIntRef> {
17-
PyNumber::from(obj).index(self)
19+
PyNumber::new(obj, self).index(self)
1820
}
1921

2022
#[inline]

0 commit comments

Comments
 (0)