Skip to content

Commit c7bca44

Browse files
youknowoneyouknowone
committed
additional optimization with MaybeInterned
1 parent dafa10f commit c7bca44

File tree

12 files changed

+148
-55
lines changed

12 files changed

+148
-55
lines changed

‎common/src/refcount.rs

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use crate::atomic::{Ordering::*, PyAtomic, Radium};
55
///
66
/// Going above this limit will abort your program (although not
77
/// necessarily) at _exactly_ `MAX_REFCOUNT + 1` references.
8-
const MAX_REFCOUNT: usize = (isize::MAX) as usize;
8+
const MAX_REFCOUNT: usize = std::usize::MAX - 0x1000;
99

1010
pub struct RefCount {
1111
strong: PyAtomic<usize>,
@@ -58,3 +58,19 @@ impl RefCount {
5858
true
5959
}
6060
}
61+
62+
impl RefCount {
63+
// move these functions out and give separated type once type range is stabilized
64+
65+
pub fn mark_intern(&self) {
66+
debug_assert!(!self.is_interned());
67+
const BIT_MARKER: usize = (std::isize::MAX as usize) + 1;
68+
debug_assert_eq!(BIT_MARKER.count_ones(), 1);
69+
debug_assert_eq!(BIT_MARKER.leading_zeros(), 0);
70+
self.strong.fetch_add(BIT_MARKER, Relaxed);
71+
}
72+
73+
pub fn is_interned(&self) -> bool {
74+
(self.strong.load(Acquire) as isize) < 0
75+
}
76+
}

‎vm/src/builtins/function.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -463,7 +463,7 @@ impl GetAttr for PyBoundMethod {
463463
fn getattro(zelf: &Py<Self>, name: PyStrRef, vm: &VirtualMachine) -> PyResult {
464464
let class_attr = vm
465465
.ctx
466-
.interned_str(name.as_str())
466+
.interned_str(&*name)
467467
.and_then(|attr_name| zelf.get_class_attr(attr_name));
468468
if let Some(obj) = class_attr {
469469
return vm.call_if_get_descriptor(obj, zelf.to_owned().into());

‎vm/src/builtins/mappingproxy.rs

Lines changed: 8 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,13 @@
11
use std::borrow::Cow;
22

3-
use super::{PyDict, PyGenericAlias, PyList, PyStr, PyStrRef, PyTuple, PyTypeRef};
3+
use super::{PyDict, PyGenericAlias, PyList, PyTuple, PyTypeRef};
44
use crate::{
55
class::PyClassImpl,
66
convert::ToPyObject,
77
function::OptionalArg,
88
protocol::{PyMapping, PyMappingMethods, PySequence, PySequenceMethods},
99
types::{AsMapping, AsSequence, Constructor, Iterable},
10-
AsObject, Context, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
11-
VirtualMachine,
10+
AsObject, Context, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
1211
};
1312

1413
#[pyclass(module = false, name = "mappingproxy")]
@@ -63,12 +62,9 @@ impl Constructor for PyMappingProxy {
6362
impl PyMappingProxy {
6463
fn get_inner(&self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult<Option<PyObjectRef>> {
6564
let opt = match &self.mapping {
66-
MappingProxyInner::Class(class) => {
67-
let key = PyStrRef::try_from_object(vm, key)?;
68-
vm.ctx
69-
.interned_str(key.as_str())
70-
.and_then(|key| class.attributes.read().get(key).cloned())
71-
}
65+
MappingProxyInner::Class(class) => key
66+
.as_interned_str(vm)
67+
.and_then(|key| class.attributes.read().get(key).cloned()),
7268
MappingProxyInner::Dict(obj) => obj.get_item(key, vm).ok(),
7369
};
7470
Ok(opt)
@@ -94,16 +90,9 @@ impl PyMappingProxy {
9490

9591
fn _contains(&self, key: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
9692
match &self.mapping {
97-
MappingProxyInner::Class(class) => {
98-
// let key = PyStrRef::try_from_object(vm, key)?;
99-
let key = key
100-
.payload::<PyStr>()
101-
.ok_or_else(|| vm.new_downcast_type_error(PyStr::class(vm), key))?;
102-
Ok(vm
103-
.ctx
104-
.interned_str(key.as_str())
105-
.map_or(false, |key| class.attributes.read().contains_key(key)))
106-
}
93+
MappingProxyInner::Class(class) => Ok(key
94+
.as_interned_str(vm)
95+
.map_or(false, |key| class.attributes.read().contains_key(key))),
10796
MappingProxyInner::Dict(obj) => PySequence::from(obj.as_ref()).contains(key, vm),
10897
}
10998
}

‎vm/src/builtins/str.rs

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ use crate::{
1616
AsMapping, AsSequence, Comparable, Constructor, Hashable, IterNext, IterNextIterable,
1717
Iterable, PyComparisonOp, Unconstructible,
1818
},
19-
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
19+
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult,
2020
TryFromBorrowedObject, VirtualMachine,
2121
};
2222
use ascii::{AsciiStr, AsciiString};
@@ -116,6 +116,12 @@ impl AsRef<str> for PyStr {
116116
}
117117
}
118118

119+
impl AsRef<str> for Py<PyStr> {
120+
fn as_ref(&self) -> &str {
121+
self.as_str()
122+
}
123+
}
124+
119125
impl AsRef<str> for PyStrRef {
120126
fn as_ref(&self) -> &str {
121127
self.as_str()
@@ -1512,6 +1518,12 @@ impl SliceableSequenceOp for PyStr {
15121518
}
15131519
}
15141520

1521+
impl AsRef<str> for PyRefExact<PyStr> {
1522+
fn as_ref(&self) -> &str {
1523+
self.as_str()
1524+
}
1525+
}
1526+
15151527
#[cfg(test)]
15161528
mod tests {
15171529
use super::*;

‎vm/src/builtins/super.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,7 @@ impl GetAttr for PySuper {
148148
return skip(zelf, name);
149149
}
150150

151-
if let Some(name) = vm.ctx.interned_str(name.as_str()) {
151+
if let Some(name) = vm.ctx.interned_str(&*name) {
152152
// skip the classes in start_type.mro up to and including zelf.typ
153153
let mro: Vec<_> = start_type
154154
.iter_mro()

‎vm/src/builtins/type.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -658,7 +658,7 @@ impl GetAttr for PyType {
658658
))
659659
}
660660

661-
let name = if let Some(name) = vm.ctx.interned_str(name_str.as_str()) {
661+
let name = if let Some(name) = vm.ctx.interned_str(&*name_str) {
662662
name
663663
} else {
664664
return Err(attribute_error(zelf, name_str.as_str(), vm));

‎vm/src/intern.rs

Lines changed: 75 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use crate::{
22
builtins::{PyStr, PyTypeRef},
33
common::lock::PyRwLock,
44
convert::ToPyObject,
5-
Py, PyObject, PyObjectRef, PyRef, PyRefExact,
5+
AsObject, Py, PyObject, PyObjectRef, PyRef, PyRefExact,
66
};
77
use std::{
88
borrow::{Borrow, ToOwned},
@@ -33,26 +33,32 @@ impl Clone for StringPool {
3333
impl StringPool {
3434
#[inline]
3535
pub unsafe fn intern<S: Internable>(&self, s: S, typ: PyTypeRef) -> &'static PyStrInterned {
36-
if let Some(found) = self.inner.read().get(s.as_str()) {
37-
return PyStrInterned::borrow_cache(found);
36+
if let Some(found) = self.interned(s.as_ref()) {
37+
return found;
3838
}
3939

4040
#[cold]
4141
fn miss(zelf: &StringPool, s: PyRefExact<PyStr>) -> &'static PyStrInterned {
4242
let cache = CachedPyStrRef { inner: s };
4343
let interned = unsafe { PyStrInterned::borrow_cache(&cache) };
44-
zelf.inner.write().insert(cache);
44+
let inserted = zelf.inner.write().insert(cache);
45+
if inserted {
46+
unsafe { interned.as_object().mark_intern() };
47+
}
4548
interned
4649
}
47-
let str_ref = s.into_pyref(typ);
50+
let str_ref = s.into_pyref_exact(typ);
4851
miss(self, str_ref)
4952
}
5053

5154
#[inline]
52-
pub fn interned(&self, s: &str) -> Option<&'static PyStrInterned> {
55+
pub fn interned<S: MaybeInterned + ?Sized>(&self, s: &S) -> Option<&'static PyStrInterned> {
56+
if let Some(interned) = s.as_interned() {
57+
return Some(interned);
58+
}
5359
self.inner
5460
.read()
55-
.get(s)
61+
.get(s.as_ref())
5662
.map(|cached| unsafe { PyStrInterned::borrow_cache(cached) })
5763
}
5864
}
@@ -171,7 +177,10 @@ impl std::fmt::Display for PyStrInterned {
171177
}
172178

173179
mod sealed {
174-
use crate::{builtins::PyStr, object::PyRefExact};
180+
use crate::{
181+
builtins::PyStr,
182+
object::{Py, PyRefExact},
183+
};
175184

176185
pub trait SealedInternable {}
177186

@@ -180,44 +189,84 @@ mod sealed {
180189
impl SealedInternable for &str {}
181190

182191
impl SealedInternable for PyRefExact<PyStr> {}
192+
193+
pub trait SealedMaybeInterned {}
194+
195+
impl SealedMaybeInterned for str {}
196+
impl SealedMaybeInterned for PyRefExact<PyStr> {}
197+
impl SealedMaybeInterned for Py<PyStr> {}
183198
}
184199

185200
/// A sealed marker trait for `DictKey` types that always become an exact instance of `str`
186-
pub trait Internable: sealed::SealedInternable + crate::dictdatatype::DictKey + ToPyObject {
187-
fn as_str(&self) -> &str;
188-
fn into_pyref(self, str_type: PyTypeRef) -> PyRefExact<PyStr>;
201+
pub trait Internable
202+
where
203+
Self: sealed::SealedInternable
204+
+ crate::dictdatatype::DictKey
205+
+ ToPyObject
206+
+ AsRef<Self::Interned>,
207+
Self::Interned: MaybeInterned,
208+
{
209+
type Interned: ?Sized;
210+
fn into_pyref_exact(self, str_type: PyTypeRef) -> PyRefExact<PyStr>;
189211
}
190212

191213
impl Internable for String {
214+
type Interned = str;
192215
#[inline]
193-
fn as_str(&self) -> &str {
194-
String::as_str(self)
195-
}
196-
#[inline]
197-
fn into_pyref(self, str_type: PyTypeRef) -> PyRefExact<PyStr> {
216+
fn into_pyref_exact(self, str_type: PyTypeRef) -> PyRefExact<PyStr> {
198217
let obj = PyRef::new_ref(PyStr::from(self), str_type, None);
199218
unsafe { PyRefExact::new_unchecked(obj) }
200219
}
201220
}
202221

203222
impl Internable for &str {
223+
type Interned = str;
204224
#[inline]
205-
fn as_str(&self) -> &str {
206-
self
207-
}
208-
#[inline]
209-
fn into_pyref(self, str_type: PyTypeRef) -> PyRefExact<PyStr> {
210-
self.to_owned().into_pyref(str_type)
225+
fn into_pyref_exact(self, str_type: PyTypeRef) -> PyRefExact<PyStr> {
226+
self.to_owned().into_pyref_exact(str_type)
211227
}
212228
}
213229

214230
impl Internable for PyRefExact<PyStr> {
231+
type Interned = Py<PyStr>;
215232
#[inline]
216-
fn as_str(&self) -> &str {
217-
self.deref().as_str()
233+
fn into_pyref_exact(self, _str_type: PyTypeRef) -> PyRefExact<PyStr> {
234+
self
235+
}
236+
}
237+
238+
pub trait MaybeInterned: AsRef<str> + sealed::SealedMaybeInterned {
239+
fn as_interned(&self) -> Option<&'static PyStrInterned>;
240+
}
241+
242+
impl MaybeInterned for str {
243+
#[inline(always)]
244+
fn as_interned(&self) -> Option<&'static PyStrInterned> {
245+
None
246+
}
247+
}
248+
249+
impl MaybeInterned for Py<PyStr> {
250+
#[inline(always)]
251+
fn as_interned(&self) -> Option<&'static PyStrInterned> {
252+
if self.as_object().is_interned() {
253+
Some(unsafe { std::mem::transmute(self) })
254+
} else {
255+
None
256+
}
218257
}
258+
}
259+
260+
impl PyObject {
219261
#[inline]
220-
fn into_pyref(self, _str_type: PyTypeRef) -> PyRefExact<PyStr> {
221-
self
262+
pub fn as_interned_str(&self, vm: &crate::VirtualMachine) -> Option<&'static PyStrInterned> {
263+
let s: Option<&Py<PyStr>> = self.downcast_ref();
264+
if self.is_interned() {
265+
s.unwrap().as_interned()
266+
} else if let Some(s) = s {
267+
vm.ctx.interned_str(s.as_str())
268+
} else {
269+
None
270+
}
222271
}
223272
}

‎vm/src/object/core.rs

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -783,6 +783,16 @@ impl PyObject {
783783
// call drop only when there are no references in scope - stacked borrows stuff
784784
drop_dealloc(ptr.as_ptr())
785785
}
786+
787+
/// # Safety
788+
/// This call will make the object live forever.
789+
pub(crate) unsafe fn mark_intern(&self) {
790+
self.0.ref_count.mark_intern();
791+
}
792+
793+
pub(crate) fn is_interned(&self) -> bool {
794+
self.0.ref_count.is_interned()
795+
}
786796
}
787797

788798
impl Borrow<PyObject> for PyObjectRef {
@@ -1001,6 +1011,16 @@ where
10011011
}
10021012
}
10031013

1014+
impl<T> AsRef<Py<T>> for PyRef<T>
1015+
where
1016+
T: PyObjectPayload,
1017+
{
1018+
#[inline(always)]
1019+
fn as_ref(&self) -> &Py<T> {
1020+
self
1021+
}
1022+
}
1023+
10041024
impl<T> Deref for PyRef<T>
10051025
where
10061026
T: PyObjectPayload,

‎vm/src/object/ext.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,13 @@ impl<T: PyPayload> Deref for PyRefExact<T> {
117117
}
118118
}
119119

120+
impl<T: PyObjectPayload> AsRef<Py<T>> for PyRefExact<T> {
121+
#[inline(always)]
122+
fn as_ref(&self) -> &Py<T> {
123+
self.inner.as_ref()
124+
}
125+
}
126+
120127
impl<T: PyPayload> ToPyObject for PyRefExact<T> {
121128
#[inline(always)]
122129
fn to_pyobject(self, _vm: &VirtualMachine) -> PyObjectRef {

‎vm/src/protocol/object.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ impl PyObject {
140140
vm_trace!("object.__setattr__({:?}, {}, {:?})", obj, attr_name, value);
141141
if let Some(attr) = vm
142142
.ctx
143-
.interned_str(attr_name.as_str())
143+
.interned_str(&*attr_name)
144144
.and_then(|attr_name| self.get_class_attr(attr_name))
145145
{
146146
let descr_set = attr.class().mro_find_map(|cls| cls.slots.descr_set.load());
@@ -189,7 +189,7 @@ impl PyObject {
189189
) -> PyResult<Option<PyObjectRef>> {
190190
let name = name_str.as_str();
191191
let obj_cls = self.class();
192-
let cls_attr_name = vm.ctx.interned_str(name_str.as_str());
192+
let cls_attr_name = vm.ctx.interned_str(&*name_str);
193193
let cls_attr = match cls_attr_name.and_then(|name| obj_cls.get_attr(name)) {
194194
Some(descr) => {
195195
let descr_cls = descr.class();

‎vm/src/vm/context.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use crate::{
1111
class::{PyClassImpl, StaticType},
1212
exceptions,
1313
function::IntoPyNativeFunc,
14-
intern::{Internable, PyStrInterned, StringPool},
14+
intern::{Internable, MaybeInterned, PyStrInterned, StringPool},
1515
interned,
1616
object::{PyObjectPayload, PyObjectRef, PyPayload, PyRef},
1717
types::{PyTypeFlags, PyTypeSlots, TypeZoo},
@@ -277,7 +277,7 @@ impl Context {
277277
unsafe { self.string_pool.intern(s, self.types.str_type.clone()) }
278278
}
279279

280-
pub fn interned_str(&self, s: &str) -> Option<&'static PyStrInterned> {
280+
pub fn interned_str<S: MaybeInterned + ?Sized>(&self, s: &S) -> Option<&'static PyStrInterned> {
281281
self.string_pool.interned(s)
282282
}
283283

‎vm/src/vm/method.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ impl PyMethod {
2727
return obj.get_attr(name, vm).map(Self::Attribute);
2828
}
2929

30-
let interned_name = vm.ctx.interned_str(name.as_str());
30+
let interned_name = vm.ctx.interned_str(&*name);
3131
let mut is_method = false;
3232

3333
let cls_attr = match interned_name.and_then(|name| cls.get_attr(name)) {

0 commit comments

Comments
 (0)