Skip to content

Commit 273ffd9

Browse files
qingshi163qingshi163
authored
Optimize PyType Creation (RustPython#4209)
1 parent 186990f commit 273ffd9

File tree

7 files changed

+114
-54
lines changed

7 files changed

+114
-54
lines changed

‎stdlib/src/ssl.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -228,7 +228,7 @@ mod _ssl {
228228
/// SSL/TLS connection terminated abruptly.
229229
#[pyattr(name = "SSLEOFError", once)]
230230
fn ssl_eof_error(vm: &VirtualMachine) -> PyTypeRef {
231-
PyType::new_simple_ref("ssl.SSLEOFError", &ssl_error(vm)).unwrap()
231+
PyType::new_simple_ref("ssl.SSLEOFError", &ssl_error(vm), &vm.ctx).unwrap()
232232
}
233233

234234
type OpensslVersionInfo = (u8, u8, u8, u8, u8);

‎vm/src/builtins/type.rs‎

Lines changed: 74 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -121,13 +121,18 @@ impl PyPayload for PyType {
121121
}
122122

123123
impl PyType {
124-
pub fn new_simple_ref(name: &str, base: &PyTypeRef) -> Result<PyRef<Self>, String> {
124+
pub fn new_simple_ref(
125+
name: &str,
126+
base: &PyTypeRef,
127+
ctx: &Context,
128+
) -> Result<PyRef<Self>, String> {
125129
Self::new_ref(
126130
name,
127131
vec![base.clone()],
128132
Default::default(),
129133
Default::default(),
130134
Self::static_type().to_owned(),
135+
ctx,
131136
)
132137
}
133138
pub fn new_ref(
@@ -136,6 +141,7 @@ impl PyType {
136141
attrs: PyAttributes,
137142
slots: PyTypeSlots,
138143
metaclass: PyRef<Self>,
144+
ctx: &Context,
139145
) -> Result<PyRef<Self>, String> {
140146
Self::new_verbose_ref(
141147
name,
@@ -145,8 +151,11 @@ impl PyType {
145151
slots,
146152
HeapTypeExt::default(),
147153
metaclass,
154+
ctx,
148155
)
149156
}
157+
158+
#[allow(clippy::too_many_arguments)]
150159
fn new_verbose_ref(
151160
name: &str,
152161
base: PyRef<Self>,
@@ -155,6 +164,7 @@ impl PyType {
155164
mut slots: PyTypeSlots,
156165
heaptype_ext: HeapTypeExt,
157166
metaclass: PyRef<Self>,
167+
ctx: &Context,
158168
) -> Result<PyRef<Self>, String> {
159169
// Check for duplicates in bases.
160170
let mut unique_bases = HashSet::new();
@@ -177,17 +187,22 @@ impl PyType {
177187
*slots.name.get_mut() = Some(String::from(name));
178188

179189
#[allow(clippy::mutable_key_type)]
180-
let mut attr_name_set = HashSet::new();
190+
let mut slot_name_set = HashSet::new();
181191

182192
for cls in mro.iter() {
183193
for &name in cls.attributes.read().keys() {
184-
if name.as_str() != "__new__" {
185-
attr_name_set.insert(name);
194+
if name != identifier!(ctx, __new__)
195+
&& name.as_str().starts_with("__")
196+
&& name.as_str().ends_with("__")
197+
{
198+
slot_name_set.insert(name);
186199
}
187200
}
188201
}
189202
for &name in attrs.keys() {
190-
attr_name_set.insert(name);
203+
if name.as_str().starts_with("__") && name.as_str().ends_with("__") {
204+
slot_name_set.insert(name);
205+
}
191206
}
192207

193208
let new_type = PyRef::new_ref(
@@ -204,11 +219,53 @@ impl PyType {
204219
None,
205220
);
206221

207-
for attr_name in attr_name_set {
208-
if attr_name.as_str().starts_with("__") && attr_name.as_str().ends_with("__") {
209-
new_type.update_slot(attr_name, true);
210-
}
222+
for attr_name in slot_name_set {
223+
new_type.update_slot::<true>(attr_name, ctx);
211224
}
225+
226+
let weakref_type = super::PyWeak::static_type();
227+
for base in &new_type.bases {
228+
base.subclasses.write().push(
229+
new_type
230+
.as_object()
231+
.downgrade_with_weakref_typ_opt(None, weakref_type.to_owned())
232+
.unwrap(),
233+
);
234+
}
235+
236+
Ok(new_type)
237+
}
238+
239+
pub fn new_bare_ref(
240+
name: &str,
241+
base: PyRef<Self>,
242+
attrs: PyAttributes,
243+
mut slots: PyTypeSlots,
244+
metaclass: PyRef<Self>,
245+
) -> Result<PyRef<Self>, String> {
246+
if base.slots.flags.has_feature(PyTypeFlags::HAS_DICT) {
247+
slots.flags |= PyTypeFlags::HAS_DICT
248+
}
249+
250+
*slots.name.get_mut() = Some(String::from(name));
251+
252+
let bases = vec![base.clone()];
253+
let mro = base.iter_mro().cloned().collect();
254+
255+
let new_type = PyRef::new_ref(
256+
PyType {
257+
base: Some(base),
258+
bases,
259+
mro,
260+
subclasses: PyRwLock::default(),
261+
attributes: PyRwLock::new(attrs),
262+
slots,
263+
heaptype_ext: None,
264+
},
265+
metaclass,
266+
None,
267+
);
268+
212269
let weakref_type = super::PyWeak::static_type();
213270
for base in &new_type.bases {
214271
base.subclasses.write().push(
@@ -693,6 +750,7 @@ impl PyType {
693750
slots,
694751
heaptype_ext,
695752
metatype,
753+
&vm.ctx,
696754
)
697755
.map_err(|e| vm.new_type_error(e))?;
698756

@@ -923,7 +981,11 @@ impl SetAttr for PyType {
923981
}
924982
}
925983
if attr_name.as_str().starts_with("__") && attr_name.as_str().ends_with("__") {
926-
zelf.update_slot(attr_name, assign);
984+
if assign {
985+
zelf.update_slot::<true>(attr_name, &vm.ctx);
986+
} else {
987+
zelf.update_slot::<false>(attr_name, &vm.ctx);
988+
}
927989
}
928990
Ok(())
929991
}
@@ -1177,6 +1239,7 @@ mod tests {
11771239
PyAttributes::default(),
11781240
Default::default(),
11791241
type_type.clone(),
1242+
&context,
11801243
)
11811244
.unwrap();
11821245
let b = PyType::new_ref(
@@ -1185,6 +1248,7 @@ mod tests {
11851248
PyAttributes::default(),
11861249
Default::default(),
11871250
type_type,
1251+
&context,
11881252
)
11891253
.unwrap();
11901254

‎vm/src/class.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,9 @@ pub trait StaticType {
4343
where
4444
Self: PyClassImpl,
4545
{
46-
PyType::new_ref(
46+
PyType::new_bare_ref(
4747
Self::NAME,
48-
vec![Self::static_baseclass().to_owned()],
48+
Self::static_baseclass().to_owned(),
4949
Default::default(),
5050
Self::make_slots(),
5151
Self::static_metaclass().to_owned(),

‎vm/src/stdlib/io.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3669,6 +3669,7 @@ mod _io {
36693669
Default::default(),
36703670
Default::default(),
36713671
ctx.types.type_type.to_owned(),
3672+
ctx,
36723673
)
36733674
.unwrap()
36743675
}

‎vm/src/types/slot.rs‎

Lines changed: 33 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -356,13 +356,13 @@ fn del_wrapper(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
356356
}
357357

358358
impl PyType {
359-
pub(crate) fn update_slot(&self, name: &'static PyStrInterned, add: bool) {
359+
pub(crate) fn update_slot<const ADD: bool>(&self, name: &'static PyStrInterned, ctx: &Context) {
360360
debug_assert!(name.as_str().starts_with("__"));
361361
debug_assert!(name.as_str().ends_with("__"));
362362

363363
macro_rules! toggle_slot {
364364
($name:ident, $func:expr) => {{
365-
self.slots.$name.store(if add { Some($func) } else { None });
365+
self.slots.$name.store(if ADD { Some($func) } else { None });
366366
}};
367367
}
368368

@@ -382,95 +382,88 @@ impl PyType {
382382

383383
macro_rules! toggle_ext_func {
384384
($n1:ident, $n2:ident, $func:expr) => {{
385-
self.heaptype_ext.as_ref().unwrap().$n1.$n2.store(if add {
385+
self.heaptype_ext.as_ref().unwrap().$n1.$n2.store(if ADD {
386386
Some($func)
387387
} else {
388388
None
389389
});
390390
}};
391391
}
392392

393-
match name.as_str() {
394-
"__len__" => {
393+
match name {
394+
_ if name == identifier!(ctx, __len__) => {
395395
update_slot!(as_mapping, slot_as_mapping);
396396
toggle_ext_func!(sequence_methods, length, |seq, vm| len_wrapper(seq.obj, vm));
397397
update_pointer_slot!(as_sequence, sequence_methods);
398398
}
399-
"__getitem__" => {
399+
_ if name == identifier!(ctx, __getitem__) => {
400400
update_slot!(as_mapping, slot_as_mapping);
401401
toggle_ext_func!(sequence_methods, item, |seq, i, vm| getitem_wrapper(
402402
seq.obj, i, vm
403403
));
404404
update_pointer_slot!(as_sequence, sequence_methods);
405405
}
406-
"__setitem__" | "__delitem__" => {
406+
_ if name == identifier!(ctx, __setitem__) || name == identifier!(ctx, __delitem__) => {
407407
update_slot!(as_mapping, slot_as_mapping);
408408
toggle_ext_func!(sequence_methods, ass_item, |seq, i, value, vm| {
409409
setitem_wrapper(seq.obj, i, value, vm)
410410
});
411411
update_pointer_slot!(as_sequence, sequence_methods);
412412
}
413-
"__hash__" => {
413+
_ if name == identifier!(ctx, __hash__) => {
414414
toggle_slot!(hash, hash_wrapper);
415415
}
416-
"__call__" => {
416+
_ if name == identifier!(ctx, __call__) => {
417417
toggle_slot!(call, call_wrapper);
418418
}
419-
"__getattr__" | "__getattribute__" => {
419+
_ if name == identifier!(ctx, __getattr__)
420+
|| name == identifier!(ctx, __getattribute__) =>
421+
{
420422
update_slot!(getattro, getattro_wrapper);
421423
}
422-
"__setattr__" | "__delattr__" => {
424+
_ if name == identifier!(ctx, __setattr__) || name == identifier!(ctx, __delattr__) => {
423425
update_slot!(setattro, setattro_wrapper);
424426
}
425-
"__eq__" | "__ne__" | "__le__" | "__lt__" | "__ge__" | "__gt__" => {
427+
_ if name == identifier!(ctx, __eq__)
428+
|| name == identifier!(ctx, __ne__)
429+
|| name == identifier!(ctx, __le__)
430+
|| name == identifier!(ctx, __lt__)
431+
|| name == identifier!(ctx, __ge__)
432+
|| name == identifier!(ctx, __gt__) =>
433+
{
426434
update_slot!(richcompare, richcompare_wrapper);
427435
}
428-
"__iter__" => {
436+
_ if name == identifier!(ctx, __iter__) => {
429437
toggle_slot!(iter, iter_wrapper);
430438
}
431-
"__next__" => {
439+
_ if name == identifier!(ctx, __next__) => {
432440
toggle_slot!(iternext, iternext_wrapper);
433441
}
434-
"__get__" => {
442+
_ if name == identifier!(ctx, __get__) => {
435443
toggle_slot!(descr_get, descr_get_wrapper);
436444
}
437-
"__set__" | "__delete__" => {
445+
_ if name == identifier!(ctx, __set__) || name == identifier!(ctx, __delete__) => {
438446
update_slot!(descr_set, descr_set_wrapper);
439447
}
440-
"__init__" => {
448+
_ if name == identifier!(ctx, __init__) => {
441449
toggle_slot!(init, init_wrapper);
442450
}
443-
"__new__" => {
451+
_ if name == identifier!(ctx, __new__) => {
444452
toggle_slot!(new, new_wrapper);
445453
}
446-
"__del__" => {
454+
_ if name == identifier!(ctx, __del__) => {
447455
toggle_slot!(del, del_wrapper);
448456
}
449-
"__int__" => {
450-
self.heaptype_ext
451-
.as_ref()
452-
.unwrap()
453-
.number_methods
454-
.int
455-
.store(Some(int_wrapper));
457+
_ if name == identifier!(ctx, __int__) => {
458+
toggle_ext_func!(number_methods, int, int_wrapper);
456459
update_pointer_slot!(as_number, number_methods);
457460
}
458-
"__index__" => {
459-
self.heaptype_ext
460-
.as_ref()
461-
.unwrap()
462-
.number_methods
463-
.index
464-
.store(Some(index_wrapper));
461+
_ if name == identifier!(ctx, __index__) => {
462+
toggle_ext_func!(number_methods, index, index_wrapper);
465463
update_pointer_slot!(as_number, number_methods);
466464
}
467-
"__float__" => {
468-
self.heaptype_ext
469-
.as_ref()
470-
.unwrap()
471-
.number_methods
472-
.float
473-
.store(Some(float_wrapper));
465+
_ if name == identifier!(ctx, __float__) => {
466+
toggle_ext_func!(number_methods, float, float_wrapper);
474467
update_pointer_slot!(as_number, number_methods);
475468
}
476469
_ => {}

‎vm/src/vm/context.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -422,6 +422,7 @@ impl Context {
422422
attrs,
423423
slots,
424424
self.types.type_type.to_owned(),
425+
self,
425426
)
426427
.unwrap()
427428
}
@@ -446,6 +447,7 @@ impl Context {
446447
attrs,
447448
PyBaseException::make_slots(),
448449
self.types.type_type.to_owned(),
450+
self,
449451
)
450452
.unwrap()
451453
}

‎wasm/lib/src/js_module.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -619,7 +619,7 @@ mod _js {
619619
fn js_error(vm: &VirtualMachine) -> PyTypeRef {
620620
let ctx = &vm.ctx;
621621
let js_error = PyRef::leak(
622-
PyType::new_simple_ref("JSError", &vm.ctx.exceptions.exception_type.to_owned())
622+
PyType::new_simple_ref("JSError", &vm.ctx.exceptions.exception_type.to_owned(), ctx)
623623
.unwrap(),
624624
);
625625
extend_class!(ctx, js_error, {

0 commit comments

Comments
 (0)