Optimize DAG-CBOR encode/decode with in-place CBOR headers, inline short copies, and a two-way key cache

.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

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

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

src/cid.rs

@@ -52,6 +52,12 @@ fn read_varint(bytes: &[u8], pos: &mut usize) -> Option<u64> {
 /// 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)?;

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"));
                     }
                 }