Implement string interning

Lib/test/test_pathlib.py

@@ -2044,8 +2044,6 @@ class _BasePathTest(object):
             pp = pickle.loads(dumped)
             self.assertEqual(pp.stat(), p.stat())
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_parts_interning(self):
         P = self.cls
         p = P('/usr/bin/foo')

Lib/test/test_sys.py

@@ -539,7 +539,6 @@ class SysModuleTest(unittest.TestCase):
         # the test runs under regrtest.
         self.assertEqual(sys.__stdout__.encoding, sys.__stderr__.encoding)
 
-    @unittest.skip("TODO: RUSTPYTHON; sys.intern() string interning")
     def test_intern(self):
         global INTERN_NUMRUNS
         INTERN_NUMRUNS += 1

examples/freeze/main.rs

@@ -15,8 +15,7 @@ fn run(vm: &vm::VirtualMachine) -> vm::pyobject::PyResult<()> {
         vm::py_freeze!(file = "examples/freeze/freeze.py");
 
     let res = vm.run_code_obj(
-        vm.ctx
-            .new_code_object(modules.get("frozen").unwrap().code.clone()),
+        vm.new_code_object(modules.get("frozen").unwrap().code.clone()),
         scope,
     );
 

examples/mini_repl.rs

@@ -17,7 +17,7 @@ macro_rules! add_python_function {
         // compile the code to bytecode
         let code = vm::py_compile!(source = $src);
         // convert the rustpython_bytecode::CodeObject to a PyCodeRef
-        let code = $vm.ctx.new_code_object(code);
+        let code = $vm.new_code_object(code);
 
         // run the python code in the scope to store the function
         $vm.run_code_obj(code, $scope.clone())

vm/src/builtins/code.rs

@@ -74,16 +74,59 @@ impl Constant for PyConstant {
     }
 }
 
-pub(crate) struct PyObjBag<'a>(pub &'a PyContext);
+pub(crate) struct PyObjBag<'a>(pub &'a VirtualMachine);
 
 impl ConstantBag for PyObjBag<'_> {
     type Constant = PyConstant;
     fn make_constant(&self, constant: bytecode::ConstantData) -> Self::Constant {
-        PyConstant(self.0.unwrap_constant(constant))
+        let vm = self.0;
+        let ctx = &vm.ctx;
+        let obj = match constant {
+            bytecode::ConstantData::Integer { value } => ctx.new_int(value),
+            bytecode::ConstantData::Float { value } => ctx.new_float(value),
+            bytecode::ConstantData::Complex { value } => ctx.new_complex(value),
+            bytecode::ConstantData::Str { value } => vm.intern_string(value).into_object(),
+            bytecode::ConstantData::Bytes { value } => ctx.new_bytes(value.to_vec()),
+            bytecode::ConstantData::Boolean { value } => ctx.new_bool(value),
+            bytecode::ConstantData::Code { code } => {
+                ctx.new_code_object(code.map_bag(self)).into_object()
+            }
+            bytecode::ConstantData::Tuple { elements } => {
+                let elements = elements
+                    .into_iter()
+                    .map(|constant| self.make_constant(constant).0)
+                    .collect();
+                ctx.new_tuple(elements)
+            }
+            bytecode::ConstantData::None => ctx.none(),
+            bytecode::ConstantData::Ellipsis => ctx.ellipsis(),
+        };
+        PyConstant(obj)
     }
     fn make_constant_borrowed<C: Constant>(&self, constant: BorrowedConstant<C>) -> Self::Constant {
-        // TODO: check if the constant is a string and try interning it without cloning
-        self.make_constant(constant.into_data())
+        let vm = self.0;
+        let ctx = &vm.ctx;
+        let obj = match constant {
+            bytecode::BorrowedConstant::Integer { value } => ctx.new_bigint(value),
+            bytecode::BorrowedConstant::Float { value } => ctx.new_float(value),
+            bytecode::BorrowedConstant::Complex { value } => ctx.new_complex(value),
+            bytecode::BorrowedConstant::Str { value } => vm.intern_string(value).into_object(),
+            bytecode::BorrowedConstant::Bytes { value } => ctx.new_bytes(value.to_vec()),
+            bytecode::BorrowedConstant::Boolean { value } => ctx.new_bool(value),
+            bytecode::BorrowedConstant::Code { code } => {
+                ctx.new_code_object(code.map_clone_bag(self)).into_object()
+            }
+            bytecode::BorrowedConstant::Tuple { elements } => {
+                let elements = elements
+                    .into_iter()
+                    .map(|constant| self.make_constant_borrowed(constant).0)
+                    .collect();
+                ctx.new_tuple(elements)
+            }
+            bytecode::BorrowedConstant::None => ctx.none(),
+            bytecode::BorrowedConstant::Ellipsis => ctx.ellipsis(),
+        };
+        PyConstant(obj)
     }
 }
 
@@ -93,16 +136,16 @@ pub type CodeObject = bytecode::CodeObject<PyConstant>;
 pub type FrozenModule = bytecode::FrozenModule<PyConstant>;
 
 pub trait IntoCodeObject {
-    fn into_codeobj(self, ctx: &PyContext) -> CodeObject;
+    fn into_codeobj(self, vm: &VirtualMachine) -> CodeObject;
 }
 impl IntoCodeObject for CodeObject {
-    fn into_codeobj(self, _ctx: &PyContext) -> CodeObject {
+    fn into_codeobj(self, _vm: &VirtualMachine) -> CodeObject {
         self
     }
 }
 impl IntoCodeObject for bytecode::CodeObject {
-    fn into_codeobj(self, ctx: &PyContext) -> CodeObject {
-        ctx.map_codeobj(self)
+    fn into_codeobj(self, vm: &VirtualMachine) -> CodeObject {
+        vm.map_codeobj(self)
     }
 }
 

vm/src/builtins/mod.rs

@@ -53,7 +53,7 @@ pub use property::PyProperty;
 pub(crate) mod pybool;
 pub use pybool::{IntoPyBool, PyBool};
 pub(crate) mod pystr;
-pub use pystr::{PyStr, PyStrRef};
+pub use pystr::{PyStr, PyStrExact, PyStrRef};
 pub(crate) mod pysuper;
 pub use pysuper::PySuper;
 pub(crate) mod pytype;

vm/src/builtins/pystr.rs

@@ -1451,3 +1451,34 @@ impl<'s> AnyStr<'s, char> for str {
         splited
     }
 }
+
+#[repr(transparent)]
+pub struct PyStrExact(PyStrRef);
+impl PyStrExact {
+    pub fn from_str(s: PyStrRef, ctx: &PyContext) -> Result<Self, PyStrRef> {
+        if s.class().is(&ctx.types.str_type) {
+            Ok(Self(s))
+        } else {
+            Err(s)
+        }
+    }
+    pub fn into_pystr(self) -> PyStrRef {
+        self.0
+    }
+}
+impl IntoPyObject for PyStrExact {
+    fn into_pyobject(self, _vm: &VirtualMachine) -> PyObjectRef {
+        self.0.into_object()
+    }
+}
+impl crate::dictdatatype::DictKey for PyStrExact {
+    fn key_hash(&self, vm: &VirtualMachine) -> PyResult<hash::PyHash> {
+        self.0.key_hash(vm)
+    }
+    fn key_is(&self, other: &PyObjectRef) -> bool {
+        self.0.key_is(other)
+    }
+    fn key_eq(&self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult<bool> {
+        self.0.key_eq(vm, other_key)
+    }
+}

vm/src/dictdatatype.rs

@@ -4,9 +4,10 @@ use crate::builtins::pystr::{PyStr, PyStrRef};
 /// And: https://www.youtube.com/watch?v=p33CVV29OG8
 /// And: http://code.activestate.com/recipes/578375/
 use crate::common::lock::{PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard};
-use crate::pyobject::{BorrowValue, IdProtocol, IntoPyObject, PyObjectRef, PyResult};
+use crate::pyobject::{BorrowValue, IdProtocol, IntoPyObject, PyObjectRef, PyResult, TypeProtocol};
 use crate::vm::VirtualMachine;
 use rustpython_common::hash;
+use std::fmt;
 use std::mem::size_of;
 
 // HashIndex is intended to be same size with hash::PyHash
@@ -24,6 +25,11 @@ type EntryIndex = usize;
 pub struct Dict<T = PyObjectRef> {
     inner: PyRwLock<DictInner<T>>,
 }
+impl<T> fmt::Debug for Dict<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Debug").finish()
+    }
+}
 
 #[derive(Debug, Copy, Clone)]
 enum IndexEntry {
@@ -336,8 +342,9 @@ impl<T: Clone> Dict<T> {
         Ok(())
     }
 
-    pub fn setdefault<F>(&self, vm: &VirtualMachine, key: PyObjectRef, default: F) -> PyResult<T>
+    pub fn setdefault<K, F>(&self, vm: &VirtualMachine, key: K, default: F) -> PyResult<T>
     where
+        K: DictKey,
         F: FnOnce() -> T,
     {
         let hash = key.key_hash(vm)?;
@@ -359,13 +366,57 @@ impl<T: Clone> Dict<T> {
             } else {
                 let value = default();
                 let mut inner = self.borrow_value_mut();
-                inner.unchecked_push(index_index, hash, key, value.clone(), entry);
+                inner.unchecked_push(
+                    index_index,
+                    hash,
+                    key.into_pyobject(vm),
+                    value.clone(),
+                    entry,
+                );
                 break value;
             }
         };
         Ok(res)
     }
 
+    pub fn setdefault_entry<K, F>(
+        &self,
+        vm: &VirtualMachine,
+        key: K,
+        default: F,
+    ) -> PyResult<(PyObjectRef, T)>
+    where
+        K: DictKey,
+        F: FnOnce() -> T,
+    {
+        let hash = key.key_hash(vm)?;
+        let res = loop {
+            let lookup = self.lookup(vm, &key, hash, None)?;
+            let (entry, index_index) = lookup;
+            if let IndexEntry::Index(index) = entry {
+                let inner = self.borrow_value();
+                if let Some(entry) = inner.entries.get(index) {
+                    if entry.index == index_index {
+                        break (entry.key.clone(), entry.value.clone());
+                    } else {
+                        // stuff shifted around, let's try again
+                    }
+                } else {
+                    // The dict was changed since we did lookup, let's try again.
+                    continue;
+                }
+            } else {
+                let value = default();
+                let key = key.into_pyobject(vm);
+                let mut inner = self.borrow_value_mut();
+                let ret = (key.clone(), value.clone());
+                inner.unchecked_push(index_index, hash, key, value, entry);
+                break ret;
+            }
+        };
+        Ok(res)
+    }
+
     pub fn len(&self) -> usize {
         self.borrow_value().used
     }
@@ -561,8 +612,8 @@ impl DictKey for PyStrRef {
     fn key_eq(&self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult<bool> {
         if self.is(other_key) {
             Ok(true)
-        } else if let Some(py_str_value) = other_key.payload::<PyStr>() {
-            Ok(py_str_value.borrow_value() == self.borrow_value())
+        } else if let Some(pystr) = str_exact(other_key, vm) {
+            Ok(pystr.borrow_value() == self.borrow_value())
         } else {
             vm.bool_eq(self.as_object(), other_key)
         }
@@ -586,8 +637,8 @@ impl DictKey for &str {
     }
 
     fn key_eq(&self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult<bool> {
-        if let Some(py_str_value) = other_key.payload::<PyStr>() {
-            Ok(py_str_value.borrow_value() == *self)
+        if let Some(pystr) = str_exact(other_key, vm) {
+            Ok(pystr.borrow_value() == *self)
         } else {
             // Fall back to PyObjectRef implementation.
             let s = vm.ctx.new_str(*self);
@@ -596,6 +647,28 @@ impl DictKey for &str {
     }
 }
 
+impl DictKey for String {
+    fn key_hash(&self, vm: &VirtualMachine) -> PyResult<HashValue> {
+        self.as_str().key_hash(vm)
+    }
+
+    fn key_is(&self, other: &PyObjectRef) -> bool {
+        self.as_str().key_is(other)
+    }
+
+    fn key_eq(&self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult<bool> {
+        self.as_str().key_eq(vm, other_key)
+    }
+}
+
+fn str_exact<'a>(obj: &'a PyObjectRef, vm: &VirtualMachine) -> Option<&'a PyStr> {
+    if obj.class().is(&vm.ctx.types.str_type) {
+        obj.payload::<PyStr>()
+    } else {
+        None
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::{Dict, DictKey};

vm/src/frame.rs

@@ -255,8 +255,11 @@ impl ExecutingFrame<'_> {
         flame_guard!(format!("Frame::run({})", self.code.obj_name));
         // Execute until return or exception:
         loop {
-            let loc = self.current_location();
-            let result = self.execute_instruction(vm);
+            let idx = self.lasti.fetch_add(1, Ordering::Relaxed);
+            let loc = self.code.locations[idx];
+            let instr = &self.code.instructions[idx];
+            vm.check_signals()?;
+            let result = self.execute_instruction(instr, vm);
             match result {
                 Ok(None) => {}
                 Ok(Some(value)) => {
@@ -329,11 +332,12 @@ impl ExecutingFrame<'_> {
     }
 
     /// Execute a single instruction.
-    fn execute_instruction(&mut self, vm: &VirtualMachine) -> FrameResult {
-        vm.check_signals()?;
-
-        let instruction = &self.code.instructions[self.lasti.fetch_add(1, Ordering::Relaxed)];
-
+    #[inline(always)]
+    fn execute_instruction(
+        &mut self,
+        instruction: &bytecode::Instruction,
+        vm: &VirtualMachine,
+    ) -> FrameResult {
         flame_guard!(format!("Frame::execute_instruction({:?})", instruction));
 
         #[cfg(feature = "vm-tracing-logging")]
@@ -1465,10 +1469,6 @@ impl ExecutingFrame<'_> {
         self.lasti.load(Ordering::Relaxed)
     }
 
-    fn current_location(&self) -> bytecode::Location {
-        self.code.locations[self.lasti()]
-    }
-
     fn push_block(&mut self, typ: BlockType) {
         self.state.blocks.push(Block {
             typ,

vm/src/frozen.rs

@@ -1,10 +1,10 @@
-use super::pyobject::PyContext;
 use crate::builtins::code;
 use crate::bytecode;
+use crate::VirtualMachine;
 use std::collections::HashMap;
 
 pub fn map_frozen<'a>(
-    ctx: &'a PyContext,
+    vm: &'a VirtualMachine,
     i: impl IntoIterator<Item = (String, bytecode::FrozenModule)> + 'a,
 ) -> impl Iterator<Item = (String, code::FrozenModule)> + 'a {
     i.into_iter()
@@ -12,19 +12,19 @@ pub fn map_frozen<'a>(
             (
                 k,
                 code::FrozenModule {
-                    code: ctx.map_codeobj(code),
+                    code: vm.map_codeobj(code),
                     package,
                 },
             )
         })
 }
 
-pub fn get_module_inits(ctx: &PyContext) -> HashMap<String, code::FrozenModule> {
+pub fn get_module_inits(vm: &VirtualMachine) -> HashMap<String, code::FrozenModule> {
     let mut modules = HashMap::new();
 
     macro_rules! ext_modules {
         ($($t:tt)*) => {
-            modules.extend(map_frozen(ctx, py_freeze!($($t)*)));
+            modules.extend(map_frozen(vm, py_freeze!($($t)*)));
         };
     }
 
@@ -48,7 +48,7 @@ pub fn get_module_inits(ctx: &PyContext) -> HashMap<String, code::FrozenModule>
     // if we're on freeze-stdlib, the core stdlib modules will be included anyway
     #[cfg(feature = "freeze-stdlib")]
     {
-        modules.extend(map_frozen(ctx, rustpython_pylib::frozen_stdlib()));
+        modules.extend(map_frozen(vm, rustpython_pylib::frozen_stdlib()));
     }
 
     modules

vm/src/import.rs

@@ -100,7 +100,7 @@ pub fn import_file(
 ) -> PyResult {
     let code_obj = compile::compile(&content, compile::Mode::Exec, file_path, vm.compile_opts())
         .map_err(|err| vm.new_syntax_error(&err))?;
-    import_codeobj(vm, module_name, vm.ctx.map_codeobj(code_obj), true)
+    import_codeobj(vm, module_name, vm.map_codeobj(code_obj), true)
 }
 
 pub fn import_codeobj(

vm/src/pyobject.rs

@@ -29,11 +29,11 @@ use crate::builtins::singletons::{PyNone, PyNoneRef, PyNotImplemented, PyNotImpl
 use crate::builtins::slice::PyEllipsis;
 use crate::builtins::staticmethod::PyStaticMethod;
 use crate::builtins::tuple::{PyTuple, PyTupleRef};
-use crate::bytecode;
 pub use crate::common::borrow::BorrowValue;
 use crate::common::lock::{PyRwLock, PyRwLockReadGuard};
 use crate::common::rc::PyRc;
 use crate::common::static_cell;
+use crate::dictdatatype::Dict;
 use crate::exceptions::{self, PyBaseExceptionRef};
 use crate::function::{IntoFuncArgs, IntoPyNativeFunc};
 use crate::iterator;
@@ -103,6 +103,8 @@ pub struct PyContext {
     pub types: TypeZoo,
     pub exceptions: exceptions::ExceptionZoo,
     pub int_cache_pool: Vec<PyIntRef>,
+    // there should only be exact objects of str in here, no non-strs and no subclasses
+    pub(crate) string_cache: Dict<()>,
     tp_new_wrapper: PyObjectRef,
 }
 
@@ -136,6 +138,8 @@ impl PyContext {
             &types.tuple_type,
         );
 
+        let string_cache = Dict::default();
+
         let tp_new_wrapper = create_object(
             PyNativeFuncDef::from(pytype::tp_new_wrapper.into_func()).into_function(),
             &types.builtin_function_or_method_type,
@@ -153,6 +157,7 @@ impl PyContext {
             types,
             exceptions,
             int_cache_pool,
+            string_cache,
             tp_new_wrapper,
         };
         TypeZoo::extend(&context);
@@ -342,12 +347,11 @@ impl PyContext {
         )
     }
 
-    pub fn new_code_object(&self, code: impl code::IntoCodeObject) -> PyCodeRef {
-        PyRef::new_ref(
-            code::PyCode::new(code.into_codeobj(self)),
-            self.types.code_type.clone(),
-            None,
-        )
+    /// Create a new `PyCodeRef` from a `code::CodeObject`. If you have a non-mapped codeobject or
+    /// this is giving you a type error even though you've passed a `CodeObject`, try
+    /// [`vm.new_code_object()`](VirtualMachine::new_code_object) instead.
+    pub fn new_code_object(&self, code: code::CodeObject) -> PyCodeRef {
+        PyRef::new_ref(code::PyCode::new(code), self.types.code_type.clone(), None)
     }
 
     pub fn new_pyfunction(
@@ -381,31 +385,6 @@ impl PyContext {
         .into_ref()
     }
 
-    pub fn unwrap_constant(&self, constant: bytecode::ConstantData) -> PyObjectRef {
-        match constant {
-            bytecode::ConstantData::Integer { value } => self.new_int(value),
-            bytecode::ConstantData::Float { value } => self.new_float(value),
-            bytecode::ConstantData::Complex { value } => self.new_complex(value),
-            bytecode::ConstantData::Str { value } => self.new_str(value),
-            bytecode::ConstantData::Bytes { value } => self.new_bytes(value),
-            bytecode::ConstantData::Boolean { value } => self.new_bool(value),
-            bytecode::ConstantData::Code { code } => self.new_code_object(*code).into_object(),
-            bytecode::ConstantData::Tuple { elements } => {
-                let elements = elements
-                    .into_iter()
-                    .map(|constant| self.unwrap_constant(constant))
-                    .collect();
-                self.new_tuple(elements)
-            }
-            bytecode::ConstantData::None => self.none(),
-            bytecode::ConstantData::Ellipsis => self.ellipsis(),
-        }
-    }
-
-    pub fn map_codeobj(&self, code: bytecode::CodeObject) -> code::CodeObject {
-        code.map_bag(&code::PyObjBag(self))
-    }
-
     pub fn add_tp_new_wrapper(&self, ty: &PyTypeRef) {
         if !ty.attributes.read().contains_key("__new__") {
             let new_wrapper =

vm/src/stdlib/marshal.rs

@@ -26,7 +26,7 @@ mod decl {
         let code = bytecode::CodeObject::from_bytes(&*code_bytes.borrow_value())
             .map_err(|_| vm.new_value_error("Couldn't deserialize python bytecode".to_owned()))?;
         Ok(PyCode {
-            code: vm.ctx.map_codeobj(code),
+            code: vm.map_codeobj(code),
         })
     }
 

vm/src/sysmodule.rs

@@ -1,14 +1,14 @@
 use num_traits::ToPrimitive;
 use std::{env, mem, path};
 
-use crate::builtins::pystr::PyStrRef;
+use crate::builtins::pystr::{PyStrExact, PyStrRef};
 use crate::builtins::pytype::PyTypeRef;
 use crate::common::hash::{PyHash, PyUHash};
 use crate::frame::FrameRef;
 use crate::function::{Args, FuncArgs, OptionalArg};
 use crate::pyobject::{
-    IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyObjectRc, PyObjectRef, PyResult,
-    PyStructSequence,
+    ItemProtocol, PyClassImpl, PyContext, PyObjectRc, PyObjectRef, PyResult, PyStructSequence,
+    TypeProtocol,
 };
 use crate::vm::{PySettings, VirtualMachine};
 use crate::{builtins, exceptions, py_io, version};
@@ -214,9 +214,10 @@ fn sys_setrecursionlimit(recursion_limit: i32, vm: &VirtualMachine) -> PyResult<
     }
 }
 
-// TODO implement string interning, this will be key for performance
-fn sys_intern(value: PyStrRef) -> PyStrRef {
-    value
+fn sys_intern(value: PyStrRef, vm: &VirtualMachine) -> PyResult<PyStrRef> {
+    PyStrExact::from_str(value, &vm.ctx)
+        .map(|s| vm.intern_string(s))
+        .map_err(|subcls| vm.new_type_error(format!("can't intern {}", subcls.class().name)))
 }
 
 fn sys_exc_info(vm: &VirtualMachine) -> (PyObjectRef, PyObjectRef, PyObjectRef) {
@@ -630,7 +631,7 @@ settrace() -- set the global debug tracing function
     module_names.push("builtins".to_owned());
     module_names.sort();
     let builtin_module_names =
-        ctx.new_tuple(module_names.iter().map(|v| v.into_pyobject(vm)).collect());
+        ctx.new_tuple(module_names.into_iter().map(|n| ctx.new_str(n)).collect());
     let modules = ctx.new_dict();
 
     let prefix = option_env!("RUSTPYTHON_PREFIX").unwrap_or("/usr/local");

vm/src/vm.rs

@@ -20,7 +20,7 @@ use crate::builtins::list::PyList;
 use crate::builtins::module::{self, PyModule};
 use crate::builtins::object;
 use crate::builtins::pybool;
-use crate::builtins::pystr::{PyStr, PyStrRef};
+use crate::builtins::pystr::{PyStr, PyStrExact, PyStrRef};
 use crate::builtins::pytype::PyTypeRef;
 use crate::builtins::tuple::PyTuple;
 use crate::common::{hash::HashSecret, lock::PyMutex, rc::PyRc};
@@ -118,15 +118,6 @@ pub struct PyGlobalState {
     pub atexit_funcs: PyMutex<Vec<(PyObjectRef, FuncArgs)>>,
 }
 
-impl PyGlobalState {
-    pub fn add_frozen<I>(&mut self, ctx: &PyContext, frozen: I)
-    where
-        I: IntoIterator<Item = (String, bytecode::FrozenModule)>,
-    {
-        self.frozen.extend(frozen::map_frozen(ctx, frozen))
-    }
-}
-
 pub const NSIG: usize = 64;
 
 #[derive(Copy, Clone, PartialEq, Eq)]
@@ -257,14 +248,13 @@ impl VirtualMachine {
         let signal_handlers = RefCell::new(arr![ctx.none(); 64]);
 
         let stdlib_inits = stdlib::get_module_inits();
-        let frozen = frozen::get_module_inits(&ctx);
 
         let hash_secret = match settings.hash_seed {
             Some(seed) => HashSecret::new(seed),
             None => rand::random(),
         };
 
-        let vm = VirtualMachine {
+        let mut vm = VirtualMachine {
             builtins,
             sys_module: sysmod,
             ctx: PyRc::new(ctx),
@@ -281,7 +271,7 @@ impl VirtualMachine {
             state: PyRc::new(PyGlobalState {
                 settings,
                 stdlib_inits,
-                frozen,
+                frozen: HashMap::new(),
                 stacksize: AtomicCell::new(0),
                 thread_count: AtomicCell::new(0),
                 hash_secret,
@@ -290,6 +280,9 @@ impl VirtualMachine {
             initialized: false,
         };
 
+        let frozen = frozen::get_module_inits(&vm);
+        PyRc::get_mut(&mut vm.state).unwrap().frozen = frozen;
+
         module::init_module_dict(
             &vm,
             &builtins_dict,
@@ -355,6 +348,24 @@ impl VirtualMachine {
         self.initialized = true;
     }
 
+    /// Can only be used in the initialization closure passed to [`Interpreter::new_with_init`]
+    pub fn add_native_module(&mut self, name: String, module: stdlib::StdlibInitFunc) {
+        let state = PyRc::get_mut(&mut self.state)
+            .expect("can't add_native_module when there are multiple threads");
+        state.stdlib_inits.insert(name, module);
+    }
+
+    /// Can only be used in the initialization closure passed to [`Interpreter::new_with_init`]
+    pub fn add_frozen<I>(&mut self, frozen: I)
+    where
+        I: IntoIterator<Item = (String, bytecode::FrozenModule)>,
+    {
+        let frozen = frozen::map_frozen(self, frozen).collect::<Vec<_>>();
+        let state = PyRc::get_mut(&mut self.state)
+            .expect("can't add_frozen when there are multiple threads");
+        state.frozen.extend(frozen);
+    }
+
     /// Start a new thread with access to the same interpreter.
     ///
     /// # Note
@@ -513,6 +524,10 @@ impl VirtualMachine {
         value.into_pyobject(self)
     }
 
+    pub fn new_code_object(&self, code: impl code::IntoCodeObject) -> PyCodeRef {
+        self.ctx.new_code_object(code.into_codeobj(self))
+    }
+
     pub fn new_module(&self, name: &str, dict: PyDictRef) -> PyObjectRef {
         module::init_module_dict(
             self,
@@ -1229,7 +1244,7 @@ impl VirtualMachine {
         opts: CompileOpts,
     ) -> Result<PyCodeRef, CompileError> {
         compile::compile(source, mode, source_path, opts)
-            .map(|code| PyCode::new(self.ctx.map_codeobj(code)).into_ref(self))
+            .map(|code| PyCode::new(self.map_codeobj(code)).into_ref(self))
     }
 
     fn call_codec_func(
@@ -1622,6 +1637,20 @@ impl VirtualMachine {
         Ok(value)
     }
 
+    pub fn map_codeobj(&self, code: bytecode::CodeObject) -> code::CodeObject {
+        code.map_bag(&code::PyObjBag(self))
+    }
+
+    pub fn intern_string<S: Internable>(&self, s: S) -> PyStrRef {
+        let (s, _) = self
+            .ctx
+            .string_cache
+            .setdefault_entry(self, s, || ())
+            .expect("string_cache lookup should never error");
+        s.downcast()
+            .expect("only strings should be in string_cache")
+    }
+
     #[doc(hidden)]
     pub fn __module_set_attr(
         &self,
@@ -1634,6 +1663,17 @@ impl VirtualMachine {
     }
 }
 
+mod sealed {
+    pub trait SealedInternable {}
+    impl SealedInternable for String {}
+    impl SealedInternable for &str {}
+    impl SealedInternable for super::PyStrExact {}
+}
+pub trait Internable: sealed::SealedInternable + crate::dictdatatype::DictKey {}
+impl Internable for String {}
+impl Internable for &str {}
+impl Internable for PyStrExact {}
+
 pub struct ReprGuard<'vm> {
     vm: &'vm VirtualMachine,
     id: usize,

wasm/lib/src/browser_module.rs

@@ -276,12 +276,6 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
 }
 
 pub fn setup_browser_module(vm: &mut VirtualMachine) {
-    let state = PyRc::get_mut(&mut vm.state).unwrap();
-    state
-        .stdlib_inits
-        .insert("_browser".to_owned(), Box::new(make_module));
-    state.add_frozen(
-        &vm.ctx,
-        py_freeze!(file = "src/browser.py", module_name = "browser"),
-    );
+    state.add_native_module("_browser".to_owned(), Box::new(make_module));
+    vm.add_frozen(py_freeze!(file = "src/browser.py", module_name = "browser"));
 }

wasm/lib/src/js_module.rs

@@ -488,8 +488,5 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
 }
 
 pub fn setup_js_module(vm: &mut VirtualMachine) {
-    let state = PyRc::get_mut(&mut vm.state).unwrap();
-    state
-        .stdlib_inits
-        .insert("_js".to_owned(), Box::new(make_module));
+    vm.add_native_module("_js".to_owned(), Box::new(make_module));
 }

wasm/lib/src/vm_class.rs

@@ -32,14 +32,12 @@ impl StoredVirtualMachine {
 
             js_module::setup_js_module(vm);
             if inject_browser_module {
-                PyRc::get_mut(&mut vm.state).unwrap().stdlib_inits.insert(
-                    "_window".to_owned(),
-                    Box::new(|vm| {
-                        py_module!(vm, "_window", {
-                            "window" => js_module::PyJsValue::new(wasm_builtins::window()).into_ref(vm),
-                        })
-                    }),
-                );
+                fn init_window_module(vm: &VirtualMachine) -> PyObjectRef {
+                    py_module!(vm, "_window", {
+                        "window" => js_module::PyJsValue::new(wasm_builtins::window()).into_ref(vm),
+                    })
+                }
+                vm.add_native_module("_window".to_owned(), Box::new(init_window_module));
                 setup_browser_module(vm);
             }