Skip to content

Commit ec7f273

Browse files
youknowoneyouknowone
committed
marshal: write and read co_consts through the runtime bag
serialize_code gains a serialize_code_with variant that writes each co_consts entry through a caller-supplied writer; serialize_code keeps the BorrowedConstant writer as its default. The VM writer passes its own write_object_depth, so a code constant that Literal holds but BorrowedConstant cannot describe reaches the stream instead of panicking in borrow_obj_constant, and a constant shared with the enclosing object takes an entry in the writer's reference table. PyMarshalBag implements constant_ref_from_value, bytes_from_value, str_from_value and tuple_elements_from_value, which deserialize_code_value_inner requires to decode code fields through the runtime bag. Assisted-by: Claude
1 parent fec66fc commit ec7f273

3 files changed

Lines changed: 71 additions & 2 deletions

File tree

‎crates/compiler-core/src/marshal.rs‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1398,6 +1398,25 @@ pub fn serialize_value<W: Write, D: Dumpable>(
13981398
/// Split varnames/cellvars/freevars are reassembled into
13991399
/// co_localsplusnames/co_localspluskinds.
14001400
pub fn serialize_code<W: Write, C: Constant>(buf: &mut W, code: &CodeObject<C>) {
1401+
serialize_code_with(buf, code, |buf, constant| {
1402+
serialize_value(buf, constant.borrow_constant().into()).unwrap_or_else(|x| match x {});
1403+
Ok::<(), core::convert::Infallible>(())
1404+
})
1405+
.unwrap_or_else(|x| match x {})
1406+
}
1407+
1408+
/// Serialize a code object, writing each `co_consts` entry through
1409+
/// `write_constant`.
1410+
///
1411+
/// A runtime caller passes its own object writer so that values its constant
1412+
/// representation carries but `BorrowedConstant` cannot describe — lists,
1413+
/// dicts, sets — reach the stream, and so a constant shared with the enclosing
1414+
/// object keeps its entry in that writer's reference table.
1415+
pub fn serialize_code_with<W: Write, C: Constant, E>(
1416+
buf: &mut W,
1417+
code: &CodeObject<C>,
1418+
mut write_constant: impl FnMut(&mut W, &C) -> core::result::Result<(), E>,
1419+
) -> core::result::Result<(), E> {
14011420
// 1–5: scalar fields
14021421
buf.write_u32(code.arg_count);
14031422
buf.write_u32(code.posonlyarg_count);
@@ -1414,7 +1433,7 @@ pub fn serialize_code<W: Write, C: Constant>(buf: &mut W, code: &CodeObject<C>)
14141433
buf.write_u8(Type::Tuple as u8);
14151434
write_len(buf, code.constants.len());
14161435
for constant in &*code.constants {
1417-
serialize_value(buf, constant.borrow_constant().into()).unwrap_or_else(|x| match x {})
1436+
write_constant(buf, constant)?;
14181437
}
14191438

14201439
// 8: co_names (tuple of strings)
@@ -1457,6 +1476,7 @@ pub fn serialize_code<W: Write, C: Constant>(buf: &mut W, code: &CodeObject<C>)
14571476
// 16: co_exceptiontable
14581477
buf.write_u8(Type::Bytes as u8);
14591478
write_vec(buf, &code.exceptiontable);
1479+
Ok(())
14601480
}
14611481

14621482
fn write_marshal_str<W: Write>(buf: &mut W, s: &str) {

‎crates/vm/src/stdlib/marshal.rs‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -324,7 +324,14 @@ mod decl {
324324
}
325325
} else if let Some(co) = obj.downcast_ref::<PyCode>() {
326326
buf.write_u8(b'c');
327-
marshal::serialize_code(buf, &co.code);
327+
// `Literal` holds the exact object a constant was built from, so
328+
// route `co_consts` back through the object writer: it reaches the
329+
// values `BorrowedConstant` cannot describe and shares the one
330+
// reference table the reader indexes against.
331+
marshal::serialize_code_with(buf, &co.code, |buf, constant| {
332+
let constant = PyObjectRef::from(constant.clone());
333+
write_object_depth(buf, &constant, refs, version, vm, depth - 1)
334+
})?;
328335
} else if let Some(sl) = obj.downcast_ref::<crate::builtins::PySlice>() {
329336
if version < 5 {
330337
return Err(vm.new_value_error("unmarshallable object"));
@@ -570,6 +577,27 @@ mod decl {
570577
fn constant_bag(self) -> Self::ConstantBag {
571578
PyVmBag(self.vm)
572579
}
580+
/// `Literal` wraps any object, so a decoded `co_consts` entry is
581+
/// already its own compiler-side constant — no placeholder is needed
582+
/// and `make_code_with_constants` keeps the default.
583+
fn constant_ref_from_value(&self, value: &Self::Value) -> Option<Literal> {
584+
Some(Literal::from(value.clone()))
585+
}
586+
fn bytes_from_value(&self, value: &Self::Value) -> Option<Vec<u8>> {
587+
value
588+
.downcast_ref::<PyBytes>()
589+
.map(|bytes| bytes.as_bytes().to_vec())
590+
}
591+
fn str_from_value(&self, value: &Self::Value) -> Option<String> {
592+
value
593+
.downcast_ref::<PyStr>()
594+
.map(|str| str.to_string_lossy().into_owned())
595+
}
596+
fn tuple_elements_from_value(&self, value: &Self::Value) -> Option<Vec<Self::Value>> {
597+
value
598+
.downcast_ref::<PyTuple>()
599+
.map(|tuple| tuple.as_slice().to_vec())
600+
}
573601
}
574602

575603
fn deserialize_value(

‎extra_tests/snippets/stdlib_marshal.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,27 @@ def test_roundtrip(self):
7474

7575
assert eval(loaded) == eval(orig)
7676

77+
def test_roundtrip_non_constant_co_consts(self):
78+
# `code.replace` accepts any marshalable object, including values the
79+
# compiler constant representation cannot describe.
80+
orig = compile("1 + 1", "", "eval").replace(
81+
co_consts=([1, 2], {"a": 3}, {4, 5}, 6)
82+
)
83+
84+
loaded = marshal.loads(marshal.dumps(orig))
85+
86+
self.assertEqual(loaded.co_consts, ([1, 2], {"a": 3}, {4, 5}, 6))
87+
88+
def test_roundtrip_shared_co_const(self):
89+
# A constant shared with the enclosing object is written once and both
90+
# readers resolve the same reference.
91+
shared = ["shared"]
92+
orig = compile("1 + 1", "", "eval").replace(co_consts=(shared,))
93+
94+
loaded_code, loaded_shared = marshal.loads(marshal.dumps((orig, shared)))
95+
96+
self.assertIs(loaded_code.co_consts[0], loaded_shared)
97+
7798

7899
if __name__ == "__main__":
79100
unittest.main()

0 commit comments

Comments
 (0)