Cache parsed struct formats per interpreter (#8903)

* Cache parsed struct formats per interpreter

Assisted-by: Codex:GPT-6

* Preserve struct cache entries on concurrent misses

Assisted-by: Codex:GPT-6

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> {
             FormatSpec::parse(self.0.as_bytes(), vm)
         }
+
+        fn cached_format_spec(&self, vm: &VirtualMachine) -> PyResult<PyRc<FormatSpec>> {
+            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<Vec<u8>> {
-        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<PyTupleRef> {
-        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<PyTupleRef> {
-        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<FormatSpec>,
         buffer: ArgBytesLike,
         #[pytraverse(skip)]
         offset: AtomicCell<usize>,
@@ -186,7 +195,7 @@ pub(crate) mod _struct {
     impl UnpackIterator {
         fn with_buffer(
             vm: &VirtualMachine,
-            format_spec: FormatSpec,
+            format_spec: PyRc<FormatSpec>,
             buffer: ArgBytesLike,
         ) -> PyResult<Self> {
             if format_spec.size == 0 {
@@ -243,13 +252,13 @@ pub(crate) mod _struct {
         buffer: ArgBytesLike,
         vm: &VirtualMachine,
     ) -> PyResult<UnpackIterator> {
-        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<usize> {
-        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<FormatSpec>,
         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 {

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<HashMap<Box<[u8]>, PyRc<FormatSpec>>>,
+}
+
+impl FormatSpecCache {
+    pub fn get_or_parse(&self, format: &[u8], vm: &VirtualMachine) -> PyResult<PyRc<FormatSpec>> {
+        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<Self> {
         let mut chars = fmt.iter().copied().peekable();
@@ -972,3 +1003,35 @@ pub fn new_struct_error<T: Into<Wtf8Buf>>(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"<IH", vm).unwrap();
+            assert!(PyRc::ptr_eq(
+                &spec,
+                &cache.get_or_parse(b"<IH", vm).unwrap()
+            ));
+            let weak = PyRc::downgrade(&spec);
+            drop(spec);
+
+            for padding in 0..100 {
+                let format = format!("{padding}x");
+                cache.get_or_parse(format.as_bytes(), vm).unwrap();
+            }
+            assert!(weak.upgrade().is_none());
+
+            let spec = cache.get_or_parse(b"<IH", vm).unwrap();
+            let weak = PyRc::downgrade(&spec);
+            drop(spec);
+            cache.clear();
+            assert!(weak.upgrade().is_none());
+        });
+    }
+}

crates/vm/src/vm/interpreter.rs

@@ -161,6 +161,7 @@ where
         atexit_funcs: PyMutex::default(),
         audit_hooks: PyMutex::default(),
         codec_registry,
+        struct_format_cache: crate::buffer::FormatSpecCache::default(),
         finalizing: AtomicBool::new(false),
         warnings,
         override_frozen_modules: AtomicCell::new(0),

crates/vm/src/vm/mod.rs

@@ -824,6 +824,7 @@ pub struct PyGlobalState {
     /// `sys.addaudithook` hooks, shared by all threads of this interpreter.
     pub(crate) audit_hooks: PyMutex<Vec<PyObjectRef>>,
     pub codec_registry: CodecsRegistry,
+    pub struct_format_cache: crate::buffer::FormatSpecCache,
     pub finalizing: AtomicBool,
     pub warnings: WarningsState,
     pub override_frozen_modules: AtomicCell<isize>,

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("<H", b"\x01\x00\x02\x00")
+struct._clearcache()
+for padding in range(110):
+    assert struct.calcsize(f"{padding}x") == padding
+assert list(iterator) == [(1,), (2,)]
+
+
+class ClearFormats:
+    def __index__(self):
+        struct._clearcache()
+        assert struct.pack("<H", 7) == b"\x07\x00"
+        return 42
+
+
+# Argument conversion may clear and repopulate the cache during an active pack.
+assert struct.pack("<H", ClearFormats()) == b"\x2a\x00"
+buffer = bytearray(2)
+struct.pack_into("<H", buffer, 0, ClearFormats())
+assert buffer == b"\x2a\x00"