diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index a2e9ffb..531b4de 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -17,7 +17,7 @@ env: jobs: benchmarks: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 # required for correct CodSpeed integration diff --git a/pytests/test_dag_cbor.py b/pytests/test_dag_cbor.py index 0e31a82..0e9d35c 100644 --- a/pytests/test_dag_cbor.py +++ b/pytests/test_dag_cbor.py @@ -363,3 +363,110 @@ def test_dag_cbor_decode_error_mid_array() -> None: # 83010263616263 - [1, 2, "abc"] truncated after the second element's header with pytest.raises(ValueError): libipld.decode_dag_cbor(bytes.fromhex('8301026361')) + + +@pytest.mark.parametrize('length', [*range(70), 255, 256, 257, 65535, 65536, 65537]) +def test_dag_cbor_roundtrip_string_and_bytes_lengths(length) -> None: + text = ('abcdefghij' * (length // 10 + 1))[:length] + obj = {'b': text.encode(), 's': text, text: length, 'l': [0] * min(length, 300)} + + encoded = libipld.encode_dag_cbor(obj) + assert obj == libipld.decode_dag_cbor(encoded) + + if length < 24: + head = bytes([0x60 | length]) + elif length < 256: + head = bytes([0x78, length]) + elif length < 65536: + head = b'\x79' + length.to_bytes(2, 'big') + else: + head = b'\x7a' + length.to_bytes(4, 'big') + assert libipld.encode_dag_cbor(text) == head + text.encode() + + +@pytest.mark.parametrize( + ('value', 'expected'), + [ + (0, '00'), + (23, '17'), + (24, '1818'), + (255, '18ff'), + (256, '190100'), + (65535, '19ffff'), + (65536, '1a00010000'), + (2**32 - 1, '1affffffff'), + (2**32, '1b0000000100000000'), + (2**64 - 1, '1bffffffffffffffff'), + (-1, '20'), + (-24, '37'), + (-25, '3818'), + (-(2**64), '3bffffffffffffffff'), + (1.5, 'fb3ff8000000000000'), + (True, 'f5'), + (False, 'f4'), + (None, 'f6'), + (b'\x01\x02', '420102'), + ([], '80'), + ({}, 'a0'), + ], +) +def test_dag_cbor_encode_canonical_heads(value, expected) -> None: + encoded = libipld.encode_dag_cbor(value) + assert encoded.hex() == expected + assert libipld.decode_dag_cbor(encoded) == value + + +def test_dag_cbor_encode_map_key_order() -> None: + keys = [ + '', + 'b', + 'a', + 'zz', + 'aaaaaaaa', + 'aaaaaaab', + 'aaaaaaaaa', + 'aaaaaaaab', + 'baaaaaaaa', + 'aaaaaaaa\x00', + 'prefix__0002', + 'prefix__0001', + 'prefix__0003', + 'é', + 'zé', + 'ab', + '\x7f', + 'Ā', + ] + expected = sorted(keys, key=lambda k: (len(k.encode()), k.encode())) + + for ordering in (keys, keys[::-1], expected): + encoded = libipld.encode_dag_cbor({k: i for i, k in enumerate(ordering)}) + # decode rejects unsorted keys and keeps their order + assert list(libipld.decode_dag_cbor(encoded)) == expected + + +def test_dag_cbor_decode_many_distinct_keys() -> None: + # more distinct keys than the key cache holds + obj = [{f'key{i}': i, f'k{i}': None} for i in range(6000)] + encoded = libipld.encode_dag_cbor(obj) + + assert libipld.decode_dag_cbor(encoded) == obj + assert libipld.decode_dag_cbor(encoded) == obj + + +def test_dag_cbor_cid_roundtrip_shapes() -> None: + cids = [ + bytes.fromhex('01711220') + bytes(range(32)), # v1 dag-cbor sha2-256 + bytes.fromhex('01551220') + bytes(range(32)), # v1 raw sha2-256 + bytes.fromhex('01701340') + bytes(range(64)), # v1 dag-pb sha2-512 + bytes.fromhex('0171000401020304'), # v1 identity, 4-byte digest + bytes.fromhex('1220') + bytes(range(32)), # v0 + ] + for cid in cids: + encoded = libipld.encode_dag_cbor({'c': cid}) + assert encoded[3:5] == b'\xd8\x2a' # tag 42 + assert libipld.decode_dag_cbor(encoded) == {'c': cid} + + # truncated digest: plain bytes, not a CID + not_cid = bytes.fromhex('01711220') + bytes(31) + assert libipld.encode_dag_cbor(not_cid)[0] == 0x58 diff --git a/pytests/test_free_threading.py b/pytests/test_free_threading.py new file mode 100644 index 0000000..d141f9f --- /dev/null +++ b/pytests/test_free_threading.py @@ -0,0 +1,96 @@ +import threading +import time + +import libipld +import pytest + +# Mutating an argument from another thread mid-call is a data race on the +# free-threaded build. A clean `ValueError` is acceptable; a crash is not. + +_DURATION = 0.3 +_READERS = 4 + + +def _race(mutate, call) -> None: + deadline = time.monotonic() + _DURATION + errors = [] + + def writer() -> None: + while time.monotonic() < deadline: + mutate() + + def reader() -> None: + while time.monotonic() < deadline: + try: + call() + except ValueError: + pass + except Exception as e: # noqa: BLE001 + errors.append(e) + + threads = [threading.Thread(target=writer)] + [threading.Thread(target=reader) for _ in range(_READERS)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors, errors + + +@pytest.mark.thread_unsafe(reason='spawns its own threads') +def test_encode_list_mutated_concurrently() -> None: + data = {'a': [{'k': [1, 2, 3]} for _ in range(200)], 'b': list(range(2000))} + + def mutate() -> None: + data['a'] = [{'k': [i]} for i in range(200)] + data['a'].clear() + data['b'].clear() + data['b'].extend(range(2000)) + + _race(mutate, lambda: libipld.encode_dag_cbor(data)) + + +@pytest.mark.thread_unsafe(reason='spawns its own threads') +def test_encode_dict_mutated_concurrently() -> None: + data = {str(i): 'x' * 10 for i in range(500)} + + def mutate() -> None: + for i in range(500): + data.pop(str(i), None) + for i in range(500): + data[str(i)] = 'y' * 10 + + _race(mutate, lambda: libipld.encode_dag_cbor(data)) + + +@pytest.mark.thread_unsafe(reason='spawns its own threads') +def test_bytearray_resized_concurrently() -> None: + cid = bytes.fromhex('01711220' + '00' * 32) + buf = bytearray(cid) + + def mutate() -> None: + buf[:] = b'' + buf[:] = cid + b'\x00' * 100_000 + buf[:] = cid + + def call() -> None: + libipld.decode_cid(buf) + libipld.encode_multibase('b', buf) + + _race(mutate, call) + + +@pytest.mark.thread_unsafe(reason='spawns its own threads') +def test_shared_data_roundtrip_concurrently() -> None: + obj = {'$type': 'app.bsky.feed.post', 'text': 'hi', 'n': [1, 2.5, None, True, b'x'] * 50} + encoded = libipld.encode_dag_cbor(obj) + mismatches = [] + + def call() -> None: + decoded = libipld.decode_dag_cbor(encoded) + if decoded != obj or libipld.encode_dag_cbor(decoded) != encoded: + mismatches.append(decoded) + + _race(lambda: time.sleep(0.01), call) + + assert not mismatches diff --git a/src/cid.rs b/src/cid.rs index abf8782..0efaf5e 100644 --- a/src/cid.rs +++ b/src/cid.rs @@ -52,6 +52,12 @@ fn read_varint(bytes: &[u8], pos: &mut usize) -> Option { /// construction and its 64-byte digest copy. #[inline] pub(crate) fn parse_cid_prefix(bytes: &[u8]) -> Option<(usize, u64)> { + // Fast path: CIDv1 with single-byte codec, hash code and digest size. + if let [1, codec @ ..=0x7f, ..=0x7f, size @ ..=64, ..] = *bytes { + let end = 4 + size as usize; + return (bytes.len() >= end).then_some((end, codec as u64)); + } + let mut pos = 0; let version = read_varint(bytes, &mut pos)?; let codec = read_varint(bytes, &mut pos)?; diff --git a/src/dag_cbor/de.rs b/src/dag_cbor/de.rs index c6b37d7..2cd2d35 100644 --- a/src/dag_cbor/de.rs +++ b/src/dag_cbor/de.rs @@ -13,15 +13,22 @@ use crate::ffi::recursion::current_recursion_limit; use crate::ffi::string::from_bytes; use crate::io::{peek_one, SliceReader}; -fn map_key_cmp(a: &[u8], b: &[u8]) -> std::cmp::Ordering { - /* The keys in every map must be sorted length-first by the byte representation of the string keys, where: - - If two keys have different lengths, the shorter one sorts earlier; - - If two keys have the same length, the one with the lower value in (byte-wise) lexical order sorts earlier. - */ +// DAG-CBOR map key order: shorter keys first, same-length keys bytewise. +// Equal keys are not "less", so duplicates are rejected too. +#[inline] +fn map_key_less(a: &[u8], b: &[u8]) -> bool { if a.len() != b.len() { - a.len().cmp(&b.len()) - } else { - a.cmp(b) + return a.len() < b.len(); + } + match (a.first_chunk::<8>(), b.first_chunk::<8>()) { + // Big-endian integers order the same way as the bytes they hold. + (Some(x), Some(y)) if x != y => u64::from_be_bytes(*x) < u64::from_be_bytes(*y), + (Some(_), Some(_)) => a[8..] < b[8..], + _ => a + .iter() + .zip(b) + .find(|(x, y)| x != y) + .is_some_and(|(x, y)| x < y), } } @@ -255,8 +262,7 @@ where let key = decode_seg(r, key_byte)?; if let Some(prev_key) = prev_key { - // it cares about duplicated keys too thanks to Ordering::Equal - if map_key_cmp(prev_key, key) != std::cmp::Ordering::Less { + if !map_key_less(prev_key, key) { return Err(anyhow!("Map keys must be sorted and unique")); } } diff --git a/src/dag_cbor/ser.rs b/src/dag_cbor/ser.rs index 9149f8f..beb5ab0 100644 --- a/src/dag_cbor/ser.rs +++ b/src/dag_cbor/ser.rs @@ -1,34 +1,22 @@ use std::cell::Cell; use anyhow::{anyhow, Result}; -use cbor4ii::core::{ - enc::{self, Encode}, - types, -}; use pyo3::sync::critical_section::with_critical_section; use pyo3::{ffi, prelude::*, types::*, Borrowed}; use crate::cid::{looks_like_cid, parse_cid_prefix}; use crate::error::{pending_or_value_error, recursion_error}; use crate::ffi::recursion::current_recursion_limit; +use crate::io::writer::{ARRAY, BYTES, MAP, NEGATIVE, STRING, TAG, UNSIGNED}; use crate::io::VecWriter; -struct PrefixedCidBytes<'a>(&'a [u8]); - -impl<'a> Encode for PrefixedCidBytes<'a> { - fn encode(&self, w: &mut W) -> Result<(), enc::Error> { - // length prefix for bytes: 1 (leading 0) + payload - types::Bytes::bounded(1 + self.0.len(), w)?; - w.push(&[0x00])?; - w.push(self.0)?; - Ok(()) - } -} - #[derive(Clone, Copy)] struct MapEntry { key: *const u8, key_len: usize, + // First 8 key bytes as a big-endian integer (zero-padded), so sorting + // rarely has to read the key itself. + prefix: u64, value: *mut ffi::PyObject, #[cfg(Py_GIL_DISABLED)] key_obj: *mut ffi::PyObject, @@ -39,6 +27,25 @@ impl MapEntry { fn key(&self) -> &[u8] { unsafe { std::slice::from_raw_parts(self.key, self.key_len) } } + + // DAG-CBOR map order: shorter keys first, then bytewise. + #[inline] + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + (self.key_len, self.prefix) + .cmp(&(other.key_len, other.prefix)) + .then_with(|| self.key().cmp(other.key())) + } +} + +#[inline] +fn key_prefix(key: &[u8]) -> u64 { + match key.first_chunk::<8>() { + Some(chunk) => u64::from_be_bytes(*chunk), + None => key + .iter() + .enumerate() + .fold(0, |prefix, (i, &b)| prefix | (b as u64) << (56 - 8 * i)), + } } const MAX_POOLED_ENTRIES: usize = 1 << 14; @@ -105,17 +112,14 @@ fn list_item<'py>(l: &Bound<'py, PyList>, i: usize) -> Result> } #[inline] -fn encode_int(obj: &Bound<'_, PyAny>, w: &mut W) -> Result<()> -where - W::Error: Send + Sync, -{ +fn encode_int(obj: &Bound<'_, PyAny>, w: &mut VecWriter) -> Result<()> { #[cfg(all(CPython, Py_3_12))] { if let Some((abs_val, neg)) = unsafe { crate::ffi::int::pylong_parts(obj.as_ptr()) } { if neg { - types::Negative(abs_val - 1).encode(w)?; + w.head(NEGATIVE, abs_val - 1); } else { - abs_val.encode(w)?; + w.head(UNSIGNED, abs_val); } return Ok(()); } @@ -126,12 +130,12 @@ where if -(i + 1) > u64::MAX as i128 { return Err(anyhow!("Number out of range")); } - types::Negative(-(i + 1) as u64).encode(w)?; + w.head(NEGATIVE, -(i + 1) as u64); } else { if i > u64::MAX as i128 { return Err(anyhow!("Number out of range")); } - (i as u64).encode(w)?; + w.head(UNSIGNED, i as u64); } Ok(()) } @@ -156,9 +160,11 @@ fn collect_map_entries(py: Python<'_>, map: &Bound<'_, PyDict>, enc: &mut Encode ffi::Py_INCREF(key); ffi::Py_INCREF(value); } + let key_bytes = std::slice::from_raw_parts(utf8.cast::(), len as usize); enc.entries.0.push(MapEntry { - key: utf8.cast(), - key_len: len as usize, + key: key_bytes.as_ptr(), + key_len: key_bytes.len(), + prefix: key_prefix(key_bytes), value, #[cfg(Py_GIL_DISABLED)] key_obj: key, @@ -180,21 +186,14 @@ fn encode_map_entries<'py>( let entries = &mut enc.entries.0[base..]; if entries.len() >= 2 { // Keys are unique, so stability buys nothing. - entries.sort_unstable_by(|a, b| { - if a.key_len != b.key_len { - a.key_len.cmp(&b.key_len) - } else { - a.key().cmp(b.key()) - } - }); + entries.sort_unstable_by(MapEntry::cmp); } let len = entries.len(); - types::Map::bounded(len, &mut enc.w)?; + enc.w.head(MAP, len as u64); for i in base..base + len { let entry = enc.entries.0[i]; - // CPython hands out the UTF-8 buffer of a valid `str`. - unsafe { std::str::from_utf8_unchecked(entry.key()) }.encode(&mut enc.w)?; + enc.w.segment(STRING, entry.key()); let value = unsafe { Borrowed::from_ptr(py, entry.value) }; encode_value(py, &value, enc)?; } @@ -227,7 +226,7 @@ fn encode_value<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>, enc: &mut Encoder unsafe { if tp == &raw mut ffi::PyUnicode_Type { let s = obj.cast_unchecked::(); - s.to_str()?.encode(&mut enc.w)?; + enc.w.segment(STRING, s.to_str()?.as_bytes()); return Ok(()); } if tp == &raw mut ffi::PyDict_Type { @@ -243,11 +242,11 @@ fn encode_value<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>, enc: &mut Encoder return encode_bytes(obj.cast_unchecked::(), &mut enc.w); } if tp == &raw mut ffi::PyBool_Type { - (obj.as_ptr() == ffi::Py_True()).encode(&mut enc.w)?; + enc.w.bool(obj.as_ptr() == ffi::Py_True()); return Ok(()); } if obj.as_ptr() == ffi::Py_None() { - types::Null.encode(&mut enc.w)?; + enc.w.null(); return Ok(()); } if tp == &raw mut ffi::PyFloat_Type { @@ -262,7 +261,7 @@ fn encode_value<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>, enc: &mut Encoder fn encode_list<'py>(py: Python<'py>, l: &Bound<'py, PyList>, enc: &mut Encoder) -> Result<()> { enc.enter(py)?; let len = l.len(); - types::Array::bounded(len, &mut enc.w)?; + enc.w.head(ARRAY, len as u64); for i in 0..len { #[cfg(not(Py_GIL_DISABLED))] let item = unsafe { list_item(l, i) }; @@ -275,32 +274,29 @@ fn encode_list<'py>(py: Python<'py>, l: &Bound<'py, PyList>, enc: &mut Encoder) } #[inline] -fn encode_bytes(b: &Bound<'_, PyBytes>, w: &mut W) -> Result<()> -where - W::Error: Send + Sync, -{ +fn encode_bytes(b: &Bound<'_, PyBytes>, w: &mut VecWriter) -> Result<()> { let bytes = b.as_bytes(); 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)?; + // Tag 42 around a byte string holding 0x00 followed by the CID. + w.head(TAG, 42); + w.head(BYTES, 1 + bytes.len() as u64); + w.raw(&[0x00]); + w.raw(bytes); } else { - types::Bytes(bytes).encode(w)?; + w.segment(BYTES, bytes); } Ok(()) } #[inline] -fn encode_float(f: &Bound<'_, PyFloat>, w: &mut W) -> Result<()> -where - W::Error: Send + Sync, -{ +fn encode_float(f: &Bound<'_, PyFloat>, w: &mut VecWriter) -> Result<()> { let v = f.value(); if !v.is_finite() { return Err(anyhow!("Number out of range")); } // DAG-CBOR forbids -0.0; it equals 0.0, so encode it as 0.0. let v = if v == 0.0 { 0.0 } else { v }; - v.encode(w)?; + w.f64(v); Ok(()) } @@ -309,7 +305,7 @@ where #[inline(never)] fn encode_subclass<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>, enc: &mut Encoder) -> Result<()> { if obj.is_instance_of::() { - (obj.as_ptr() == unsafe { ffi::Py_True() }).encode(&mut enc.w)?; + enc.w.bool(obj.as_ptr() == unsafe { ffi::Py_True() }); Ok(()) } else if obj.is_instance_of::() { encode_int(obj, &mut enc.w) @@ -318,7 +314,7 @@ fn encode_subclass<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>, enc: &mut Enco } else if let Ok(map) = obj.cast::() { encode_map(py, map, enc) } else if let Ok(s) = obj.cast::() { - s.to_str()?.encode(&mut enc.w)?; + enc.w.segment(STRING, s.to_str()?.as_bytes()); Ok(()) } else if let Ok(b) = obj.cast::() { encode_bytes(b, &mut enc.w) diff --git a/src/ffi/key_cache.rs b/src/ffi/key_cache.rs index 6076317..e57d6cc 100644 --- a/src/ffi/key_cache.rs +++ b/src/ffi/key_cache.rs @@ -1,4 +1,4 @@ -//! Direct-mapped intern cache for short map keys. atproto-shape payloads +//! Two-way intern cache for short map keys. atproto-shape payloads //! reuse a small vocabulary (`$type`, `did`, `cid`, `uri`, `text`, ...) per //! record; caching the constructed `PyUnicode`, hash already computed and //! stored inside it, skips both the rebuild and the rehash on dict insert. @@ -9,6 +9,7 @@ mod cached { use pyo3::{ffi, prelude::*}; use crate::ffi::string::from_bytes; + use crate::io::short::eq_short; const CAP: usize = 2048; const MAX_KEY_LEN: usize = 64; @@ -31,14 +32,24 @@ mod cached { static mut SLOTS: [Entry; CAP] = [const { Entry::empty() }; CAP]; + // Hashes the length plus the first and last (up to) 8 bytes. Longer keys + // differing only in the middle share a slot: more evictions, same result. #[inline] - fn fx_hash(bytes: &[u8]) -> usize { - const K: u64 = 0x517c_c1b7_2722_0a95; - let mut h: u64 = 0; - for &b in bytes { - h = (h.rotate_left(5) ^ b as u64).wrapping_mul(K); - } - h as usize + fn slot_index(bytes: &[u8]) -> usize { + const K: u64 = 0x9e37_79b9_7f4a_7c15; + let (head, tail) = match (bytes.first_chunk::<8>(), bytes.last_chunk::<8>()) { + (Some(head), Some(tail)) => (u64::from_le_bytes(*head), u64::from_le_bytes(*tail)), + _ => match (bytes.first_chunk::<4>(), bytes.last_chunk::<4>()) { + (Some(head), Some(tail)) => ( + u32::from_le_bytes(*head) as u64, + u32::from_le_bytes(*tail) as u64, + ), + _ => (bytes.iter().fold(0, |acc, &b| (acc << 8) | b as u64), 0), + }, + }; + let len = bytes.len() as u64; + let hash = ((head ^ len.wrapping_mul(K)).wrapping_mul(K) ^ tail).wrapping_mul(K); + (hash >> (64 - CAP.trailing_zeros())) as usize } /// Returns a strong-ref `PyUnicode*`; the caller owns one ref. @@ -49,23 +60,29 @@ mod cached { return build(py, bytes); } - let slot_idx = fx_hash(bytes) & (CAP - 1); - // `&raw mut` is the supported path to a `static mut`; the explicit - // re-borrow keeps the field accesses readable. Clippy's `deref_addrof` - // suggestion would re-introduce `static_mut_refs`. - #[allow(clippy::deref_addrof)] - let slot = &mut *(&raw mut SLOTS[slot_idx]); - - if slot.len as usize == bytes.len() - && !slot.obj.is_null() - && slot.bytes[..bytes.len()] == *bytes - { - ffi::Py_INCREF(slot.obj); - return Ok(slot.obj); + // A key lives in one of two adjacent slots, so two hot keys that hash + // alike don't keep evicting each other. + let idx = slot_index(bytes); + let first = &raw mut SLOTS[idx]; + let second = &raw mut SLOTS[idx ^ 1]; + for slot in [first, second] { + let slot = &*slot; + if !slot.obj.is_null() + && slot.len as usize == bytes.len() + && eq_short(&slot.bytes[..bytes.len()], bytes) + { + ffi::Py_INCREF(slot.obj); + return Ok(slot.obj); + } } let obj = build(py, bytes)?; - // Evict the previous occupant before claiming the slot. + // Prefer an empty slot; otherwise evict the first. + let slot = if !(*first).obj.is_null() && (*second).obj.is_null() { + &mut *second + } else { + &mut *first + }; if !slot.obj.is_null() { ffi::Py_DECREF(slot.obj); } diff --git a/src/ffi/string.rs b/src/ffi/string.rs index c5d9ac7..4582fe4 100644 --- a/src/ffi/string.rs +++ b/src/ffi/string.rs @@ -21,7 +21,7 @@ pub(crate) fn from_bytes<'py>(py: Python<'py>, bytes: &[u8]) -> PyResult().offset(1).cast::(); - std::ptr::copy_nonoverlapping(bytes.as_ptr(), data, bytes.len()); + crate::io::short::copy_short(bytes.as_ptr(), data, bytes.len()); *data.add(bytes.len()) = 0; Ok(Bound::from_owned_ptr(py, obj).cast_into_unchecked::()) diff --git a/src/io.rs b/src/io.rs index b8f02d8..473c8c9 100644 --- a/src/io.rs +++ b/src/io.rs @@ -3,6 +3,7 @@ pub(crate) mod leb128; pub(crate) mod reader; +pub(crate) mod short; pub(crate) mod writer; pub(crate) use reader::{peek_one, SliceReader}; diff --git a/src/io/short.rs b/src/io/short.rs new file mode 100644 index 0000000..d4fac5b --- /dev/null +++ b/src/io/short.rs @@ -0,0 +1,51 @@ +//! Copy and compare for short byte runs, where a `memcpy`/`memcmp` call costs +//! more than the work itself. + +/// `copy_nonoverlapping` that handles up to 32 bytes inline with two +/// overlapping fixed-width moves. +/// +/// # Safety +/// `src` must be readable and `dst` writable for `len` bytes, non-overlapping. +#[inline(always)] +pub(crate) unsafe fn copy_short(src: *const u8, dst: *mut u8, len: usize) { + macro_rules! overlapping { + ($t:ty) => {{ + let tail = len - std::mem::size_of::<$t>(); + let first = src.cast::<$t>().read_unaligned(); + let last = src.add(tail).cast::<$t>().read_unaligned(); + dst.cast::<$t>().write_unaligned(first); + dst.add(tail).cast::<$t>().write_unaligned(last); + }}; + } + // An `if` chain keeps this branch order; a `match` measured 3-6% slower. + if len > 32 { + std::ptr::copy_nonoverlapping(src, dst, len); + } else if len >= 16 { + overlapping!(u128) + } else if len >= 8 { + overlapping!(u64) + } else if len >= 4 { + overlapping!(u32) + } else if len >= 2 { + overlapping!(u16) + } else if len == 1 { + *dst = *src; + } +} + +/// Equality of two same-length byte runs, compared a word at a time. +#[cfg(all(CPython, not(Py_GIL_DISABLED)))] +#[inline(always)] +pub(crate) fn eq_short(a: &[u8], b: &[u8]) -> bool { + debug_assert_eq!(a.len(), b.len()); + if a.len() >= 8 { + a.chunks_exact(8) + .zip(b.chunks_exact(8)) + .all(|(x, y)| x == y) + && a.last_chunk::<8>() == b.last_chunk::<8>() + } else if a.len() >= 4 { + a.first_chunk::<4>() == b.first_chunk::<4>() && a.last_chunk::<4>() == b.last_chunk::<4>() + } else { + a.iter().zip(b).all(|(x, y)| x == y) + } +} diff --git a/src/io/writer.rs b/src/io/writer.rs index e500c1b..2ba7f24 100644 --- a/src/io/writer.rs +++ b/src/io/writer.rs @@ -1,7 +1,6 @@ use std::cell::Cell; -use std::convert::Infallible; -use cbor4ii::core::enc; +use crate::io::short::copy_short; // Retaining bigger buffers would pin worst-case memory per thread forever. const MAX_POOLED_CAPACITY: usize = 1 << 20; @@ -10,7 +9,7 @@ thread_local! { static POOL: Cell> = const { Cell::new(Vec::new()) }; } -// `enc::Write` over a `Vec` recycled through a thread-local pool: encoding +// CBOR writer over a `Vec` 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. @@ -39,12 +38,93 @@ impl Drop for VecWriter { } } -impl enc::Write for VecWriter { - type Error = Infallible; +// Major types, already shifted into the top 3 bits of the header byte. +pub(crate) const UNSIGNED: u8 = 0x00; +pub(crate) const NEGATIVE: u8 = 0x20; +pub(crate) const BYTES: u8 = 0x40; +pub(crate) const STRING: u8 = 0x60; +pub(crate) const ARRAY: u8 = 0x80; +pub(crate) const MAP: u8 = 0xa0; +pub(crate) const TAG: u8 = 0xc0; + +const FALSE: u8 = 0xf4; +const TRUE: u8 = 0xf5; +const NULL: u8 = 0xf6; +const F64: u8 = 0xfb; + +// Longest header: the type byte plus an 8-byte argument. +const MAX_HEAD: usize = 9; + +// Writes the shortest header for `major` with argument `n` and returns its +// length. `p` must be writable for `MAX_HEAD` bytes. +#[inline(always)] +unsafe fn put_head(p: *mut u8, major: u8, n: u64) -> usize { + macro_rules! put { + ($info:literal, $t:ty) => {{ + *p = major | $info; + p.add(1).cast::<$t>().write_unaligned((n as $t).to_be()); + 1 + std::mem::size_of::<$t>() + }}; + } + if n < 24 { + *p = major | n as u8; + 1 + } else if n <= u8::MAX as u64 { + put!(24, u8) + } else if n <= u16::MAX as u64 { + put!(25, u16) + } else if n <= u32::MAX as u64 { + put!(26, u32) + } else { + put!(27, u64) + } +} + +impl VecWriter { + /// Header of `major` with argument `n`: an integer value, a container + /// length or a tag number. + #[inline] + pub(crate) fn head(&mut self, major: u8, n: u64) { + self.0.reserve(MAX_HEAD); + unsafe { + let len = self.0.len(); + let head = put_head(self.0.as_mut_ptr().add(len), major, n); + self.0.set_len(len + head); + } + } + + /// Length header followed by `data`: a whole byte or text string. + #[inline] + pub(crate) fn segment(&mut self, major: u8, data: &[u8]) { + self.0.reserve(MAX_HEAD + data.len()); + unsafe { + let len = self.0.len(); + let p = self.0.as_mut_ptr().add(len); + let head = put_head(p, major, data.len() as u64); + copy_short(data.as_ptr(), p.add(head), data.len()); + self.0.set_len(len + head + data.len()); + } + } + + #[inline] + pub(crate) fn raw(&mut self, data: &[u8]) { + self.0.extend_from_slice(data); + } + + #[inline] + pub(crate) fn bool(&mut self, v: bool) { + self.0.push(if v { TRUE } else { FALSE }); + } + + #[inline] + pub(crate) fn null(&mut self) { + self.0.push(NULL); + } #[inline] - fn push(&mut self, input: &[u8]) -> Result<(), Self::Error> { - self.0.extend_from_slice(input); - Ok(()) + pub(crate) fn f64(&mut self, v: f64) { + let mut buf = [F64; 9]; + buf[1..].copy_from_slice(&v.to_be_bytes()); + self.0.extend_from_slice(&buf); } }