didbot_data/mst/

tree.rs

//! The tree: nodes, the layering, and the blocks they encode to.

use std::collections::BTreeMap;

use super::{common_prefix_len, key_height};
use crate::cid::Cid;
use crate::dag_cbor;
use crate::value::Value;

/// One entry in a node: a key, the record it names, and what hangs to its
/// right.
///
/// The key is stored split, because that is how it is encoded. A node's
/// entries are sorted, so each one records how many leading **bytes** it
/// shares with its predecessor and then only the rest — `p` and `k` in the
/// node encoding. The first entry in a node shares nothing, because it has no
/// predecessor *in that node*; the key of the last entry in the subtree to its
/// left is not what the prefix is counted against.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Entry {
    /// How many leading bytes this key shares with the previous entry's.
    pub prefix: usize,
    /// The rest of the key, from `prefix` onwards, as bytes.
    pub suffix: Vec<u8>,
    /// The record this key names.
    pub value: Cid,
    /// The subtree of keys that sort after this one and before the next.
    pub tree: Option<Cid>,
}

/// One node of the tree.
///
/// A node holds every key at its own layer within some range, and a pointer
/// to a subtree between each pair of them. `left` is the subtree before the
/// first entry. A node with no entries is legal and common — it is what the
/// layering rule produces when a run of low keys hangs under a much taller
/// one — and it carries `left` alone.
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct Node {
    /// Keys below this node's layer that sort before every entry.
    pub left: Option<Cid>,
    /// This node's own keys, in order.
    pub entries: Vec<Entry>,
}

impl Node {
    /// The node as a data-model value, which is what it is hashed as.
    ///
    /// The field names are the protocol's and are one letter each: `l` for the
    /// left subtree, `e` for the entries, and within an entry `p` for the
    /// shared prefix length, `k` for the rest of the key, `v` for the record
    /// and `t` for the subtree to its right. `k` is a **byte string** and not
    /// text — a node is keyed on bytes, and a prefix length may land inside a
    /// multi-byte character.
    ///
    /// Ordering is not decided here. [`dag_cbor`] sorts map keys length-first
    /// and then bytewise, so a node built with the fields in any order encodes
    /// to the same bytes.
    #[must_use]
    pub fn to_value(&self) -> Value {
        let entries = self
            .entries
            .iter()
            .map(|entry| {
                Value::Map(BTreeMap::from([
                    ("p".to_owned(), Value::Integer(entry.prefix as i64)),
                    ("k".to_owned(), Value::Bytes(entry.suffix.clone())),
                    ("v".to_owned(), Value::Link(entry.value.clone())),
                    ("t".to_owned(), link_or_null(entry.tree.as_ref())),
                ]))
            })
            .collect();
        Value::Map(BTreeMap::from([
            ("l".to_owned(), link_or_null(self.left.as_ref())),
            ("e".to_owned(), Value::List(entries)),
        ]))
    }
}

/// A link, or the null a missing subtree is written as.
///
/// The field is always present. Omitting it would be a different map with a
/// different encoding and therefore a different CID.
fn link_or_null(cid: Option<&Cid>) -> Value {
    cid.map_or(Value::Null, |cid| Value::Link(cid.clone()))
}

/// A built Merkle search tree: a root CID and every block under it.
///
/// Built from a leaf set in one pass, then changed one key at a time by
/// [`Tree::insert`] and [`Tree::remove`], which leave it exactly the tree
/// [`Tree::build`] makes from the new leaf set. The blocks are handed
/// back rather than written anywhere, because what wants them differs — a CAR
/// export wants all of them, a proof wants the ones on one key's path — and a
/// tree that also owned a block store would have to be told which.
#[derive(Debug, Clone)]
pub struct Tree {
    pub(super) root: Cid,
    /// The layer the root sits on: the height of the tallest key, or zero.
    pub(super) layer: usize,
    pub(super) blocks: BTreeMap<Cid, Vec<u8>>,
    /// The same nodes, still as nodes.
    ///
    /// Kept alongside the encoded blocks rather than instead of them because
    /// [`Tree::path`] has to *read* a node to say which way to go, and this
    /// crate has no DAG-CBOR decoder to read one back out of its block.
    /// The duplication is bounded by the tree's node count, which is a small
    /// fraction of its record count.
    pub(super) nodes: BTreeMap<Cid, Node>,
}

/// Which blocks prove what a tree says about one key.
///
/// The nodes are the walk from the root down, in that order, and each one
/// names the next: a verifier hashes each block to check it is the node it
/// was told to expect, reads the entry range it lands in, and follows the
/// link. `value` is the record's CID when the key is in the tree and `None`
/// when it is not — and the *absence* is proved by the same nodes, because
/// the last one on the path is where the key would have to be if it existed.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Path {
    /// The nodes walked, root first.
    pub nodes: Vec<Cid>,
    /// The record the key names, if the tree holds it.
    pub value: Option<Cid>,
}

impl Tree {
    /// Builds the tree over a leaf set.
    ///
    /// The keys are the repository's, which atproto writes as
    /// `<collection>/<rkey>`, and the values are the CIDs of the records they
    /// name. A `BTreeMap<String, _>` is already in the order the tree needs:
    /// `str` compares bytewise, and bytewise over the keys is exactly the
    /// ordering a node's entries are in.
    ///
    /// An empty repository is not an absent tree. It is one node holding no
    /// entries and no left pointer, and it has a CID like any other, which is
    /// what an account's first commit signs.
    ///
    /// ```
    /// use std::collections::BTreeMap;
    /// use didbot_data::{mst::Tree, Cid};
    ///
    /// let empty = Tree::build(&BTreeMap::new());
    /// assert_eq!(
    ///     empty.root().to_string(),
    ///     "bafyreie5737gdxlw5i64vzichcalba3z2v5n6icifvx5xytvske7mr3hpm"
    /// );
    /// assert_eq!(empty.blocks().len(), 1);
    /// ```
    #[must_use]
    pub fn build(leaves: &BTreeMap<String, Cid>) -> Self {
        let entries: Vec<(&str, &Cid)> = leaves
            .iter()
            .map(|(key, value)| (key.as_str(), value))
            .collect();
        // The root sits at the height of the tallest key, with nothing above
        // it. An empty repository has no tallest key and roots at layer zero.
        let top = entries
            .iter()
            .map(|(key, _)| key_height(key))
            .max()
            .unwrap_or(0);
        let mut blocks = BTreeMap::new();
        let mut nodes = BTreeMap::new();
        let root = build_layer(&entries, top, &mut blocks, &mut nodes);
        Self {
            root,
            layer: top,
            blocks,
            nodes,
        }
    }

    /// The path through the tree to `key`, whether or not it is there.
    ///
    /// This is what an inclusion proof is made of, and what makes it a proof
    /// rather than an export: the nodes on one root-to-key walk and nothing
    /// else. A repository of ten thousand records proves any one of them with
    /// a handful of blocks, and a verifier that is handed the whole
    /// repository instead has been given an export with a different name.
    ///
    /// A key the tree does not hold gets the same walk with no `value`. That
    /// is the exclusion proof: the last node on it is the one the key would
    /// have to be in, and a verifier that reaches it and does not find the key
    /// has been shown the key is not in the tree at all.
    #[must_use]
    pub fn path(&self, key: &str) -> Path {
        let mut walked = Vec::new();
        let mut current = self.root.clone();
        loop {
            walked.push(current.clone());
            // Unreachable for a tree this crate built: every link in it names
            // a node this crate stored. Answering with the walk so far rather
            // than panicking keeps a corrupted tree a bad proof instead of a
            // dead server.
            let Some(node) = self.nodes.get(&current) else {
                return Path {
                    nodes: walked,
                    value: None,
                };
            };
            match step(node, key.as_bytes()) {
                Step::Found(value) => {
                    return Path {
                        nodes: walked,
                        value: Some(value),
                    }
                }
                Step::Descend(next) => current = next,
                Step::Absent => {
                    return Path {
                        nodes: walked,
                        value: None,
                    }
                }
            }
        }
    }

    /// Every key in the tree, in the order the tree holds them.
    ///
    /// Reconstructed from the nodes rather than remembered from the leaf set,
    /// so it says what the tree *is* rather than what it was built from. A
    /// node stores each key as a shared prefix length and a suffix, and the
    /// prefix is counted against the previous entry **in that node**, so the
    /// walk has to rebuild the keys node by node rather than across the whole
    /// tree.
    #[must_use]
    pub fn keys(&self) -> Vec<String> {
        let mut found = Vec::new();
        self.walk_keys(&self.root, &mut found);
        found
    }

    /// Appends every key under one node, in order.
    fn walk_keys(&self, node: &Cid, found: &mut Vec<String>) {
        // A link this crate did not store is a corrupted tree, and answering
        // with the keys found so far keeps that a short answer rather than a
        // panic — the same choice [`Tree::path`] makes.
        let Some(node) = self.nodes.get(node) else {
            return;
        };
        if let Some(left) = &node.left {
            self.walk_keys(left, found);
        }
        let mut previous: Vec<u8> = Vec::new();
        for entry in &node.entries {
            let mut full = previous[..entry.prefix.min(previous.len())].to_vec();
            full.extend_from_slice(&entry.suffix);
            // A repository key is a collection and a record key, both of which
            // are ASCII, so this is lossless in practice. It is written as a
            // lossy conversion rather than an unwrap because a node is bytes
            // and nothing in the encoding promises otherwise.
            found.push(String::from_utf8_lossy(&full).into_owned());
            if let Some(tree) = &entry.tree {
                self.walk_keys(tree, found);
            }
            previous = full;
        }
    }

    /// The blocks a consumer needs to verify one key changing.
    ///
    /// atproto calls this a *covering proof*, and it is what a firehose commit
    /// carries so that a subscriber can check an operation without holding the
    /// repository. It is more than [`Tree::path`]: adding or removing a key
    /// can split a node or merge two, so the nodes that move are the ones on
    /// the key's own walk **and** on the walks to the keys either side of it.
    /// A proof of the key alone verifies the lookup and cannot reproduce the
    /// root the operation started from.
    ///
    /// The neighbours are found by listing the keys, which is a walk of the
    /// whole tree. A caller that already keeps the keys in order finds them
    /// itself and calls [`Tree::covering_proof_between`].
    #[must_use]
    pub fn covering_proof(&self, key: &str) -> BTreeMap<Cid, Vec<u8>> {
        let keys = self.keys();
        let before = keys.iter().rev().find(|held| held.as_str() < key);
        let after = keys.iter().find(|held| held.as_str() > key);
        self.covering_proof_between(key, before.map(String::as_str), after.map(String::as_str))
    }

    /// [`Tree::covering_proof`], given the keys either side of `key`.
    ///
    /// `before` and `after` are the nearest keys the tree holds below and
    /// above `key`, or `None` where there is none.
    #[must_use]
    pub fn covering_proof_between(
        &self,
        key: &str,
        before: Option<&str>,
        after: Option<&str>,
    ) -> BTreeMap<Cid, Vec<u8>> {
        let mut proof = BTreeMap::new();
        for walked in std::iter::once(key).chain(before).chain(after) {
            for node in self.path(walked).nodes {
                if let Some(block) = self.blocks.get(&node) {
                    proof.insert(node, block.clone());
                }
            }
        }
        proof
    }

    /// The CID a commit signs over.
    #[must_use]
    pub fn root(&self) -> &Cid {
        &self.root
    }

    /// Every node block in the tree, keyed by its CID.
    #[must_use]
    pub fn blocks(&self) -> &BTreeMap<Cid, Vec<u8>> {
        &self.blocks
    }

    /// The nodes, still as nodes, keyed by CID.
    #[must_use]
    pub fn nodes(&self) -> &BTreeMap<Cid, Node> {
        &self.nodes
    }

    /// How many nodes the tree has.
    #[must_use]
    pub fn len(&self) -> usize {
        self.blocks.len()
    }

    /// Never true: a tree always has at least its root node.
    ///
    /// Present because `len` without `is_empty` is a lint, and answering it
    /// honestly is better than suppressing the lint — an empty repository
    /// still has a root, which is the whole reason a first commit can be
    /// signed.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        false
    }
}

/// Builds one node at `layer` and everything under it.
///
/// Every entry handed in has a height at or below `layer`, and the caller has
/// established that this layer is one somebody asked for. Keys at the layer
/// become entries; runs of lower keys between them become subtrees, each built
/// at exactly one layer down whether or not a key sits there. That last clause
/// is the load-bearing one — see the module documentation.
fn build_layer(
    entries: &[(&str, &Cid)],
    layer: usize,
    blocks: &mut BTreeMap<Cid, Vec<u8>>,
    nodes: &mut BTreeMap<Cid, Node>,
) -> Cid {
    let mut node = Node::default();
    let mut below: Vec<(&str, &Cid)> = Vec::new();
    let mut previous: Option<&str> = None;

    for (key, value) in entries {
        if key_height(key) == layer {
            attach(&mut node, &mut below, layer, blocks, nodes);
            let prefix = previous.map_or(0, |previous| common_prefix_len(previous, key));
            node.entries.push(Entry {
                prefix,
                suffix: key.as_bytes()[prefix..].to_vec(),
                value: (*value).clone(),
                tree: None,
            });
            previous = Some(key);
        } else {
            below.push((key, value));
        }
    }
    attach(&mut node, &mut below, layer, blocks, nodes);

    store(node, blocks, nodes)
}

/// Where one node says to go next for a key.
enum Step {
    /// The key is an entry in this node, naming this record.
    Found(Cid),
    /// The key, if it exists, is in this subtree.
    Descend(Cid),
    /// The key would be in this node and is not, and there is no subtree
    /// under it to look in. The walk ends and the key is not in the tree.
    Absent,
}

/// Reads one node and says which way `key` lies.
///
/// A node's entries are stored with their keys split against the previous
/// one, so the full keys are rebuilt as the entries are walked — the same
/// reconstruction any other implementation does on the way down.
fn step(node: &Node, key: &[u8]) -> Step {
    let mut previous: Vec<u8> = Vec::new();
    // The subtree to the left of the entry being considered. It starts as the
    // node's own left pointer and becomes each entry's right-hand subtree as
    // the walk passes it, which is the same statement as "a subtree sits
    // between the entry before it and the entry after it".
    let mut before = node.left.clone();
    for entry in &node.entries {
        let mut full = previous[..entry.prefix.min(previous.len())].to_vec();
        full.extend_from_slice(&entry.suffix);
        match key.cmp(full.as_slice()) {
            std::cmp::Ordering::Less => {
                return before.map_or(Step::Absent, Step::Descend);
            }
            std::cmp::Ordering::Equal => return Step::Found(entry.value.clone()),
            std::cmp::Ordering::Greater => {
                before = entry.tree.clone();
                previous = full;
            }
        }
    }
    before.map_or(Step::Absent, Step::Descend)
}

/// Builds the pending run of lower keys and hangs it off the node.
///
/// It goes on the last entry if there is one and on `left` if there is not,
/// which is the same statement as "the subtree sits between the entry before
/// it and the entry after it".
fn attach<'a>(
    node: &mut Node,
    below: &mut Vec<(&'a str, &'a Cid)>,
    layer: usize,
    blocks: &mut BTreeMap<Cid, Vec<u8>>,
    nodes: &mut BTreeMap<Cid, Node>,
) {
    if below.is_empty() {
        return;
    }
    // Unreachable: at layer zero every key has height zero, so nothing is
    // ever below it. Written as a saturating step rather than a subtraction
    // so that a future caller's mistake is a wrong tree and not a panic in a
    // server.
    debug_assert!(layer > 0, "a key cannot sit below layer zero");
    let subtree = build_layer(below, layer.saturating_sub(1), blocks, nodes);
    match node.entries.last_mut() {
        Some(last) => last.tree = Some(subtree),
        None => node.left = Some(subtree),
    }
    below.clear();
}

/// Encodes a node, records its block and the node itself, and answers its CID.
fn store(node: Node, blocks: &mut BTreeMap<Cid, Vec<u8>>, nodes: &mut BTreeMap<Cid, Node>) -> Cid {
    let block = dag_cbor::encode(&node.to_value());
    let cid = Cid::of_dag_cbor(&block);
    blocks.insert(cid.clone(), block);
    nodes.insert(cid.clone(), node);
    cid
}

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

    /// A CID to hang off a key when the test is about the tree and not about
    /// what the leaves hold.
    fn leaf(seed: &str) -> Cid {
        Cid::of_dag_cbor(seed.as_bytes())
    }

    fn tree_of(keys: &[&str]) -> Tree {
        let leaves: BTreeMap<String, Cid> = keys
            .iter()
            .map(|key| ((*key).to_owned(), leaf(key)))
            .collect();
        Tree::build(&leaves)
    }

    /// The shape is a function of the key set alone, so the order the leaves
    /// were collected in cannot reach it. This is the property that makes two
    /// servers agree, and it is worth asserting even though a `BTreeMap`
    /// makes it hard to violate.
    #[test]
    fn the_root_does_not_depend_on_insertion_order() {
        let forwards = tree_of(&["a", "b", "c", "d", "e", "f", "g", "h"]);
        let backwards = tree_of(&["h", "g", "f", "e", "d", "c", "b", "a"]);
        assert_eq!(forwards.root(), backwards.root());
        assert_eq!(forwards.blocks(), backwards.blocks());
    }

    /// One record's tree is one node holding one entry, and it is not the
    /// empty tree.
    #[test]
    fn a_single_record_roots_at_its_own_node() {
        let single = tree_of(&["app.bsky.feed.post/asdf"]);
        let empty = Tree::build(&BTreeMap::new());
        assert_eq!(single.len(), 1);
        assert_ne!(single.root(), empty.root());
    }

    /// Changing one leaf changes the root. Trivially true of a Merkle
    /// structure and the reason one is used, so it is stated.
    #[test]
    fn a_changed_leaf_moves_the_root() {
        let mut leaves: BTreeMap<String, Cid> = ["a", "b", "c"]
            .iter()
            .map(|k| ((*k).to_owned(), leaf(k)))
            .collect();
        let before = Tree::build(&leaves).root().clone();
        leaves.insert("b".to_owned(), leaf("something else"));
        assert_ne!(Tree::build(&leaves).root(), &before);
    }

    /// The first entry of a node shares no prefix, because it has no
    /// predecessor in that node — even when the key immediately before it in
    /// the repository shares most of its bytes.
    #[test]
    fn a_nodes_first_entry_carries_its_whole_key() {
        // Both at layer 0, so both land in one node.
        let tree = tree_of(&["app.bsky.feed.post/aaaa", "app.bsky.feed.post/aaab"]);
        let block = tree.blocks().get(tree.root()).expect("the root block");
        // The whole first key appears in the block; the second is stored as
        // its last byte alone.
        let needle = b"app.bsky.feed.post/aaaa";
        assert!(
            block.windows(needle.len()).any(|window| window == needle),
            "the first entry should carry its whole key"
        );
    }

    /// Every block a tree hands back is the encoding of the CID it is filed
    /// under. A block store that disagreed with its own keys would produce a
    /// CAR nothing could walk.
    #[test]
    fn every_block_hashes_to_the_cid_it_is_filed_under() {
        let tree = tree_of(&[
            "app.bsky.feed.post/3jt6wh3ckcc2y",
            "app.bsky.feed.like/3jtdufrp6md2l",
            "com.example.thing/3lqmnopqrstuv",
            "app.bsky.actor.profile/self",
        ]);
        for (cid, block) in tree.blocks() {
            assert_eq!(&Cid::of_dag_cbor(block), cid);
        }
    }

    /// The path to a key is the walk that reaches it, and the last node on it
    /// holds the record's CID.
    #[test]
    fn a_path_ends_at_the_node_holding_the_key() {
        let tree = tree_of(&["a", "b", "c", "d", "e", "f", "g", "h"]);
        let path = tree.path("d");
        assert_eq!(path.value.as_ref(), Some(&leaf("d")));
        assert_eq!(path.nodes.first(), Some(tree.root()));
        // Every node on the walk is one the tree actually holds.
        for cid in &path.nodes {
            assert!(tree.blocks().contains_key(cid), "{cid} is not in the tree");
        }
    }

    /// A key the tree does not hold gets the same walk and no value. That is
    /// the exclusion proof, and it is why `path` does not return an `Option`.
    #[test]
    fn a_missing_key_still_has_a_path() {
        let tree = tree_of(&["a", "b", "c", "d"]);
        let path = tree.path("zzzz");
        assert_eq!(path.value, None);
        assert!(!path.nodes.is_empty());
        assert_eq!(path.nodes.first(), Some(tree.root()));
    }

    /// The point of a proof: it is the path and not the repository. A tree
    /// with hundreds of records proves one of them with a handful of blocks.
    #[test]
    fn a_path_is_shorter_than_the_tree() {
        let keys: Vec<String> = (0..400)
            .map(|i| format!("com.example.thing/3lqm{i:04}"))
            .collect();
        let leaves: BTreeMap<String, Cid> =
            keys.iter().map(|key| (key.clone(), leaf(key))).collect();
        let tree = Tree::build(&leaves);
        let path = tree.path(&keys[137]);
        assert_eq!(path.value.as_ref(), Some(&leaf(&keys[137])));
        assert!(
            path.nodes.len() * 4 < tree.len(),
            "a path of {} nodes out of {} is not a proof",
            path.nodes.len(),
            tree.len()
        );
    }

    /// Every key in the tree is reachable by a walk, which is the property a
    /// wrong `step` would break for exactly the keys that sit in a subtree.
    #[test]
    fn every_key_is_reachable_by_its_own_path() {
        let keys: Vec<String> = (0..200)
            .map(|i| format!("app.bsky.feed.post/3jt6wh3ck{i:03}"))
            .collect();
        let leaves: BTreeMap<String, Cid> =
            keys.iter().map(|key| (key.clone(), leaf(key))).collect();
        let tree = Tree::build(&leaves);
        for key in &keys {
            assert_eq!(
                tree.path(key).value.as_ref(),
                Some(&leaf(key)),
                "{key} is in the tree and its path did not find it"
            );
        }
    }

    /// A prefix length may land inside a multi-byte character, because the
    /// tree is keyed on bytes. Not a key atproto permits, but the encoding
    /// has to be right for the same reason `common_prefix_len` counts bytes.
    #[test]
    fn a_prefix_may_split_a_character() {
        let tree = tree_of(&["ααx", "ααy"]);
        for (cid, block) in tree.blocks() {
            assert_eq!(&Cid::of_dag_cbor(block), cid);
        }
    }
}