uqa-storage 0.5.3

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

//! Point lookups into the writes of a prepared commit.

use std::cell::RefCell;
use std::collections::BTreeMap;
use std::sync::Arc;

use uqa_core::memory::{MemoryError, MemoryReservation};

use crate::mvcc::overlay::run::{RunCacheReader, RunEntry, SpilledRun};
use crate::mvcc::VersionResult;
use crate::read_control::StorageReadControl;

use super::writes::PreparedWrites;
use super::{PreparedRecordCommit, PreparedRecordWrite, PreparedWriteMetadata, RecordWriteKind};

enum Source<'a> {
    Resident {
        writes: BTreeMap<&'a [u8], &'a PreparedRecordWrite>,
        _memory: MemoryReservation,
    },
    Spilled {
        run: &'a Arc<SpilledRun>,
        _reader: RunCacheReader,
        metadata: RefCell<Option<RunEntry>>,
    },
}

/// Finds the write of a key in a prepared commit: through an ordered index of a batch in memory, or through the key filter and block index of a spilled batch.
pub(in crate::mvcc) struct PreparedLookup<'a> {
    source: Source<'a>,
}

impl<'a> PreparedLookup<'a> {
    pub(in crate::mvcc) fn new(
        prepared: &'a PreparedRecordCommit,
        control: &StorageReadControl,
    ) -> VersionResult<Self> {
        let source = match &prepared.writes {
            PreparedWrites::Resident(writes) => {
                let size = writes
                    .len()
                    .checked_mul(size_of::<(&[u8], &PreparedRecordWrite)>())
                    .ok_or(MemoryError::SizeOverflow)?;
                // Charge logical entries; allocator-specific tree-node bookkeeping follows the other record maps.
                let memory = control.memory().reserve(size)?;
                let mut index = BTreeMap::new();
                for write in writes.iter() {
                    control.cancellation().check()?;
                    index.insert(write.key(), write);
                }
                Source::Resident {
                    writes: index,
                    _memory: memory,
                }
            }
            PreparedWrites::Spilled(run) => Source::Spilled {
                run,
                _reader: run.cache_reader(),
                metadata: RefCell::new(None),
            },
        };
        Ok(Self { source })
    }

    /// The write's identity, condition and kind without reading its payload.
    pub(in crate::mvcc) fn metadata(
        &self,
        key: &[u8],
        control: &StorageReadControl,
    ) -> VersionResult<Option<PreparedWriteMetadata>> {
        control.check()?;
        match &self.source {
            Source::Resident { writes, .. } => Ok(writes.get(key).map(|write| {
                PreparedWriteMetadata::new(
                    write.shared_key(),
                    write.expected(),
                    write.kind(),
                    write.value().map(|value| value.len() as u64),
                )
            })),
            Source::Spilled { run, metadata, .. } => {
                let mut cached = metadata.borrow_mut();
                if cached.as_ref().is_none_or(|entry| entry.key.bytes() != key) {
                    // Drop the previous charged key before reading another. The prepared run is immutable for this lookup's lifetime.
                    *cached = None;
                    *cached = run.get(key, control)?;
                }
                Ok(cached.as_ref().map(|entry| {
                    let value_len = entry.value.map(|location| location.len);
                    PreparedWriteMetadata::new(
                        entry.key.clone(),
                        entry.expected,
                        entry.kind,
                        value_len,
                    )
                }))
            }
        }
    }

    /// The write of `key`, with its value.
    pub(in crate::mvcc) fn get(
        &self,
        key: &[u8],
        control: &StorageReadControl,
    ) -> VersionResult<Option<PreparedRecordWrite>> {
        match &self.source {
            Source::Resident { writes, .. } => Ok(writes.get(key).map(|write| (*write).clone())),
            Source::Spilled { run, .. } => run
                .get(key, control)?
                .map(|entry| {
                    let value = entry
                        .value
                        .map(|location| run.load_value(location, control))
                        .transpose()?;
                    Ok(
                        PreparedRecordWrite::from_shared(entry.key, entry.expected, value)
                            .with_kind(entry.kind),
                    )
                })
                .transpose(),
        }
    }

    /// Visit the writes whose keys start with `prefix` in key order, until `visit` returns false; the batch must be ordered by key.
    pub(in crate::mvcc) fn visit_prefix(
        &self,
        prefix: &[u8],
        control: &StorageReadControl,
        visit: &mut dyn FnMut(&PreparedRecordWrite) -> VersionResult<bool>,
    ) -> VersionResult<()> {
        self.visit_matching_prefix(prefix, None, control, visit)
    }

    /// Select one write kind inside an exact key range without loading other kinds' payloads.
    pub(in crate::mvcc) fn visit_prefix_with_kind(
        &self,
        prefix: &[u8],
        kind: RecordWriteKind,
        control: &StorageReadControl,
        visit: &mut dyn FnMut(&PreparedRecordWrite) -> VersionResult<bool>,
    ) -> VersionResult<()> {
        self.visit_matching_prefix(prefix, Some(kind), control, visit)
    }

    fn visit_matching_prefix(
        &self,
        prefix: &[u8],
        kind: Option<RecordWriteKind>,
        control: &StorageReadControl,
        visit: &mut dyn FnMut(&PreparedRecordWrite) -> VersionResult<bool>,
    ) -> VersionResult<()> {
        match &self.source {
            Source::Resident { writes, .. } => {
                for (key, write) in writes.range::<[u8], _>((
                    std::ops::Bound::Included(prefix),
                    std::ops::Bound::Unbounded,
                )) {
                    control.cancellation().check()?;
                    if !key.starts_with(prefix) {
                        break;
                    }
                    if kind.is_none_or(|kind| write.kind() == kind) && !visit(write)? {
                        break;
                    }
                }
            }
            Source::Spilled { run, .. } => {
                let mut cursor = run.cursor(std::ops::Bound::Included(prefix));
                while let Some(entry) = cursor.next(control)? {
                    if !entry.key.bytes().starts_with(prefix) {
                        break;
                    }
                    if kind.is_some_and(|kind| entry.kind != kind) {
                        continue;
                    }
                    let value = entry
                        .value
                        .map(|location| cursor.load_value(location, control))
                        .transpose()?;
                    let write = PreparedRecordWrite::from_shared(entry.key, entry.expected, value)
                        .with_kind(entry.kind);
                    if !visit(&write)? {
                        break;
                    }
                }
            }
        }
        Ok(())
    }

    pub(in crate::mvcc) fn contains(
        &self,
        key: &[u8],
        control: &StorageReadControl,
    ) -> VersionResult<bool> {
        match &self.source {
            Source::Resident { writes, .. } => Ok(writes.contains_key(key)),
            Source::Spilled { run, .. } => Ok(run.get(key, control)?.is_some()),
        }
    }
}