didbot_data/mst/

store.rs

//! A tree somebody else built: nodes by CID, and the walks over them.
//!
//! [`Tree`](super::Tree) is this crate's own tree — built from a whole leaf
//! set, immutable, holding its own blocks. A [`Store`] is the opposite
//! situation: a pile of blocks that arrived from elsewhere, one of which is
//! claimed to be a root. Nothing about it is trusted, so every walk here can
//! fail — a link may name a block that is not in the pile, and a subtree may
//! claim a layer it is not on.
//!
//! Three things are built on the walks:
//!
//! * [`Store::leaves`] — what the tree says the repository holds.
//! * [`Store::find`] — the walk to one path, which is an inclusion proof when
//!   the path is there and an exclusion proof when it is not.
//! * [`diff`] — what changed between two roots, as the nodes that came and
//!   went, the records that came and went, and the blocks that prove it.
//!
//! # Why the walk stops early
//!
//! [`Store::find`] gives up as soon as it is below the layer the path's own
//! height puts it on. A path can only ever be an entry at that one layer, so
//! a walk that continued would be reading blocks that cannot hold the answer
//! — and an exclusion proof made of them would be larger than it needs to be
//! while proving exactly the same thing. This is the rule that decides how
//! big `proof_nodes` is in the Merkle search tree conformance vectors, and
//! getting it wrong is not a wrong answer, only a wasteful one.

use std::collections::{BTreeMap, BTreeSet};

use super::read::NodeError;
use super::{key_height, Node};
use crate::cid::Cid;

/// Why a tree could not be walked.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum StoreError {
    /// A link naming a block the store does not hold.
    #[error("the tree links to {0}, which is not in the store")]
    Missing(Box<Cid>),
    /// A subtree that is not one layer below the node pointing at it.
    #[error("{cid} says it is at layer {found}, and its parent puts it at {expected}")]
    Layer {
        /// The subtree that disagreed.
        cid: Box<Cid>,
        /// The layer the node's own keys put it on.
        found: usize,
        /// The layer its parent puts it on.
        expected: usize,
    },
    /// A root whose own layer cannot be told, which every walk needs to know.
    ///
    /// A node with no keys and a subtree below it is one layer under whatever
    /// pointed at it, and a root has nothing pointing at it. The tree that
    /// produced it is malformed: [`Tree::build`](super::Tree::build) strips
    /// those from the top, and so does every other implementation.
    #[error("the root {0} has no keys, so nothing says which layer it is on")]
    RootLayer(Box<Cid>),
}

/// Blocks that are claimed to be a Merkle search tree.
///
/// Nodes are held decoded, because every operation reads them and decoding is
/// where the structural checks live. The blocks are held too: a proof is
/// answered as a set of CIDs, and whatever sends it needs the bytes.
#[derive(Debug, Clone, Default)]
pub struct Store {
    /// Every node, by the CID it was filed under.
    nodes: BTreeMap<Cid, Node>,
    /// The blocks those nodes were decoded from.
    blocks: BTreeMap<Cid, Vec<u8>>,
}

/// What a walk to one path found, and what proves it.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Found {
    /// The record the path names, if the tree holds it.
    pub value: Option<Cid>,
    /// The nodes walked, which are the proof of either answer.
    ///
    /// A set rather than a list: it is combined with other proofs far more
    /// often than it is followed, and the order is recoverable from the
    /// blocks themselves.
    pub proof: BTreeSet<Cid>,
}

impl Store {
    /// Builds a store over blocks that are all claimed to be nodes.
    ///
    /// Every block is decoded here rather than lazily, so a malformed node is
    /// a failure to load a tree rather than a failure part-way through
    /// answering a query about it. A CAR file that also holds records is not
    /// what this takes; see [`Store::insert`] for adding them one at a time.
    pub fn load(
        blocks: impl IntoIterator<Item = (Cid, Vec<u8>)>,
    ) -> Result<Self, (Box<Cid>, NodeError)> {
        let mut store = Self::default();
        for (cid, block) in blocks {
            store
                .insert(&cid, block)
                .map_err(|err| (Box::new(cid.clone()), err))?;
        }
        Ok(store)
    }

    /// Decodes one block and files it under `cid`.
    ///
    /// The CID is taken on trust: whether the block hashes to it is a
    /// question about a *container* — see `didbot_repo::car::read`,
    /// which checks it — and re-hashing here would make every store pay for
    /// it twice.
    pub fn insert(&mut self, cid: &Cid, block: Vec<u8>) -> Result<(), NodeError> {
        let node = Node::decode(&block)?;
        self.nodes.insert(cid.clone(), node);
        self.blocks.insert(cid.clone(), block);
        Ok(())
    }

    /// Files a node that was built rather than decoded.
    ///
    /// Used by anything that *writes* a tree into a store — a rewrite, a
    /// proof being replayed — where the node is already in hand and encoding
    /// it only to decode it again would prove nothing.
    pub fn put(&mut self, node: Node) -> Cid {
        let (cid, block) = node.block();
        self.blocks.insert(cid.clone(), block);
        self.nodes.insert(cid.clone(), node);
        cid
    }

    /// The node under `cid`, or the error naming what is missing.
    pub fn node(&self, cid: &Cid) -> Result<&Node, StoreError> {
        self.nodes
            .get(cid)
            .ok_or_else(|| StoreError::Missing(Box::new(cid.clone())))
    }

    /// The block under `cid`, if the store holds it.
    #[must_use]
    pub fn block(&self, cid: &Cid) -> Option<&[u8]> {
        self.blocks.get(cid).map(Vec::as_slice)
    }

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

    /// Whether the store holds nothing at all.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.nodes.is_empty()
    }

    /// Which layer a root is on, which is the one thing a walk needs before
    /// it starts.
    pub fn root_layer(&self, root: &Cid) -> Result<usize, StoreError> {
        self.node(root)?
            .layer()
            .ok_or_else(|| StoreError::RootLayer(Box::new(root.clone())))
    }

    /// Every node reachable from `root`, itself included.
    ///
    /// This is what a tree *is*, as a set of blocks, and the two directions
    /// of [`diff`] are the two differences of these sets. Subtree layers are
    /// checked on the way down: a node claiming the wrong one is a tree that
    /// two implementations would walk differently.
    pub fn nodes_under(&self, root: &Cid) -> Result<BTreeSet<Cid>, StoreError> {
        let mut found = BTreeSet::new();
        let mut pending = vec![(root.clone(), self.root_layer(root)?)];
        while let Some((cid, layer)) = pending.pop() {
            if !found.insert(cid.clone()) {
                continue;
            }
            let node = self.node(&cid)?;
            self.check_layer(&cid, node, layer)?;
            for subtree in node.parts().subtrees.into_iter().flatten() {
                pending.push((subtree, layer.saturating_sub(1)));
            }
        }
        Ok(found)
    }

    /// Every path in the tree and the record it names, in key order.
    pub fn leaves(&self, root: &Cid) -> Result<BTreeMap<String, Cid>, StoreError> {
        let mut leaves = BTreeMap::new();
        self.collect(root, self.root_layer(root)?, &mut leaves)?;
        Ok(leaves)
    }

    /// Appends every path under one node.
    ///
    /// The descent is an explicit stack, as [`Store::nodes_under`]'s is, and
    /// for the same reason: the depth of a tree that arrived from elsewhere
    /// is whatever its blocks say. A chain of empty nodes is a few bytes per
    /// link and is walked clean, so recursing once per link would let a small
    /// pile of blocks end the process.
    fn collect(
        &self,
        cid: &Cid,
        layer: usize,
        leaves: &mut BTreeMap<String, Cid>,
    ) -> Result<(), StoreError> {
        let mut pending = vec![(cid.clone(), layer)];
        while let Some((cid, layer)) = pending.pop() {
            let node = self.node(&cid)?;
            self.check_layer(&cid, node, layer)?;
            let parts = node.parts();
            for (key, value) in parts.keys.into_iter().zip(parts.values) {
                leaves.insert(key, value);
            }
            // Pushed right to left so they pop left to right, which is the
            // order the keys are in.
            let below = layer.saturating_sub(1);
            for subtree in parts.subtrees.into_iter().rev().flatten() {
                pending.push((subtree, below));
            }
        }
        Ok(())
    }

    /// Refuses a subtree that is not on the layer its parent puts it on.
    ///
    /// An empty intermediate has no layer of its own and is skipped: it
    /// exists precisely to carry a pointer down a layer nothing sits on.
    fn check_layer(&self, cid: &Cid, node: &Node, expected: usize) -> Result<(), StoreError> {
        match node.layer() {
            Some(found) if found != expected && !node.is_empty() => Err(StoreError::Layer {
                cid: Box::new(cid.clone()),
                found,
                expected,
            }),
            _ => Ok(()),
        }
    }

    /// Walks to one path, and hands back the walk as well as the answer.
    ///
    /// The walk is the proof: a verifier holding these blocks and the root
    /// CID can repeat it, and reach the same answer or find that the blocks
    /// do not hash to what they were named. That is true of a path that is
    /// *not* in the tree as well — the walk ends where the path would have to
    /// be, and its absence there is its absence everywhere.
    ///
    /// The descent stops as soon as it is below the path's own layer; see the
    /// module documentation for why that is not merely an optimisation.
    pub fn find(&self, root: &Cid, path: &str) -> Result<Found, StoreError> {
        let wanted = key_height(path);
        let mut proof = BTreeSet::new();
        let mut cid = root.clone();
        let mut layer = self.root_layer(root)?;
        proof.insert(cid.clone());
        loop {
            // Below the layer the path belongs on, so the tree does not hold
            // it however far down the blocks go.
            if wanted > layer {
                return Ok(Found { value: None, proof });
            }
            let node = self.node(&cid)?;
            self.check_layer(&cid, node, layer)?;
            let parts = node.parts();
            let mut index = 0;
            while index < parts.keys.len() && parts.keys[index].as_str() < path {
                index += 1;
            }
            if parts.keys.get(index).is_some_and(|key| key == path) {
                return Ok(Found {
                    value: Some(parts.values[index].clone()),
                    proof,
                });
            }
            let Some(subtree) = parts.subtrees[index].clone() else {
                return Ok(Found { value: None, proof });
            };
            proof.insert(subtree.clone());
            cid = subtree;
            layer = layer.saturating_sub(1);
        }
    }
}

/// One record that changed between two trees.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RecordOp {
    /// The repository path: `<collection>/<record key>`.
    pub path: String,
    /// What the record was before, or `None` if it was created.
    pub old_value: Option<Cid>,
    /// What it is after, or `None` if it was deleted.
    pub new_value: Option<Cid>,
}

/// What changed between two roots.
///
/// The two node sets are the difference of the trees as *blocks*, and they
/// are not symmetric with the records: a node can be created without any
/// record changing inside it, because a sibling changed and the node above
/// both had to be rewritten.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Diff {
    /// Nodes the second tree has and the first does not.
    pub created_nodes: BTreeSet<Cid>,
    /// Nodes the first tree has and the second does not.
    pub deleted_nodes: BTreeSet<Cid>,
    /// The records that changed, in path order.
    pub ops: Vec<RecordOp>,
    /// The blocks that prove those record changes against the second root.
    ///
    /// One walk per changed path, over the tree *after* the change: an
    /// inclusion proof for a record that is now there, an exclusion proof for
    /// one that is now gone. It is usually the created nodes over again, and
    /// it is neither a subset nor a superset of them in general — a deletion
    /// can be proved by a node that was already there, and a node can be
    /// created that no changed path walks through.
    pub proof_nodes: BTreeSet<Cid>,
}

/// Diffs two trees held in one store.
///
/// Both roots are walked whole, which is the obvious implementation and not
/// the clever one. A production diff would descend the two trees together and
/// skip the subtrees whose CIDs already match; this reads every node of both.
/// For a repository-sized tree that is the difference between a proof and an
/// export, and it is the right trade here — the operation is checked against
/// sixteen thousand vectors, and the version that is obviously correct is the
/// one worth checking them against.
pub fn diff(store: &Store, before: &Cid, after: &Cid) -> Result<Diff, StoreError> {
    let old_nodes = store.nodes_under(before)?;
    let new_nodes = store.nodes_under(after)?;
    let old_leaves = store.leaves(before)?;
    let new_leaves = store.leaves(after)?;

    let mut ops = Vec::new();
    for (path, old_value) in &old_leaves {
        match new_leaves.get(path) {
            Some(new_value) if new_value == old_value => {}
            new_value => ops.push(RecordOp {
                path: path.clone(),
                old_value: Some(old_value.clone()),
                new_value: new_value.cloned(),
            }),
        }
    }
    for (path, new_value) in &new_leaves {
        if !old_leaves.contains_key(path) {
            ops.push(RecordOp {
                path: path.clone(),
                old_value: None,
                new_value: Some(new_value.clone()),
            });
        }
    }
    ops.sort_by(|left, right| left.path.cmp(&right.path));

    let mut proof_nodes = BTreeSet::new();
    for op in &ops {
        proof_nodes.extend(store.find(after, &op.path)?.proof);
    }

    Ok(Diff {
        created_nodes: new_nodes.difference(&old_nodes).cloned().collect(),
        deleted_nodes: old_nodes.difference(&new_nodes).cloned().collect(),
        ops,
        proof_nodes,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mst::{Parts, Tree};

    /// A record CID to hang off a path when the test is about the tree.
    fn record(seed: &str) -> Cid {
        Cid::of_dag_cbor(seed.as_bytes())
    }

    /// Builds a tree over `paths` and loads it into a store.
    fn tree(paths: &[&str]) -> (Store, Cid) {
        let leaves: BTreeMap<String, Cid> = paths
            .iter()
            .map(|path| ((*path).to_owned(), record(path)))
            .collect();
        let built = Tree::build(&leaves);
        let store = Store::load(
            built
                .blocks()
                .iter()
                .map(|(cid, block)| (cid.clone(), block.clone())),
        )
        .expect("a tree this crate built");
        (store, built.root().clone())
    }

    /// A tree that went out as blocks comes back as the same leaves. This is
    /// the round trip the whole file rests on.
    #[test]
    fn a_tree_reads_back_the_leaves_it_was_built_from() {
        let paths: Vec<String> = (0..250)
            .map(|i| format!("app.bsky.feed.post/3jt6wh3ck{i:03}"))
            .collect();
        let leaves: BTreeMap<String, Cid> = paths
            .iter()
            .map(|path| (path.clone(), record(path)))
            .collect();
        let built = Tree::build(&leaves);
        let store = Store::load(
            built
                .blocks()
                .iter()
                .map(|(cid, block)| (cid.clone(), block.clone())),
        )
        .expect("a tree this crate built");
        assert_eq!(store.leaves(built.root()).unwrap(), leaves);
        assert_eq!(store.nodes_under(built.root()).unwrap().len(), built.len());
    }

    /// A path in the tree is found, and the walk that found it is shorter
    /// than the tree.
    #[test]
    fn a_walk_finds_a_path_and_is_smaller_than_the_tree() {
        let paths: Vec<String> = (0..400)
            .map(|i| format!("com.example.thing/3lqm{i:04}"))
            .collect();
        let borrowed: Vec<&str> = paths.iter().map(String::as_str).collect();
        let (store, root) = tree(&borrowed);
        let found = store.find(&root, &paths[137]).unwrap();
        assert_eq!(found.value, Some(record(&paths[137])));
        assert!(found.proof.len() * 4 < store.len());
    }

    /// A path the tree does not hold gets a walk and no value, and the walk
    /// stops at the layer the path would have to be on rather than running to
    /// a leaf.
    #[test]
    fn an_absent_path_stops_at_its_own_layer() {
        let paths: Vec<String> = (0..400)
            .map(|i| format!("com.example.thing/3lqm{i:04}"))
            .collect();
        let borrowed: Vec<&str> = paths.iter().map(String::as_str).collect();
        let (store, root) = tree(&borrowed);
        let missing = "com.example.thing/zzzzzzzz";
        let found = store.find(&root, missing).unwrap();
        assert_eq!(found.value, None);
        let deepest = found
            .proof
            .iter()
            .map(|cid| store.node(cid).unwrap().layer().unwrap_or(usize::MAX))
            .min()
            .expect("the walk visited a node");
        assert!(
            deepest >= key_height(missing),
            "the walk went below the layer the path belongs on"
        );
    }

    /// A link the store does not hold is an error naming it, rather than a
    /// tree that quietly reports fewer records than it has.
    #[test]
    fn a_missing_block_is_named() {
        let (mut store, root) = tree(&["a", "b", "c", "d", "e", "f", "g", "h"]);
        let under = store.nodes_under(&root).unwrap();
        let child = under
            .iter()
            .find(|cid| *cid != &root)
            .expect("the tree has more than one node")
            .clone();
        store.nodes.remove(&child);
        store.blocks.remove(&child);
        assert_eq!(
            store.leaves(&root),
            Err(StoreError::Missing(Box::new(child)))
        );
    }

    /// One record added: the record shows up as a create, and the nodes that
    /// had to be rewritten show up as created.
    #[test]
    fn adding_one_record_is_one_op() {
        let (mut store, before) = tree(&["a", "b", "c", "d"]);
        let (after_store, after) = tree(&["a", "b", "c", "d", "e"]);
        for (cid, block) in after_store.blocks {
            store.insert(&cid, block).unwrap();
        }
        let diff = diff(&store, &before, &after).unwrap();
        assert_eq!(
            diff.ops,
            vec![RecordOp {
                path: "e".to_owned(),
                old_value: None,
                new_value: Some(record("e")),
            }]
        );
        assert!(!diff.created_nodes.is_empty());
        assert!(diff.created_nodes.contains(&after));
        assert!(diff.deleted_nodes.contains(&before));
        // The proof of a created record is a walk through the new tree, and
        // every node on it is one the new tree holds.
        for cid in &diff.proof_nodes {
            assert!(after_store_holds(&store, &after, cid));
        }
    }

    /// Whether a node is in the tree under `root`.
    fn after_store_holds(store: &Store, root: &Cid, cid: &Cid) -> bool {
        store.nodes_under(root).unwrap().contains(cid)
    }

    /// A tree diffed against itself changed nothing, which is the case a
    /// firehose consumer sees most often and the one an off-by-one in the
    /// set arithmetic breaks first.
    #[test]
    fn a_tree_against_itself_is_empty() {
        let (store, root) = tree(&["a", "b", "c", "d"]);
        let diff = diff(&store, &root, &root).unwrap();
        assert!(diff.created_nodes.is_empty());
        assert!(diff.deleted_nodes.is_empty());
        assert!(diff.ops.is_empty());
        assert!(diff.proof_nodes.is_empty());
    }

    /// A record whose value changed is an update, with both sides named.
    #[test]
    fn a_changed_value_is_an_update() {
        let mut leaves: BTreeMap<String, Cid> = ["a", "b", "c"]
            .iter()
            .map(|path| ((*path).to_owned(), record(path)))
            .collect();
        let before_tree = Tree::build(&leaves);
        leaves.insert("b".to_owned(), record("something else"));
        let after_tree = Tree::build(&leaves);
        let mut store = Store::default();
        for tree in [&before_tree, &after_tree] {
            for (cid, block) in tree.blocks() {
                store.insert(cid, block.clone()).unwrap();
            }
        }
        let diff = diff(&store, before_tree.root(), after_tree.root()).unwrap();
        assert_eq!(
            diff.ops,
            vec![RecordOp {
                path: "b".to_owned(),
                old_value: Some(record("b")),
                new_value: Some(record("something else")),
            }]
        );
    }

    /// A deep chain of blocks is walked without ending the process.
    ///
    /// An empty intermediate is a handful of bytes and is walked clean —
    /// [`Store::check_layer`] passes it whatever layer its parent claims —
    /// so a chain of them is the cheapest deep tree a sender can build. A
    /// walk that took one stack frame per link would reach the guard page
    /// here, and a stack overflow is not something the caller of a read can
    /// catch.
    #[test]
    fn a_deep_chain_of_empty_nodes_is_walked_rather_than_overflowing() {
        const LINKS: usize = 100_000;

        let mut store = Store::default();
        // The bottom of the chain: an empty node with nothing under it.
        let mut below = store.put(Node::from_parts(&Parts {
            keys: Vec::new(),
            values: Vec::new(),
            subtrees: vec![None],
        }));
        for _ in 0..LINKS {
            below = store.put(Node::from_parts(&Parts {
                keys: Vec::new(),
                values: Vec::new(),
                subtrees: vec![Some(below)],
            }));
        }
        // A root has to carry a key for its layer to be readable at all, so
        // the chain hangs off one entry and that entry is the only leaf.
        let path = "com.example.thing/3lqm0000";
        let root = store.put(Node::from_parts(&Parts {
            keys: vec![path.to_owned()],
            values: vec![record(path)],
            subtrees: vec![Some(below), None],
        }));

        assert_eq!(
            store.leaves(&root).unwrap(),
            BTreeMap::from([(path.to_owned(), record(path))])
        );
        assert_eq!(store.nodes_under(&root).unwrap().len(), LINKS + 2);
    }
}