Merge pull request RustPython#3700 from youknowone/attr-intern

PyStrInterned for PyAttributes/Constants

Author: youknowone

ast/asdl_rs.py

@@ -389,9 +389,9 @@ def gen_classdef(self, name, fields, attrs, depth, base="AstNode"):
         self.emit(f"#[extend_class]", depth + 1)
         self.emit("fn extend_class_with_fields(ctx: &Context, class: &PyTypeRef) {", depth + 1)
         fields = ",".join(f"ctx.new_str(ascii!({json.dumps(f.name)})).into()" for f in fields)
-        self.emit(f'class.set_str_attr("_fields", ctx.new_list(vec![{fields}]));', depth + 2)
+        self.emit(f'class.set_attr(interned!(ctx, _fields), ctx.new_list(vec![{fields}]).into());', depth + 2)
         attrs = ",".join(f"ctx.new_str(ascii!({json.dumps(attr.name)})).into()" for attr in attrs)
-        self.emit(f'class.set_str_attr("_attributes", ctx.new_list(vec![{attrs}]));', depth + 2)
+        self.emit(f'class.set_attr(interned!(ctx, _attributes), ctx.new_list(vec![{attrs}]).into());', depth + 2)
         self.emit("}", depth + 1)
         self.emit("}", depth)
 

benches/execution.rs

@@ -116,7 +116,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
         .map(|entry| {
             let path = entry.unwrap().path();
             (
-                path.file_name().unwrap().to_str().unwrap().to_owned(),
+                path.file_name().unwrap().to_owned().unwrap().to_owned(),
                 std::fs::read_to_string(path).unwrap(),
             )
         })

common/src/hash.rs

@@ -54,8 +54,9 @@ impl HashSecret {
     pub fn new(seed: u32) -> Self {
         let mut buf = [0u8; 16];
         lcg_urandom(seed, &mut buf);
-        let k0 = u64::from_le_bytes(buf[..8].try_into().unwrap());
-        let k1 = u64::from_le_bytes(buf[8..].try_into().unwrap());
+        let (left, right) = buf.split_at(8);
+        let k0 = u64::from_le_bytes(left.try_into().unwrap());
+        let k1 = u64::from_le_bytes(right.try_into().unwrap());
         Self { k0, k1 }
     }
 }

common/src/refcount.rs

@@ -5,7 +5,7 @@ use crate::atomic::{Ordering::*, PyAtomic, Radium};
 ///
 /// Going above this limit will abort your program (although not
 /// necessarily) at _exactly_ `MAX_REFCOUNT + 1` references.
-const MAX_REFCOUNT: usize = (isize::MAX) as usize;
+const MAX_REFCOUNT: usize = isize::MAX as usize;
 
 pub struct RefCount {
     strong: PyAtomic<usize>,
@@ -18,6 +18,8 @@ impl Default for RefCount {
 }
 
 impl RefCount {
+    const MASK: usize = MAX_REFCOUNT;
+
     pub fn new() -> Self {
         RefCount {
             strong: Radium::new(1),
@@ -33,7 +35,7 @@ impl RefCount {
     pub fn inc(&self) {
         let old_size = self.strong.fetch_add(1, Relaxed);
 
-        if old_size > MAX_REFCOUNT {
+        if old_size & Self::MASK == Self::MASK {
             std::process::abort();
         }
     }
@@ -58,3 +60,19 @@ impl RefCount {
         true
     }
 }
+
+impl RefCount {
+    // move these functions out and give separated type once type range is stabilized
+
+    pub fn leak(&self) {
+        debug_assert!(!self.is_leaked());
+        const BIT_MARKER: usize = (std::isize::MAX as usize) + 1;
+        debug_assert_eq!(BIT_MARKER.count_ones(), 1);
+        debug_assert_eq!(BIT_MARKER.leading_zeros(), 0);
+        self.strong.fetch_add(BIT_MARKER, Relaxed);
+    }
+
+    pub fn is_leaked(&self) -> bool {
+        (self.strong.load(Acquire) as isize) < 0
+    }
+}

derive/src/pyclass.rs

@@ -503,7 +503,8 @@ where
                     #py_name,
                     ctx.make_funcdef(#py_name, Self::#ident)
                         #doc
-                        #build_func
+                        #build_func,
+                    ctx,
                 );
             }
         };
@@ -611,7 +612,7 @@ where
                     ident,
                     py_name.clone(),
                     quote! {
-                        class.set_str_attr(#py_name, #value);
+                        class.set_str_attr(#py_name, #value, ctx);
                     },
                 )
             } else {
@@ -736,7 +737,8 @@ impl ToTokens for GetSetNursery {
                             ::rustpython_vm::builtins::PyGetSet::new(#name.into(), class.clone())
                                 .with_get(&Self::#getter)
                                 #setter #deleter,
-                            ctx.types.getset_type.clone(), None)
+                            ctx.types.getset_type.clone(), None),
+                        ctx
                     );
                 }
             });

derive/src/pymodule.rs

@@ -433,7 +433,7 @@ impl ModuleItem for ClassItem {
             };
             let class_new = quote_spanned!(ident.span() =>
                 let new_class = <#ident as ::rustpython_vm::class::PyClassImpl>::make_class(&vm.ctx);
-                new_class.set_str_attr("__module__", vm.new_pyobj(#module_name));
+                new_class.set_attr(rustpython_vm::identifier!(vm, __module__), vm.new_pyobj(#module_name));
             );
             (class_name, class_new)
         };

src/shell/helper.rs

@@ -2,7 +2,7 @@
 use rustpython_vm::{
     builtins::{PyDictRef, PyStrRef},
     function::ArgIterable,
-    PyResult, TryFromObject, VirtualMachine,
+    identifier, PyResult, TryFromObject, VirtualMachine,
 };
 
 pub struct ShellHelper<'vm> {
@@ -81,14 +81,19 @@ impl<'vm> ShellHelper<'vm> {
                 current = current.get_attr(attr.as_str(), self.vm).ok()?;
             }
 
-            let current_iter = str_iter_method(current, "__dir__").ok()?;
+            let current_iter = str_iter_method(current, identifier!(self.vm, __dir__)).ok()?;
 
             (last, current_iter, None)
         } else {
             // we need to get a variable based off of globals/builtins
 
-            let globals = str_iter_method(self.globals.clone().into(), "keys").ok()?;
-            let builtins = str_iter_method(self.vm.builtins.clone().into(), "__dir__").ok()?;
+            let globals =
+                str_iter_method(self.globals.clone().into(), identifier!(self.vm, keys)).ok()?;
+            let builtins = str_iter_method(
+                self.vm.builtins.clone().into(),
+                identifier!(self.vm, __dir__),
+            )
+            .ok()?;
             (first, globals, Some(builtins))
         };
         Some((word_start, iter1.chain(iter2.into_iter().flatten())))

stdlib/src/array.rs

@@ -683,7 +683,7 @@ mod array {
         fn typecode(&self, vm: &VirtualMachine) -> PyStrRef {
             vm.ctx
                 .intern_str(self.read().typecode().to_string())
-                .to_str()
+                .to_owned()
         }
 
         #[pyproperty]

stdlib/src/math.rs

@@ -3,9 +3,9 @@ pub(crate) use math::make_module;
 #[pymodule]
 mod math {
     use crate::vm::{
-        builtins::{try_bigint_to_f64, try_f64_to_bigint, PyFloat, PyInt, PyIntRef},
+        builtins::{try_bigint_to_f64, try_f64_to_bigint, PyFloat, PyInt, PyIntRef, PyStrInterned},
         function::{ArgIntoFloat, ArgIterable, Either, OptionalArg, PosArgs},
-        AsObject, PyObject, PyObjectRef, PyRef, PyResult, VirtualMachine,
+        identifier, AsObject, PyObject, PyObjectRef, PyRef, PyResult, VirtualMachine,
     };
     use num_bigint::BigInt;
     use num_traits::{One, Signed, Zero};
@@ -430,25 +430,29 @@ mod math {
         }
     }
 
-    fn try_magic_method(func_name: &str, vm: &VirtualMachine, value: &PyObject) -> PyResult {
+    fn try_magic_method(
+        func_name: &'static PyStrInterned,
+        vm: &VirtualMachine,
+        value: &PyObject,
+    ) -> PyResult {
         let method = vm.get_method_or_type_error(value.to_owned(), func_name, || {
             format!(
                 "type '{}' doesn't define '{}' method",
                 value.class().name(),
-                func_name,
+                func_name.as_str(),
             )
         })?;
         vm.invoke(&method, ())
     }
 
     #[pyfunction]
     fn trunc(x: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        try_magic_method("__trunc__", vm, &x)
+        try_magic_method(identifier!(vm, __trunc__), vm, &x)
     }
 
     #[pyfunction]
     fn ceil(x: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        let result_or_err = try_magic_method("__ceil__", vm, &x);
+        let result_or_err = try_magic_method(identifier!(vm, __ceil__), vm, &x);
         if result_or_err.is_err() {
             if let Ok(Some(v)) = x.try_to_f64(vm) {
                 let v = try_f64_to_bigint(v.ceil(), vm)?;
@@ -460,7 +464,7 @@ mod math {
 
     #[pyfunction]
     fn floor(x: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        let result_or_err = try_magic_method("__floor__", vm, &x);
+        let result_or_err = try_magic_method(identifier!(vm, __floor__), vm, &x);
         if result_or_err.is_err() {
             if let Ok(Some(v)) = x.try_to_f64(vm) {
                 let v = try_f64_to_bigint(v.floor(), vm)?;

stdlib/src/pyexpat.rs

@@ -25,7 +25,7 @@ macro_rules! create_property {
             move |this: &PyExpatLikeXmlParser, func: PyObjectRef| *this.$element.write() = func,
         );
 
-        $attributes.insert($name.to_owned(), attr.into());
+        $attributes.insert($ctx.intern_str($name), attr.into());
     };
 }