perf(msgpack-bounds): walk shallow documents without allocating the nesting stack (LAB-12134)

README.md

@@ -274,7 +274,10 @@ them, so every input within the size cap gets their result and error. The pre-sc
 MessagePack themselves; it returns a `MsgpackStructureError` whose `Display` is
 the bare reason, with no prefix. `msgpack_document_len(bytes, max_depth)` runs
 the same walk and returns the document's length, so a caller that must reject
-trailing bytes compares it with `bytes.len()` instead of walking twice. See
+trailing bytes compares it with `bytes.len()` instead of walking twice. The walk
+allocates nothing for a document that nests at most 8 levels deep, and for a
+deeper one allocates only for the levels past 8; a rejection allocates its error
+message. See
 [`SECURITY.md`](SECURITY.md#envelope-decode-bounds).
 
 </details>

benches/perf_ir_baselines.json

@@ -15,9 +15,9 @@
         "keyring_decrypt/1024": 34739,
         "keyring_decrypt/64": 31754,
         "keyring_decrypt/65536": 264387,
-        "prescan/1024": 985,
-        "prescan/64": 883,
-        "prescan/65536": 935,
+        "prescan/1024": 910,
+        "prescan/64": 815,
+        "prescan/65536": 864,
         "retrieve/1024": 7409,
         "retrieve/1048576": 6679353,
         "retrieve/40": 2637,

src/msgpack_bounds.rs

@@ -27,7 +27,10 @@ fn reject(reason: impl Into<String>) -> MsgpackStructureError {
 /// A header-only walk: str/bin/ext payloads are skipped by offset, never read,
 /// and nothing is decoded. It keeps one `u64` per open non-empty collection, so
 /// at most `min(max_depth, bytes.len())` of them (8 KiB at a depth of 1024), and
-/// allocates nothing proportional to any declared length.
+/// allocates nothing proportional to any declared length. The first 8 live on
+/// the call stack, so a document that nests no deeper than that (an envelope
+/// `ByteStorage::store` writes nests 2 deep) is walked with no allocation. A
+/// rejection allocates its error message.
 ///
 /// Bytes after the document are neither read nor rejected. To find where the
 /// document ends, call [`msgpack_document_len`], which runs this same walk.
@@ -76,6 +79,57 @@ pub fn check_msgpack_structure(
     msgpack_document_len(bytes, max_depth).map(|_| ())
 }
 
+/// Open collections the walk keeps on the call stack before it spills to the heap.
+const INLINE_LEVELS: usize = 8;
+
+/// Elements still owed per open collection, innermost last; `len` is the
+/// depth. The outermost [`INLINE_LEVELS`] live inline and deeper ones in
+/// `spill`, so a shallow document never allocates. `spill` holds exactly the
+/// levels past the inline ones.
+struct OpenCollections {
+    inline: [u64; INLINE_LEVELS],
+    len: usize,
+    spill: Vec<u64>,
+}
+
+impl OpenCollections {
+    fn new() -> Self {
+        Self {
+            inline: [0; INLINE_LEVELS],
+            len: 0,
+            spill: Vec::new(),
+        }
+    }
+
+    fn len(&self) -> usize {
+        self.len
+    }
+
+    fn last_mut(&mut self) -> Option<&mut u64> {
+        match self.len {
+            0 => None,
+            n @ 1..=INLINE_LEVELS => Some(&mut self.inline[n - 1]),
+            _ => self.spill.last_mut(),
+        }
+    }
+
+    fn push(&mut self, owed: u64) {
+        if self.len < INLINE_LEVELS {
+            self.inline[self.len] = owed;
+        } else {
+            self.spill.push(owed);
+        }
+        self.len += 1;
+    }
+
+    fn pop(&mut self) {
+        if self.len > INLINE_LEVELS {
+            self.spill.pop();
+        }
+        self.len -= 1;
+    }
+}
+
 /// Runs the walk [`check_msgpack_structure`] describes and returns the length,
 /// in bytes, of the document at the start of `bytes`.
 ///
@@ -119,9 +173,9 @@ pub fn msgpack_document_len(
 
     let mut pos = 0usize;
     let mut pending: u64 = 1; // elements owed across all open collections (the root is one)
-    let mut open: Vec<u64> = Vec::new(); // elements still owed per open collection = depth
+    let mut open = OpenCollections::new();
     while pending > 0 {
-        while open.last() == Some(&0) {
+        while open.last_mut() == Some(&mut 0) {
             open.pop();
         }
         let marker = *bytes
@@ -319,6 +373,52 @@ mod tests {
         }
     }
 
+    #[test]
+    fn depth_bound_is_exact_where_the_stack_spills_to_the_heap() {
+        for max_depth in INLINE_LEVELS - 1..=INLINE_LEVELS + 1 {
+            let too_deep = Err(format!("nests deeper than {max_depth} levels"));
+            for (header, empty) in LEVELS {
+                let at = |levels, tail: &[u8]| {
+                    check_msgpack_structure(&nested(header, levels, tail), max_depth)
+                        .map_err(|e| e.to_string())
+                };
+                assert_eq!(at(max_depth, &[0xc0]), Ok(()), "max_depth {max_depth}");
+                assert_eq!(
+                    at(max_depth + 1, &[0xc0]),
+                    too_deep,
+                    "max_depth {max_depth}"
+                );
+                assert_eq!(at(max_depth - 1, &[empty]), Ok(()), "max_depth {max_depth}");
+                assert_eq!(at(max_depth, &[empty]), too_deep, "max_depth {max_depth}");
+            }
+        }
+    }
+
+    #[test]
+    fn a_sibling_after_a_spilled_subtree_is_measured_from_its_own_parent() {
+        // [a levels of [..[nil]..], b levels of the same]: depth is 1 + max(a, b).
+        // An `a` past the inline slots spills the stack to the heap and drains
+        // it again before `b` starts.
+        let max_depth = INLINE_LEVELS + 2;
+        for a in 0..=max_depth {
+            for b in 0..=max_depth {
+                let doc = [
+                    vec![0x92],
+                    nested(&[0x91], a, &[0xc0]),
+                    nested(&[0x91], b, &[0xc0]),
+                ]
+                .concat();
+                let expected = if a.max(b) < max_depth {
+                    Ok(doc.len())
+                } else {
+                    Err(format!("nests deeper than {max_depth} levels"))
+                };
+                let got = msgpack_document_len(&doc, max_depth).map_err(|e| e.to_string());
+                assert_eq!(got, expected, "a {a}, b {b}");
+            }
+        }
+    }
+
     #[test]
     fn map_pairs_cost_two_slots() {
         assert_eq!(check(&[0x81, 0xa0, 0xc0]), Ok(()));

src/read_allocation_probe.rs

@@ -25,7 +25,10 @@
 //!
 //! The same count, [`bytes_requested`], also pins the write side: the
 //! `byte_storage` unit test `store_allocates_only_its_output_and_block` checks
-//! that `ByteStorage::store` requests only its output and its LZ4 block.
+//! that `ByteStorage::store` requests only its output and its LZ4 block, and
+//! `accepted_envelope_prescan_allocates_nothing` below checks that the
+//! structural pre-scan, which the fallback decode and SDK payload walks run,
+//! requests nothing for an envelope a writer emits.
 //!
 //! This lives in the crate's unit tests, compiled only under `cfg(test)`,
 //! because the positive control needs a hook inside `extract_payload`
@@ -151,6 +154,44 @@ mod tests {
         }
     }
 
+    /// The structural pre-scan keeps its nesting stack on the call stack for
+    /// an envelope, current or legacy, at every payload size, so walking one
+    /// requests nothing from the allocator. `retrieve` runs the walk only in
+    /// its fallback decode (`decode_envelope_full`); SDK bindings run it over
+    /// every payload they unpack.
+    #[test]
+    fn accepted_envelope_prescan_allocates_nothing() {
+        // The envelope decode's bound (`decode_envelope_full`).
+        const MAX_DEPTH: usize = 100;
+        let storage = ByteStorage::new(None);
+        for len in [0, 40, 64, 1024, 64 * 1024] {
+            let data: Vec<u8> = (0..len).map(|i| (i * 31 % 251) as u8).collect();
+            let envelope = storage.store(&data, None).unwrap();
+            let e: crate::StorageEnvelope = rmp_serde::from_slice(&envelope).unwrap();
+            // Legacy writers encoded `compressed_data` as an array of ints.
+            let legacy =
+                rmp_serde::to_vec(&(&e.compressed_data, e.checksum, e.original_size, &e.format))
+                    .unwrap();
+            for (shape, doc) in [("current", envelope), ("legacy", legacy)] {
+                assert!(
+                    storage.retrieve(&doc).is_ok(),
+                    "[{shape} {len}] not accepted"
+                );
+                let (result, requested) =
+                    bytes_requested(|| crate::check_msgpack_structure(&doc, MAX_DEPTH));
+                assert_eq!(result, Ok(()), "[{shape} {len}]");
+                assert_eq!(requested, 0, "[{shape} {len}] pre-scan allocated");
+            }
+        }
+        // Positive control: past the inline levels the stack spills to the
+        // heap, and this probe sees it.
+        let deep = [vec![0x91; 9], vec![0xc0]].concat();
+        let (result, requested) =
+            bytes_requested(|| crate::check_msgpack_structure(&deep, MAX_DEPTH));
+        assert_eq!(result, Ok(()));
+        assert!(requested > 0, "probe missed the spilled nesting stack");
+    }
+
     /// On every `msgpack bin` vector, `retrieve_data` allocates what the LZ4
     /// decoder does on its own and nothing else, since the one-pass decode
     /// takes every such vector without allocating, and `retrieve` adds only