Optimize DAG-CBOR encode by reusing a thread-local buffer and validating CIDs without Multihash construction (#116)

Author: MarshalX

src/car.rs

@@ -5,6 +5,7 @@ use cbor4ii::core::dec::Read;
 use pyo3::prelude::*;
 use pyo3::types::*;
 
+use crate::cid::parse_cid_prefix;
 use crate::dag_cbor::de::to_pyobject;
 use crate::error::value_error;
 use crate::ffi::recursion::current_recursion_limit;
@@ -68,25 +69,20 @@ pub fn decode_car<'py>(py: Python<'py>, data: &[u8]) -> PyResult<(Py<PyAny>, Bou
         }
 
         let cid_bytes_before = buf.buf;
-        // `&[u8]` is itself an `io::Read`, so we hand it to `Cid::read_bytes`
-        // directly and recover the consumed length from the slice shrink.
-        let mut slice: &[u8] = cid_bytes_before;
-        let cid_result = ::cid::Cid::read_bytes(&mut slice);
-        let Ok(cid) = cid_result else {
+        let Some((consumed, codec)) = parse_cid_prefix(cid_bytes_before) else {
             return Err(value_error(
                 "Failed to read CID of block",
-                cid_result.unwrap_err().to_string(),
+                "Invalid CID".to_string(),
             ));
         };
 
-        if cid.codec() != 0x71 {
+        if codec != 0x71 {
             return Err(value_error(
                 "Failed to read CAR block",
                 "Unsupported codec. For now we support only DAG-CBOR (0x71)".to_string(),
             ));
         }
 
-        let consumed = cid_bytes_before.len() - slice.len();
         buf.advance(consumed);
         let cid_raw = &cid_bytes_before[..consumed];
 

src/cid.rs

@@ -27,6 +27,53 @@ pub(crate) fn looks_like_cid(bytes: &[u8]) -> bool {
     bytes.len() == 34 && bytes[0] == 0x12 && bytes[1] == 0x20
 }
 
+// Minimal-encoding unsigned varint, as `unsigned_varint::decode::u64`:
+// ≤10 bytes, last byte of a multi-byte varint must be non-zero.
+#[inline]
+fn read_varint(bytes: &[u8], pos: &mut usize) -> Option<u64> {
+    let mut n: u64 = 0;
+    for i in 0..10 {
+        let &b = bytes.get(*pos + i)?;
+        n |= ((b & 0x7f) as u64) << (i * 7);
+        if b & 0x80 == 0 {
+            if b == 0 && i > 0 {
+                return None;
+            }
+            *pos += i + 1;
+            return Some(n);
+        }
+    }
+    None
+}
+
+/// Structural check that `bytes` starts with a valid binary CID; returns
+/// `(consumed, codec)`. Accepts exactly what `::cid::Cid::try_from` accepts
+/// (trailing bytes are the caller's concern) but skips the `Multihash`
+/// construction and its 64-byte digest copy.
+#[inline]
+pub(crate) fn parse_cid_prefix(bytes: &[u8]) -> Option<(usize, u64)> {
+    let mut pos = 0;
+    let version = read_varint(bytes, &mut pos)?;
+    let codec = read_varint(bytes, &mut pos)?;
+
+    // CIDv0: `0x12 0x20` + 32-byte sha2-256 digest; codec is implicitly dag-pb
+    if (version, codec) == (0x12, 0x20) {
+        let end = pos + 32;
+        return (bytes.len() >= end).then_some((end, 0x70));
+    }
+    if version != 1 {
+        return None;
+    }
+
+    read_varint(bytes, &mut pos)?; // multihash code
+    let hash_size = read_varint(bytes, &mut pos)?;
+    if hash_size > 64 {
+        return None;
+    }
+    let end = pos + hash_size as usize;
+    (bytes.len() >= end).then_some((end, codec))
+}
+
 pub(crate) fn extract_cid(data: &Bound<PyAny>) -> PyResult<::cid::Cid> {
     let cid = if let Ok(s) = data.cast::<PyString>() {
         ::cid::Cid::try_from(s.to_str()?)

src/dag_cbor/de.rs

@@ -5,6 +5,7 @@ use cbor4ii::core::{
 };
 use pyo3::{ffi, prelude::*, types::*, BoundObject};
 
+use crate::cid::parse_cid_prefix;
 use crate::error::value_error;
 use crate::ffi::dict::new_presized;
 use crate::ffi::key_cache::intern;
@@ -158,7 +159,7 @@ where
             }
 
             let cid_without_prefix = &cid[1..];
-            if ::cid::Cid::try_from(cid_without_prefix).is_err() {
+            if parse_cid_prefix(cid_without_prefix).is_none() {
                 return Err(anyhow!("Invalid CID"));
             }
 

src/dag_cbor/ser.rs

@@ -6,7 +6,7 @@ use cbor4ii::core::{
 use pyo3::pybacked::PyBackedStr;
 use pyo3::{ffi, prelude::*, types::*};
 
-use crate::cid::looks_like_cid;
+use crate::cid::{looks_like_cid, parse_cid_prefix};
 use crate::error::value_error;
 use crate::io::VecWriter;
 
@@ -133,7 +133,7 @@ where
         if tp == &raw mut ffi::PyBytes_Type {
             let b = obj.cast_unchecked::<PyBytes>();
             let bytes = b.as_bytes();
-            if looks_like_cid(bytes) && ::cid::Cid::try_from(bytes).is_ok() {
+            if looks_like_cid(bytes) && parse_cid_prefix(bytes).is_some() {
                 // by providing custom encoding we avoid extra allocation
                 types::Tag(42, PrefixedCidBytes(bytes)).encode(w)?;
             } else {
@@ -187,7 +187,7 @@ where
         Ok(())
     } else if let Ok(b) = obj.cast::<PyBytes>() {
         let bytes = b.as_bytes();
-        if looks_like_cid(bytes) && ::cid::Cid::try_from(bytes).is_ok() {
+        if looks_like_cid(bytes) && parse_cid_prefix(bytes).is_some() {
             types::Tag(42, PrefixedCidBytes(bytes)).encode(w)?;
         } else {
             types::Bytes(bytes).encode(w)?;

src/io/writer.rs

@@ -1,15 +1,27 @@
+use std::cell::Cell;
 use std::convert::Infallible;
 
 use cbor4ii::core::enc;
 
-// `enc::Write` over a raw `Vec<u8>`: no syscalls behind it, so a `BufWriter`
-// wrapper would just add a memcpy per push for no benefit.
+// Retaining bigger buffers would pin worst-case memory per thread forever.
+const MAX_POOLED_CAPACITY: usize = 1 << 20;
+
+thread_local! {
+    static POOL: Cell<Vec<u8>> = const { Cell::new(Vec::new()) };
+}
+
+// `enc::Write` over a `Vec<u8>` recycled through a thread-local pool: encoding
+// many small records reuses one grown allocation instead of re-growing from
+// zero on every call. `Cell::take` leaves an empty `Vec` behind, so a
+// re-entrant encode just falls back to a fresh buffer.
 pub(crate) struct VecWriter(Vec<u8>);
 
 impl VecWriter {
     #[inline]
     pub(crate) fn new() -> Self {
-        VecWriter(Vec::new())
+        let mut buf = POOL.take();
+        buf.clear();
+        VecWriter(buf)
     }
 
     #[inline]
@@ -18,6 +30,15 @@ impl VecWriter {
     }
 }
 
+impl Drop for VecWriter {
+    fn drop(&mut self) {
+        let buf = std::mem::take(&mut self.0);
+        if buf.capacity() <= MAX_POOLED_CAPACITY {
+            POOL.set(buf);
+        }
+    }
+}
+
 impl enc::Write for VecWriter {
     type Error = Infallible;