didbot_data/mst/

shape.rs

//! The two pure functions a Merkle search tree's *shape* comes from.
//!
//! These sit below [`Tree`](super::Tree) and are checked without one, which is
//! the whole argument: a tree built on a wrong height function is not
//! merely self-consistent, it is *correct* — every lookup succeeds and every
//! proof it makes verifies against itself — and the only symptom is a root CID
//! no other implementation reproduces.
//!
//! * [`key_height`] decides which layer a key sits on. It is derived from the
//!   key's own hash, so every implementation puts a key on the same layer
//!   without coordinating.
//! * [`common_prefix_len`] decides how much of a key a node has to store. A
//!   node's entries are sorted, so each one is written as the number of bytes
//!   it shares with its predecessor plus the rest.

use sha2::{Digest, Sha256};

/// atproto's Merkle search tree has a fanout of four, so two zero bits of the
/// key's hash buy one layer.
const FANOUT_BITS: u32 = 2;

/// Which layer of the tree a key belongs on.
///
/// The rule is the whole of it: hash the key with sha-256, count the leading
/// zero *bits* of the digest, and divide by two. Every part of that sentence
/// is a way to get it wrong — bytes instead of bits gives a height that is
/// almost always zero, and forgetting the divide doubles every layer — and a
/// tree built either way still works, which is why this is checked against
/// `mst/key_heights.json` rather than against a tree.
///
/// The key is hashed as its UTF-8 bytes. atproto record keys are ASCII by
/// rule, so this is not a choice that shows up in the vectors, but it is the
/// one every implementation makes and it is written down here rather than
/// left to the caller.
///
/// The height is unbounded in principle — a digest of 256 zero bits would
/// answer 128 — and in practice a key above layer four is already rare. There
/// is no error case: every string has a hash and every hash has a height.
///
/// ```
/// use didbot_data::mst::key_height;
///
/// assert_eq!(key_height("asdf"), 0);
/// assert_eq!(key_height("blue"), 1);
/// assert_eq!(key_height("app.bsky.feed.post/9adeb165882c"), 8);
/// ```
#[must_use]
pub fn key_height(key: &str) -> usize {
    let digest = Sha256::digest(key.as_bytes());

    let mut zeros: u32 = 0;
    for byte in digest {
        zeros += byte.leading_zeros();
        if byte != 0 {
            break;
        }
    }

    (zeros / FANOUT_BITS) as usize
}

/// How many leading **bytes** two keys share.
///
/// This is what lets a node store a suffix rather than a whole key: entries
/// are sorted, so each one records how much of its predecessor it repeats and
/// then only the rest. Get it wrong and the node encodes to different bytes,
/// which means a different CID, which means a different root — the error does
/// not show up as a lookup that fails, it shows up as a repository nobody else
/// can verify.
///
/// The count is in bytes, not characters, and that is the load-bearing detail.
/// Two keys sharing a three-character prefix may share four bytes or nine, and
/// the number a node carries is the one that indexes into the UTF-8. A
/// character count would agree with a byte count on every ASCII key — which is
/// every key upstream's vectors contain, and every record key atproto permits
/// — and would disagree the first time a tree held a key that was neither.
///
/// A returned length can therefore land inside a multi-byte character. That is
/// correct: the tree is keyed on byte strings and the encoding of a node
/// splits on the byte offset this returns.
///
/// ```
/// use didbot_data::mst::common_prefix_len;
///
/// assert_eq!(common_prefix_len("abcde", "abc1"), 3);
/// assert_eq!(common_prefix_len("abc", "abc\0"), 3);
/// // Two characters shared, four bytes shared.
/// assert_eq!(common_prefix_len("ααx", "ααy"), 4);
/// ```
#[must_use]
pub fn common_prefix_len(left: &str, right: &str) -> usize {
    left.as_bytes()
        .iter()
        .zip(right.as_bytes())
        .take_while(|(left, right)| left == right)
        .count()
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The empty key has a height like any other: it is the height of the
    /// sha-256 of nothing, which happens to be zero.
    #[test]
    fn the_empty_key_has_a_height() {
        assert_eq!(key_height(""), 0);
    }

    /// Counting zero *bytes* rather than zero bits is the mistake this
    /// guards. `88bfafc7` hashes to a digest beginning `0x0c`: no zero byte at
    /// all, four leading zero bits, height two.
    #[test]
    fn height_counts_bits_not_bytes() {
        assert_eq!(key_height("88bfafc7"), 2);
        assert_eq!(
            Sha256::digest("88bfafc7".as_bytes())[0],
            0x0c,
            "the digest this test reasons about has changed"
        );
    }

    /// A key whose digest starts with a whole zero byte forces the loop past
    /// its first iteration, which the low-height keys never do.
    #[test]
    fn height_carries_across_a_whole_zero_byte() {
        // `884976f5` is upstream's height-6 vector: twelve leading zero bits,
        // so a whole zero byte and then four more.
        assert_eq!(key_height("884976f5"), 6);
        assert_eq!(Sha256::digest("884976f5".as_bytes())[0], 0x00);
    }

    /// Neither argument's length bounds the answer on its own.
    #[test]
    fn a_prefix_is_symmetric_and_bounded_by_the_shorter_key() {
        assert_eq!(common_prefix_len("abcde", "abc"), 3);
        assert_eq!(common_prefix_len("abc", "abcde"), 3);
        assert_eq!(common_prefix_len("", "abc"), 0);
        assert_eq!(common_prefix_len("abc", ""), 0);
        assert_eq!(common_prefix_len("", ""), 0);
    }

    /// Bytes, not characters. Upstream's vectors are all ASCII, so this is the
    /// one claim in this module the interop suite cannot make on its own.
    #[test]
    fn a_prefix_is_counted_in_bytes_not_characters() {
        // Two shared characters, four shared bytes.
        assert_eq!(common_prefix_len("ααx", "ααy"), 4);
        // A shared first byte inside characters that differ: 'α' is 0xCE 0xB1
        // and 'ω' is 0xCF 0x89, so nothing is shared at all here...
        assert_eq!(common_prefix_len("α", "ω"), 0);
        // ...while 'α' (0xCE 0xB1) and 'β' (0xCE 0xB2) share their lead byte.
        assert_eq!(common_prefix_len("α", "β"), 1);
    }
}