didbot_data/mst/

edit.rs

//! Changing one key of a built tree without building it again.
//!
//! A node covers the keys between two keys taller than it. Changing one key
//! rewrites exactly the nodes whose range holds that key: one per layer, from
//! the root down. Below the key's own layer, an insert splits each of those
//! nodes in two around the key, and a remove joins the two nodes either side
//! of it. Every other node is left where it is, shared with the tree before.
//!
//! Nothing here is written into the tree until the edit is done. An edit
//! reads the nodes it needs, stores the ones it writes beside the tree, and
//! then applies both at once, so a node it writes and then replaces again
//! never reaches the tree at all.

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

use super::key_height;
use super::tree::{Entry, Node, Tree};
use crate::cid::Cid;
use crate::dag_cbor;

/// The blocks an edit added and took away.
///
/// For [`Tree::insert`] and [`Tree::remove`] these are tree nodes. A
/// repository commit reports its record and commit blocks the same way.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Rewritten {
    /// Blocks held now and not held before.
    pub created: BTreeSet<Cid>,
    /// Blocks held before and not held now.
    pub removed: BTreeSet<Cid>,
}

impl Rewritten {
    /// Folds a later edit into this one, so the two read as a single edit.
    ///
    /// A block the later edit creates that this one removed was there all
    /// along, and a block it removes that this one created never was.
    pub fn then(&mut self, later: Rewritten) {
        for cid in later.created {
            if !self.removed.remove(&cid) {
                self.created.insert(cid);
            }
        }
        for cid in later.removed {
            if !self.created.remove(&cid) {
                self.removed.insert(cid);
            }
        }
    }
}

impl Tree {
    /// Points `key` at `value`, adding the key if the tree does not hold it.
    ///
    /// The tree afterwards is the one [`Tree::build`] makes from the new leaf
    /// set, node for node. Setting a key to the value it already names
    /// changes nothing and answers an empty [`Rewritten`].
    pub fn insert(&mut self, key: &str, value: Cid) -> Rewritten {
        let height = key_height(key);
        let key = key.as_bytes();
        let mut edit = Edit::new(self);
        let (top, layer) = if height > self.layer {
            // Taller than every key held: the key becomes the root's only
            // entry, and the old tree is cut in two either side of it. Each
            // half is lifted to sit exactly one layer below the new root.
            let empty = self
                .nodes
                .get(&self.root)
                .is_some_and(|root| root.entries.is_empty() && root.left.is_none());
            let old = if empty {
                // The empty tree's root is replaced, not split.
                edit.open(&self.root);
                None
            } else {
                Some(self.root.clone())
            };
            let (low, high) = edit.split(old, self.layer, key);
            let low = edit.lift(low, self.layer, height - 1);
            let high = edit.lift(high, self.layer, height - 1);
            let top = Open {
                left: low,
                entries: vec![(key.to_vec(), value, high)],
            };
            (top, height)
        } else {
            let root = edit.open(&self.root);
            (
                edit.insert(root, self.layer, key, value, height),
                self.layer,
            )
        };
        let root = edit.store(top);
        let staged = edit.finish(root, layer);
        self.apply(staged)
    }

    /// Takes `key` out of the tree, if the tree holds it.
    ///
    /// The tree afterwards is the one [`Tree::build`] makes from the leaf set
    /// without the key. A key the tree does not hold changes nothing and
    /// answers an empty [`Rewritten`].
    pub fn remove(&mut self, key: &str) -> Rewritten {
        if self.path(key).value.is_none() {
            return Rewritten::default();
        }
        let height = key_height(key);
        let mut edit = Edit::new(self);
        let root = edit.open(&self.root);
        let mut top = edit.remove(root, self.layer, key.as_bytes(), height);
        // The root sits at the height of the tallest key left. A root with
        // no entries is a layer nothing holds up any more, so the node under
        // it takes its place, until one has entries or nothing is left.
        let mut layer = self.layer;
        while top.entries.is_empty() {
            let Some(below) = top.left.take() else {
                layer = 0;
                break;
            };
            top = edit.open(&below);
            layer = layer.saturating_sub(1);
        }
        let root = edit.store(top);
        let staged = edit.finish(root, layer);
        self.apply(staged)
    }

    /// Writes an edit's nodes into the tree and takes its removed ones out.
    fn apply(&mut self, staged: Staged) -> Rewritten {
        let Staged {
            root,
            layer,
            mut written,
            present,
        } = staged;
        let mut rewritten = Rewritten::default();
        for (cid, present) in present {
            if present {
                if self.nodes.contains_key(&cid) {
                    continue;
                }
                if let Some((node, block)) = written.remove(&cid) {
                    self.nodes.insert(cid.clone(), node);
                    self.blocks.insert(cid.clone(), block);
                    rewritten.created.insert(cid);
                }
            } else if self.nodes.remove(&cid).is_some() {
                self.blocks.remove(&cid);
                rewritten.removed.insert(cid);
            }
        }
        self.root = root;
        self.layer = layer;
        rewritten
    }
}

/// A node with its keys written out whole, which is how an edit works on it.
///
/// A stored node writes each key as a prefix of the one before it, which is
/// only meaningful while the entries stay where they are. An edit moves
/// entries between nodes, so it opens a node into this, works on it, and
/// closes it back into a [`Node`] that computes the prefixes afresh.
#[derive(Debug, Default)]
struct Open {
    /// The subtree before the first entry.
    left: Option<Cid>,
    /// Each entry's whole key, its value, and the subtree to its right.
    entries: Vec<(Vec<u8>, Cid, Option<Cid>)>,
}

impl Open {
    /// Opens a stored node.
    fn of(node: &Node) -> Self {
        let mut entries: Vec<(Vec<u8>, Cid, Option<Cid>)> = Vec::with_capacity(node.entries.len());
        for entry in &node.entries {
            let previous: &[u8] = entries.last().map_or(&[], |(key, _, _)| key.as_slice());
            let mut key = previous[..entry.prefix.min(previous.len())].to_vec();
            key.extend_from_slice(&entry.suffix);
            entries.push((key, entry.value.clone(), entry.tree.clone()));
        }
        Self {
            left: node.left.clone(),
            entries,
        }
    }

    /// Closes it back into a node, each key written against the one before.
    fn close(&self) -> Node {
        let mut previous: &[u8] = &[];
        let entries = self
            .entries
            .iter()
            .map(|(key, value, tree)| {
                let prefix = previous
                    .iter()
                    .zip(key)
                    .take_while(|(left, right)| left == right)
                    .count();
                previous = key;
                Entry {
                    prefix,
                    suffix: key[prefix..].to_vec(),
                    value: value.clone(),
                    tree: tree.clone(),
                }
            })
            .collect();
        Node {
            left: self.left.clone(),
            entries,
        }
    }

    /// Whether it holds nothing at all, which is no node.
    fn is_empty(&self) -> bool {
        self.entries.is_empty() && self.left.is_none()
    }

    /// How many entries sort before `key`.
    fn position(&self, key: &[u8]) -> usize {
        self.entries
            .partition_point(|(held, _, _)| held.as_slice() < key)
    }

    /// The subtree just before entry `at`: `left` for the first, and the
    /// previous entry's right-hand subtree for every other.
    fn slot(&self, at: usize) -> Option<Cid> {
        match at.checked_sub(1) {
            None => self.left.clone(),
            Some(previous) => self.entries[previous].2.clone(),
        }
    }

    /// Replaces the subtree just before entry `at`.
    fn set_slot(&mut self, at: usize, subtree: Option<Cid>) {
        match at.checked_sub(1) {
            None => self.left = subtree,
            Some(previous) => self.entries[previous].2 = subtree,
        }
    }
}

/// An edit in progress: the tree as it was, and what the edit has done.
struct Edit<'t> {
    tree: &'t Tree,
    /// Every node the edit has stored, by CID, with its block.
    written: BTreeMap<Cid, (Node, Vec<u8>)>,
    /// Whether each node the edit touched is in the tree once it is done.
    ///
    /// Opening a node takes it out and storing one puts it in, in the order
    /// they happen. A node taken out and then stored again identically is
    /// still there, which is what an edit that changes nothing amounts to.
    present: BTreeMap<Cid, bool>,
}

/// A finished edit, ready to be written into the tree.
struct Staged {
    root: Cid,
    layer: usize,
    written: BTreeMap<Cid, (Node, Vec<u8>)>,
    present: BTreeMap<Cid, bool>,
}

impl<'t> Edit<'t> {
    fn new(tree: &'t Tree) -> Self {
        Self {
            tree,
            written: BTreeMap::new(),
            present: BTreeMap::new(),
        }
    }

    /// Takes a node out of the tree to be rewritten.
    ///
    /// # Panics
    ///
    /// If the tree links to a node it does not hold. Every link in a tree
    /// names a node the same tree stored, and nothing outside this module
    /// can change the nodes a tree holds.
    fn open(&mut self, cid: &Cid) -> Open {
        let node = self
            .written
            .get(cid)
            .map(|(node, _)| node)
            .or_else(|| self.tree.nodes.get(cid))
            .expect("a tree links only to nodes it holds");
        let open = Open::of(node);
        self.present.insert(cid.clone(), false);
        open
    }

    /// Encodes a node and puts it in the tree.
    fn store(&mut self, open: Open) -> Cid {
        let node = open.close();
        let block = dag_cbor::encode(&node.to_value());
        let cid = Cid::of_dag_cbor(&block);
        self.present.insert(cid.clone(), true);
        self.written.insert(cid.clone(), (node, block));
        cid
    }

    /// [`Self::store`], or no subtree at all for a node that holds nothing.
    fn store_unless_empty(&mut self, open: Open) -> Option<Cid> {
        (!open.is_empty()).then(|| self.store(open))
    }

    /// Puts `key` into the node at `layer`, and answers the node.
    fn insert(
        &mut self,
        mut node: Open,
        layer: usize,
        key: &[u8],
        value: Cid,
        height: usize,
    ) -> Open {
        let at = node.position(key);
        if layer == height {
            if let Some(held) = node.entries.get_mut(at).filter(|(held, _, _)| held == key) {
                held.1 = value;
                return node;
            }
            // The subtree that ran across the key's place is cut in two: the
            // half before it stays in its slot, and the half after it hangs
            // off the new entry.
            let (low, high) = self.split(node.slot(at), layer.saturating_sub(1), key);
            node.set_slot(at, low);
            node.entries.insert(at, (key.to_vec(), value, high));
            return node;
        }
        // Below this layer. A missing subtree is a node that holds nothing
        // yet, one layer down, whether or not the key lands on that layer.
        let below = match node.slot(at) {
            Some(cid) => self.open(&cid),
            None => Open::default(),
        };
        let below = self.insert(below, layer - 1, key, value, height);
        let below = self.store(below);
        node.set_slot(at, Some(below));
        node
    }

    /// Cuts the subtree at `layer` into the keys before `key` and the keys
    /// after it. `key` is taller than `layer`, so no entry here is `key`.
    fn split(
        &mut self,
        subtree: Option<Cid>,
        layer: usize,
        key: &[u8],
    ) -> (Option<Cid>, Option<Cid>) {
        let Some(cid) = subtree else {
            return (None, None);
        };
        let node = self.open(&cid);
        let at = node.position(key);
        let (low, high) = self.split(node.slot(at), layer.saturating_sub(1), key);
        let mut after = node.entries;
        let before: Vec<_> = after.drain(..at).collect();
        let mut left = Open {
            left: node.left,
            entries: before,
        };
        left.set_slot(at, low);
        let right = Open {
            left: high,
            entries: after,
        };
        (
            self.store_unless_empty(left),
            self.store_unless_empty(right),
        )
    }

    /// Hangs a subtree at `from` under entry-less nodes until it sits at `to`.
    fn lift(&mut self, subtree: Option<Cid>, from: usize, to: usize) -> Option<Cid> {
        let mut subtree = subtree?;
        for _ in from..to {
            subtree = self.store(Open {
                left: Some(subtree),
                entries: Vec::new(),
            });
        }
        Some(subtree)
    }

    /// Takes `key` out of the node at `layer`, and answers the node.
    ///
    /// The caller has established that the tree holds `key`.
    fn remove(&mut self, mut node: Open, layer: usize, key: &[u8], height: usize) -> Open {
        let at = node.position(key);
        if layer == height {
            debug_assert!(
                node.entries.get(at).is_some_and(|(held, _, _)| held == key),
                "the key is on its own layer"
            );
            // The two subtrees either side of the key become one.
            let (_, _, after) = node.entries.remove(at);
            let joined = self.merge(node.slot(at), after, layer.saturating_sub(1));
            node.set_slot(at, joined);
            return node;
        }
        let Some(below) = node.slot(at) else {
            return node;
        };
        let below = self.open(&below);
        let below = self.remove(below, layer - 1, key, height);
        let below = self.store_unless_empty(below);
        node.set_slot(at, below);
        node
    }

    /// Joins two neighbouring subtrees at `layer`, every key of `low` sorting
    /// before every key of `high`.
    fn merge(&mut self, low: Option<Cid>, high: Option<Cid>, layer: usize) -> Option<Cid> {
        let (low, high) = match (low, high) {
            (None, other) | (other, None) => return other,
            (Some(low), Some(high)) => (low, high),
        };
        let mut low = self.open(&low);
        let high = self.open(&high);
        // The last subtree of one and the first of the other meet in the
        // middle, and are joined one layer down in turn.
        let seam = low.entries.len();
        let joined = self.merge(low.slot(seam), high.left, layer.saturating_sub(1));
        low.set_slot(seam, joined);
        low.entries.extend(high.entries);
        Some(self.store(low))
    }

    /// Ends the edit with `root` at the top.
    fn finish(self, root: Cid, layer: usize) -> Staged {
        Staged {
            root,
            layer,
            written: self.written,
            present: self.present,
        }
    }
}

#[cfg(test)]
mod tests {
    use rand::rngs::StdRng;
    use rand::{Rng, SeedableRng};

    use super::*;

    fn leaf(seed: &str) -> Cid {
        Cid::of_dag_cbor(seed.as_bytes())
    }

    /// Asserts that `tree` is the tree [`Tree::build`] makes from `leaves`,
    /// node for node, and that `rewritten` is the difference from `before`.
    fn assert_built(
        tree: &Tree,
        leaves: &BTreeMap<String, Cid>,
        before: &BTreeSet<Cid>,
        rewritten: &Rewritten,
        step: &str,
    ) {
        let built = Tree::build(leaves);
        assert_eq!(tree.root(), built.root(), "{step}: the roots differ");
        assert_eq!(tree.layer, built.layer, "{step}: the root layers differ");
        assert_eq!(tree.blocks(), built.blocks(), "{step}: the blocks differ");
        assert_eq!(tree.nodes(), built.nodes(), "{step}: the nodes differ");
        let after: BTreeSet<Cid> = built.blocks().keys().cloned().collect();
        let created: BTreeSet<Cid> = after.difference(before).cloned().collect();
        let removed: BTreeSet<Cid> = before.difference(&after).cloned().collect();
        assert_eq!(
            rewritten.created, created,
            "{step}: the created nodes differ"
        );
        assert_eq!(
            rewritten.removed, removed,
            "{step}: the removed nodes differ"
        );
    }

    /// Keys whose heights run from zero to five, so a sequence of edits grows
    /// the root, shrinks it, and splits and joins nodes on every layer.
    ///
    /// Random keys are almost all on layer zero: a quarter reach layer one
    /// and one in a thousand layer five. The tall ones are mined from a
    /// counter rather than waited for.
    fn pool() -> Vec<String> {
        let collections = [
            "app.bsky.feed.post",
            "app.bsky.feed.like",
            "com.example.thing",
        ];
        let mut keys: Vec<String> = Vec::new();
        let mut per_height = [0usize; 6];
        let wanted = [40, 20, 10, 6, 4, 2];
        let mut n = 0u64;
        while per_height
            .iter()
            .zip(&wanted)
            .any(|(have, want)| have < want)
        {
            let collection = collections[(n % 3) as usize];
            let key = format!("{collection}/3l{n:011}");
            let height = key_height(&key).min(5);
            if per_height[height] < wanted[height] {
                per_height[height] += 1;
                keys.push(key);
            }
            n += 1;
        }
        // Literal keys, which sort among the TIDs differently.
        keys.extend([
            "app.bsky.actor.profile/self".to_owned(),
            "bot.did.registration/self".to_owned(),
        ]);
        keys
    }

    /// Random inserts, overwrites and removes, compared against a rebuild
    /// after every one.
    ///
    /// `DIDBOT_MST_SEEDS` runs more seeds than the default.
    #[test]
    fn every_edit_lands_on_the_tree_a_rebuild_makes() {
        let seeds: u64 = std::env::var("DIDBOT_MST_SEEDS")
            .ok()
            .and_then(|seeds| seeds.parse().ok())
            .unwrap_or(24);
        let pool = pool();
        let mut layers = BTreeSet::new();
        for seed in 0..seeds {
            let mut rng = StdRng::seed_from_u64(seed);
            let mut leaves: BTreeMap<String, Cid> = BTreeMap::new();
            let mut tree = Tree::build(&leaves);
            // Fill towards the whole pool, then drain back to nothing, with
            // the bias flipped halfway so both ends are reached.
            for step in 0..600 {
                let filling = step < 300;
                let key = &pool[rng.random_range(0..pool.len())];
                let before: BTreeSet<Cid> = tree.blocks().keys().cloned().collect();
                let insert = if filling {
                    rng.random_bool(0.8)
                } else {
                    rng.random_bool(0.2)
                };
                let rewritten = if insert {
                    let value = leaf(&format!("{key}@{step}"));
                    leaves.insert(key.clone(), value.clone());
                    tree.insert(key, value)
                } else {
                    leaves.remove(key);
                    tree.remove(key)
                };
                layers.insert(tree.layer);
                let what = if insert { "insert" } else { "remove" };
                assert_built(
                    &tree,
                    &leaves,
                    &before,
                    &rewritten,
                    &format!("seed {seed} step {step}: {what} {key}"),
                );
            }
            // And everything that is left, down to the empty tree.
            for key in leaves.keys().cloned().collect::<Vec<_>>() {
                let before: BTreeSet<Cid> = tree.blocks().keys().cloned().collect();
                leaves.remove(&key);
                let rewritten = tree.remove(&key);
                assert_built(
                    &tree,
                    &leaves,
                    &before,
                    &rewritten,
                    &format!("seed {seed}: draining {key}"),
                );
            }
            assert_eq!(tree.root(), Tree::build(&BTreeMap::new()).root());
        }
        // The root moved through every layer the pool reaches, so growing
        // and shrinking it were both exercised rather than assumed.
        assert_eq!(
            layers,
            (0..=5).collect(),
            "the root layers the edits reached"
        );
    }

    /// Every key of the pool, inserted in every order a seed picks, gives one
    /// tree. The order the keys arrive in cannot reach the shape.
    #[test]
    fn insertion_order_does_not_reach_the_root() {
        let pool = pool();
        let leaves: BTreeMap<String, Cid> =
            pool.iter().map(|key| (key.clone(), leaf(key))).collect();
        let built = Tree::build(&leaves);
        for seed in 0..8 {
            let mut rng = StdRng::seed_from_u64(seed);
            let mut order = pool.clone();
            for i in (1..order.len()).rev() {
                order.swap(i, rng.random_range(0..=i));
            }
            let mut tree = Tree::build(&BTreeMap::new());
            for key in &order {
                tree.insert(key, leaf(key));
            }
            assert_eq!(tree.blocks(), built.blocks(), "seed {seed}");
            assert_eq!(tree.root(), built.root(), "seed {seed}");
        }
    }

    /// Setting a key to what it already names, or removing a key the tree
    /// does not hold, is no edit at all.
    #[test]
    fn an_edit_that_changes_nothing_rewrites_nothing() {
        let pool = pool();
        let leaves: BTreeMap<String, Cid> =
            pool.iter().map(|key| (key.clone(), leaf(key))).collect();
        let mut tree = Tree::build(&leaves);
        let root = tree.root().clone();
        assert_eq!(tree.insert(&pool[7], leaf(&pool[7])), Rewritten::default());
        assert_eq!(
            tree.remove("com.example.thing/absent"),
            Rewritten::default()
        );
        assert_eq!(tree.root(), &root);
    }
}