diff --git a/crates/compiler-core/src/marshal.rs b/crates/compiler-core/src/marshal.rs index fee6a111c98..5587c6314a6 100644 --- a/crates/compiler-core/src/marshal.rs +++ b/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 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, ) -> Result; + fn freeze_dict(&self, _dict: Self::Value) -> Result { + 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( 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( 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( 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))??; diff --git a/crates/vm/src/stdlib/marshal.rs b/crates/vm/src/stdlib/marshal.rs index 9d7d6458eec..7e6d1700007 100644 --- a/crates/vm/src/stdlib/marshal.rs +++ b/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::>(); f(DumpableValue::Dict(&entries)) } + ref pyfrozen @ PyFrozenDict => { + let entries = DictIter::new(&pyfrozen.dict).collect::>(); + 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::().is_some() - || obj.downcast_ref::().is_some(); + || obj.downcast_ref::().is_some() + || obj.downcast_ref::().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::() { - 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 { Some(self.vm.ctx.new_dict().into()) } + fn freeze_dict(&self, source: Self::Value) -> Result { + 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,