haematite_cli/serve/storage/

counts.rs

//! Counts move with semantic writes in one durable branch version.
//! Served document writes carry no expiry, so their presence at a version is fixed.
//! Unready branches keep ordinary writes and never invent count metadata.

use std::collections::{BTreeMap, btree_map::Entry};

use haematite::branch::{CommitDurability, CommitRequest, current_timestamp};
use haematite::db::{ReadOnlyDatabase, branches::BranchDatabase};
use haematite::tree::Hash;

use super::mutations::PreparedWrite;
use super::{
    BranchRefRecord, BranchRefStore, BranchRegistry, Captured, Failure, MergePlan, MergeState,
    Receipt, SERVICE, Storage, branch_name, collection_key, edges, identifier, open_branch, reads,
    receipt_key, revision_key, version,
};
use crate::serve::access::Access;
use crate::serve::protocol::{
    AuditVerdict, AuditedBucket, Change, CountBucket, CountReason, CountsAudit, CountsRebuild,
    KeptCount, ReceiptQuery, Version,
};

pub(super) struct View<'a> {
    reader: &'a ReadOnlyDatabase,
    at: &'a Captured,
    ready: bool,
}

impl<'a> View<'a> {
    pub(super) fn new(reader: &'a ReadOnlyDatabase, at: &'a Captured) -> Result<Self, Failure> {
        Ok(Self {
            reader,
            at,
            ready: is_ready(reader, at)?,
        })
    }

    pub(super) const fn unavailable(&self, allowed: bool) -> Option<CountReason> {
        if !self.ready {
            Some(CountReason::CountsNotReady)
        } else if !allowed {
            Some(CountReason::Withheld)
        } else {
            None
        }
    }

    pub(super) fn documents(
        &self,
        access: &Access,
        collection: &str,
    ) -> Result<KeptCount, Failure> {
        if let Some(reason) = self.unavailable(access.can_read_collection(collection)) {
            return Ok(KeptCount {
                value: None,
                reason: Some(reason),
            });
        }
        Ok(KeptCount {
            value: Some(required_documents(self.reader, self.at, collection)?),
            reason: None,
        })
    }

    pub(super) fn relationships(
        &self,
        access: &Access,
        relation: &str,
    ) -> Result<KeptCount, Failure> {
        if let Some(reason) = self.unavailable(access.can_read_store()) {
            return Ok(KeptCount {
                value: None,
                reason: Some(reason),
            });
        }
        let key = format!("k\0e\0{relation}");
        let value = read(self.reader, self.at, key.as_bytes())?
            .ok_or_else(|| Failure::new("counts_corrupt", "declared relation has no count"))?;
        Ok(KeptCount {
            value: Some(decode(&value)?),
            reason: None,
        })
    }
}

const READY_KEY: &[u8] = b"k\0ready";
const READY_VALUE: &[u8] = b"1";

#[derive(Default)]
struct BucketCount {
    counted: u64,
    kept: Option<u64>,
    present: bool,
}

impl CountBucket {
    fn key(&self) -> Vec<u8> {
        match self {
            Self::Documents { collection } => document_count_key(collection),
            Self::Relationships { relation } => format!("k\0e\0{relation}").into_bytes(),
            Self::Vectors { collection, field } => {
                format!("k\0v\0{collection}\0{field}").into_bytes()
            }
        }
    }
}

fn counted_buckets(
    reader: &ReadOnlyDatabase,
    at: &Captured,
) -> Result<BTreeMap<CountBucket, BucketCount>, Failure> {
    let mut counts = BTreeMap::new();
    for (prefix, parts, kind) in [
        (b"c\0".as_slice(), 1, 0),
        (b"n\0".as_slice(), 1, 1),
        (b"w\0".as_slice(), 2, 2),
        (b"d\0".as_slice(), 2, 3),
        (b"e\0".as_slice(), 5, 4),
        (b"v\0".as_slice(), 3, 5),
    ] {
        reads::scan_rows(reader, at, prefix, parts, |parts, bytes| {
            let bucket = match kind {
                0 | 3 => CountBucket::Documents {
                    collection: parts[0].to_owned(),
                },
                1 => {
                    super::relations::decode(parts[0], bytes)?;
                    CountBucket::Relationships {
                        relation: parts[0].to_owned(),
                    }
                }
                4 => CountBucket::Relationships {
                    relation: parts[2].to_owned(),
                },
                _ => {
                    if kind == 2
                        && matches!(
                            super::embedded::read_field(bytes)?,
                            super::embedded::Field::Embedded(_)
                        )
                    {
                        return Ok(());
                    }
                    CountBucket::Vectors {
                        collection: parts[0].to_owned(),
                        field: parts[1].to_owned(),
                    }
                }
            };
            let counter: &mut BucketCount = counts.entry(bucket).or_default();
            counter.present = true;
            if kind >= 3 {
                counter.counted = added_count(counter.counted, 1)?;
            }
            Ok(())
        })?;
    }
    Ok(counts)
}

fn count_rows(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    counts: &mut BTreeMap<CountBucket, BucketCount>,
    audit: bool,
) -> Result<(), Failure> {
    for (prefix, parts, kind) in [
        (b"k\0d\0".as_slice(), 1, 0),
        (b"k\0e\0".as_slice(), 1, 1),
        (b"k\0v\0".as_slice(), 2, 2),
    ] {
        reads::scan_rows(reader, at, prefix, parts, |parts, bytes| {
            let bucket = match kind {
                0 => CountBucket::Documents {
                    collection: parts[0].to_owned(),
                },
                1 => CountBucket::Relationships {
                    relation: parts[0].to_owned(),
                },
                _ => CountBucket::Vectors {
                    collection: parts[0].to_owned(),
                    field: parts[1].to_owned(),
                },
            };
            let counter = counts.entry(bucket).or_default();
            if audit && counter.kept.replace(decode(bytes)?).is_some() {
                return Err(Failure::new("counts_corrupt", "duplicate count bucket"));
            }
            Ok(())
        })?;
    }
    Ok(())
}

impl Storage {
    /// Counts visible content and installs all counters with one branch version.
    ///
    /// # Errors
    /// Refuses incomplete operations and scan errors before staging; commit errors
    /// preserve the common mutation outcome and receipt reconciliation rules.
    pub fn rebuild_counts(&mut self, request: &CountsRebuild) -> Result<Receipt, Failure> {
        identifier(&request.branch)?;
        identifier(&request.request_id)?;
        let name = branch_name(&request.branch);
        if self.pending_discards.contains(&name) {
            return Err(Failure::new(
                "counts_discard_pending",
                "branch discard is incomplete",
            ));
        }
        let record = self
            .refs
            .get(&name)
            .ok_or_else(|| Failure::new("counts_branch_missing", &request.branch))?
            .clone();
        let fingerprint =
            Hash::of(&serde_json::to_vec(&("counts.rebuild", request)).map_err(Failure::storage)?)
                .to_string();
        if let Some(receipt) = self.receipt(&ReceiptQuery {
            branch: request.branch.clone(),
            request_id: request.request_id.clone(),
        })? {
            return if receipt.fingerprint == fingerprint {
                Ok(receipt)
            } else {
                Err(Failure::new(
                    "request_conflict",
                    "request identity has different content",
                ))
            };
        }
        let at = Captured {
            branch: request.branch.clone(),
            version: version(&record),
            shards: record.shards.clone(),
        };
        self.refuse_count_merge(&at).map_err(|error| {
            if error.code == "counts_merge_pending" {
                error
            } else {
                count_scan_failure(&error)
            }
        })?;
        let next = Version {
            generation: record.created,
            sequence: record
                .seq
                .checked_add(1)
                .ok_or_else(|| Failure::new("sequence_exhausted", "branch sequence exhausted"))?,
        };
        let counts = rebuild_buckets(&self.reader, &at)?;
        #[cfg(test)]
        rebuild_probe::pause()?;
        let receipt = Receipt {
            retention: None,
            versions_released: None,
            branch: request.branch.clone(),
            request_id: request.request_id.clone(),
            fingerprint,
            version: next,
        };
        self.commit_rebuilt_counts(&name, &record, &at, receipt, counts)
    }

    fn commit_rebuilt_counts(
        &mut self,
        name: &str,
        record: &BranchRefRecord,
        at: &Captured,
        receipt: Receipt,
        counts: BTreeMap<CountBucket, BucketCount>,
    ) -> Result<Receipt, Failure> {
        self.rows_before_advance(&at.branch, record)?;
        let branch =
            open_branch(name, &self.refs, &self.registry).map_err(Failure::branch_commit)?;
        for (bucket, counter) in counts {
            let key = bucket.key();
            let shard = self.writer.shard_for(&key);
            if counter.present {
                branch.put(shard, key, counter.counted.to_string().into_bytes())
            } else {
                branch.delete(shard, key)
            }
            .map_err(Failure::branch)?;
        }
        branch
            .put(self.writer.shard_for(READY_KEY), READY_KEY, READY_VALUE)
            .map_err(Failure::branch)?;
        let key = receipt_key(
            &receipt.branch,
            receipt.version.generation,
            &receipt.request_id,
        );
        branch
            .put(
                self.writer.shard_for(&key),
                key,
                serde_json::to_vec(&receipt).map_err(Failure::storage)?,
            )
            .map_err(Failure::branch)?;
        #[cfg(test)]
        rebuild_probe::before_commit();
        self.writer
            .commit(
                &branch,
                &self.registry,
                CommitRequest {
                    durability: CommitDurability::Durable {
                        refs: &mut self.refs,
                    },
                    extra_parents: &[],
                    timestamp: current_timestamp(),
                    note: None,
                },
            )
            .map_err(|error| {
                Failure::new(
                    if self
                        .refs
                        .get(name)
                        .is_some_and(|now| version(now) != at.version)
                    {
                        "durability_unknown"
                    } else {
                        "write_refused"
                    },
                    error.to_string(),
                )
            })?;
        self.publish_record(name)?;
        Ok(receipt)
    }

    fn refuse_count_merge(&self, at: &Captured) -> Result<(), Failure> {
        if reads::get_bytes(&self.reader, at, b"b\0merge-candidate-progress")?.is_some() {
            return Err(Failure::new(
                "counts_merge_pending",
                "branch holds a merge candidate",
            ));
        }
        let service = self.merge_service_selection()?;
        reads::scan_rows(&self.reader, &service, b"j\0merge\0", 1, |parts, bytes| {
            let plan: MergePlan = serde_json::from_slice(bytes).map_err(Failure::storage)?;
            identifier(&plan.id)?;
            for pin in [&plan.source, &plan.target, &plan.origin] {
                identifier(&pin.branch)?;
            }
            if Hash::of(plan.id.as_bytes()).to_string() != parts[0]
                || plan.origin.branch != plan.target.branch
            {
                return Err(Failure::new(
                    "merge_corrupt",
                    "merge plan key differs from identity",
                ));
            }
            if plan.state == MergeState::Open
                && [&plan.source, &plan.target, &plan.origin]
                    .iter()
                    .any(|pin| pin.branch == at.branch)
            {
                return Err(Failure::new(
                    "counts_merge_pending",
                    "branch has an open merge",
                ));
            }
            Ok(())
        })
    }

    /// Read-only comparison; withheld callers receive no bucket identities or numbers.
    ///
    /// # Errors
    /// Propagates corrupt metadata, malformed keys and content read failures.
    pub fn audit_counts(
        access: &Access,
        reader: &ReadOnlyDatabase,
        at: &Captured,
    ) -> Result<CountsAudit, Failure> {
        if !access.can_read_store() {
            return Ok(CountsAudit::CannotTell {
                reason: CountReason::Withheld,
            });
        }
        if !is_ready(reader, at)? {
            return Ok(CountsAudit::CannotTell {
                reason: CountReason::CountsNotReady,
            });
        }
        let mut counts = counted_buckets(reader, at)?;
        count_rows(reader, at, &mut counts, true)?;
        Ok(CountsAudit::Compared {
            buckets: counts
                .into_iter()
                .map(|(bucket, count)| {
                    let verdict = if count.kept == Some(count.counted) {
                        AuditVerdict::Agree
                    } else {
                        AuditVerdict::Differ {
                            kept: count.kept,
                            counted: count.counted,
                        }
                    };
                    AuditedBucket { bucket, verdict }
                })
                .collect(),
        })
    }

    #[cfg(test)]
    pub(in crate::serve) fn arm_count_rebuild_pause(
        paused: std::sync::mpsc::Sender<()>,
        release: std::sync::mpsc::Receiver<()>,
    ) {
        rebuild_probe::arm(paused, release);
    }
}

fn rebuild_buckets(
    reader: &ReadOnlyDatabase,
    at: &Captured,
) -> Result<BTreeMap<CountBucket, BucketCount>, Failure> {
    let mut counts = counted_buckets(reader, at).map_err(|error| count_scan_failure(&error))?;
    count_rows(reader, at, &mut counts, false).map_err(|error| count_scan_failure(&error))?;
    Ok(counts)
}

fn count_scan_failure(error: &Failure) -> Failure {
    Failure::new(
        "counts_scan_refused",
        format!("{}: {}", error.code, error.message),
    )
}

#[cfg(test)]
pub(super) mod rebuild_probe {
    use super::Failure;
    use std::cell::{Cell, RefCell};
    use std::sync::mpsc::{Receiver, Sender};
    thread_local! {
        static CRASH: Cell<bool> = const { Cell::new(false) };
        static PAUSE: RefCell<Option<(Sender<()>, Receiver<()>)>> = const { RefCell::new(None) };
    }
    pub fn crash_next() {
        CRASH.with(|slot| slot.set(true));
    }
    pub fn arm(paused: Sender<()>, release: Receiver<()>) {
        PAUSE.with(|slot| *slot.borrow_mut() = Some((paused, release)));
    }
    pub fn pause() -> Result<(), Failure> {
        if let Some((paused, release)) = PAUSE.with(|slot| slot.borrow_mut().take()) {
            paused.send(()).map_err(Failure::storage)?;
            release.recv().map_err(Failure::storage)?;
        }
        Ok(())
    }
    pub fn before_commit() {
        assert!(
            !CRASH.with(|slot| slot.replace(false)),
            "injected rebuild interruption"
        );
    }
}

pub(in crate::serve::storage) fn initialize_new_store(
    writer: &mut BranchDatabase,
    refs: &mut BranchRefStore,
    registry: &BranchRegistry,
) -> Result<(), Failure> {
    let service = open_branch(SERVICE, refs, registry).map_err(Failure::branch_commit)?;
    service
        .put(writer.shard_for(READY_KEY), READY_KEY, READY_VALUE)
        .map_err(Failure::branch)?;
    writer
        .commit(
            &service,
            registry,
            CommitRequest {
                durability: CommitDurability::Durable { refs },
                extra_parents: &[],
                timestamp: current_timestamp(),
                note: None,
            },
        )
        .map_err(Failure::branch_commit)?;
    Ok(())
}

fn document_count_key(collection: &str) -> Vec<u8> {
    format!("k\0d\0{collection}").into_bytes()
}

fn read(reader: &ReadOnlyDatabase, at: &Captured, key: &[u8]) -> Result<Option<Vec<u8>>, Failure> {
    #[cfg(test)]
    probe::read(key);
    reads::get_bytes(reader, at, key)
}

fn is_ready(reader: &ReadOnlyDatabase, at: &Captured) -> Result<bool, Failure> {
    match read(reader, at, READY_KEY)? {
        None => Ok(false),
        Some(value) if value.as_slice() == READY_VALUE => Ok(true),
        Some(_) => Err(Failure::new(
            "counts_corrupt",
            "count readiness format is invalid",
        )),
    }
}

fn decode(value: &[u8]) -> Result<u64, Failure> {
    if value.is_empty()
        || value.len() > 20
        || (value.len() > 1 && value.first() == Some(&b'0'))
        || !value.iter().all(u8::is_ascii_digit)
    {
        return Err(Failure::new(
            "counts_corrupt",
            "count is not a canonical decimal u64",
        ));
    }
    std::str::from_utf8(value)
        .map_err(|error| Failure::new("counts_corrupt", error.to_string()))?
        .parse()
        .map_err(|error: std::num::ParseIntError| Failure::new("counts_corrupt", error.to_string()))
}

fn required_documents(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    collection: &str,
) -> Result<u64, Failure> {
    let value = read(reader, at, &document_count_key(collection))?
        .ok_or_else(|| Failure::new("counts_corrupt", "ready collection has no document count"))?;
    decode(&value)
}

fn document_exists(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    collection: &str,
    id: &str,
) -> Result<bool, Failure> {
    let key = revision_key(collection, id);
    Ok(read(reader, at, &key)?.is_some())
}

fn increased(count: u64, added: u64) -> Result<Option<u64>, Failure> {
    if added == 0 {
        return Ok(None);
    }
    added_count(count, added).map(Some)
}

fn added_count(count: u64, added: u64) -> Result<u64, Failure> {
    count
        .checked_add(added)
        .ok_or_else(|| Failure::new("count_overflow", "object count exceeds u64"))
}

fn removed_count(count: u64, removed: u64) -> Result<u64, Failure> {
    count
        .checked_sub(removed)
        .ok_or_else(|| Failure::new("count_underflow", "removed objects exceed count"))
}

pub(super) fn append(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    change: &Change,
    writes: &mut Vec<PreparedWrite>,
) -> Result<(), Failure> {
    let collection = match change {
        Change::Collection(write) => Some(&write.collection),
        Change::Put(write) => Some(&write.collection),
        Change::AppendPrune(write) => Some(&write.collection),
        Change::Delete(write) => Some(&write.collection),
        Change::DeleteMany(delete) => Some(&delete.collection),
        Change::PutMany(import) => Some(&import.collection),
        Change::EdgePut(_)
        | Change::EdgeDelete(_)
        | Change::EdgePutMany(_)
        | Change::VectorPut(_)
        | Change::VectorPutMany(_)
        | Change::VectorDelete(_)
        | Change::RelationPut(_) => None,
        _ => return Ok(()),
    };
    if !is_ready(reader, at)? {
        return Ok(());
    }
    if let Some(collection) = collection {
        append_documents(reader, at, change, collection, writes)?;
    }
    let mut counters = declared_counters(reader, at, writes)?;
    append_objects(reader, at, writes, &mut counters)?;
    for (key, counter) in counters {
        if counter.write {
            writes.push((key, Some(counter.value.to_string().into_bytes())));
        }
    }
    #[cfg(test)]
    probe::write_count(writes);
    Ok(())
}

fn append_documents(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    change: &Change,
    collection: &str,
    writes: &mut Vec<PreparedWrite>,
) -> Result<(), Failure> {
    let count = match change {
        Change::Collection(_) => {
            if read(reader, at, &collection_key(collection))?.is_some() {
                required_documents(reader, at, collection)?;
                None
            } else {
                if read(reader, at, &document_count_key(collection))?.is_some() {
                    return Err(Failure::new(
                        "counts_corrupt",
                        "document count has no collection",
                    ));
                }
                Some(0)
            }
        }
        Change::Put(write) => {
            let present = document_exists(reader, at, collection, &write.document_id)?;
            increased(
                required_documents(reader, at, collection)?,
                u64::from(!present),
            )?
        }
        Change::Delete(write) => {
            if document_exists(reader, at, collection, &write.document_id)? {
                Some(removed_count(
                    required_documents(reader, at, collection)?,
                    1,
                )?)
            } else {
                None
            }
        }
        Change::DeleteMany(delete) => Some(removed_count(
            required_documents(reader, at, collection)?,
            u64::try_from(delete.document_ids.len()).map_err(Failure::storage)?,
        )?),
        Change::AppendPrune(_) => {
            let prefix = revision_key(collection, "");
            let mut added = 0_u64;
            let mut removed = 0_u64;
            for (key, value) in writes.iter() {
                let Some(id) = key.strip_prefix(prefix.as_slice()) else {
                    continue;
                };
                let id = std::str::from_utf8(id).map_err(Failure::storage)?;
                let present = document_exists(reader, at, collection, id)?;
                if value.is_some() && !present {
                    added = added_count(added, 1)?;
                }
                if value.is_none() && present {
                    removed = added_count(removed, 1)?;
                }
            }
            Some(added_count(
                removed_count(required_documents(reader, at, collection)?, removed)?,
                added,
            )?)
        }
        Change::PutMany(import) => {
            let mut added = 0_u64;
            for document in &import.documents {
                if !document_exists(reader, at, collection, &document.document_id)? {
                    added = added.checked_add(1).ok_or_else(|| {
                        Failure::new("count_overflow", "import delta exceeds u64")
                    })?;
                }
            }
            increased(required_documents(reader, at, collection)?, added)?
        }
        _ => return Ok(()),
    };
    if let Some(count) = count {
        writes.push((
            document_count_key(collection),
            Some(count.to_string().into_bytes()),
        ));
    }
    Ok(())
}

#[derive(Default)]
struct Delta {
    added: u64,
    removed: u64,
    existing: bool,
}

struct Counter {
    value: u64,
    write: bool,
    initialized: bool,
}

pub(super) struct CandidateCounts {
    ready: bool,
    changes: BTreeMap<Vec<u8>, (bool, bool)>,
}

impl CandidateCounts {
    pub(super) fn new(reader: &ReadOnlyDatabase, at: &Captured) -> Result<Self, Failure> {
        Ok(Self {
            ready: is_ready(reader, at)?,
            changes: BTreeMap::new(),
        })
    }

    pub(super) fn record(
        &mut self,
        reader: &ReadOnlyDatabase,
        at: &Captured,
        key: &[u8],
        presence_key: &[u8],
        after: bool,
    ) -> Result<(), Failure> {
        if !self.ready {
            return Ok(());
        }
        if let Some(change) = self.changes.get_mut(key) {
            change.1 = after;
        } else {
            let before = read(reader, at, presence_key)?.is_some();
            self.changes.insert(key.to_vec(), (before, after));
        }
        Ok(())
    }

    pub(super) fn append(
        self,
        reader: &ReadOnlyDatabase,
        at: &Captured,
        rows: &mut BTreeMap<Vec<u8>, Option<Vec<u8>>>,
    ) -> Result<(), Failure> {
        if !self.ready {
            return Ok(());
        }
        let mut counters = BTreeMap::new();
        for (key, value) in &*rows {
            let Some(value) = value else { continue };
            if key.starts_with(b"w\0")
                && matches!(
                    super::embedded::read_field(value)?,
                    super::embedded::Field::Embedded(_)
                )
            {
                continue;
            }
            let count_key = if let Some(name) = key.strip_prefix(b"c\0") {
                let collection = std::str::from_utf8(name).map_err(Failure::storage)?;
                identifier(collection)?;
                Some(document_count_key(collection))
            } else {
                declaration_count_key(key)?
            };
            let Some(count_key) = count_key else { continue };
            let stored = read(reader, at, &count_key)?;
            let count = match &stored {
                Some(bytes) => decode(bytes)?,
                None if read(reader, at, key)?.is_none() => 0,
                None => return Err(Failure::new("counts_corrupt", "declared kind has no count")),
            };
            counters.insert(
                count_key,
                Counter {
                    value: count,
                    write: stored.is_none(),
                    initialized: stored.is_none(),
                },
            );
        }
        let deltas = candidate_deltas(self.changes)?;
        for (key, delta) in deltas {
            let counter = match counters.entry(key) {
                Entry::Occupied(entry) => entry.into_mut(),
                Entry::Vacant(entry) => {
                    let value = match read(reader, at, entry.key())? {
                        Some(bytes) => decode(&bytes)?,
                        None if !delta.existing => {
                            candidate_missing_count(reader, at, entry.key())?
                        }
                        None => {
                            return Err(Failure::new(
                                "counts_corrupt",
                                "present object has no count",
                            ));
                        }
                    };
                    entry.insert(Counter {
                        value,
                        write: false,
                        initialized: false,
                    })
                }
            };
            if counter.initialized && delta.existing {
                return Err(Failure::new(
                    "counts_corrupt",
                    "present object has no previous count",
                ));
            }
            counter.value = added_count(removed_count(counter.value, delta.removed)?, delta.added)?;
            counter.write |= delta.added != 0 || delta.removed != 0;
        }
        for (key, counter) in counters {
            if counter.write {
                rows.insert(key, Some(counter.value.to_string().into_bytes()));
            }
        }
        Ok(())
    }
}

fn candidate_deltas(
    changes: BTreeMap<Vec<u8>, (bool, bool)>,
) -> Result<BTreeMap<Vec<u8>, Delta>, Failure> {
    let mut deltas: BTreeMap<Vec<u8>, Delta> = BTreeMap::new();
    for (key, (before, after)) in changes {
        let count_key = if let Some(suffix) = key.strip_prefix(b"d\0") {
            let text = std::str::from_utf8(suffix).map_err(Failure::storage)?;
            let (collection, document) = text
                .split_once('\0')
                .ok_or_else(|| Failure::new("corrupt_key", "document owner has no identifier"))?;
            identifier(collection)?;
            identifier(document)?;
            document_count_key(collection)
        } else {
            object_count_key(&key)?
                .ok_or_else(|| Failure::new("counts_corrupt", "candidate object has no bucket"))?
        };
        let delta = deltas.entry(count_key).or_default();
        delta.existing |= before;
        if after && !before {
            delta.added = added_count(delta.added, 1)?;
        }
        if before && !after {
            delta.removed = added_count(delta.removed, 1)?;
        }
    }
    Ok(deltas)
}

fn candidate_missing_count(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    key: &[u8],
) -> Result<u64, Failure> {
    if let Some(collection) = key.strip_prefix(b"k\0d\0") {
        let collection = std::str::from_utf8(collection).map_err(Failure::storage)?;
        identifier(collection)?;
        if read(reader, at, &collection_key(collection))?.is_some() {
            return Err(Failure::new(
                "counts_corrupt",
                "declared collection has no count",
            ));
        }
        Ok(0)
    } else {
        unregistered_count(reader, at, key)
    }
}

fn declaration_count_key(key: &[u8]) -> Result<Option<Vec<u8>>, Failure> {
    if let Some(name) = key.strip_prefix(b"n\0") {
        let name = std::str::from_utf8(name).map_err(Failure::storage)?;
        identifier(name)?;
        return Ok(Some(format!("k\0e\0{name}").into_bytes()));
    }
    if !key.starts_with(b"w\0") {
        return Ok(None);
    }
    let text = std::str::from_utf8(key).map_err(Failure::storage)?;
    let parts: Vec<_> = text.split('\0').collect();
    let ["w", collection, field] = parts.as_slice() else {
        return Err(invalid_vector_key());
    };
    identifier(collection)?;
    identifier(field)?;
    Ok(Some(format!("k\0v\0{collection}\0{field}").into_bytes()))
}

fn declared_counters(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    writes: &[PreparedWrite],
) -> Result<BTreeMap<Vec<u8>, Counter>, Failure> {
    let mut counters = BTreeMap::new();
    for (key, value) in writes {
        if value.is_none() {
            continue;
        }
        let Some(count_key) = declaration_count_key(key)? else {
            continue;
        };
        if counters.contains_key(&count_key) {
            continue;
        }
        let declared = read(reader, at, key)?.is_some();
        let counter = match read(reader, at, &count_key)? {
            Some(value) => Counter {
                value: decode(&value)?,
                write: false,
                initialized: false,
            },
            None if !declared => Counter {
                value: 0,
                write: true,
                initialized: true,
            },
            None => return Err(Failure::new("counts_corrupt", "declared kind has no count")),
        };
        counters.insert(count_key, counter);
    }
    Ok(counters)
}

fn object_count_key(key: &[u8]) -> Result<Option<Vec<u8>>, Failure> {
    if key.starts_with(b"e\0") {
        let edge = edges::decode_key(key)?;
        return Ok(Some(format!("k\0e\0{}", edge.relation).into_bytes()));
    }
    if !key.starts_with(b"v\0") {
        return Ok(None);
    }
    let text = std::str::from_utf8(key).map_err(Failure::storage)?;
    let mut parts = text.split('\0');
    if parts.next() != Some("v") {
        return Err(invalid_vector_key());
    }
    let collection = parts.next().ok_or_else(invalid_vector_key)?;
    let field = parts.next().ok_or_else(invalid_vector_key)?;
    let document = parts.next().ok_or_else(invalid_vector_key)?;
    if parts.next().is_some() {
        return Err(invalid_vector_key());
    }
    identifier(collection)?;
    identifier(field)?;
    identifier(document)?;
    Ok(Some(format!("k\0v\0{collection}\0{field}").into_bytes()))
}

fn invalid_vector_key() -> Failure {
    Failure::new("vector_key", "vector key must have three identifier parts")
}

fn append_objects(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    writes: &[PreparedWrite],
    counters: &mut BTreeMap<Vec<u8>, Counter>,
) -> Result<(), Failure> {
    // Only the final effect of each canonical key changes a bucket count.
    // Incoming edge mirrors and field declarations are not counted objects.
    let mut final_presence = BTreeMap::new();
    for (key, value) in writes {
        if key.starts_with(b"e\0") || key.starts_with(b"v\0") {
            final_presence.insert(key.as_slice(), value.is_some());
        }
    }
    let mut deltas: BTreeMap<Vec<u8>, Delta> = BTreeMap::new();
    for (key, present) in final_presence {
        let count_key = object_count_key(key)?
            .ok_or_else(|| Failure::new("counts_corrupt", "counted object has no bucket"))?;
        let existed = read(reader, at, key)?.is_some();
        let delta = deltas.entry(count_key).or_default();
        delta.existing |= existed;
        if present && !existed {
            delta.added = delta
                .added
                .checked_add(1)
                .ok_or_else(|| Failure::new("count_overflow", "object delta exceeds u64"))?;
        } else if existed && !present {
            delta.removed = delta
                .removed
                .checked_add(1)
                .ok_or_else(|| Failure::new("count_overflow", "object delta exceeds u64"))?;
        }
    }
    for (key, delta) in deltas {
        if let Some(counter) = counters.get_mut(&key) {
            if counter.initialized && delta.existing {
                return Err(Failure::new(
                    "counts_corrupt",
                    "unregistered kind has stored objects",
                ));
            }
            counter.value = added_count(removed_count(counter.value, delta.removed)?, delta.added)?;
            counter.write |= delta.added != 0 || delta.removed != 0;
            continue;
        }
        let count = match read(reader, at, &key)? {
            Some(value) => decode(&value)?,
            None if !delta.existing => unregistered_count(reader, at, &key)?,
            None => {
                return Err(Failure::new(
                    "counts_corrupt",
                    "present object has no count",
                ));
            }
        };
        let remaining = removed_count(count, delta.removed)?;
        let next = added_count(remaining, delta.added)?;
        if delta.added != 0 || delta.removed != 0 {
            counters.insert(
                key,
                Counter {
                    value: next,
                    write: true,
                    initialized: false,
                },
            );
        }
    }
    Ok(())
}

fn unregistered_count(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    key: &[u8],
) -> Result<u64, Failure> {
    let declaration = if let Some(name) = key.strip_prefix(b"k\0e\0") {
        let mut declaration = b"n\0".to_vec();
        declaration.extend_from_slice(name);
        declaration
    } else if let Some(field) = key.strip_prefix(b"k\0v\0") {
        let mut declaration = b"w\0".to_vec();
        declaration.extend_from_slice(field);
        declaration
    } else {
        return Err(Failure::new(
            "counts_corrupt",
            "count key has no declaration family",
        ));
    };
    if read(reader, at, &declaration)?.is_some() {
        return Err(Failure::new("counts_corrupt", "declared kind has no count"));
    }
    Ok(0)
}

pub(super) fn sum_rows(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    prefix: &[u8],
    parts: usize,
) -> Result<(u64, std::collections::BTreeSet<String>), Failure> {
    let mut total = 0;
    let mut names = std::collections::BTreeSet::new();
    reads::scan_rows(reader, at, prefix, parts, |parts, value| {
        let name = parts.join("\0");
        if !names.insert(name) {
            return Err(Failure::new(
                "counts_corrupt",
                "count exists in multiple shards",
            ));
        }
        total = added_count(total, decode(value)?)?;
        Ok(())
    })?;
    Ok((total, names))
}

#[cfg(test)]
mod probe {
    use std::cell::Cell;

    #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
    pub(super) struct Reads {
        pub(super) ready: usize,
        pub(super) revisions: usize,
        pub(super) bodies: usize,
        pub(super) objects: usize,
        pub(super) counts: usize,
        pub(super) count_writes: usize,
    }

    std::thread_local! {
        static READS: Cell<Reads> = const { Cell::new(Reads {
            ready: 0, revisions: 0, bodies: 0, objects: 0,
            counts: 0, count_writes: 0,
        }) };
    }

    pub(super) fn read(key: &[u8]) {
        READS.with(|slot| {
            let mut value = slot.get();
            if key == super::READY_KEY {
                value.ready += 1;
            } else if key.starts_with(b"m\0") {
                value.revisions += 1;
            } else if key.starts_with(b"d\0") {
                value.bodies += 1;
            } else if key.starts_with(b"e\0") || key.starts_with(b"v\0") {
                value.objects += 1;
            } else if key.starts_with(b"k\0") {
                value.counts += 1;
            }
            slot.set(value);
        });
    }

    pub(super) fn write_count(writes: &[super::PreparedWrite]) {
        READS.with(|slot| {
            let mut value = slot.get();
            value.count_writes += writes
                .iter()
                .filter(|(key, _)| key.starts_with(b"k\0"))
                .count();
            slot.set(value);
        });
    }

    pub(super) fn take() -> Reads {
        READS.with(Cell::take)
    }
}

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