Skip to content

Commit fd7d107

Browse files
youknowoneyouknowone
authored
marshal: round-trip code constants through the runtime bag (#8516)
* Keep abnormal marshal loop test aligned with CPython Assisted-by: OpenAI Codex:gpt-5.5 * marshal: retain runtime code constants while loading Decode code fields through the runtime MarshalBag so co_consts values that do not fit the compiler constant enum remain available to VM code wrappers while the compiler table receives shape placeholders. Assisted-by: OpenAI Codex:GPT-5.4 * 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 2ed082a commit fd7d107

4 files changed

Lines changed: 267 additions & 22 deletions

File tree

‎Lib/test/test_marshal.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -410,13 +410,13 @@ def test_loads_abnormal_reference_loops(self):
410410
self.assertIsInstance(a[0], dict)
411411
self.assertIs(a[0][None], a)
412412

413-
# Direct self-reference which cannot be created in Python. CPython
414-
# leaves this disabled because its reference counting cannot collect
415-
# the resulting cycle; RustPython's tracing collector can.
416-
data = b'\xa8\x01\x00\x00\x00r\x00\x00\x00\x00' # (<R>,)
417-
a = marshal.loads(data)
418-
self.assertIsInstance(a, tuple)
419-
self.assertIs(a[0], a)
413+
# Direct self-reference which cannot be created in Python.
414+
# This creates a reference loop which cannot be collected.
415+
if False:
416+
data = b'\xa8\x01\x00\x00\x00r\x00\x00\x00\x00' # (<R>,)
417+
a = marshal.loads(data)
418+
self.assertIsInstance(a, tuple)
419+
self.assertIs(a[0], a)
420420

421421
# Direct self-references which cannot be created in Python
422422
# because of unhashability.

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

Lines changed: 210 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -555,6 +555,18 @@ pub trait MarshalBag: Copy {
555555
code: CodeObject<<Self::ConstantBag as ConstantBag>::Constant>,
556556
) -> Self::Value;
557557

558+
/// Construct a runtime code object while retaining the exact values read
559+
/// from ``co_consts``. Compiler bags ignore this second channel; runtime
560+
/// bags use it for marshalable values (lists, dicts, sets, recursive
561+
/// containers) that their compiler constant representation cannot hold.
562+
fn make_code_with_constants(
563+
&self,
564+
code: CodeObject<<Self::ConstantBag as ConstantBag>::Constant>,
565+
_constants: Vec<Self::Value>,
566+
) -> Self::Value {
567+
self.make_code(code)
568+
}
569+
558570
fn make_stop_iter(&self) -> Result<Self::Value>;
559571

560572
fn make_list(&self, it: impl Iterator<Item = Self::Value>) -> Result<Self::Value>;
@@ -630,6 +642,30 @@ pub trait MarshalBag: Copy {
630642
) -> Option<<Self::ConstantBag as ConstantBag>::Constant> {
631643
None
632644
}
645+
646+
/// Convert a runtime constant to the compiler-side shape stored in
647+
/// ``CodeObject``. Runtime implementations may return a semantically
648+
/// unused placeholder when the exact value is carried by
649+
/// `make_code_with_constants` instead.
650+
fn code_constant_from_value(
651+
&self,
652+
value: &Self::Value,
653+
) -> Result<<Self::ConstantBag as ConstantBag>::Constant> {
654+
self.constant_ref_from_value(value)
655+
.ok_or(MarshalError::BadType)
656+
}
657+
658+
fn bytes_from_value(&self, _value: &Self::Value) -> Option<Vec<u8>> {
659+
None
660+
}
661+
662+
fn str_from_value(&self, _value: &Self::Value) -> Option<alloc::string::String> {
663+
None
664+
}
665+
666+
fn tuple_elements_from_value(&self, _value: &Self::Value) -> Option<Vec<Self::Value>> {
667+
None
668+
}
633669
}
634670

635671
impl<Bag: ConstantBag> MarshalBag for Bag {
@@ -731,6 +767,27 @@ impl<Bag: ConstantBag> MarshalBag for Bag {
731767
) -> Option<<Self::ConstantBag as ConstantBag>::Constant> {
732768
Some(value.clone())
733769
}
770+
771+
fn bytes_from_value(&self, value: &Self::Value) -> Option<Vec<u8>> {
772+
match value.borrow_constant() {
773+
BorrowedConstant::Bytes { value } => Some(value.to_vec()),
774+
_ => None,
775+
}
776+
}
777+
778+
fn str_from_value(&self, value: &Self::Value) -> Option<alloc::string::String> {
779+
match value.borrow_constant() {
780+
BorrowedConstant::Str { value } => Some(value.to_string_lossy().into_owned()),
781+
_ => None,
782+
}
783+
}
784+
785+
fn tuple_elements_from_value(&self, value: &Self::Value) -> Option<Vec<Self::Value>> {
786+
match value.borrow_constant() {
787+
BorrowedConstant::Tuple { elements } => Some(elements.to_vec()),
788+
_ => None,
789+
}
790+
}
734791
}
735792

736793
pub const MAX_MARSHAL_STACK_DEPTH: usize = 2000;
@@ -789,20 +846,8 @@ fn deserialize_value_after_header<R: Read, Bag: MarshalBag>(
789846
};
790847

791848
let typ = Type::try_from(type_code)?;
792-
// CPython's r_object() uses one global ref table: TYPE_CODE reserves its
793-
// slot before reading code fields, and those fields may use later TYPE_REF
794-
// indexes. Keep the same indexes even when Bag::Value and Constant differ.
795849
let value = if matches!(typ, Type::Code) {
796-
let mut inner_refs: Vec<Option<<Bag::ConstantBag as ConstantBag>::Constant>> = refs
797-
.iter()
798-
.map(|value| {
799-
value
800-
.as_ref()
801-
.and_then(|value| bag.constant_ref_from_value(value))
802-
})
803-
.collect();
804-
let code = deserialize_code_inner(rdr, bag.constant_bag(), depth - 1, &mut inner_refs)?;
805-
bag.make_code(code)
850+
deserialize_code_value_inner(rdr, bag, depth - 1, refs)?
806851
} else {
807852
deserialize_value_typed(rdr, bag, depth, refs, typ, slot)?
808853
};
@@ -813,6 +858,137 @@ fn deserialize_value_after_header<R: Read, Bag: MarshalBag>(
813858
Ok(value)
814859
}
815860

861+
/// Decode a code object through the runtime bag. CPython's marshal reader
862+
/// keeps one reference table for the code fields and `co_consts`; using
863+
/// `Bag::Value` here preserves that index space and lets runtime-only
864+
/// constants survive alongside the compiler representation.
865+
fn deserialize_code_value_inner<R: Read, Bag: MarshalBag>(
866+
rdr: &mut R,
867+
bag: Bag,
868+
depth: usize,
869+
refs: &mut Vec<Option<Bag::Value>>,
870+
) -> Result<Bag::Value> {
871+
if depth == 0 {
872+
return Err(MarshalError::InvalidBytecode);
873+
}
874+
let arg_count = rdr.read_u32()?;
875+
let posonlyarg_count = rdr.read_u32()?;
876+
let kwonlyarg_count = rdr.read_u32()?;
877+
let max_stackdepth = rdr.read_u32()?;
878+
let flags = CodeFlags::from_bits_truncate(rdr.read_u32()?);
879+
let child_depth = depth - 1;
880+
881+
let code_value = deserialize_value_depth(rdr, bag, child_depth, refs)?;
882+
let code_bytes = bag
883+
.bytes_from_value(&code_value)
884+
.ok_or(MarshalError::BadType)?;
885+
886+
let consts_value = deserialize_value_depth(rdr, bag, child_depth, refs)?;
887+
let constant_values = bag
888+
.tuple_elements_from_value(&consts_value)
889+
.ok_or(MarshalError::BadType)?;
890+
let constants = constant_values
891+
.iter()
892+
.map(|value| bag.code_constant_from_value(value))
893+
.collect::<Result<Vec<_>>>()?
894+
.into_iter()
895+
.collect();
896+
897+
let read_strings =
898+
|rdr: &mut R, refs: &mut Vec<Option<Bag::Value>>| -> Result<Vec<alloc::string::String>> {
899+
let tuple = deserialize_value_depth(rdr, bag, child_depth, refs)?;
900+
bag.tuple_elements_from_value(&tuple)
901+
.ok_or(MarshalError::BadType)?
902+
.iter()
903+
.map(|value| bag.str_from_value(value).ok_or(MarshalError::BadType))
904+
.collect()
905+
};
906+
let names_raw = read_strings(rdr, refs)?;
907+
let localsplusnames = read_strings(rdr, refs)?;
908+
909+
let kinds_value = deserialize_value_depth(rdr, bag, child_depth, refs)?;
910+
let localspluskinds = bag
911+
.bytes_from_value(&kinds_value)
912+
.ok_or(MarshalError::BadType)?;
913+
914+
let read_string =
915+
|rdr: &mut R, refs: &mut Vec<Option<Bag::Value>>| -> Result<alloc::string::String> {
916+
let value = deserialize_value_depth(rdr, bag, child_depth, refs)?;
917+
bag.str_from_value(&value).ok_or(MarshalError::BadType)
918+
};
919+
let source_path_raw = read_string(rdr, refs)?;
920+
let obj_name_raw = read_string(rdr, refs)?;
921+
let qualname_raw = read_string(rdr, refs)?;
922+
923+
let first_line_raw = rdr.read_u32()? as i32;
924+
let first_line_number = if first_line_raw > 0 {
925+
OneIndexed::new(first_line_raw as usize)
926+
} else {
927+
None
928+
};
929+
let linetable_value = deserialize_value_depth(rdr, bag, child_depth, refs)?;
930+
let linetable = bag
931+
.bytes_from_value(&linetable_value)
932+
.ok_or(MarshalError::BadType)?
933+
.into_boxed_slice();
934+
let exceptiontable_value = deserialize_value_depth(rdr, bag, child_depth, refs)?;
935+
let exceptiontable = bag
936+
.bytes_from_value(&exceptiontable_value)
937+
.ok_or(MarshalError::BadType)?
938+
.into_boxed_slice();
939+
940+
let lp = split_localplus(
941+
&localsplusnames
942+
.iter()
943+
.map(|s| s.as_str())
944+
.collect::<Vec<_>>(),
945+
&localspluskinds,
946+
arg_count,
947+
kwonlyarg_count,
948+
flags,
949+
)?;
950+
let instructions = CodeUnits::try_from(code_bytes.as_slice())?;
951+
let locations = linetable_to_locations(&linetable, first_line_raw, instructions.len());
952+
let constant_bag = bag.constant_bag();
953+
let code = CodeObject {
954+
instructions,
955+
locations,
956+
flags,
957+
posonlyarg_count,
958+
arg_count,
959+
kwonlyarg_count,
960+
source_path: constant_bag.make_name(&source_path_raw),
961+
first_line_number,
962+
max_stackdepth,
963+
obj_name: constant_bag.make_name(&obj_name_raw),
964+
qualname: constant_bag.make_name(&qualname_raw),
965+
constants,
966+
names: names_raw
967+
.iter()
968+
.map(|name| constant_bag.make_name(name))
969+
.collect(),
970+
varnames: lp
971+
.varnames
972+
.iter()
973+
.map(|name| constant_bag.make_name(name))
974+
.collect(),
975+
cellvars: lp
976+
.cellvars
977+
.iter()
978+
.map(|name| constant_bag.make_name(name))
979+
.collect(),
980+
freevars: lp
981+
.freevars
982+
.iter()
983+
.map(|name| constant_bag.make_name(name))
984+
.collect(),
985+
localspluskinds: localspluskinds.into_boxed_slice(),
986+
linetable,
987+
exceptiontable,
988+
};
989+
Ok(bag.make_code_with_constants(code, constant_values))
990+
}
991+
816992
fn deserialize_value_typed<R: Read, Bag: MarshalBag>(
817993
rdr: &mut R,
818994
bag: Bag,
@@ -1222,6 +1398,25 @@ pub fn serialize_value<W: Write, D: Dumpable>(
12221398
/// Split varnames/cellvars/freevars are reassembled into
12231399
/// co_localsplusnames/co_localspluskinds.
12241400
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> {
12251420
// 1–5: scalar fields
12261421
buf.write_u32(code.arg_count);
12271422
buf.write_u32(code.posonlyarg_count);
@@ -1238,7 +1433,7 @@ pub fn serialize_code<W: Write, C: Constant>(buf: &mut W, code: &CodeObject<C>)
12381433
buf.write_u8(Type::Tuple as u8);
12391434
write_len(buf, code.constants.len());
12401435
for constant in &*code.constants {
1241-
serialize_value(buf, constant.borrow_constant().into()).unwrap_or_else(|x| match x {})
1436+
write_constant(buf, constant)?;
12421437
}
12431438

12441439
// 8: co_names (tuple of strings)
@@ -1281,6 +1476,7 @@ pub fn serialize_code<W: Write, C: Constant>(buf: &mut W, code: &CodeObject<C>)
12811476
// 16: co_exceptiontable
12821477
buf.write_u8(Type::Bytes as u8);
12831478
write_vec(buf, &code.exceptiontable);
1479+
Ok(())
12841480
}
12851481

12861482
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)