uqa-storage 0.4.6

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

//! Incremental live payload sizes keep shared snapshot admission independent of corpus size.

use std::sync::{Arc, Weak};

use parking_lot::Mutex;
use uqa_core::memory::{MemoryError, MemoryReservation, OwnedMap};

use super::{
    DocId, FieldName, IndexedFieldMetadata, MemoryDocument, MemoryFieldCounters, MemoryIndexState,
    MemoryPosting, PostingKey,
};
use crate::{read_control::StorageReadControl, StorageBackendResult};

#[derive(Debug, Default)]
pub(super) struct RetainedPayload {
    // These sizes describe existing allocations, not requested capacities. A wide sum lets writes preserve their existing error boundary; admission checks the target-sized allowance.
    bytes: u128,
    cached: Mutex<Weak<MemoryReservation>>,
}

impl RetainedPayload {
    fn add(&mut self, bytes: u128) {
        self.bytes += bytes;
        *self.cached.get_mut() = Weak::new();
    }

    fn remove(&mut self, bytes: u128) {
        self.bytes = self
            .bytes
            .checked_sub(bytes)
            .expect("removed payload belongs to the tracked corpus");
        *self.cached.get_mut() = Weak::new();
    }

    pub(super) fn with_retained<T>(
        &self,
        control: &StorageReadControl,
        build: impl FnOnce(Arc<MemoryReservation>) -> StorageBackendResult<T>,
    ) -> StorageBackendResult<T> {
        control.check()?;
        let mut cached = self.cached.lock();
        control.check()?;
        let memory = if let Some(memory) = cached
            .upgrade()
            .filter(|memory| memory.budget().shares_allowance(control.memory()))
        {
            memory
        } else {
            let bytes = usize::try_from(
                self.bytes
                    + size_of::<MemoryIndexState>() as u128
                    + size_of::<MemoryReservation>() as u128,
            )
            .map_err(|_| MemoryError::SizeOverflow)?;
            Arc::new(control.memory().reserve(bytes)?)
        };
        let result = build(Arc::clone(&memory))?;
        control.check()?;
        // Publish only after reader construction succeeds; a late quota or cancellation failure keeps the prior cached allowance.
        *cached = Arc::downgrade(&memory);
        Ok(result)
    }
}

pub(super) fn posting_buffers(posting: &MemoryPosting) -> u128 {
    posting.occurrences.capacity() as u128 * size_of::<uqa_core::TokenOccurrence>() as u128
}

fn posting_size(posting: &MemoryPosting) -> u128 {
    OwnedMap::<DocId, MemoryPosting>::entry_bytes() as u128 + posting_buffers(posting)
}

fn term_payload(key: &PostingKey) -> u128 {
    key.0.capacity() as u128 + key.1.allocated_bytes() as u128
}

fn index_key_size(key: &PostingKey) -> u128 {
    OwnedMap::<PostingKey, OwnedMap<DocId, MemoryPosting>>::entry_bytes() as u128
        + term_payload(key)
}

fn terms_size(terms: &Vec<PostingKey>) -> u128 {
    terms.capacity() as u128 * size_of::<PostingKey>() as u128
        + terms.iter().map(term_payload).sum::<u128>()
}

fn fields_size(fields: &OwnedMap<FieldName, IndexedFieldMetadata>) -> u128 {
    fields
        .keys()
        .map(|field| {
            OwnedMap::<FieldName, IndexedFieldMetadata>::entry_bytes() as u128
                + field.capacity() as u128
        })
        .sum::<u128>()
}

fn document_size(document: &MemoryDocument) -> u128 {
    OwnedMap::<DocId, MemoryDocument>::entry_bytes() as u128
        + fields_size(&document.fields)
        + terms_size(&document.terms)
}

fn counter_size(capacity: usize) -> u128 {
    OwnedMap::<FieldName, MemoryFieldCounters>::entry_bytes() as u128 + capacity as u128
}

pub(super) fn set_counter(
    target: &mut OwnedMap<FieldName, MemoryFieldCounters>,
    field: FieldName,
    value: Option<MemoryFieldCounters>,
    retention: &mut RetainedPayload,
) {
    if let Some(value) = value {
        if let Some(previous) = target.get_mut(&field) {
            *previous = value;
        } else {
            retention.add(counter_size(field.capacity()));
            target.insert(field, value);
        }
    } else if let Some((field, _)) = target.remove_entry(&field) {
        retention.remove(counter_size(field.capacity()));
    }
}

impl MemoryIndexState {
    pub(super) fn insert_posting(
        &mut self,
        doc_id: DocId,
        key: PostingKey,
        posting: MemoryPosting,
    ) {
        self.retention.add(posting_size(&posting));
        if let Some(postings) = self.index.get_mut(&key) {
            if let Some(previous) = postings.insert(doc_id, posting) {
                self.retention.remove(posting_size(&previous));
            }
        } else {
            self.retention.add(index_key_size(&key));
            let mut postings = OwnedMap::new();
            postings.insert(doc_id, posting);
            self.index.insert(key, postings);
        }
    }

    pub(super) fn remove_posting(
        &mut self,
        doc_id: DocId,
        key: &PostingKey,
    ) -> StorageBackendResult<()> {
        let postings = self.index.get_mut(key).ok_or_else(|| {
            super::StorageBackendError::Other(format!(
                "inverted-index document {doc_id} lost a validated posting before removal"
            ))
        })?;
        if let Some(posting) = postings.remove(&doc_id) {
            self.retention.remove(posting_size(&posting));
        }
        if postings.is_empty() {
            let (key, _) = self.index.remove_entry(key).expect("validated posting key");
            self.retention.remove(index_key_size(&key));
        }
        Ok(())
    }

    pub(super) fn remove_document_metadata(&mut self, doc_id: DocId) {
        drop(self.take_document_metadata(doc_id));
    }

    pub(super) fn take_document_metadata(&mut self, doc_id: DocId) -> Option<MemoryDocument> {
        let document = self.documents.remove(&doc_id);
        if let Some(document) = &document {
            self.retention.remove(document_size(document));
        }
        document
    }

    pub(super) fn insert_postings(
        &mut self,
        key: PostingKey,
        postings: OwnedMap<DocId, MemoryPosting>,
    ) {
        if let Some(previous) = self.index.get_mut(&key) {
            for (id, posting) in &postings {
                self.retention.add(posting_size(posting));
                if let Some(replaced) = previous.get(id) {
                    self.retention.remove(posting_size(replaced));
                }
            }
            previous.append(postings);
        } else {
            self.retention
                .add(index_key_size(&key) + postings.values().map(posting_size).sum::<u128>());
            self.index.insert(key, postings);
        }
    }

    pub(super) fn insert_document_metadata(
        &mut self,
        doc_id: DocId,
        fields: OwnedMap<FieldName, IndexedFieldMetadata>,
        terms: Vec<PostingKey>,
    ) {
        self.remove_document_metadata(doc_id);
        let document = MemoryDocument { fields, terms };
        self.retention.add(document_size(&document));
        self.documents.insert(doc_id, document);
    }

    fn payload_size(&self) -> u128 {
        self.index
            .iter()
            .map(|(key, postings)| {
                index_key_size(key) + postings.values().map(posting_size).sum::<u128>()
            })
            .sum::<u128>()
            + self.documents.values().map(document_size).sum::<u128>()
            + self
                .field_counters
                .keys()
                .map(|field| counter_size(field.capacity()))
                .sum::<u128>()
    }
}

impl Clone for MemoryIndexState {
    fn clone(&self) -> Self {
        let mut state = Self {
            index: self.index.clone(),
            documents: self.documents.clone(),
            field_counters: self.field_counters.clone(),
            doc_count: self.doc_count,
            retention: RetainedPayload::default(),
        };
        // Cloned strings/vectors may have smaller capacities. Account the new allocations during this existing deep-copy boundary, never during shared capture.
        state.retention.bytes = state.payload_size();
        state
    }
}