leaked pystr and doc memory usage optimization

youknowone · youknowone · commit 10f08e085009 · 2023-06-14T17:25:36.000+09:00
diff --git a/vm/src/builtins/str.rs b/vm/src/builtins/str.rs
@@ -1802,6 +1802,30 @@ impl PyStrInterned {
     pub fn to_exact(&'static self) -> PyRefExact<PyStr> {
         unsafe { PyRefExact::new_unchecked(self.to_owned()) }
     }
+
+    pub(crate) fn new_static_str(s: &'static str, ctx: &Context) -> &'static PyStrInterned {
+        if let Some(interned) = ctx.interned_str(s) {
+            return interned;
+        }
+        struct PyStrView {
+            _bytes: &'static [u8],
+            _kind: PyStrKindData,
+            _hash: PyAtomic<hash::PyHash>,
+        }
+        let kind = PyStrKind::of_str(s);
+        let fake_payload = PyStrView {
+            _bytes: s.as_bytes(),
+            _kind: kind.new_data(),
+            _hash: Radium::new(hash::SENTINEL),
+        };
+        let payload: PyStr = unsafe {
+            // Safety:
+            // the layout of Box<[u8]> == &<[u8]>
+            // and Box::drop will never be called by leaking
+            std::mem::transmute(fake_payload)
+        };
+        ctx.intern_str(payload.into_exact_ref(ctx))
+    }
 }
 
 impl std::fmt::Display for PyStrInterned {
diff --git a/vm/src/builtins/super.rs b/vm/src/builtins/super.rs
@@ -257,22 +257,4 @@ fn supercheck(ty: PyTypeRef, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<
 pub fn init(context: &Context) {
     let super_type = &context.types.super_type;
     PySuper::extend_class(context, super_type);
-
-    let super_doc = "super() -> same as super(__class__, <first argument>)\n\
-                     super(type) -> unbound super object\n\
-                     super(type, obj) -> bound super object; requires isinstance(obj, type)\n\
-                     super(type, type2) -> bound super object; requires issubclass(type2, type)\n\
-                     Typical use to call a cooperative superclass method:\n\
-                     class C(B):\n    \
-                     def meth(self, arg):\n        \
-                     super().meth(arg)\n\
-                     This works for class methods too:\n\
-                     class C(B):\n    \
-                     @classmethod\n    \
-                     def cmeth(cls, arg):\n        \
-                     super().cmeth(arg)\n";
-
-    extend_class!(context, super_type, {
-        "__doc__" => context.new_str(super_doc),
-    });
 }
diff --git a/vm/src/class.rs b/vm/src/class.rs
@@ -76,7 +76,7 @@ pub trait PyClassImpl: PyClassDef {
             assert!(class.slots.flags.is_created_with_flags());
         }
 
-        let _ = ctx.intern_str(Self::NAME); // intern type name
+        let _ = ctx.intern_static_str(Self::NAME); // intern type name
 
         if Self::TP_FLAGS.has_feature(PyTypeFlags::HAS_DICT) {
             let __dict__ = identifier!(ctx, __dict__);
@@ -93,12 +93,15 @@ pub trait PyClassImpl: PyClassDef {
         }
         Self::impl_extend_class(ctx, class);
         if let Some(doc) = Self::DOC {
-            class.set_attr(identifier!(ctx, __doc__), ctx.new_str(doc).into());
+            class.set_attr(
+                identifier!(ctx, __doc__),
+                ctx.intern_static_str(doc).to_owned().into(),
+            );
         }
         if let Some(module_name) = Self::MODULE_NAME {
             class.set_attr(
                 identifier!(ctx, __module__),
-                ctx.new_str(module_name).into(),
+                ctx.intern_static_str(module_name).to_owned().into(),
             );
         }
 
diff --git a/vm/src/stdlib/sys.rs b/vm/src/stdlib/sys.rs
@@ -1,4 +1,4 @@
-use crate::{builtins::PyModule, convert::ToPyObject, Py, PyResult, VirtualMachine};
+use crate::{builtins::PyModule, convert::IntoObject, Py, PyResult, VirtualMachine};
 
 pub(crate) use sys::{UnraisableHookArgs, __module_def, DOC, MAXSIZE, MULTIARCH};
 
@@ -927,7 +927,7 @@ pub(crate) fn init_module(vm: &VirtualMachine, module: &Py<PyModule>, builtins:
         .set_item("builtins", builtins.to_owned().into(), vm)
         .unwrap();
     extend_module!(vm, module, {
-        "__doc__" => sys::DOC.to_owned().to_pyobject(vm),
+        "__doc__" => vm.ctx.intern_static_str(sys::DOC.unwrap()).to_owned().into_object(),
         "modules" => modules,
     });
 }
diff --git a/vm/src/vm/context.rs b/vm/src/vm/context.rs
@@ -323,6 +323,10 @@ impl Context {
         unsafe { self.string_pool.intern(s, self.types.str_type.to_owned()) }
     }
 
+    pub(crate) fn intern_static_str(&self, s: &'static str) -> &'static PyStrInterned {
+        PyStrInterned::new_static_str(s, self)
+    }
+
     pub fn interned_str<S: MaybeInternedString + ?Sized>(
         &self,
         s: &S,

Original file line number	Diff line number	Diff line change
`@@ -76,7 +76,7 @@ pub trait PyClassImpl: PyClassDef {`
`76`	`76`	`assert!(class.slots.flags.is_created_with_flags());`
`77`	`77`	`}`
`78`	`78`
`79`		`- let _ = ctx.intern_str(Self::NAME); // intern type name`
	`79`	`+ let _ = ctx.intern_static_str(Self::NAME); // intern type name`
`80`	`80`
`81`	`81`	`if Self::TP_FLAGS.has_feature(PyTypeFlags::HAS_DICT) {`
`82`	`82`	`let __dict__ = identifier!(ctx, __dict__);`
`@@ -93,12 +93,15 @@ pub trait PyClassImpl: PyClassDef {`
`93`	`93`	`}`
`94`	`94`	`Self::impl_extend_class(ctx, class);`
`95`	`95`	`if let Some(doc) = Self::DOC {`
`96`		`- class.set_attr(identifier!(ctx, __doc__), ctx.new_str(doc).into());`
	`96`	`+ class.set_attr(`
	`97`	`+ identifier!(ctx, __doc__),`
	`98`	`+ ctx.intern_static_str(doc).to_owned().into(),`
	`99`	`+ );`
`97`	`100`	`}`
`98`	`101`	`if let Some(module_name) = Self::MODULE_NAME {`
`99`	`102`	`class.set_attr(`
`100`	`103`	`identifier!(ctx, __module__),`
`101`		`- ctx.new_str(module_name).into(),`
	`104`	`+ ctx.intern_static_str(module_name).to_owned().into(),`
`102`	`105`	`);`
`103`	`106`	`}`
`104`	`107`