Materialize code constants through the VM (#8867)

PyObjBag built code object constants from a Context alone, and its
frozenset arm was unimplemented!(), so ctx.new_code() panicked on code
containing a frozenset constant such as `x in {1, 2, 3}`.

IntoCodeObject now takes a VirtualMachine and uses PyVmBag for bytecode
and frozen code objects. Context::new_code, PyObjBag and
AsBag for &Context are removed; callers use VirtualMachine::new_code.

Assisted-by: Claude:claude-opus-5-5
Assisted-by: Grok:grok-4.7

crates/vm/src/builtins/code.rs

@@ -8,7 +8,7 @@ use crate::convert::ToPyException;
 use crate::{
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
     builtins::PyStrInterned,
-    bytecode::{self, AsBag, BorrowedConstant, CodeFlags, Constant, ConstantBag, Instruction},
+    bytecode::{self, BorrowedConstant, CodeFlags, Constant, ConstantBag, Instruction},
     class::{PyClassImpl, StaticType},
     convert::ToPyObject,
     frozen,
@@ -263,16 +263,6 @@ impl Constant for Literal {
     }
 }
 
-impl<'a> AsBag for &'a Context {
-    type Bag = PyObjBag<'a>;
-    fn as_bag(self) -> PyObjBag<'a> {
-        PyObjBag(self)
-    }
-}
-
-#[derive(Clone, Copy)]
-pub struct PyObjBag<'a>(pub &'a Context);
-
 /// Whether a string constant reads as a name. Those are the ones interned,
 /// the way `all_name_chars` picks them out.
 fn is_name_chars(value: &crate::common::wtf8::Wtf8) -> bool {
@@ -282,77 +272,6 @@ fn is_name_chars(value: &crate::common::wtf8::Wtf8) -> bool {
         .all(|&b| b.is_ascii_alphanumeric() || b == b'_')
 }
 
-impl ConstantBag for PyObjBag<'_> {
-    type Constant = Literal;
-
-    fn make_constant<C: Constant>(&self, constant: BorrowedConstant<'_, C>) -> Self::Constant {
-        let ctx = self.0;
-        let obj = match constant {
-            BorrowedConstant::Integer { value } => ctx.new_bigint(value).into(),
-            BorrowedConstant::Float { value } => ctx.new_float(value).into(),
-            BorrowedConstant::Complex { value } => ctx.new_complex(value).into(),
-            BorrowedConstant::Str { value } if is_name_chars(value) => {
-                ctx.intern_str(value).to_object()
-            }
-            BorrowedConstant::Str { value } => ctx.new_str(value).into(),
-            BorrowedConstant::Bytes { value } => ctx.new_bytes(value.to_vec()).into(),
-            BorrowedConstant::Boolean { value } => ctx.new_bool(value).into(),
-            BorrowedConstant::Code { code } => ctx.new_code(code.map_clone_bag(self)).into(),
-            BorrowedConstant::Tuple { elements } => {
-                let elements = elements
-                    .iter()
-                    .map(|constant| self.make_constant(constant.borrow_constant()).0)
-                    .collect();
-                ctx.new_tuple(elements).into()
-            }
-            BorrowedConstant::Slice { elements } => {
-                let [start, stop, step] = elements;
-                let start_obj = self.make_constant(start.borrow_constant()).0;
-                let stop_obj = self.make_constant(stop.borrow_constant()).0;
-                let step_obj = self.make_constant(step.borrow_constant()).0;
-                // Store as PySlice with Some() for all fields (even None values)
-                // so borrow_obj_constant can reference them.
-                use crate::builtins::PySlice;
-                PySlice {
-                    start: Some(start_obj),
-                    stop: stop_obj,
-                    step: Some(step_obj),
-                }
-                .into_ref(ctx)
-                .into()
-            }
-            BorrowedConstant::Frozenset { elements: _ } => {
-                // Creating a frozenset requires VirtualMachine for element hashing.
-                // PyObjBag only has Context, so we cannot construct PyFrozenSet here.
-                // Frozenset constants from .pyc are handled by PyMarshalBag which has VM access.
-                unimplemented!(
-                    "frozenset constant in PyObjBag::make_constant requires VirtualMachine"
-                )
-            }
-            BorrowedConstant::None => ctx.none(),
-            BorrowedConstant::Ellipsis => ctx.ellipsis.clone().into(),
-        };
-
-        Literal(obj)
-    }
-
-    fn make_name(&self, name: &str) -> &'static PyStrInterned {
-        self.0.intern_str(name)
-    }
-
-    fn make_int(&self, value: BigInt) -> Self::Constant {
-        Literal(self.0.new_int(value).into())
-    }
-
-    fn make_tuple(&self, elements: impl Iterator<Item = Self::Constant>) -> Self::Constant {
-        Literal(self.0.new_tuple(elements.map(|lit| lit.0).collect()).into())
-    }
-
-    fn make_code(&self, code: CodeObject) -> Self::Constant {
-        Literal(self.0.new_code(code).into())
-    }
-}
-
 #[derive(Clone, Copy)]
 pub(crate) struct PyVmBag<'a>(pub &'a VirtualMachine);
 
@@ -437,24 +356,24 @@ impl ConstantBag for PyVmBag<'_> {
 pub(crate) type CodeObject = bytecode::CodeObject<Literal>;
 
 pub trait IntoCodeObject {
-    fn into_code_object(self, ctx: &Context) -> CodeObject;
+    fn into_code_object(self, vm: &VirtualMachine) -> CodeObject;
 }
 
 impl IntoCodeObject for CodeObject {
-    fn into_code_object(self, _ctx: &Context) -> Self {
+    fn into_code_object(self, _vm: &VirtualMachine) -> Self {
         self
     }
 }
 
 impl IntoCodeObject for bytecode::CodeObject {
-    fn into_code_object(self, ctx: &Context) -> CodeObject {
-        self.map_bag(PyObjBag(ctx))
+    fn into_code_object(self, vm: &VirtualMachine) -> CodeObject {
+        self.map_bag(PyVmBag(vm))
     }
 }
 
 impl<B: AsRef<[u8]>> IntoCodeObject for frozen::FrozenCodeObject<B> {
-    fn into_code_object(self, ctx: &Context) -> CodeObject {
-        self.decode(ctx)
+    fn into_code_object(self, vm: &VirtualMachine) -> CodeObject {
+        self.decode(PyVmBag(vm))
     }
 }
 
@@ -1645,7 +1564,7 @@ impl Py<PyCode> {
 
 impl ToPyObject for CodeObject {
     fn to_pyobject(self, vm: &VirtualMachine) -> PyObjectRef {
-        vm.ctx.new_code(self).into()
+        vm.new_code(self).into()
     }
 }
 

crates/vm/src/stdlib/_testinternalcapi.rs

@@ -931,7 +931,7 @@ mod _testinternalcapi {
         let unit = super::metadata_to_code_unit(&metadata, vm)?;
         let code = rustpython_codegen::ir::assemble_for_tests(filename, rust_seq, unit, true)
             .map_err(|err| super::internal_error_to_py(err, vm))?;
-        Ok(vm.ctx.new_code(code))
+        Ok(vm.new_code(code))
     }
 
     #[cfg(feature = "codegen")]
@@ -1736,7 +1736,7 @@ fn constant_data_to_py(
         )
         .map_or_else(|_| vm.ctx.none(), |set| set.to_pyobject(vm)),
         ConstantData::Ellipsis => vm.ctx.ellipsis.clone().into(),
-        ConstantData::Code { code } => vm.ctx.new_code(*code).into(),
+        ConstantData::Code { code } => vm.new_code(*code).into(),
         ConstantData::Slice { elements } => {
             let [start, stop, step] = *elements;
             crate::builtins::PySlice {

crates/vm/src/vm/context.rs

@@ -5,7 +5,6 @@ use crate::{
         PyFrozenSet, PyInt, PyIntRef, PyList, PyListRef, PyNone, PyNotImplemented, PyStr,
         PyStrInterned, PyTuple, PyTupleRef, PyType, PyTypeRef, PyUtf8Str,
         bool_::PyBool,
-        code::{self, PyCode},
         descriptor::{
             MemberAccess, MemberKind, PyDescriptorOwned, PyMemberDef, PyMemberDescriptor,
             PyMemberFlags,
@@ -781,11 +780,6 @@ impl Context {
         PyRef::new_ref(object::PyBaseObject, class, dict).into()
     }
 
-    pub fn new_code(&self, code: impl code::IntoCodeObject) -> PyRef<PyCode> {
-        let code = code.into_code_object(self);
-        PyRef::new_ref(PyCode::new(code), self.types.code_type.to_owned(), None)
-    }
-
     pub fn new_capsule(
         &self,
         ptr: *mut c_void,

crates/vm/src/vm/interpreter.rs

@@ -2163,4 +2163,34 @@ for _ in range(40):
             assert!(result.try_to_bool(vm).unwrap());
         });
     }
+
+    /// Set and `for` literals become frozenset constants. Materializing them
+    /// needs the VM so element hashing can run.
+    #[cfg(feature = "compiler")]
+    #[test]
+    fn new_code_materializes_frozenset_constants() {
+        Interpreter::without_stdlib(Default::default()).enter(|vm| {
+            let source = "\
+a = 2 in {1, 2, 3}
+n = 0
+for e in {1, 2, 3}:
+    n += e
+";
+            let code = crate::compiler::compile(
+                source,
+                crate::compiler::Mode::Exec,
+                "<test>",
+                vm.compile_opts(),
+            )
+            .unwrap();
+            let code = vm.new_code(code);
+            let scope = vm.new_scope_with_builtins();
+            vm.run_code_obj(code, scope.clone()).unwrap();
+
+            let a = scope.globals.get_item("a", vm).unwrap();
+            assert!(a.is(&vm.ctx.true_value));
+            let n = scope.globals.get_item("n", vm).unwrap();
+            assert_eq!(*int::get_value(&n), 6_i32.to_bigint().unwrap());
+        })
+    }
 }

crates/vm/src/vm/vm_new.rs

@@ -18,6 +18,7 @@ use crate::{
         PyImportError, PyMemoryError, PyModule, PyNameError, PyOSError, PyStopIteration, PyStrRef,
         PySyntaxError, PySystemExit, PyType, PyTypeRef,
         builtin_func::PyNativeFunction,
+        code::{IntoCodeObject, PyCode},
         descriptor::PyMethodDescriptor,
         tuple::{IntoPyTuple, PyTupleRef},
     },
@@ -311,6 +312,11 @@ impl VirtualMachine {
         value.to_pyobject(self)
     }
 
+    pub fn new_code(&self, code: impl IntoCodeObject) -> PyRef<PyCode> {
+        let code = code.into_code_object(self);
+        PyRef::new_ref(PyCode::new(code), self.ctx.types.code_type.to_owned(), None)
+    }
+
     pub fn new_tuple(&self, value: impl IntoPyTuple) -> PyTupleRef {
         value.into_pytuple(self)
     }

examples/freeze/main.rs

@@ -11,7 +11,7 @@ fn run(vm: &vm::VirtualMachine) -> vm::PyResult<()> {
     // the file parameter is relative to the current file.
     let module = vm::py_compile!(file = "freeze.py");
 
-    let res = vm.run_code_obj(vm.ctx.new_code(module), scope);
+    let res = vm.run_code_obj(vm.new_code(module), scope);
 
     if let Err(exc) = res {
         vm.print_exception(&exc);

examples/mini_repl.rs

@@ -16,7 +16,7 @@ macro_rules! add_python_function {
         // compile the code to bytecode
         let code = vm::py_compile!(source = $src);
         // convert the rustpython_compiler_core::CodeObject to a PyRef<PyCode>
-        let code = $vm.ctx.new_code(code);
+        let code = $vm.new_code(code);
 
         // run the python code in the scope to store the function
         $vm.run_code_obj(code, $scope.clone())

src/lib.rs

@@ -143,7 +143,7 @@ __import__("io").TextIOWrapper(
         mode = "eval"
     );
     eprintln!("downloading get-pip.py...");
-    let getpip_code = vm.run_code_obj(vm.ctx.new_code(get_getpip), vm.new_scope_with_builtins())?;
+    let getpip_code = vm.run_code_obj(vm.new_code(get_getpip), vm.new_scope_with_builtins())?;
     let getpip_code: rustpython_vm::builtins::PyStrRef = getpip_code
         .downcast()
         .expect("TextIOWrapper.read() should return str");