uqa-storage 0.5.2

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Sparse document selections compose actual retained canonical sources without replacing the physical reader.

mod populations;
mod source;
pub use source::DiskANNReadSnapshot;

use super::invalid;
use crate::diskann_index::{
    format::{
        DiskANNCanonicalOrigin, DiskANNChangeIdentity, DiskANNGeneration, DiskANNVectorVersion,
    },
    pages::DiskANNOriginReader,
    DiskANNCanonicalCounts, DiskANNCanonicalRead, DiskANNCanonicalVectorVisitor, DiskANNQueryRead,
};
use crate::{read_control::StorageReadControl, StorageBackendResult};
use std::sync::Arc;
use uqa_core::{
    memory::{Budgeted, BudgetedVec},
    DocId,
};

type SelectedSources = Arc<Budgeted<BudgetedVec<(DocId, Option<DiskANNReadSnapshot>)>>>;

/// Immutable, ordered replacements selected from actual `DiskANN` snapshots. None selects an evaluated deletion independently of retained source membership. A selected document absent from its source is also deleted; a zero-ordinal origin remains an explicit empty replacement. None of these cases falls through to the original tensor. Clones share the sparse directory and its original allowance.
#[derive(Clone)]
pub struct DiskANNReadChanges {
    sources: SelectedSources,
    control: StorageReadControl,
    complete: bool,
}

impl DiskANNReadChanges {
    /// Capture a strictly increasing document selection. Callers supply the already evaluated private source for each replacement and None for each deletion; unselected documents remain on the original view.
    pub fn capture(
        sources: impl IntoIterator<Item = StorageBackendResult<(DocId, Option<DiskANNReadSnapshot>)>>,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Self> {
        control.check()?;
        let mut selected = BudgetedVec::new(control.memory());
        let mut after = None;
        for entry in sources {
            control.check()?;
            let (document, source) = entry?;
            if after.is_some_and(|after| document <= after) {
                return Err(invalid(
                    "private document selection is not strictly increasing",
                ));
            }
            if let Some(source) = &source {
                source.check_control(control)?;
            }
            selected.push((document, source))?;
            after = Some(document);
        }
        let sources =
            Budgeted::new(selected, control.memory().empty_reservation()).into_shared()?;
        control.check()?;
        Ok(Self {
            sources,
            control: control.clone(),
            complete: false,
        })
    }

    /// Select only these evaluated documents, masking every unselected document in the receiver. The caller must establish that this is the complete visible row selection; source lineage and width are still validated against the retained physical index.
    pub fn without_unselected_documents(mut self) -> Self {
        self.complete = true;
        self
    }

    fn source(&self, document: DocId) -> Option<&Option<DiskANNReadSnapshot>> {
        self.sources
            .binary_search_by_key(&document, |(id, _)| *id)
            .ok()
            .map(|position| &self.sources[position].1)
    }

    fn next(&self, after: Option<DocId>) -> Option<DocId> {
        let position = self
            .sources
            .partition_point(|(document, _)| after.is_some_and(|after| *document <= after));
        self.sources.get(position).map(|(document, _)| *document)
    }
}

pub(super) struct SelectedCanonical {
    base: DiskANNReadSnapshot,
    changes: DiskANNReadChanges,
    generation: DiskANNGeneration,
    counts: Option<DiskANNCanonicalCounts>,
}

impl SelectedCanonical {
    pub(super) fn new(
        base: DiskANNReadSnapshot,
        changes: &DiskANNReadChanges,
        built: &DiskANNOriginReader,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Self> {
        base.check_control(control)?;
        changes.control.check()?;
        for source in changes
            .sources
            .iter()
            .filter_map(|(_, source)| source.as_ref())
        {
            source.check_control(control)?;
            if !base.same_lineage(source) || base.dimensions() != source.dimensions() {
                return Err(invalid(
                    "private canonical source belongs to a different index lineage or width",
                ));
            }
        }
        let counts = populations::capture(&base, changes, built, control)?;
        Ok(Self {
            base,
            changes: changes.clone(),
            generation: built.manifest().input().generation,
            counts,
        })
    }

    fn source(&self, document: DocId) -> Option<&DiskANNReadSnapshot> {
        self.changes.source(document).map_or_else(
            || (!self.changes.complete).then_some(&self.base),
            Option::as_ref,
        )
    }
}

impl DiskANNCanonicalRead for SelectedCanonical {
    fn population_counts(
        &self,
        generation: DiskANNGeneration,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<DiskANNCanonicalCounts>> {
        self.check_control(control)?;
        for source in self
            .changes
            .sources
            .iter()
            .filter_map(|(_, source)| source.as_ref())
        {
            source.check_control(control)?;
        }
        Ok((generation == self.generation)
            .then_some(self.counts)
            .flatten())
    }

    fn corpus_fingerprint(
        &self,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<[u8; 32]>> {
        self.check_control(control)?;
        crate::vector_index::selected_fingerprint(
            self.dimensions(),
            (!self.changes.complete).then_some(&self.base as &dyn crate::vector_index::VectorRead),
            self.changes.sources.iter().map(|(document, source)| {
                (
                    *document,
                    source
                        .as_ref()
                        .map(|source| source as &dyn crate::vector_index::VectorRead),
                )
            }),
            control,
        )
    }
    fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
        self.changes.control.check()?;
        self.base.check_control(control)
    }

    fn dimensions(&self) -> u32 {
        self.base.dimensions()
    }

    fn next_document_after(
        &self,
        mut after: Option<DocId>,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<DocId>> {
        loop {
            self.check_control(control)?;
            let base = if self.changes.complete {
                None
            } else {
                self.base.next_document_after(after, control)?
            };
            let selected = self.changes.next(after);
            let Some(document) = first(base, selected) else {
                return Ok(None);
            };
            if after.is_some_and(|after| document <= after) {
                return Err(invalid("private canonical document cursor did not advance"));
            }
            // Preserve the base cursor's integrity checks, including unstamped values. Only a selected, genuinely absent replacement removes an enumerated identity.
            if self.changes.source(document).is_none() || self.origin(document, control)?.is_some()
            {
                return Ok(Some(document));
            }
            after = Some(document);
        }
    }

    fn origin(
        &self,
        document: DocId,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<DiskANNVectorVersion>> {
        self.check_control(control)?;
        self.source(document)
            .map(|source| source.origin(document, control))
            .transpose()
            .map(Option::flatten)
    }

    fn visit_document(
        &self,
        document: DocId,
        control: &StorageReadControl,
        visit: &mut DiskANNCanonicalVectorVisitor<'_>,
    ) -> StorageBackendResult<Option<DiskANNVectorVersion>> {
        self.check_control(control)?;
        self.source(document)
            .map(|source| source.visit_document(document, control, visit))
            .transpose()
            .map(Option::flatten)
    }

    fn read_vector(
        &self,
        document: DocId,
        ordinal: u32,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<uqa_core::memory::BudgetedVec<f32>>> {
        self.check_control(control)?;
        self.source(document)
            .map(|source| source.read_vector(document, ordinal, control))
            .transpose()
            .map(Option::flatten)
    }
}

impl DiskANNQueryRead for SelectedCanonical {
    fn document_origin(
        &self,
        document: DocId,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<DiskANNCanonicalOrigin>> {
        self.check_control(control)?;
        self.source(document)
            .map(|source| source.document_origin(document, control))
            .transpose()
            .map(Option::flatten)
    }

    fn next_change_after(
        &self,
        mut after: Option<DocId>,
        control: &StorageReadControl,
    ) -> StorageBackendResult<Option<DiskANNChangeIdentity>> {
        loop {
            self.check_control(control)?;
            let base = if self.changes.complete {
                None
            } else {
                self.base.next_change_after(after, control)?
            };
            let selected = self.changes.next(after);
            let Some(document) = first(base.map(DiskANNChangeIdentity::document), selected) else {
                return Ok(None);
            };
            if after.is_some_and(|after| document <= after) {
                return Err(invalid("private canonical change cursor did not advance"));
            }
            if let Some(source) = self.changes.source(document) {
                // A private source may already have compacted its own journal. Its actual origin still marks coverage relative to the retained base generation.
                if let Some(version) = source
                    .as_ref()
                    .map(|source| source.origin(document, control))
                    .transpose()?
                    .flatten()
                {
                    return Ok(Some(DiskANNChangeIdentity::new(document, version)));
                }
            } else {
                return Ok(base);
            }
            after = Some(document);
        }
    }
}

fn first(left: Option<DocId>, right: Option<DocId>) -> Option<DocId> {
    match (left, right) {
        (Some(left), Some(right)) => Some(left.min(right)),
        (left, right) => left.or(right),
    }
}