diff --git a/Cargo.toml b/Cargo.toml index 574a664..9890470 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,6 +19,9 @@ anyhow = "1.0.102" cid = "0.11.3" cbor4ii = { version = "1.2.2" } +[build-dependencies] +pyo3-build-config = "0.28.3" + [workspace] members = [ "profiling" ] diff --git a/build.rs b/build.rs new file mode 100644 index 0000000..40d12b4 --- /dev/null +++ b/build.rs @@ -0,0 +1,9 @@ +fn main() { + if matches!( + pyo3_build_config::get().implementation, + pyo3_build_config::PythonImplementation::CPython + ) { + println!("cargo:rustc-cfg=CPython"); + } + println!("cargo:rustc-check-cfg=cfg(CPython)"); +} diff --git a/src/lib.rs b/src/lib.rs index 0dfbff2..00ec94d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,4 @@ -use std::io::{self, Cursor, Write}; +use std::convert::Infallible; use anyhow::{anyhow, Result}; use cbor4ii::core::{ @@ -10,28 +10,53 @@ use cid::{multibase, Cid}; use pyo3::pybacked::PyBackedStr; use pyo3::{ffi, prelude::*, types::*, BoundObject, Python}; -struct BufWriter(io::BufWriter); +// Private CPython symbol; not provided by pyo3-ffi and CPython-only. +#[cfg(CPython)] +extern "C" { + fn _PyDict_NewPresized(minused: ffi::Py_ssize_t) -> *mut ffi::PyObject; +} -impl BufWriter { - pub fn new(inner: W) -> Self { - BufWriter(io::BufWriter::new(inner)) +// Empty CPython dicts already have 8 slots, so presizing below that buys +// nothing and lets us stay on the public `PyDict_New` path. +#[inline] +unsafe fn new_presized_dict(len: usize) -> *mut ffi::PyObject { + #[cfg(CPython)] + { + if len > 8 { + _PyDict_NewPresized(len as ffi::Py_ssize_t) + } else { + ffi::PyDict_New() + } + } + #[cfg(not(CPython))] + { + let _ = len; + ffi::PyDict_New() } +} - pub fn flush(&mut self) -> io::Result<()> { - self.0.flush() +// `enc::Write` over a raw `Vec`: no syscalls behind it, so a `BufWriter` +// wrapper would just add a memcpy per push for no benefit. +struct VecWriter(Vec); + +impl VecWriter { + #[inline] + fn new() -> Self { + VecWriter(Vec::new()) } - pub fn get_ref(&self) -> &W { - self.0.get_ref() + #[inline] + fn as_slice(&self) -> &[u8] { + &self.0 } } -impl enc::Write for BufWriter { - type Error = io::Error; +impl enc::Write for VecWriter { + type Error = Infallible; #[inline] fn push(&mut self, input: &[u8]) -> Result<(), Self::Error> { - self.0.write_all(input)?; + self.0.extend_from_slice(input); Ok(()) } } @@ -108,40 +133,39 @@ fn map_key_cmp(a: &[u8], b: &[u8]) -> std::cmp::Ordering { } } -fn sort_map_keys(keys: &Bound, len: usize) -> Result> { - // Returns key and index. - let mut keys_str = Vec::with_capacity(len); - for i in 0..len { - let item = unsafe { keys.get_item_unchecked(i) }; - let key = match item.cast::() { - Ok(k) => k.to_owned(), +// One dict walk collects (key, value) pairs together; sorting by-index and +// re-fetching values through `map.values()` would materialize two extra +// PyLists and walk the dict three times. +fn collect_and_sort_map_entries<'py>( + map: &Bound<'py, PyDict>, +) -> Result)>> { + let len = map.len(); + let mut entries: Vec<(PyBackedStr, Bound<'py, PyAny>)> = Vec::with_capacity(len); + + for (key, value) in map.iter() { + let key_str = match key.cast_into::() { + Ok(k) => k, Err(_) => return Err(anyhow!("Map keys must be strings")), }; - let backed_str = match PyBackedStr::try_from(key) { - Ok(bs) => bs, - Err(_) => return Err(anyhow!("Failed to convert PyString to PyBackedStr")), - }; - keys_str.push((backed_str, i)); + let backed = PyBackedStr::try_from(key_str) + .map_err(|_| anyhow!("Failed to convert PyString to PyBackedStr"))?; + entries.push((backed, value)); } - if keys_str.len() < 2 { - return Ok(keys_str); + if entries.len() >= 2 { + entries.sort_by(|a, b| { + // sort_unstable_by performs bad in past benchmarks; revisit if data shape changes. + let (s1, _) = a; + let (s2, _) = b; + if s1.len() != s2.len() { + s1.len().cmp(&s2.len()) + } else { + s1.as_bytes().cmp(s2.as_bytes()) + } + }); } - keys_str.sort_by(|a, b| { - // sort_unstable_by performs bad - let (s1, _) = a; - let (s2, _) = b; - - // sorted length-first by the byte representation of the string keys - if s1.len() != s2.len() { - s1.len().cmp(&s2.len()) - } else { - s1.cmp(s2) - } - }); - - Ok(keys_str) + Ok(entries) } fn get_bytes_from_py_any<'py>(obj: &'py Bound<'py, PyAny>) -> PyResult<&'py [u8]> { @@ -229,7 +253,14 @@ where } major::MAP => { let len = types::Map::len(r)?.ok_or_else(|| anyhow!("Map must contain length"))?; - let dict = PyDict::new(py); + // Length is known up front; presize to avoid rehashes as we fill. + let dict = unsafe { + let ptr = new_presized_dict(len); + if ptr.is_null() { + return Err(anyhow!(PyErr::fetch(py))); + } + Bound::from_owned_ptr(py, ptr).cast_into_unchecked::() + }; let mut prev_key: Option<&[u8]> = None; for _ in 0..len { @@ -315,6 +346,20 @@ where }) } +// `Cid::try_from` parses two varints + a multihash on every call; this O(1) +// shape check rejects payloads that can't be a CID without paying for it. +// CIDv1 starts with `0x01`; CIDv0 is exactly 34 bytes starting `0x12 0x20`. +#[inline] +fn looks_like_cid(bytes: &[u8]) -> bool { + if bytes.len() < 4 { + return false; + } + if bytes[0] == 0x01 { + return true; + } + bytes.len() == 34 && bytes[0] == 0x12 && bytes[1] == 0x20 +} + fn encode_dag_cbor_from_pyobject<'py, W: enc::Write>( _py: Python<'py>, obj: &Bound<'py, PyAny>, @@ -323,103 +368,134 @@ fn encode_dag_cbor_from_pyobject<'py, W: enc::Write>( where W::Error: Send + Sync, { - /* Order is important for performance! - - Fast checks go first: - - None - - bool - - int - - list - - dict - - str - Then slow checks: - - bytes - - float - */ - - if obj.is_none() { - types::Null.encode(w)?; - - Ok(()) - } else if obj.is_instance_of::() { - obj.is_truthy()?.encode(w)?; - - Ok(()) - } else if obj.is_instance_of::() { - let i: i128 = obj.extract()?; - - if i.is_negative() { - if -(i + 1) > u64::MAX as i128 { - return Err(anyhow!("Number out of range")); + // Exact-type pointer compare per branch avoids the MRO walk that + // `is_instance_of` / `cast` perform. Order tuned for typical ATProto + // record shapes; subclasses fall through to the slow path below. + let tp = unsafe { ffi::Py_TYPE(obj.as_ptr()) }; + unsafe { + if tp == &raw mut ffi::PyUnicode_Type { + let s = obj.cast_unchecked::(); + s.to_str()?.encode(w)?; + return Ok(()); + } + if tp == &raw mut ffi::PyDict_Type { + let map = obj.cast_unchecked::(); + let entries = collect_and_sort_map_entries(map)?; + types::Map::bounded(entries.len(), w)?; + for (key, value) in &entries { + (&**key).encode(w)?; + encode_dag_cbor_from_pyobject(_py, value, w)?; } - - types::Negative(-(i + 1) as u64).encode(w)?; - } else { - if i > u64::MAX as i128 { + return Ok(()); + } + if tp == &raw mut ffi::PyList_Type { + let l = obj.cast_unchecked::(); + let len = l.len(); + types::Array::bounded(len, w)?; + for i in 0..len { + let item = l.get_item_unchecked(i); + encode_dag_cbor_from_pyobject(_py, &item, w)?; + } + return Ok(()); + } + if tp == &raw mut ffi::PyLong_Type { + return encode_int(obj, w); + } + if tp == &raw mut ffi::PyBytes_Type { + let b = obj.cast_unchecked::(); + let bytes = b.as_bytes(); + if looks_like_cid(bytes) && Cid::try_from(bytes).is_ok() { + // by providing custom encoding we avoid extra allocation + types::Tag(42, PrefixedCidBytes(bytes)).encode(w)?; + } else { + types::Bytes(bytes).encode(w)?; + } + return Ok(()); + } + if tp == &raw mut ffi::PyBool_Type { + (obj.as_ptr() == ffi::Py_True()).encode(w)?; + return Ok(()); + } + if obj.as_ptr() == ffi::Py_None() { + types::Null.encode(w)?; + return Ok(()); + } + if tp == &raw mut ffi::PyFloat_Type { + let f = obj.cast_unchecked::(); + let v = f.value(); + if !v.is_finite() { return Err(anyhow!("Number out of range")); } - - (i as u64).encode(w)?; + v.encode(w)?; + return Ok(()); } + } + // Slow path: subclasses of supported types (rare in DAG-CBOR usage). + if obj.is_instance_of::() { + (obj.as_ptr() == unsafe { ffi::Py_True() }).encode(w)?; Ok(()) + } else if obj.is_instance_of::() { + encode_int(obj, w) } else if let Ok(l) = obj.cast::() { let len = l.len(); - types::Array::bounded(len, w)?; - for i in 0..len { let item = unsafe { l.get_item_unchecked(i) }; encode_dag_cbor_from_pyobject(_py, &item, w)?; } - Ok(()) } else if let Ok(map) = obj.cast::() { - let len = map.len(); - let keys = sort_map_keys(&map.keys(), len)?; - let values = map.values(); - - types::Map::bounded(len, w)?; - - for (key, i) in keys { - key.get(..) - .ok_or_else(|| anyhow!("Invalid string range while encoding map key"))? - .encode(w)?; - let value = unsafe { values.get_item_unchecked(i) }; - encode_dag_cbor_from_pyobject(_py, &value, w)?; + let entries = collect_and_sort_map_entries(map)?; + types::Map::bounded(entries.len(), w)?; + for (key, value) in &entries { + (&**key).encode(w)?; + encode_dag_cbor_from_pyobject(_py, value, w)?; } - Ok(()) - } else if let Ok(f) = obj.cast::() { - let v = f.value(); - if !v.is_finite() { - return Err(anyhow!("Number out of range")); - } - - v.encode(w)?; - + } else if let Ok(s) = obj.cast::() { + s.to_str()?.encode(w)?; Ok(()) } else if let Ok(b) = obj.cast::() { - // FIXME (MarshalX): it's not efficient to try to parse it as CID let bytes = b.as_bytes(); - let cid = Cid::try_from(bytes); - if cid.is_ok() { - // by providing custom encoding we avoid extra allocation + if looks_like_cid(bytes) && Cid::try_from(bytes).is_ok() { types::Tag(42, PrefixedCidBytes(bytes)).encode(w)?; } else { types::Bytes(bytes).encode(w)?; } - Ok(()) - } else if let Ok(s) = obj.cast::() { - s.to_str()?.encode(w)?; - + } else if let Ok(f) = obj.cast::() { + let v = f.value(); + if !v.is_finite() { + return Err(anyhow!("Number out of range")); + } + v.encode(w)?; Ok(()) } else { Err(anyhow!("Unknown tag")) } } +#[inline] +fn encode_int(obj: &Bound<'_, PyAny>, w: &mut W) -> Result<()> +where + W::Error: Send + Sync, +{ + let i: i128 = obj.extract()?; + if i.is_negative() { + if -(i + 1) > u64::MAX as i128 { + return Err(anyhow!("Number out of range")); + } + types::Negative(-(i + 1) as u64).encode(w)?; + } else { + if i > u64::MAX as i128 { + return Err(anyhow!("Number out of range")); + } + (i as u64).encode(w)?; + } + Ok(()) +} + #[pyfunction] fn decode_dag_cbor_multi<'py>(py: Python<'py>, data: &[u8]) -> PyResult> { let mut reader = SliceReader::new(data); @@ -529,9 +605,10 @@ pub fn decode_car<'py>(py: Python<'py>, data: &[u8]) -> PyResult<(Py, Bou } let cid_bytes_before = buf.buf; - let mut cursor = Cursor::new(cid_bytes_before); - - let cid_result = Cid::read_bytes(&mut cursor); + // `&[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::read_bytes(&mut slice); let Ok(cid) = cid_result else { return Err(get_err( "Failed to read CID of block", @@ -546,7 +623,7 @@ pub fn decode_car<'py>(py: Python<'py>, data: &[u8]) -> PyResult<(Py, Bou )); } - let consumed = cursor.position() as usize; + let consumed = cid_bytes_before.len() - slice.len(); buf.advance(consumed); let cid_raw = &cid_bytes_before[..consumed]; @@ -602,14 +679,11 @@ pub fn encode_dag_cbor<'py>( py: Python<'py>, data: &Bound<'py, PyAny>, ) -> PyResult> { - let mut buf = &mut BufWriter::new(Vec::new()); + let mut buf = VecWriter::new(); if let Err(e) = encode_dag_cbor_from_pyobject(py, data, &mut buf) { return Err(get_err("Failed to encode DAG-CBOR", e.to_string())); } - if let Err(e) = buf.flush() { - return Err(get_err("Failed to flush buffer", e.to_string())); - } - Ok(PyBytes::new(py, buf.get_ref())) + Ok(PyBytes::new(py, buf.as_slice())) } fn get_cid_from_py_any(data: &Bound) -> PyResult {