haematite-cli 0.21.1

//! Private candidate pages share the server-owned affected-owner walk.

use std::collections::BTreeMap;
use std::num::NonZeroUsize;

use haematite::branch::history::Key as HistoryKey;
use haematite::branch::{
    CommitDurability, CommitRequest, ForkWithChanges, ForkWrite, current_timestamp,
    open as open_branch, remove as remove_branch,
};
use haematite::tree::Hash;
use serde::{Deserialize, Serialize};

use super::{HeldInputs, MergeChoice, MergeOwner, MissingChoice, Storage, Triple, coverage};
use crate::serve::storage::{
    Captured, DocumentRevision, Failure, MergePlan, SERVICE, counts, edges, identifier,
    merge_holders::holder_name, reads, review::choice, revision_key, version,
};

const MARKER: &[u8] = b"b\0merge-candidate-progress";

#[derive(Debug, Clone, Copy, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum CandidateAction {
    Provisioned,
    Confirmed,
    Page,
    End,
}

#[derive(Debug, Serialize)]
pub struct CandidateAdvance {
    pub action: CandidateAction,
    pub page: u64,
    pub end: bool,
}

#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct Marker {
    plan: MergePlan,
    revision: u64,
    page: u64,
    end: bool,
    position: Option<coverage::Position>,
}

#[derive(Clone, Copy)]
enum Side {
    Source,
    Target,
    Removed,
}

struct Resolved {
    side: Side,
    value: Option<Vec<u8>>,
}

type Writes = BTreeMap<Vec<u8>, Option<Vec<u8>>>;

impl Storage {
    /// Release a private candidate without changing its inputs or choice journal.
    ///
    /// A later page call provisions again at the journal's current revision.
    ///
    /// # Errors
    /// Refuses invalid identities and reports uncertain reference removal explicitly.
    pub fn merge_candidate_restart(&mut self, id: &str) -> Result<bool, Failure> {
        identifier(id)?;
        let name = holder_name(id, "candidate");
        self.confirmed_candidates.remove(&name);
        remove_branch(&name, &mut self.refs)
            .map(|removed| removed.is_some())
            .map_err(|error| Failure::new("candidate_restart_unknown", error.to_string()))
    }

    /// Advance a private candidate from its durable marker, never a caller cursor.
    ///
    /// End means the affected-owner walk ended; it is not a seal or a graph proof.
    /// Supplied vector changes require the separate field-wide rebuild and refuse here.
    ///
    /// # Errors
    /// Refuses moved inputs or choices, missing choices, invalid markers and uncertain writes.
    pub fn merge_candidate_page(&mut self, id: &str) -> Result<CandidateAdvance, Failure> {
        let plan = self.open_merge_plan(id)?;
        let inputs = self.choice_inputs(&plan)?;
        let service = self.merge_service_selection()?;
        let revision = self.choice_revision(&service, id)?;
        let name = holder_name(id, "candidate");
        let Some(at) = self.candidate_selection(&name) else {
            return self.provision_candidate(&name, plan, revision);
        };
        let mut marker = load_marker(self, &at)?;
        validate_marker(&marker, &plan, revision, &at, &inputs)?;
        if self.confirmed_candidates.get(&name) != Some(&at.version) {
            self.refs
                .confirm_version(&name, at.version.generation, at.version.sequence)
                .map_err(|error| Failure::new("durability_unknown", error.to_string()))?;
            self.confirmed_candidates.insert(name, at.version);
            return Ok(answer(CandidateAction::Confirmed, &marker));
        }
        if marker.end {
            return Ok(answer(CandidateAction::End, &marker));
        }
        let step = self.coverage_step(
            id,
            marker.position.take(),
            NonZeroUsize::MIN.saturating_add(63),
        )?;
        if let Some(missing) = step.missing.first() {
            return Err(missing_choice(missing)?);
        }
        if step.keys.len() > coverage::PAGE_OWNER_LIMIT {
            return Err(Failure::new(
                "candidate_page_budget",
                "affected-owner batch exceeds its bound",
            ));
        }
        let rows = materialise(self, id, &inputs, &service, &at, &step.keys)?;
        marker.position = step.next;
        marker.end = marker.position.is_none();
        marker.page = marker.page.checked_add(1).ok_or_else(|| {
            Failure::new("candidate_page_exhausted", "page counter cannot advance")
        })?;
        self.commit_candidate(&name, &at, &marker, rows)?;
        Ok(answer(CandidateAction::Page, &marker))
    }

    fn candidate_selection(&self, name: &str) -> Option<Captured> {
        self.refs.get(name).map(|record| Captured {
            branch: name.to_owned(),
            version: version(record),
            shards: record.shards.clone(),
        })
    }

    fn provision_candidate(
        &mut self,
        name: &str,
        plan: MergePlan,
        revision: u64,
    ) -> Result<CandidateAdvance, Failure> {
        let marker = Marker {
            plan,
            revision,
            page: 0,
            end: false,
            position: None,
        };
        let service =
            open_branch(SERVICE, &self.refs, &self.registry).map_err(Failure::branch_commit)?;
        let writes = [ForkWrite {
            shard_id: self.writer.shard_for(MARKER),
            key: MARKER.to_vec(),
            value: Some(serde_json::to_vec(&marker).map_err(Failure::storage)?),
        }];
        self.writer
            .fork_kept(
                &service,
                HistoryKey {
                    branch: Hash::of(marker.plan.target.branch.as_bytes()),
                    generation: marker.plan.target.version.generation,
                    sequence: marker.plan.target.version.sequence,
                },
                &mut self.refs,
                &self.registry,
                ForkWithChanges {
                    name,
                    timestamp: current_timestamp(),
                    writes: &writes,
                },
            )
            .map_err(|error| {
                Failure::new(
                    if self.refs.get(name).is_some() {
                        "durability_unknown"
                    } else {
                        "candidate_provision_refused"
                    },
                    error.to_string(),
                )
            })?;
        let at = self.candidate_selection(name).ok_or_else(|| {
            Failure::new(
                "candidate_reference_missing",
                "fork did not install its reference",
            )
        })?;
        self.confirmed_candidates
            .insert(name.to_owned(), at.version);
        Ok(answer(CandidateAction::Provisioned, &marker))
    }

    fn commit_candidate(
        &mut self,
        name: &str,
        before: &Captured,
        marker: &Marker,
        mut rows: Writes,
    ) -> Result<(), Failure> {
        rows.insert(
            MARKER.to_vec(),
            Some(serde_json::to_vec(marker).map_err(Failure::storage)?),
        );
        let branch =
            open_branch(name, &self.refs, &self.registry).map_err(Failure::branch_commit)?;
        for (key, value) in rows {
            let shard = self.writer.shard_for(&key);
            match value {
                Some(value) => branch.put(shard, key, value).map_err(Failure::branch)?,
                None => branch.delete(shard, key).map_err(Failure::branch)?,
            }
        }
        self.confirmed_candidates.remove(name);
        self.writer
            .commit(
                &branch,
                &self.registry,
                CommitRequest {
                    durability: CommitDurability::Durable {
                        refs: &mut self.refs,
                    },
                    extra_parents: &[],
                    timestamp: current_timestamp(),
                },
            )
            .map_err(|error| {
                Failure::new(
                    if self
                        .refs
                        .get(name)
                        .is_some_and(|record| version(record) != before.version)
                    {
                        "durability_unknown"
                    } else {
                        "candidate_write_refused"
                    },
                    error.to_string(),
                )
            })?;
        let at = self.candidate_selection(name).ok_or_else(|| {
            Failure::new(
                "candidate_reference_missing",
                "committed reference disappeared",
            )
        })?;
        self.confirmed_candidates
            .insert(name.to_owned(), at.version);
        Ok(())
    }
}

const fn answer(action: CandidateAction, marker: &Marker) -> CandidateAdvance {
    CandidateAdvance {
        action,
        page: marker.page,
        end: marker.end,
    }
}

fn load_marker(storage: &Storage, at: &Captured) -> Result<Marker, Failure> {
    let bytes = reads::get_bytes(&storage.reader, at, MARKER)?.ok_or_else(|| {
        Failure::new(
            "candidate_marker_missing",
            "candidate has no progress marker",
        )
    })?;
    serde_json::from_slice(&bytes)
        .map_err(|error| Failure::new("candidate_marker_corrupt", error.to_string()))
}

fn validate_marker(
    marker: &Marker,
    plan: &MergePlan,
    revision: u64,
    at: &Captured,
    inputs: &HeldInputs,
) -> Result<(), Failure> {
    if marker.plan != *plan
        || marker.page != at.version.sequence
        || (marker.end && marker.position.is_some())
        || (marker.page == 0 && (marker.end || marker.position.is_some()))
        || (!marker.end && marker.page > 0 && marker.position.is_none())
        || !at
            .shards
            .iter()
            .map(|shard| (shard.shard_id, shard.fork_anchor))
            .eq(inputs
                .target
                .shards
                .iter()
                .map(|shard| (shard.shard_id, shard.head)))
    {
        return Err(Failure::new(
            "candidate_marker_corrupt",
            "candidate identity, anchors or progress differs",
        ));
    }
    if marker.revision != revision {
        return Err(Failure::new(
            "choice_journal_moved",
            "candidate is bound to another choice revision",
        ));
    }
    Ok(())
}

fn missing_choice(missing: &MissingChoice) -> Result<Failure, Failure> {
    Ok(Failure::new(
        "choice_missing",
        serde_json::to_string(missing).map_err(Failure::storage)?,
    ))
}

fn embedded_choice(
    storage: &Storage,
    id: &str,
    inputs: &HeldInputs,
    service: &Captured,
    key: &[u8],
    values: &Triple,
) -> Result<choice::EmbeddedSide, Failure> {
    let Some(side) = choice::embedded_field(
        key,
        values.ancestor.as_deref(),
        values.target.as_deref(),
        values.source.as_deref(),
    )?
    else {
        return Err(Failure::new(
            "candidate_vector_rebuild_required",
            "supplied vector fields require field-wide validation and dimensions",
        ));
    };
    let token = values.token(&inputs.pins, key)?;
    if storage.read_choice_row(service, id, &token)?.is_some() {
        return Err(Failure::new(
            "choice_owner_unsupported",
            "embedded declarations do not admit a choice",
        ));
    }
    Ok(side)
}

fn materialise(
    storage: &Storage,
    id: &str,
    inputs: &HeldInputs,
    service: &Captured,
    candidate: &Captured,
    keys: &[Vec<u8>],
) -> Result<Writes, Failure> {
    let mut rows = BTreeMap::new();
    let mut kept_counts = counts::CandidateCounts::new(&storage.reader, candidate)?;
    for key in keys {
        if key.starts_with(b"k\0") {
            continue;
        }
        // Supplied vectors require field-wide validation before they can be carried.
        if key.starts_with(b"v\0") {
            return Err(Failure::new(
                "candidate_vector_rebuild_required",
                "vector changes require field-wide validation and dimensions",
            ));
        }
        let values = storage.choice_values(inputs, key)?;
        if key.starts_with(b"w\0") {
            let side = embedded_choice(storage, id, inputs, service, key, &values)?;
            if side == choice::EmbeddedSide::Source {
                rows.insert(key.clone(), values.source);
            }
            continue;
        }
        let owner = coverage::owner_key(key)?;
        let token = values.token(&inputs.pins, key)?;
        let row = if owner.is_some() {
            storage.read_choice_row(service, id, &token)?
        } else {
            None
        };
        if row
            .as_ref()
            .is_some_and(|row| Some(&row.owner) != owner.as_ref())
        {
            return Err(Failure::new(
                "choice_row_corrupt",
                "candidate choice names another owner",
            ));
        }
        let side = selected(&values, row.as_ref().map(|row| row.choice)).ok_or_else(|| {
            Failure::new(
                "choice_missing",
                format!(
                    "owner {} token {token}",
                    String::from_utf8_lossy(key).replace('\0', "/")
                ),
            )
        })?;
        let value = match side {
            Side::Source => values.source,
            Side::Target => values.target,
            Side::Removed => None,
        };
        let resolved = Resolved { side, value };
        match owner {
            Some(MergeOwner::Document {
                collection,
                document_id,
            }) => {
                let (revision, bytes) =
                    revision_row(storage, inputs, &resolved, &collection, &document_id)?;
                if !matches!(resolved.side, Side::Target) {
                    kept_counts.record(
                        &storage.reader,
                        candidate,
                        key,
                        &revision,
                        bytes.is_some(),
                    )?;
                }
                rows.insert(revision, bytes);
            }
            Some(MergeOwner::Edge { relation, from, to }) => {
                if let Some(value) = &resolved.value {
                    let text = std::str::from_utf8(key).map_err(Failure::storage)?;
                    edges::decode_edge(b"", text, value)?;
                }
                let incoming = edges::keys(&from, &relation, &to)?.1;
                if !matches!(resolved.side, Side::Target) {
                    kept_counts.record(
                        &storage.reader,
                        candidate,
                        key,
                        key,
                        resolved.value.is_some(),
                    )?;
                }
                rows.insert(incoming, resolved.value.clone());
            }
            None => {}
        }
        rows.insert(key.clone(), resolved.value);
    }
    kept_counts.append(&storage.reader, candidate, &mut rows)?;
    Ok(rows)
}

fn revision_row(
    storage: &Storage,
    inputs: &HeldInputs,
    resolved: &Resolved,
    collection: &str,
    document_id: &str,
) -> Result<(Vec<u8>, Option<Vec<u8>>), Failure> {
    let at = match resolved.side {
        Side::Source => &inputs.source,
        Side::Target => &inputs.target,
        Side::Removed => {
            return Err(Failure::new(
                "choice_row_corrupt",
                "an edge removal is recorded for a document",
            ));
        }
    };
    let revision = revision_key(collection, document_id);
    let bytes = if resolved.value.is_some() {
        let bytes = reads::get_bytes(&storage.reader, at, &revision)?.ok_or_else(|| {
            Failure::new(
                "document_revision_missing",
                "resolved document has no revision",
            )
        })?;
        serde_json::from_slice::<DocumentRevision>(&bytes).map_err(Failure::storage)?;
        Some(bytes)
    } else {
        None
    };
    Ok((revision, bytes))
}

fn selected(values: &Triple, choice: Option<MergeChoice>) -> Option<Side> {
    match choice {
        Some(MergeChoice::TakeSource) => Some(Side::Source),
        Some(MergeChoice::TakeTarget) => Some(Side::Target),
        Some(MergeChoice::RemoveEdge) => Some(Side::Removed),
        None if values.ancestor == values.target => Some(Side::Source),
        None if values.ancestor == values.source || values.target == values.source => {
            Some(Side::Target)
        }
        None => None,
    }
}

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