diff --git a/crates/stdlib/src/pystruct.rs b/crates/stdlib/src/pystruct.rs index b217e2005..7425b3c33 100644 --- a/crates/stdlib/src/pystruct.rs +++ b/crates/stdlib/src/pystruct.rs @@ -13,7 +13,10 @@ pub(crate) mod _struct { AsObject, Py, PyObjectRef, PyPayload, PyResult, TryFromObject, VirtualMachine, buffer::{FormatSpec, new_struct_error, struct_error_type}, builtins::{PyBytes, PyStr, PyStrRef, PyTupleRef, PyType, PyTypeRef}, - common::lock::{PyMappedRwLockReadGuard, PyRwLock, PyRwLockReadGuard}, + common::{ + lock::{PyMappedRwLockReadGuard, PyRwLock, PyRwLockReadGuard}, + rc::PyRc, + }, function::{ArgBytesLike, ArgMemoryBuffer, FuncArgs, PosArgs}, match_class, protocol::PyIterReturn, @@ -70,6 +73,12 @@ pub(crate) mod _struct { fn format_spec(&self, vm: &VirtualMachine) -> PyResult { FormatSpec::parse(self.0.as_bytes(), vm) } + + fn cached_format_spec(&self, vm: &VirtualMachine) -> PyResult> { + vm.state + .struct_format_cache + .get_or_parse(self.0.as_bytes(), vm) + } } fn get_buffer_offset( @@ -126,7 +135,7 @@ pub(crate) mod _struct { #[pyfunction] fn pack(fmt: IntoStructFormatBytes, args: PosArgs, vm: &VirtualMachine) -> PyResult> { - fmt.format_spec(vm)?.pack(args.into_vec(), vm) + fmt.cached_format_spec(vm)?.pack(args.into_vec(), vm) } #[pyfunction] @@ -137,7 +146,7 @@ pub(crate) mod _struct { args: PosArgs, vm: &VirtualMachine, ) -> PyResult<()> { - let format_spec = fmt.format_spec(vm)?; + let format_spec = fmt.cached_format_spec(vm)?; let offset = get_buffer_offset(buffer.len(), offset, format_spec.size, true, vm)?; buffer.with_ref(|data| format_spec.pack_into(&mut data[offset..], args.into_vec(), vm)) } @@ -148,7 +157,7 @@ pub(crate) mod _struct { buffer: ArgBytesLike, vm: &VirtualMachine, ) -> PyResult { - let format_spec = format.format_spec(vm)?; + let format_spec = format.cached_format_spec(vm)?; buffer.with_ref(|buf| format_spec.unpack(buf, vm)) } @@ -165,7 +174,7 @@ pub(crate) mod _struct { args: UpdateFromArgs, vm: &VirtualMachine, ) -> PyResult { - let format_spec = format.format_spec(vm)?; + let format_spec = format.cached_format_spec(vm)?; let offset = get_buffer_offset(args.buffer.len(), args.offset, format_spec.size, false, vm)?; args.buffer @@ -177,7 +186,7 @@ pub(crate) mod _struct { #[derive(Debug, PyPayload)] struct UnpackIterator { #[pytraverse(skip)] - format_spec: FormatSpec, + format_spec: PyRc, buffer: ArgBytesLike, #[pytraverse(skip)] offset: AtomicCell, @@ -186,7 +195,7 @@ pub(crate) mod _struct { impl UnpackIterator { fn with_buffer( vm: &VirtualMachine, - format_spec: FormatSpec, + format_spec: PyRc, buffer: ArgBytesLike, ) -> PyResult { if format_spec.size == 0 { @@ -243,13 +252,13 @@ pub(crate) mod _struct { buffer: ArgBytesLike, vm: &VirtualMachine, ) -> PyResult { - let format_spec = format.format_spec(vm)?; + let format_spec = format.cached_format_spec(vm)?; UnpackIterator::with_buffer(vm, format_spec, buffer) } #[pyfunction] fn calcsize(format: IntoStructFormatBytes, vm: &VirtualMachine) -> PyResult { - Ok(format.format_spec(vm)?.size) + Ok(format.cached_format_spec(vm)?.size) } /// What a `Struct` is once a format has been read into it. Held apart @@ -258,7 +267,7 @@ pub(crate) mod _struct { /// already holds a format. #[derive(Debug)] struct StructSpec { - spec: FormatSpec, + spec: PyRc, format: PyStrRef, } @@ -288,7 +297,7 @@ pub(crate) mod _struct { // cannot be read leaves the object as it was. let spec = fmt.format_spec(vm)?; *zelf.inner.write() = Some(StructSpec { - spec, + spec: PyRc::new(spec), format: fmt.0, }); Ok(()) @@ -385,9 +394,10 @@ pub(crate) mod _struct { } // seems weird that this is part of the "public" API, but whatever - // TODO: implement a format code->spec cache like CPython does? #[pyfunction] - const fn _clearcache() {} + fn _clearcache(vm: &VirtualMachine) { + vm.state.struct_format_cache.clear(); + } #[pyattr(name = "error")] fn error_type(vm: &VirtualMachine) -> PyTypeRef { diff --git a/crates/vm/src/buffer.rs b/crates/vm/src/buffer.rs index 957d9c525..15c0050cb 100644 --- a/crates/vm/src/buffer.rs +++ b/crates/vm/src/buffer.rs @@ -1,7 +1,7 @@ use crate::{ AsObject, Py, PyObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine, builtins::{PyBaseExceptionRef, PyBytesRef, PyComplex, PyTuple, PyTupleRef, PyType, PyTypeRef}, - common::{static_cell, str::wchar_t}, + common::{lock::PyRwLock, rc::PyRc, static_cell, str::wchar_t}, convert::ToPyObject, exceptions, function::{ArgBytesLike, ArgIntoBool, ArgIntoComplex, ArgIntoFloat}, @@ -15,7 +15,7 @@ use itertools::Itertools; use malachite_bigint::BigInt; use num_complex::Complex64; use num_traits::{PrimInt, ToPrimitive}; -use std::os::raw; +use std::{collections::HashMap, os::raw}; type PackFunc = fn(&VirtualMachine, FormatType, PyObjectRef, &mut [u8]) -> Result<(), PackError>; type UnpackFunc = fn(&VirtualMachine, &[u8]) -> PyObjectRef; @@ -478,6 +478,37 @@ pub struct FormatSpec { pub arg_count: usize, } +/// Bounded, interpreter-local cache for the module-level `struct` functions. +/// Keys and specifications contain no Python objects or user callbacks. +#[derive(Default)] +pub struct FormatSpecCache { + entries: PyRwLock, PyRc>>, +} + +impl FormatSpecCache { + pub fn get_or_parse(&self, format: &[u8], vm: &VirtualMachine) -> PyResult> { + if let Some(spec) = self.entries.read().get(format) { + return Ok(spec.clone()); + } + + // Parsing can raise; neither errors nor Python conversions run under the cache lock. + let spec = PyRc::new(FormatSpec::parse(format, vm)?); + let mut entries = self.entries.write(); + if let Some(spec) = entries.get(format) { + return Ok(spec.clone()); + } + if entries.len() >= 100 { + entries.clear(); + } + entries.insert(format.into(), spec.clone()); + Ok(spec) + } + + pub fn clear(&self) { + *self.entries.write() = HashMap::new(); + } +} + impl FormatSpec { pub fn parse(fmt: &[u8], vm: &VirtualMachine) -> PyResult { let mut chars = fmt.iter().copied().peekable(); @@ -972,3 +1003,35 @@ pub fn new_struct_error>(vm: &VirtualMachine, msg: T) -> PyBase // machinery, independent of whether _struct was ever imported vm.new_exception_msg(struct_error_type(vm).to_owned(), msg.into()) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::Interpreter; + + #[test] + fn format_cache_reuses_specs_and_releases_evicted_entries() { + Interpreter::without_stdlib(Default::default()).enter(|vm| { + let cache = &vm.state.struct_format_cache; + let spec = cache.get_or_parse(b">, pub codec_registry: CodecsRegistry, + pub struct_format_cache: crate::buffer::FormatSpecCache, pub finalizing: AtomicBool, pub warnings: WarningsState, pub override_frozen_modules: AtomicCell, diff --git a/extra_tests/snippets/stdlib_struct.py b/extra_tests/snippets/stdlib_struct.py index 213059482..c80f55117 100644 --- a/extra_tests/snippets/stdlib_struct.py +++ b/extra_tests/snippets/stdlib_struct.py @@ -162,3 +162,25 @@ for call in ( # allocate must raise instead of aborting. with assert_raises(MemoryError): struct.pack("%dx" % (2**60)) + + +# Clearing or evicting cached formats must not invalidate a live iterator. +iterator = struct.iter_unpack("