JSON perf improvements (#314)

Author: samuelcolvin

.gitignore

@@ -19,3 +19,4 @@
 /crates/fuzz/corpus/
 /.worktrees/
 /.zed/
+/scripts/json_files/

crates/monty/src/bytecode/vm/mod.rs

@@ -31,7 +31,7 @@ use crate::{
     heap_data::{Closure, FunctionDefaults},
     intern::{FunctionId, Interns, StringId},
     io::PrintWriter,
-    modules::StandardLib,
+    modules::{StandardLib, json::JsonStringCache},
     os::OsFunction,
     parse::CodeRange,
     resource::ResourceTracker,
@@ -583,6 +583,13 @@ pub struct VM<'h, 'a, T: ResourceTracker> {
     /// back to a `NameError`, so the traceback points to the name reference rather than
     /// the call expression.
     ext_function_load_ip: Option<usize>,
+
+    /// Per-run string cache for `json.loads()`.
+    ///
+    /// Deduplicates heap allocations for repeated strings (especially dict keys)
+    /// across multiple `json.loads()` calls within a single execution. Lazily
+    /// initialized on first use, cleaned up in [`cleanup()`](Self::cleanup).
+    pub(crate) json_string_cache: JsonStringCache,
 }
 
 impl<'h, 'a, T: ResourceTracker> VM<'h, 'a, T> {
@@ -605,6 +612,7 @@ impl<'h, 'a, T: ResourceTracker> VM<'h, 'a, T> {
             scheduler: Scheduler::new(),
             ext_function_load_ip: None, // Set by LoadGlobalCallable/LoadLocalCallable
             module_code: None,
+            json_string_cache: JsonStringCache::default(),
         }
     }
 
@@ -666,8 +674,10 @@ impl<'h, 'a, T: ResourceTracker> VM<'h, 'a, T> {
             scheduler: snapshot.scheduler,
             module_code: Some(module_code),
             ext_function_load_ip: None,
+            json_string_cache: JsonStringCache::default(),
         }
     }
+
     /// Consumes the VM and creates a snapshot for pause/resume.
     ///
     /// **Ownership transfer:** This method takes `self` by value, consuming the VM.
@@ -676,7 +686,11 @@ impl<'h, 'a, T: ResourceTracker> VM<'h, 'a, T> {
     ///
     /// This is NOT a clone - it's a transfer. After calling this, the original VM
     /// is gone and only the snapshot (+ serialized heap/namespaces) represents the state.
-    pub fn snapshot(self) -> VMSnapshot {
+    pub fn snapshot(mut self) -> VMSnapshot {
+        // Drop cached JSON strings before consuming the VM — they are not
+        // included in the snapshot and their refcounts must be decremented.
+        self.json_string_cache.drop_all(self.heap);
+
         VMSnapshot {
             // Move values directly — no clone, no refcount increment needed
             // (the VM owned them, now the snapshot owns them)
@@ -709,6 +723,8 @@ impl<'h, 'a, T: ResourceTracker> VM<'h, 'a, T> {
         // Clean up scheduler state (task stacks, pending calls, resolved values, frame cells)
         self.scheduler.cleanup(self.heap);
         self.globals.drain(..).drop_with_heap(self.heap);
+        // Release cached JSON string values
+        self.json_string_cache.drop_all(self.heap);
     }
 
     /// Returns the `stack_base` of the current (topmost) call frame.
@@ -1710,9 +1726,14 @@ impl<'h, 'a, T: ResourceTracker> VM<'h, 'a, T> {
         let stack_roots = self.stack.iter().filter_map(Value::ref_id);
         let globals_roots = self.globals.iter().filter_map(Value::ref_id);
         let exc_roots = self.exception_stack.iter().filter_map(Value::ref_id);
+        let json_cache_roots = self.json_string_cache.gc_roots();
 
         // Collect all roots into a vec to avoid lifetime issues
-        let roots: Vec<HeapId> = stack_roots.chain(globals_roots).chain(exc_roots).collect();
+        let roots: Vec<HeapId> = stack_roots
+            .chain(globals_roots)
+            .chain(exc_roots)
+            .chain(json_cache_roots)
+            .collect();
 
         self.heap.collect_garbage(roots);
     }

crates/monty/src/modules/json/dump.rs

@@ -6,7 +6,6 @@
 use std::{
     cmp::Ordering,
     fmt::{Display, Write},
-    mem,
 };
 
 use crate::{
@@ -17,7 +16,7 @@ use crate::{
     heap::{DropWithHeap, HeapData, HeapGuard, HeapId, HeapReadOutput},
     intern::StaticStrings,
     resource::ResourceTracker,
-    sorting::sort_indices,
+    sorting::{apply_permutation, sort_indices},
     types::{PyTrait, long_int::check_bigint_str_digits_limit, str::allocate_string},
     value::Value,
 };
@@ -521,18 +520,7 @@ fn serialize_dict(
     vm: &mut VM<'_, '_, impl ResourceTracker>,
 ) -> RunResult<()> {
     if config.skipkeys() {
-        // Cannot use `retain` here because removed `Value::Ref` entries need
-        // `drop_with_heap` to decrement reference counts properly.
-        let mut i = 0;
-        while i < entries.len() {
-            if is_json_key_allowed(&entries[i].0, vm) {
-                i += 1;
-            } else {
-                let (key, value) = entries.remove(i);
-                key.drop_with_heap(vm);
-                value.drop_with_heap(vm);
-            }
-        }
+        skip_disallowed_dict_keys(entries, vm);
     } else if let Some((key, _)) = entries.iter().find(|(key, _)| !is_json_key_allowed(key, vm)) {
         return Err(ExcType::json_invalid_key_error(key.py_type(vm)));
     }
@@ -575,18 +563,31 @@ fn sort_dict_entries(entries: &mut Vec<(Value, Value)>, vm: &mut VM<'_, '_, impl
     let mut compare_values_guard = HeapGuard::new(compare_values, vm);
     let (compare_values, vm) = compare_values_guard.as_parts_mut();
     sort_indices(&mut indices, compare_values.as_slice(), false, vm)?;
+    apply_permutation(entries.as_mut_slice(), &mut indices);
+    Ok(())
+}
 
-    let mut ordered: Vec<(Value, Value)> = Vec::with_capacity(entries.len());
-    for index in indices {
-        ordered.push((
-            entries[index].0.clone_with_heap(vm),
-            entries[index].1.clone_with_heap(vm),
-        ));
+/// Removes dict entries whose keys are not JSON-serializable, preserving order.
+///
+/// `skipkeys=True` must drop invalid entries without disturbing the relative
+/// order of the retained pairs. A two-pointer compaction avoids the repeated
+/// shifting cost of `Vec::remove(i)` while still cleaning up skipped `Value`
+/// references with `drop_with_heap`.
+fn skip_disallowed_dict_keys(entries: &mut Vec<(Value, Value)>, vm: &mut VM<'_, '_, impl ResourceTracker>) {
+    let mut write = 0;
+    for read in 0..entries.len() {
+        if is_json_key_allowed(&entries[read].0, vm) {
+            if write != read {
+                entries.swap(write, read);
+            }
+            write += 1;
+        }
     }
 
-    let old_entries = mem::replace(entries, ordered);
-    old_entries.drop_with_heap(vm);
-    Ok(())
+    for (key, value) in entries.drain(write..) {
+        key.drop_with_heap(vm);
+        value.drop_with_heap(vm);
+    }
 }
 
 /// Returns whether a value is an allowed JSON object key type.
@@ -788,30 +789,105 @@ fn with_entered_container<R>(
 
 /// Writes a Rust string as a JSON string token.
 ///
-/// The writer escapes control characters, quotes, and backslashes in all modes.
-/// When `ensure_ascii` is enabled, non-ASCII code points are emitted as `\uXXXX`
-/// escapes using surrogate pairs for supplementary-plane characters.
+/// Uses a byte-oriented batch strategy inspired by serde_json: a 256-entry
+/// lookup table classifies each byte in O(1), and contiguous runs of safe bytes
+/// are flushed with a single `push_str` rather than character-by-character.
+///
+/// When `ensure_ascii` is enabled, non-ASCII code points (bytes >= 0x80) are
+/// emitted as `\uXXXX` escapes using surrogate pairs for supplementary-plane
+/// characters.
 fn write_json_string(value: &str, out: &mut String, ensure_ascii: bool) {
     out.push('"');
-    for ch in value.chars() {
-        match ch {
-            '"' => out.push_str("\\\""),
-            '\\' => out.push_str("\\\\"),
-            '\u{08}' => out.push_str("\\b"),
-            '\u{0C}' => out.push_str("\\f"),
-            '\n' => out.push_str("\\n"),
-            '\r' => out.push_str("\\r"),
-            '\t' => out.push_str("\\t"),
-            ch if ch <= '\u{1F}' => {
-                write!(out, "\\u{:04x}", ch as u32).expect("writing to String cannot fail");
+    let bytes = value.as_bytes();
+    let mut start = 0;
+    let mut i = 0;
+
+    while i < bytes.len() {
+        let byte = bytes[i];
+
+        if ensure_ascii && byte >= 0x7F {
+            // Flush the safe ASCII run accumulated so far.
+            out.push_str(&value[start..i]);
+            if byte == 0x7F {
+                // DEL (0x7F) is a control character that CPython escapes.
+                out.push_str("\\u007f");
+                i += 1;
+            } else {
+                // Decode the full character at this position and emit \uXXXX escapes.
+                let ch = value[i..].chars().next().expect("valid UTF-8");
+                write_json_escape_for_non_ascii(ch, out);
+                i += ch.len_utf8();
+            }
+            start = i;
+            continue;
+        }
+
+        let escape = ESCAPE_TABLE[byte as usize];
+        if escape == 0 {
+            // Safe byte — keep scanning.
+            i += 1;
+            continue;
+        }
+
+        // Flush the safe run before this byte.
+        out.push_str(&value[start..i]);
+
+        // Write the escape sequence.
+        match escape {
+            b'b' => out.push_str("\\b"),
+            b't' => out.push_str("\\t"),
+            b'n' => out.push_str("\\n"),
+            b'f' => out.push_str("\\f"),
+            b'r' => out.push_str("\\r"),
+            b'"' => out.push_str("\\\""),
+            b'\\' => out.push_str("\\\\"),
+            b'u' => {
+                write!(out, "\\u{:04x}", u32::from(byte)).expect("writing to String cannot fail");
             }
-            ch if ensure_ascii && (ch as u32) > 0x7E => write_json_escape_for_non_ascii(ch, out),
-            ch => out.push(ch),
+            _ => unreachable!(),
         }
+
+        i += 1;
+        start = i;
     }
+
+    // Flush the final safe run.
+    out.push_str(&value[start..]);
     out.push('"');
 }
 
+/// Byte lookup table for JSON string escaping.
+///
+/// Each entry is either 0 (byte is safe, no escaping needed) or a shorthand
+/// character that indicates which escape to emit:
+/// - `b'"'`  → `\"`
+/// - `b'\\'` → `\\`
+/// - `b'b'`  → `\b` (backspace, 0x08)
+/// - `b't'`  → `\t` (tab, 0x09)
+/// - `b'n'`  → `\n` (newline, 0x0A)
+/// - `b'f'`  → `\f` (form feed, 0x0C)
+/// - `b'r'`  → `\r` (carriage return, 0x0D)
+/// - `b'u'`  → `\u00XX` (other control characters, 0x00–0x1F)
+#[rustfmt::skip]
+static ESCAPE_TABLE: [u8; 256] = {
+    let mut table = [0u8; 256];
+    // Control characters 0x00–0x1F default to \u00XX escapes.
+    let mut i = 0;
+    while i < 0x20 {
+        table[i] = b'u';
+        i += 1;
+    }
+    // Override the named escapes.
+    table[0x08] = b'b';  // backspace
+    table[0x09] = b't';  // tab
+    table[0x0A] = b'n';  // newline
+    table[0x0C] = b'f';  // form feed
+    table[0x0D] = b'r';  // carriage return
+    table[0x22] = b'"';  // quote
+    table[0x5C] = b'\\'; // backslash
+    table
+};
+
 /// Writes a non-ASCII character using JSON `\uXXXX` escapes.
 ///
 /// Code points above `U+FFFF` are encoded as UTF-16 surrogate pairs to match