haematite-cli 0.21.0

//! Removes companion vectors of fields no longer declared at the branch head.
//!
//! Forks inherit their parent's vectors. Once the parent rewrites an undeclared
//! field, the fork's copies are unreachable and need removal.
//!
//! Fields are found by skipping, never by reading every vector: the first
//! vector key at or after a point gives a field, and a field the branch
//! declares is passed over whole, so the reads are in proportion to the
//! fields present, plus one page.
//!
//! A field with a position is never pruned, declared or not: a fork's
//! inheriting removes the position of an undeclared field first (F1, step
//! two), and until it has, the field is left alone, so the two never race.
//!
//! Durable commits: at most one a call, removing at most the page setting's
//! vector keys of one field; none when nothing is found.
//!
//! EMB-E2: a field's windows (`vw\0`) are pruned as its vectors are, by the
//! same rule, once no vector key is left to prune.

use std::num::NonZeroUsize;

use super::embedded::{self, Field};
use super::{BranchRefRecord, Captured, Failure, Storage, branch_name, identifier, reads, version};

/// What one call removed, and whether a later call may find more.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Pruned {
    pub removed: usize,
    pub more: bool,
}

/// Where each family of vector keys starts and ends: one vector a record,
/// then EMB-E2's windows.
const FAMILIES: [(&[u8], &[u8]); 2] = [(b"v\0", b"v\x01"), (b"vw\0", b"vw\x01")];

impl Storage {
    /// Removes, in one commit, up to `page_records` vector keys of the first
    /// field in `branch`'s companion that the branch does not declare
    /// embedded at its head and that has no position (F7).
    ///
    /// # Errors
    /// `branch_missing`; `embed_vector_key_corrupt` for a vector key that is
    /// not UTF-8 or names no field; whatever a read or the commit answers.
    pub fn embed_prune(
        &mut self,
        branch: &str,
        page_records: NonZeroUsize,
    ) -> Result<Pruned, Failure> {
        identifier(branch)?;
        let record = self
            .refs
            .get(&branch_name(branch))
            .ok_or_else(|| Failure::new("branch_missing", branch))?;
        let head = Captured {
            branch: branch.to_owned(),
            version: version(record),
            shards: record.shards.clone(),
        };
        let name = embedded::companion_name(branch);
        let nothing = Pruned {
            removed: 0,
            more: false,
        };
        let Some(companion) = self.refs.get(&name) else {
            return Ok(nothing);
        };
        let positions = self.companion_positions(companion)?;
        let kept = |collection: &str, field: &str| -> Result<bool, Failure> {
            Ok(
                positions.contains_key(&(collection.to_owned(), field.to_owned()))
                    || self.declares(&head, collection, field)?,
            )
        };
        let mut found = None;
        for (index, family) in FAMILIES.into_iter().enumerate() {
            let scan = Scan {
                name: &name,
                companion,
                family,
            };
            if let Some(FamilyPage { removals, more }) =
                family_page(self, &scan, page_records, &kept)?
            {
                // Windows are looked at once no vector is left to prune.
                let (from, to) = FAMILIES[1];
                let windows = index == 0 && first_key(self, companion, from, to)?.is_some();
                found = Some((removals, more || windows));
                break;
            }
        }
        let Some((removals, more)) = found else {
            return Ok(nothing);
        };
        let writes: Vec<(Vec<u8>, Option<Vec<u8>>)> =
            removals.into_iter().map(|key| (key, None)).collect();
        let removed = writes.len();
        if removed == 0 {
            return Ok(nothing);
        }
        self.commit_namespace(&name, &writes)?;
        Ok(Pruned { removed, more })
    }

    /// Whether `head` declares `collection.field` embedded.
    fn declares(&self, head: &Captured, collection: &str, field: &str) -> Result<bool, Failure> {
        let key = super::embed_job::field_key(collection, field);
        match reads::get_bytes(&self.reader, head, &key)? {
            Some(bytes) => Ok(matches!(embedded::read_field(&bytes)?, Field::Embedded(_))),
            None => Ok(false),
        }
    }
}

/// One family of a companion's vector keys, as [`FAMILIES`] names it.
struct Scan<'a> {
    name: &'a str,
    companion: &'a BranchRefRecord,
    family: (&'static [u8], &'static [u8]),
}

/// One page of a family's keys to remove.
struct FamilyPage {
    removals: Vec<Vec<u8>>,
    /// Whether any key is left in the family after this page.
    more: bool,
}

/// The first page of keys of the first field in `scan`'s family that `kept`
/// does not keep, and whether any is left in the family after them; `None`
/// when it keeps every field there.
fn family_page(
    storage: &Storage,
    scan: &Scan<'_>,
    page_records: NonZeroUsize,
    kept: &impl Fn(&str, &str) -> Result<bool, Failure>,
) -> Result<Option<FamilyPage>, Failure> {
    let (family, family_end) = scan.family;
    let mut from = family.to_vec();
    loop {
        let Some(key) = first_key(storage, scan.companion, &from, family_end)? else {
            return Ok(None);
        };
        let (collection, field) = vector_field(scan.name, family, &key)?;
        let mut prefix = family.to_vec();
        prefix.extend_from_slice(format!("{collection}\0{field}\0").as_bytes());
        let end = past(&prefix);
        if kept(&collection, &field)? {
            from = end;
            continue;
        }
        let (removals, more) = field_page(storage, scan.companion, &prefix, &end, page_records)?;
        let later = first_key(storage, scan.companion, &end, family_end)?.is_some();
        return Ok(Some(FamilyPage {
            removals,
            more: more || later,
        }));
    }
}

/// The smallest key in `[from, to)` across the companion's shards: one read
/// of one key on each.
fn first_key(
    storage: &Storage,
    companion: &BranchRefRecord,
    from: &[u8],
    to: &[u8],
) -> Result<Option<Vec<u8>>, Failure> {
    let mut first: Option<Vec<u8>> = None;
    for shard in &companion.shards {
        let page = storage
            .reader
            .range_page_at(shard.shard_id, shard.head, from, to, NonZeroUsize::MIN)
            .map_err(Failure::observer)?;
        if let Some((key, _value)) = page.entries.into_iter().next()
            && first.as_ref().is_none_or(|known| key < *known)
        {
            first = Some(key);
        }
    }
    Ok(first)
}

/// The first `limit` keys in `[prefix, end)` across the companion's shards,
/// and whether any is left after them.
fn field_page(
    storage: &Storage,
    companion: &BranchRefRecord,
    prefix: &[u8],
    end: &[u8],
    limit: NonZeroUsize,
) -> Result<(Vec<Vec<u8>>, bool), Failure> {
    let mut keys = Vec::new();
    let mut more = false;
    for shard in &companion.shards {
        let page = storage
            .reader
            .range_page_at(shard.shard_id, shard.head, prefix, end, limit)
            .map_err(Failure::observer)?;
        more |= page.next_from.is_some();
        keys.extend(page.entries.into_iter().map(|(key, _value)| key));
    }
    keys.sort();
    if keys.len() > limit.get() {
        keys.truncate(limit.get());
        more = true;
    }
    Ok((keys, more))
}

/// The (collection, field) a key of `family` names, refused by name, never
/// guessed, when they are not UTF-8 or it names no field and record. A
/// window key's record part ends in its index, which need not be UTF-8. The
/// key's bytes are never printed.
fn vector_field(companion: &str, family: &[u8], key: &[u8]) -> Result<(String, String), Failure> {
    let corrupt = || {
        Failure::new(
            "embed_vector_key_corrupt",
            format!("{companion} holds a vector key that names no collection, field and record"),
        )
    };
    let named = key.strip_prefix(family).ok_or_else(corrupt)?;
    let mut parts = named.splitn(3, |byte| *byte == 0);
    let (Some(collection), Some(field), Some(_record)) = (parts.next(), parts.next(), parts.next())
    else {
        return Err(corrupt());
    };
    let text = |part: &[u8]| {
        std::str::from_utf8(part)
            .map(str::to_owned)
            .map_err(|_error| corrupt())
    };
    Ok((text(collection)?, text(field)?))
}

/// The first key after every key that starts with `prefix`, whose last byte
/// is the NUL separator.
fn past(prefix: &[u8]) -> Vec<u8> {
    let mut end = prefix.to_vec();
    if let Some(last) = end.last_mut() {
        *last = 1;
    }
    end
}

#[cfg(test)]
#[path = "embed_prune_tests.rs"]
mod tests;