Support marshal format 6 and frozendict serialization (#9130)

Extract the Python 3.15 final mapping format while preserving current code-object field identities and owned writer references.

Assisted-by: OpenAI Codex:model version unavailable

Co-authored-by: youknowdot <dot@youknowone.org>

crates/compiler-core/src/marshal.rs

@@ -5,7 +5,7 @@ use malachite_bigint::{BigInt, Sign};
 use num_complex::Complex64;
 use rustpython_wtf8::Wtf8;
 
-pub const FORMAT_VERSION: u32 = 5;
+pub const FORMAT_VERSION: u32 = 6;
 
 #[derive(Clone, Copy, Debug)]
 pub enum MarshalError {
@@ -84,6 +84,7 @@ enum Type {
     SmallTuple = b')',
     List = b'[',
     Dict = b'{',
+    FrozenDict = b'}',
     Code = b'c',
     Unicode = b'u',
     Set = b'<',
@@ -120,6 +121,7 @@ impl TryFrom<u8> for Type {
             b')' => Self::SmallTuple,
             b'[' => Self::List,
             b'{' => Self::Dict,
+            b'}' => Self::FrozenDict,
             b'c' => Self::Code,
             b'u' => Self::Unicode,
             b'<' => Self::Set,
@@ -622,6 +624,10 @@ pub trait MarshalBag: Copy {
         it: impl Iterator<Item = (Self::Value, Self::Value)>,
     ) -> Result<Self::Value>;
 
+    fn freeze_dict(&self, _dict: Self::Value) -> Result<Self::Value> {
+        Err(MarshalError::BadType)
+    }
+
     /// Install partially-built containers in the marshal reference table
     /// before reading their children, as CPython's `r_object()` does.
     /// Runtime bags can opt in; constant bags retain collect-then-construct.
@@ -1196,22 +1202,30 @@ fn deserialize_value_typed<R: Read, Bag: MarshalBag>(
             let it = (0..len).map(|_| deserialize_value_depth(rdr, bag, d, refs));
             itertools::process_results(it, |it| bag.make_frozenset(it))??
         }
-        Type::Dict => {
+        Type::Dict | Type::FrozenDict => {
             let d = depth - 1;
-            if let Some(index) = slot
+            // Immutable mappings are published only after all entries are read.
+            let frozen = matches!(typ, Type::FrozenDict);
+            if (frozen || slot.is_some())
                 && let Some(dict) = bag.make_dict_placeholder()
             {
-                refs[index] = Some(dict.clone());
+                if !frozen && let Some(index) = slot {
+                    refs[index] = Some(dict.clone());
+                }
                 loop {
                     let raw = read_object_type(rdr)?;
                     if raw & !FLAG_REF == b'0' {
                         break;
                     }
                     let key = deserialize_value_after_header(rdr, bag, d, refs, raw)?;
-                    let value = deserialize_value_depth(rdr, bag, d, refs)?;
+                    let raw = read_object_type(rdr)?;
+                    if raw & !FLAG_REF == b'0' {
+                        break;
+                    }
+                    let value = deserialize_value_after_header(rdr, bag, d, refs, raw)?;
                     bag.insert_dict_item(&dict, key, value)?;
                 }
-                dict
+                if frozen { bag.freeze_dict(dict)? } else { dict }
             } else {
                 let mut pairs = Vec::new();
                 loop {
@@ -1220,10 +1234,15 @@ fn deserialize_value_typed<R: Read, Bag: MarshalBag>(
                         break;
                     }
                     let key = deserialize_value_after_header(rdr, bag, d, refs, raw)?;
-                    let value = deserialize_value_depth(rdr, bag, d, refs)?;
+                    let raw = read_object_type(rdr)?;
+                    if raw & !FLAG_REF == b'0' {
+                        break;
+                    }
+                    let value = deserialize_value_after_header(rdr, bag, d, refs, raw)?;
                     pairs.push((key, value));
                 }
-                bag.make_dict(pairs.into_iter())?
+                let dict = bag.make_dict(pairs.into_iter())?;
+                if frozen { bag.freeze_dict(dict)? } else { dict }
             }
         }
         Type::Bytes => {
@@ -1267,6 +1286,7 @@ pub enum DumpableValue<'a, D: Dumpable> {
     Set(&'a [D]),
     Frozenset(&'a [D]),
     Dict(&'a [(D, D)]),
+    Frozendict(&'a [(D, D)]),
     Slice(&'a D, &'a D, &'a D),
 }
 
@@ -1438,8 +1458,13 @@ pub fn serialize_value<W: Write, D: Dumpable>(
                 val.with_dump(|val| serialize_value(buf, val))??
             }
         }
-        DumpableValue::Dict(d) => {
-            buf.write_u8(Type::Dict as u8);
+        DumpableValue::Dict(d) | DumpableValue::Frozendict(d) => {
+            let typ = if matches!(constant, DumpableValue::Frozendict(_)) {
+                Type::FrozenDict
+            } else {
+                Type::Dict
+            };
+            buf.write_u8(typ as u8);
             for (k, v) in d {
                 k.with_dump(|val| serialize_value(buf, val))??;
                 v.with_dump(|val| serialize_value(buf, val))??;

crates/vm/src/stdlib/marshal.rs

@@ -4,14 +4,15 @@ pub(crate) use decl::module_def;
 #[pymodule(name = "marshal")]
 mod decl {
     use crate::builtins::code::{CodeObject, Literal, PyVmBag};
+    use crate::builtins::dict::DictIter;
     use crate::class::StaticType;
     use crate::common::wtf8::Wtf8;
     use crate::{
         PyObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine,
         builtins::{
-            PyBaseExceptionRef, PyBool, PyByteArray, PyBytes, PyCode, PyComplex, PyDict,
-            PyEllipsis, PyFloat, PyFrozenSet, PyInt, PyList, PyMemoryView, PyNone, PySet,
-            PyStopIteration, PyStr, PyTuple,
+            PyAnyDict, PyBaseExceptionRef, PyBool, PyByteArray, PyBytes, PyCode, PyComplex, PyDict,
+            PyEllipsis, PyFloat, PyFrozenDict, PyFrozenSet, PyInt, PyList, PyMemoryView, PyNone,
+            PySet, PyStopIteration, PyStr, PyTuple,
         },
         convert::ToPyObject,
         function::ArgBytesLike,
@@ -76,6 +77,10 @@ mod decl {
                     let entries = pydict.into_iter().collect::<Vec<_>>();
                     f(DumpableValue::Dict(&entries))
                 }
+                ref pyfrozen @ PyFrozenDict => {
+                    let entries = DictIter::new(&pyfrozen.dict).collect::<Vec<_>>();
+                    f(DumpableValue::Frozendict(&entries))
+                }
                 ref bytes @ PyBytes => {
                     f(DumpableValue::Bytes(bytes.as_bytes()))
                 }
@@ -95,7 +100,7 @@ mod decl {
     struct DumpsArgs {
         #[pyarg(positional)]
         value: PyObjectRef,
-        #[pyarg(positional, default = 5)]
+        #[pyarg(positional, default = 6)]
         version: i32,
         #[pyarg(named, default = true)]
         allow_code: bool,
@@ -261,12 +266,13 @@ mod decl {
             }
         }
         let type_pos = buf.len();
-        // A code or slice entry stays incomplete until its contents are
-        // written: the reader rebuilds both from their fields, so a
+        // A code, slice or frozendict entry stays incomplete until its contents are
+        // written: the reader rebuilds them from their fields, so a
         // back-reference issued while those fields are still being emitted
         // would name an object that does not exist yet.
         let requires_completion = obj.downcast_ref::<PyCode>().is_some()
-            || obj.downcast_ref::<crate::builtins::PySlice>().is_some();
+            || obj.downcast_ref::<crate::builtins::PySlice>().is_some()
+            || obj.downcast_ref::<PyFrozenDict>().is_some();
         if use_ref {
             refs.as_mut().unwrap().reserve(obj, requires_completion);
         }
@@ -385,9 +391,17 @@ mod decl {
             for elem in items.iter() {
                 write_object_depth(buf, elem, refs, version, allow_code, vm, depth - 1)?;
             }
-        } else if let Some(d) = obj.downcast_ref::<PyDict>() {
-            buf.write_u8(b'{');
-            for (k, v) in d {
+        } else if let Some(d) = PyAnyDict::from_object(obj) {
+            let typ = if d.is_frozen() {
+                if version < 6 || !obj.class().is(vm.ctx.types.frozendict_type) {
+                    return Err(vm.new_value_error("unmarshallable object"));
+                }
+                b'}'
+            } else {
+                b'{'
+            };
+            buf.write_u8(typ);
+            for (k, v) in DictIter::new(d.as_dict()) {
                 write_object_depth(buf, &k, refs, version, allow_code, vm, depth - 1)?;
                 write_object_depth(buf, &v, refs, version, allow_code, vm, depth - 1)?;
             }
@@ -507,7 +521,7 @@ mod decl {
         value: PyObjectRef,
         #[pyarg(positional)]
         file: PyObjectRef,
-        #[pyarg(positional, default = 5)]
+        #[pyarg(positional, default = 6)]
         version: i32,
         #[pyarg(named, default = true)]
         allow_code: bool,
@@ -716,6 +730,12 @@ mod decl {
         fn make_dict_placeholder(&self) -> Option<Self::Value> {
             Some(self.vm.ctx.new_dict().into())
         }
+        fn freeze_dict(&self, source: Self::Value) -> Result<Self::Value, marshal::MarshalError> {
+            let dict = PyDict::default();
+            dict.update_arg(source, true, "frozendict", self.vm)
+                .map_err(|error| self.remember_python_error(error))?;
+            Ok(PyFrozenDict::from_dict(dict).into_ref(&self.vm.ctx).into())
+        }
         fn insert_dict_item(
             &self,
             dict: &Self::Value,