moraine 0.7.0

//! Index entry maintenance: turning writer-supplied entries into staged
//! `index` writes, with commit-time uniqueness enforcement. A unique
//! entry's store key is the value, so two commits inserting the same value
//! collide on write-write detection.

use std::{
    collections::{HashMap, HashSet, VecDeque},
    ops::Bound,
    sync::atomic::{AtomicU64, Ordering},
    time::{Duration, Instant},
};

use bytes::Bytes;
use futures::{Stream, StreamExt, TryStreamExt, stream};
use slatedb::DbTransaction;
use tracing::warn;

use crate::{
    data_file::ScopedReadTally,
    error::{Error, Result},
    store::{
        StagedBytes,
        handle::{ReadHandle, ScanOrder, ScanShape},
        index_encoding::{CanonicalKey, NullOrder, non_null_flag_key},
        key::{
            IndexKey, IndexKind, Key, encode_index_entry, index_index_prefix,
            index_multi_value_prefix, index_value_above, index_value_body, index_value_suffix,
        },
    },
};

/// One index-entry mutation accumulated during a commit closure, resolved
/// against the store when the batch is staged.
#[derive(Debug, Clone, PartialEq)]
#[expect(
    clippy::struct_excessive_bools,
    reason = "each flag is an independent property of one staged entry"
)]
pub(crate) struct StagedIndexEntry {
    /// The index this entry belongs to.
    pub(crate) index_id: u64,
    /// Whether the index is unique.
    pub(crate) unique: bool,
    /// The final encoded SlateDB entry key.
    pub(crate) key: Bytes,
    /// The row the entry points at.
    pub(crate) row_id: u64,
    /// Whether this removes the entry rather than adds it.
    pub(crate) delete: bool,
    /// Whether the index is still building; a collision then poisons it
    /// instead of failing the commit.
    pub(crate) building: bool,
    /// The caller established the key is absent from committed state, so
    /// a unique put needs no read before it.
    pub(crate) known_absent: bool,
    /// The caller established the entry names this row in committed state,
    /// so a unique delete needs no read before it.
    pub(crate) known_held: bool,
}

/// A unique entry's value is the holding row id, big-endian.
fn decode_row_id(bytes: &[u8]) -> Result<u64> {
    let array: [u8; 8] = bytes.try_into().map_err(|_| {
        Error::Corruption(format!(
            "index entry value is {} bytes, expected 8",
            bytes.len()
        ))
    })?;
    Ok(u64::from_be_bytes(array))
}

/// Uniqueness probes in flight at once; sized for a remote object store.
const UNIQUENESS_PROBE_CONCURRENCY: usize = 1024;

/// Ready keys sharing one transactional scan; pending sources flush partial
/// batches.
const PROBE_BATCH_SIZE: usize = 128;

/// Additions derived while the deletion phase is still draining. Once full,
/// backpressure pauses addition sources without holding up deletion staging.
const ADDITION_PREFETCH: usize = 512;

/// The most equality-index entries one commit may derive.
const MAX_INDEX_ENTRIES_PER_COMMIT: usize = 1_000_000;

/// The most encoded equality-index key and value bytes one commit may stage.
const MAX_INDEX_STAGED_BYTES_PER_COMMIT: u64 = 64 * 1024 * 1024;

#[derive(Default)]
struct IndexCommitBudget {
    entries: usize,
    staged_bytes: u64,
}

impl IndexCommitBudget {
    fn with_entries(entries: usize) -> Result<Self> {
        let mut budget = Self::default();
        budget.observe_entry(entries)?;
        Ok(budget)
    }

    fn observe_entry(&mut self, additional: usize) -> Result<()> {
        self.entries = self.entries.saturating_add(additional);
        if self.entries > MAX_INDEX_ENTRIES_PER_COMMIT {
            return Err(oversized_commit(self.entries));
        }
        Ok(())
    }

    fn stage_bytes(&mut self, additional: u64) -> Result<()> {
        let staged = self.staged_bytes.saturating_add(additional);
        if staged > MAX_INDEX_STAGED_BYTES_PER_COMMIT {
            return Err(oversized_staged_bytes(staged));
        }
        self.staged_bytes = staged;
        Ok(())
    }
}

/// What staging a batch's entries produced: the indexes a collision
/// poisoned, and the bytes the entries put on the batch.
pub(crate) struct StagedEntries {
    /// Definitions a duplicate poisoned, sorted and deduplicated.
    pub(crate) poisoned: Vec<u64>,
    /// Deferred definitions whose SQL additions committed without entries.
    pub(crate) deferred: Vec<u64>,
    /// Key and value bytes staged, deletes included.
    pub(crate) bytes: u64,
    /// Inline chunk-directory repairs derived along the way, handed back
    /// rather than staged so the caller can order them below the batch's
    /// own directory writes.
    pub(crate) locator_writes: Vec<crate::transaction::commit::StagedWrite>,
    pub(crate) metrics: IndexMaintenanceMetrics,
}

/// Per-commit equality-index work. Derivation and probe windows overlap.
#[derive(Debug, Default, Clone, Copy)]
pub(crate) struct IndexMaintenanceMetrics {
    pub(crate) deletion_derivation: Duration,
    pub(crate) addition_derivation: Duration,
    pub(crate) probe_window: Duration,
    pub(crate) probe_service: Duration,
    pub(crate) staging: Duration,
    pub(crate) scoped_read: ScopedReadTally,
    pub(crate) additions: u64,
    pub(crate) deletions: u64,
    pub(crate) unique_probes: u64,
    pub(crate) probe_hits: u64,
    pub(crate) probe_misses: u64,
    pub(crate) probe_peak_in_flight: u64,
    pub(crate) probes_completed_during_deletions: u64,
    /// Unique puts staged without a read, their keys known absent.
    pub(crate) known_absent: u64,
    /// Unique deletes that read their entry to check which row holds it.
    pub(crate) guard_reads: u64,
}

/// One unique put awaiting its committed-state probe.
struct PendingProbe {
    /// The entry's encoded store key.
    key: Bytes,
    /// The row claiming the value.
    row_id: u64,
    /// The index the claim belongs to.
    index_id: u64,
    /// Whether that index is still building.
    building: bool,
}

/// Work one streamed index entry requires.
enum ProbePlan {
    /// Stage a non-unique put without a read.
    Put(Bytes),
    /// Stage a unique put whose key is known absent, without a read.
    Claim { key: Bytes, row_id: u64 },
    /// Read committed state before deciding whether to stage the put.
    Probe(PendingProbe),
    /// A repeated claim by the same row needs no write.
    Noop,
    /// Two rows in this batch claim one unique value.
    Collision {
        index_id: u64,
        building: bool,
        claimant: u64,
        holder: u64,
    },
    /// Input, limit, or derivation failure at this stream position.
    Failure(Error),
}

/// A completed plan, carrying a probe's value when it needed a read.
enum ReadyAddition {
    Put { key: Bytes, row_id: Option<u64> },
    Probed(CompletedProbe),
    Poison(u64),
}

struct CompletedProbe {
    probe: PendingProbe,
    present: Option<Bytes>,
    service: Duration,
    completed: Instant,
}

/// Sequential planning state. Claims are recorded before probes start, so
/// concurrent reads cannot admit duplicate values from the same commit.
#[derive(Default)]
struct ProbePlanner {
    claimed: HashMap<Bytes, u64>,
}

impl ProbePlanner {
    fn plan(
        &mut self,
        entry: Result<StagedIndexEntry>,
        budget: &mut IndexCommitBudget,
    ) -> ProbePlan {
        let entry = match entry {
            Ok(entry) => entry,
            Err(error) => return ProbePlan::Failure(error),
        };
        if let Err(error) = budget.observe_entry(1) {
            return ProbePlan::Failure(error);
        }
        let value_bytes = usize::from(entry.unique) * size_of::<u64>();
        if let Err(error) = budget.stage_bytes(
            u64::try_from(entry.key.len().saturating_add(value_bytes)).unwrap_or(u64::MAX),
        ) {
            return ProbePlan::Failure(error);
        }
        if entry.delete {
            return ProbePlan::Failure(Error::Corruption(
                "index put stream contains a deletion".to_owned(),
            ));
        }
        if !entry.unique {
            return ProbePlan::Put(entry.key);
        }

        if let Some(&holder) = self.claimed.get(&entry.key) {
            return if holder == entry.row_id {
                ProbePlan::Noop
            } else {
                ProbePlan::Collision {
                    index_id: entry.index_id,
                    building: entry.building,
                    claimant: entry.row_id,
                    holder,
                }
            };
        }
        self.claimed.insert(entry.key.clone(), entry.row_id);
        if entry.known_absent {
            return ProbePlan::Claim {
                key: entry.key,
                row_id: entry.row_id,
            };
        }

        ProbePlan::Probe(PendingProbe {
            key: entry.key,
            row_id: entry.row_id,
            index_id: entry.index_id,
            building: entry.building,
        })
    }
}

/// Which state a rejected claim collided with: the commit's own additions,
/// or a row the index already held.
#[derive(Clone, Copy)]
enum Collided {
    InBatch,
    Committed,
}

/// A collision's verdict: fail the commit, or poison the building index.
fn collision(
    index_id: u64,
    building: bool,
    claimant: u64,
    holder: u64,
    collided: Collided,
) -> Result<Option<u64>> {
    if building {
        Ok(Some(index_id))
    } else {
        Err(unique_violation(index_id, claimant, holder, collided))
    }
}

async fn resolve_probe(transaction: &DbTransaction, probe: PendingProbe) -> Result<CompletedProbe> {
    let started = Instant::now();
    let present = transaction.get(&probe.key).await?;
    Ok(CompletedProbe {
        probe,
        present,
        service: started.elapsed(),
        completed: Instant::now(),
    })
}

/// Resolves a batch by concurrent point reads. Each read is answered by
/// the SST filters when the key is absent, so its cost is bounded by the
/// sources the key could be in rather than by the batch's key span.
async fn resolve_probes(
    transaction: &DbTransaction,
    probes: Vec<PendingProbe>,
) -> Result<Vec<CompletedProbe>> {
    let reads: Vec<_> = probes
        .into_iter()
        .map(|probe| resolve_probe(transaction, probe))
        .collect();
    stream::iter(reads)
        .buffer_unordered(PROBE_BATCH_SIZE)
        .try_collect()
        .await
}

fn plan_probe_batch(
    additions: Vec<Result<StagedIndexEntry>>,
    planner: &mut ProbePlanner,
    budget: &mut IndexCommitBudget,
    ready: &mut VecDeque<ReadyAddition>,
    metrics: &mut IndexMaintenanceMetrics,
    in_flight: &mut usize,
    first_probe: &mut Option<Instant>,
) -> Result<Vec<PendingProbe>> {
    let capacity = additions.len();
    let mut probes = Vec::new();
    for addition in additions {
        metrics.additions = metrics.additions.saturating_add(1);
        match planner.plan(addition, budget) {
            ProbePlan::Put(key) => ready.push_back(ReadyAddition::Put { key, row_id: None }),
            ProbePlan::Claim { key, row_id } => {
                metrics.known_absent = metrics.known_absent.saturating_add(1);
                ready.push_back(ReadyAddition::Put {
                    key,
                    row_id: Some(row_id),
                });
            }
            ProbePlan::Probe(probe) => {
                if probes.is_empty() {
                    probes.reserve_exact(capacity);
                }
                probes.push(probe);
            }
            ProbePlan::Noop => {}
            ProbePlan::Collision {
                index_id,
                building,
                claimant,
                holder,
            } => {
                if let Some(index_id) =
                    collision(index_id, building, claimant, holder, Collided::InBatch)?
                {
                    ready.push_back(ReadyAddition::Poison(index_id));
                }
            }
            ProbePlan::Failure(error) => return Err(error),
        }
    }
    if !probes.is_empty() {
        first_probe.get_or_insert_with(Instant::now);
        let count = u64::try_from(probes.len()).unwrap_or(u64::MAX);
        metrics.unique_probes = metrics.unique_probes.saturating_add(count);
        *in_flight += probes.len();
        metrics.probe_peak_in_flight = metrics
            .probe_peak_in_flight
            .max(u64::try_from(*in_flight).unwrap_or(u64::MAX));
    }
    Ok(probes)
}

fn complete_probes(
    resolutions: Result<Vec<CompletedProbe>>,
    ready: &mut VecDeque<ReadyAddition>,
    in_flight: &mut usize,
) -> Result<usize> {
    let resolutions = resolutions?;
    let count = resolutions.len();
    *in_flight -= count;
    ready.extend(resolutions.into_iter().map(ReadyAddition::Probed));
    Ok(count)
}

/// Drops a unique deletion whose entry is held by some other row, counting
/// each read in `guard_reads`.
///
/// A unique entry's key is the value alone, so deleting it by key removes
/// whichever row holds that value — not necessarily the row being removed.
/// The stored row id is what says the two are the same, and a derivation
/// that names an already-removed row must not take the entry a later
/// insert of the same value put there. A deletion the caller established
/// names its own row is staged without the read.
fn guarded_deletions<'a, D>(
    reader: ReadHandle<'a>,
    deletes: D,
    guard_reads: &'a AtomicU64,
) -> impl Stream<Item = Result<StagedIndexEntry>> + 'a
where
    D: Stream<Item = Result<StagedIndexEntry>> + 'a,
{
    deletes
        .map(move |entry| async move {
            let entry = entry?;
            if !entry.unique || entry.known_held {
                return Ok(Some(entry));
            }
            guard_reads.fetch_add(1, Ordering::Relaxed);
            let Some(bytes) = reader.get(entry.key.clone()).await.map_err(Error::from)? else {
                return Ok(None);
            };
            Ok((decode_row_id(&bytes)? == entry.row_id).then_some(entry))
        })
        .buffer_unordered(UNIQUENESS_PROBE_CONCURRENCY)
        .try_filter_map(|entry| std::future::ready(Ok(entry)))
}

fn stage_probe_put(
    db_tx: &DbTransaction,
    staged: &mut StagedBytes,
    key: Bytes,
    row_id: Option<u64>,
) -> Result<()> {
    let value_bytes = row_id.map_or(0, |_| size_of::<u64>());
    staged.add(key.len(), value_bytes);
    match row_id {
        Some(row_id) => db_tx.put(key, row_id.to_be_bytes()),
        None => db_tx.put(key, []),
    }
    .map_err(Error::from)
}

/// Resolves accumulated entries onto `db_tx`, enforcing uniqueness at
/// commit. Deletes are staged first. For each unique put: present with a
/// different row id → [`Error::Constraint`]; present with the same row id
/// → no-op; absent → staged. Duplicates within the commit are caught in
/// memory.
pub(crate) async fn stage_index_entries(
    db_tx: &DbTransaction,
    entries: Vec<StagedIndexEntry>,
) -> Result<StagedEntries> {
    let entry_count = entries.len();
    if entry_count > MAX_INDEX_ENTRIES_PER_COMMIT {
        warn!(
            staged = entry_count,
            limit = MAX_INDEX_ENTRIES_PER_COMMIT,
            "refusing an oversized commit"
        );
        return Err(oversized_commit(entry_count));
    }
    let encoded_bytes = entries.iter().fold(0_u64, |total, entry| {
        let value_bytes = usize::from(entry.unique && !entry.delete) * size_of::<u64>();
        total.saturating_add(
            u64::try_from(entry.key.len().saturating_add(value_bytes)).unwrap_or(u64::MAX),
        )
    });
    if encoded_bytes > MAX_INDEX_STAGED_BYTES_PER_COMMIT {
        return Err(oversized_staged_bytes(encoded_bytes));
    }

    let (deletes, puts): (Vec<_>, Vec<_>) = entries.into_iter().partition(|entry| entry.delete);
    let deletes = stream::iter(deletes.into_iter().map(Ok));
    let puts = stream::iter(puts.into_iter().map(Ok));
    stage_index_entry_stream(db_tx, deletes, puts, 0).await
}

/// Consumes deletion entries, then additions, as streams. Unique probes form
/// one continuously replenished window; successful reads stage in
/// completion order, and the first fatal result aborts.
#[allow(clippy::too_many_lines)]
pub(crate) async fn stage_index_entry_stream<D, S>(
    db_tx: &DbTransaction,
    deletes: D,
    entries: S,
    prior_entry_count: usize,
) -> Result<StagedEntries>
where
    D: Stream<Item = Result<StagedIndexEntry>>,
    S: Stream<Item = Result<StagedIndexEntry>>,
{
    let started = Instant::now();
    let mut staged = StagedBytes::default();
    let mut deleted_unique = HashSet::new();
    let mut budget = IndexCommitBudget::with_entries(prior_entry_count)?;
    let reader = ReadHandle::Tx(db_tx);
    let guard_reads = AtomicU64::new(0);
    let mut deletes = std::pin::pin!(guarded_deletions(reader, deletes, &guard_reads));
    let mut entries = std::pin::pin!(entries.ready_chunks(PROBE_BATCH_SIZE));
    let mut ready = VecDeque::with_capacity(ADDITION_PREFETCH);
    let mut additions_done = false;
    let mut planner = ProbePlanner::default();

    let mut probes = futures::stream::FuturesUnordered::new();
    let mut in_flight = 0;
    let mut metrics = IndexMaintenanceMetrics::default();
    let mut first_probe = None;
    let mut last_probe_completion = None;

    loop {
        let buffered = ready.len().saturating_add(in_flight);
        let deletion = if buffered > ADDITION_PREFETCH - PROBE_BATCH_SIZE {
            deletes.next().await
        } else {
            tokio::select! {
                biased;
                deletion = deletes.next() => deletion,
                resolution = probes.next(), if !probes.is_empty() => {
                    if let Some(resolution) = resolution {
                        let count = complete_probes(resolution, &mut ready, &mut in_flight)?;
                        metrics.probes_completed_during_deletions = metrics
                            .probes_completed_during_deletions
                            .saturating_add(u64::try_from(count).unwrap_or(u64::MAX));
                    }
                    continue;
                },
                addition = entries.next(), if !additions_done => if let Some(addition) = addition {
                    let pending = plan_probe_batch(
                        addition, &mut planner, &mut budget, &mut ready,
                        &mut metrics, &mut in_flight, &mut first_probe,
                    )?;
                    if !pending.is_empty() {
                        probes.push(resolve_probes(db_tx, pending));
                    }
                    continue;
                } else {
                    additions_done = true;
                    metrics.addition_derivation = started.elapsed();
                        continue;
                }
            }
        };
        let Some(entry) = deletion else {
            break;
        };
        let entry = entry?;
        budget.observe_entry(1)?;
        budget.stage_bytes(u64::try_from(entry.key.len()).unwrap_or(u64::MAX))?;
        if !entry.delete {
            return Err(Error::Corruption(
                "index deletion stream contains a put".to_owned(),
            ));
        }
        if entry.unique {
            deleted_unique.insert(entry.key.clone());
        }
        metrics.deletions = metrics.deletions.saturating_add(1);
        let stage_started = Instant::now();
        staged.add(entry.key.len(), 0);
        db_tx.delete(entry.key).map_err(Error::from)?;
        metrics.staging = metrics.staging.saturating_add(stage_started.elapsed());
    }
    metrics.deletion_derivation = started.elapsed();
    metrics.guard_reads = guard_reads.load(Ordering::Relaxed);

    let mut poisoned = Vec::new();
    loop {
        let resolution = if let Some(resolution) = ready.pop_front() {
            Some(resolution)
        } else if additions_done || in_flight > UNIQUENESS_PROBE_CONCURRENCY - PROBE_BATCH_SIZE {
            if let Some(resolutions) = probes.next().await {
                complete_probes(resolutions, &mut ready, &mut in_flight)?;
                continue;
            }
            None
        } else {
            tokio::select! {
                biased;
                resolution = probes.next(), if !probes.is_empty() => {
                    if let Some(resolution) = resolution {
                        complete_probes(resolution, &mut ready, &mut in_flight)?;
                    }
                    continue;
                },
                addition = entries.next() => if let Some(addition) = addition {
                    let pending = plan_probe_batch(
                        addition, &mut planner, &mut budget, &mut ready,
                        &mut metrics, &mut in_flight, &mut first_probe,
                    )?;
                    if !pending.is_empty() {
                        probes.push(resolve_probes(db_tx, pending));
                    }
                    continue;
                } else {
                    additions_done = true;
                    metrics.addition_derivation = started.elapsed();
                    continue;
                }
            }
        };

        let Some(resolution) = resolution else {
            break;
        };
        match resolution {
            ReadyAddition::Put { key, row_id } => {
                let stage_started = Instant::now();
                stage_probe_put(db_tx, &mut staged, key, row_id)?;
                metrics.staging = metrics.staging.saturating_add(stage_started.elapsed());
            }
            ReadyAddition::Probed(CompletedProbe {
                probe,
                present,
                service,
                completed,
            }) => {
                metrics.probe_service = metrics.probe_service.saturating_add(service);
                last_probe_completion = Some(completed);
                if present.is_some() {
                    metrics.probe_hits = metrics.probe_hits.saturating_add(1);
                } else {
                    metrics.probe_misses = metrics.probe_misses.saturating_add(1);
                }
                if deleted_unique.contains(&probe.key) {
                    let stage_started = Instant::now();
                    stage_probe_put(db_tx, &mut staged, probe.key, Some(probe.row_id))?;
                    metrics.staging = metrics.staging.saturating_add(stage_started.elapsed());
                } else if let Some(bytes) = present {
                    match decode_row_id(&bytes) {
                        Ok(holder) if holder == probe.row_id => {}
                        Ok(holder) => {
                            if let Some(index_id) = collision(
                                probe.index_id,
                                probe.building,
                                probe.row_id,
                                holder,
                                Collided::Committed,
                            )? {
                                poisoned.push(index_id);
                            }
                        }
                        Err(error) => return Err(error),
                    }
                } else {
                    let stage_started = Instant::now();
                    stage_probe_put(db_tx, &mut staged, probe.key, Some(probe.row_id))?;
                    metrics.staging = metrics.staging.saturating_add(stage_started.elapsed());
                }
            }
            ReadyAddition::Poison(index_id) => poisoned.push(index_id),
        }
    }

    if let (Some(first_probe), Some(last_probe_completion)) = (first_probe, last_probe_completion) {
        metrics.probe_window = last_probe_completion.saturating_duration_since(first_probe);
    }
    poisoned.sort_unstable();
    poisoned.dedup();

    Ok(StagedEntries {
        poisoned,
        deferred: Vec::new(),
        bytes: staged.0,
        locator_writes: Vec::new(),
        metrics,
    })
}

/// Records `poisoned` on the working state's definitions, so the commit's
/// ordinary entity diff stages the flag. Poisoning is terminal.
pub(crate) fn apply_poison(state: &mut crate::catalog::CatalogSnapshot, poisoned: &[u64]) {
    for index_id in poisoned {
        let Some(table_id) = state.index_tables.get(index_id).copied() else {
            continue;
        };
        if let Some(value) = state
            .indexes
            .get_mut(&table_id)
            .and_then(|indexes| indexes.get_mut(index_id))
        {
            value.poisoned = Some(true);
        }
    }
}

/// Marks ready deferred definitions as awaiting repair. A definition already
/// being repaired keeps its cursor.
pub(crate) fn apply_deferred_maintenance(
    base: &crate::catalog::CatalogSnapshot,
    state: &mut crate::catalog::CatalogSnapshot,
    deferred: &[u64],
    new_snapshot: u64,
) {
    for index_id in deferred {
        let Some(table_id) = state.index_tables.get(index_id).copied() else {
            continue;
        };
        let Some(value) = state
            .indexes
            .get_mut(&table_id)
            .and_then(|indexes| indexes.get_mut(index_id))
        else {
            continue;
        };
        if value.build_state.is_none() {
            value.begin_snapshot = new_snapshot;
            value.build_state = Some("maintaining".to_owned());
            value.build_cursor_row_id = base
                .table_stats
                .get(&table_id)
                .and_then(|stats| stats.next_row_id.checked_sub(1));
            let tail = base
                .data_files
                .get(&table_id)
                .and_then(|files| files.get_max())
                .map(|(_, file)| file);
            value.build_cursor_file = Some(tail.map_or(0, |file| file.data_file_id));
            value.build_cursor_position =
                Some(tail.map_or(0, |file| file.record_count.saturating_sub(1)));
        }
    }
}

/// The refusal for a commit deriving too many equality-index entries. The
/// text must avoid DuckLake's retry
/// substrings (`conflict`, `concurrent`, `unique`, `primary key`).
fn oversized_commit(staged: usize) -> Error {
    Error::Constraint(format!(
        "commit stages {staged} equality-index entries, above the \
         {MAX_INDEX_ENTRIES_PER_COMMIT} a single commit allows. Split the work into \
         several smaller commits."
    ))
}

fn oversized_staged_bytes(staged: u64) -> Error {
    Error::Constraint(format!(
        "commit stages {staged} encoded equality-index bytes, above the \
         {MAX_INDEX_STAGED_BYTES_PER_COMMIT} a single commit allows. Split the work into \
         several smaller commits."
    ))
}

/// A uniqueness error, naming both rows and which state the claim collided
/// with. The text must avoid DuckLake's retry substrings (`conflict`,
/// `concurrent`, `unique`, `primary key`), and names row ids rather than the
/// value, which no error on this path carries.
fn unique_violation(index_id: u64, claimant: u64, holder: u64, collided: Collided) -> Error {
    Error::Constraint(match collided {
        Collided::InBatch => format!(
            "duplicate value violates equality index {index_id}: rows {claimant} and {holder} in \
             one commit claim the same value"
        ),
        Collided::Committed => format!(
            "duplicate value violates equality index {index_id}: row {claimant} claims a value \
             held by committed row {holder}"
        ),
    })
}

/// Deletes up to `limit` orphaned entries of one dropped index inside an
/// open transaction, returning how many deletes were staged.
pub(crate) async fn reclaim_entries(
    tx: &slatedb::DbTransaction,
    index_id: u64,
    limit: usize,
    staged: &mut StagedBytes,
) -> Result<usize> {
    let mut deleted = 0;
    for kind in [IndexKind::Unique, IndexKind::Multi] {
        if deleted >= limit {
            break;
        }
        let (batch, _) =
            reclaim_entries_from(tx, kind, index_id, limit - deleted, None, staged).await?;
        deleted += batch;
    }

    Ok(deleted)
}

/// Deletes up to `limit` entries of one dropped index of one kind,
/// resuming at `start_from` when given, and returns how many were staged
/// alongside the last key deleted.
pub(crate) async fn reclaim_entries_from(
    tx: &slatedb::DbTransaction,
    kind: IndexKind,
    index_id: u64,
    limit: usize,
    start_from: Option<&[u8]>,
    staged: &mut StagedBytes,
) -> Result<(usize, Option<Vec<u8>>)> {
    let prefix = index_index_prefix(kind, index_id);
    // `scan_prefix` takes its bounds as a suffix of the prefix.
    let suffix = match start_from {
        Some(key) if key.len() >= prefix.len() => key[prefix.len()..].to_vec(),
        _ => Vec::new(),
    };

    let mut iter = ReadHandle::Tx(tx)
        .scan_prefix(prefix, suffix.., ScanShape::Bulk)
        .await?;
    let mut deleted = 0;
    let mut last: Option<Bytes> = None;
    while deleted < limit {
        match iter.next().await? {
            Some(entry) => {
                staged.add(entry.key.len(), 0);
                last = Some(entry.key.clone());
                tx.delete(entry.key)?;
                deleted += 1;
            }
            None => break,
        }
    }

    Ok((deleted, last.map(|key| key.to_vec())))
}

/// Drains a scan over `multi` entries into their row ids, which live in
/// the key. `context` names the scan in a corruption error.
async fn collect_multi_row_ids(iter: &mut slatedb::DbIterator, context: &str) -> Result<Vec<u64>> {
    let mut row_ids = Vec::new();
    while let Some(entry) = iter.next().await.map_err(Error::from)? {
        match Key::decode(&entry.key)? {
            Key::Index(IndexKey::Multi { row_id, .. }) => row_ids.push(row_id),
            other => {
                return Err(Error::Corruption(format!(
                    "non-multi key in {context}: {other:?}"
                )));
            }
        }
    }

    Ok(row_ids)
}

/// The row ids holding one indexed value: a point-get for a unique index,
/// an ascending prefix scan for a non-unique one.
pub(crate) async fn lookup_row_ids(
    reader: ReadHandle<'_>,
    index_id: u64,
    unique: bool,
    key: &CanonicalKey,
) -> Result<Vec<u64>> {
    if unique {
        let entry_key = encode_index_entry(index_id, true, key, 0);
        return match reader.get(entry_key).await.map_err(Error::from)? {
            Some(bytes) => Ok(vec![decode_row_id(&bytes)?]),
            None => Ok(Vec::new()),
        };
    }

    let prefix = index_multi_value_prefix(index_id, key);
    let mut iter = reader
        .scan_prefix(prefix, .., ScanShape::Equality)
        .await
        .map_err(Error::from)?;

    collect_multi_row_ids(&mut iter, "index scan").await
}

/// The row ids whose indexed value falls between `lower` and `upper`, in the
/// requested scan order. The bounds are canonical values, already encoded
/// in the columns' declared directions.
pub(crate) async fn range_row_ids(
    reader: ReadHandle<'_>,
    index_id: u64,
    unique: bool,
    leading_nulls: NullOrder,
    lower: Bound<CanonicalKey>,
    upper: Bound<CanonicalKey>,
    order: ScanOrder,
) -> Result<Vec<u64>> {
    let kind = if unique {
        IndexKind::Unique
    } else {
        IndexKind::Multi
    };
    let prefix = index_index_prefix(kind, index_id);

    let suffix = |canon: &CanonicalKey| index_value_suffix(kind, index_id, canon);
    let above = |canon: &CanonicalKey| index_value_above(kind, index_id, canon);

    // A comparison never matches a NULL, so an open side stops at the
    // leading column's non-null region.
    let non_null = non_null_flag_key(leading_nulls);
    let past_non_null = above(&non_null).map_or(Bound::Unbounded, Bound::Excluded);

    let start = match lower {
        Bound::Included(canon) => Bound::Included(suffix(&canon)),
        // Skip every entry sharing the bound value: start above them all.
        Bound::Excluded(canon) => match above(&canon) {
            Some(above) => Bound::Included(above),
            None => Bound::Excluded(suffix(&canon)),
        },
        Bound::Unbounded => Bound::Included(index_value_body(kind, index_id, &non_null)),
    };
    let end = match upper {
        // Include every entry sharing the bound value: end above them all.
        Bound::Included(canon) => match above(&canon) {
            Some(above) => Bound::Excluded(above),
            None => past_non_null,
        },
        Bound::Excluded(canon) => Bound::Excluded(suffix(&canon)),
        Bound::Unbounded => past_non_null,
    };

    let mut iter = reader
        .scan_prefix_ordered(prefix, (start, end), ScanShape::Probe, order)
        .await
        .map_err(Error::from)?;
    if !unique {
        return collect_multi_row_ids(&mut iter, "index range scan").await;
    }
    let mut row_ids = Vec::new();
    while let Some(entry) = iter.next().await.map_err(Error::from)? {
        row_ids.push(decode_row_id(entry.value.as_ref())?);
    }

    Ok(row_ids)
}

/// The row ids whose leading indexed columns match `prefix`, a canonical
/// key over a leading run of `= value` and `IS NULL` predicates. A row with
/// any NULL indexed column is stored multi-shaped, so the `multi` subrange
/// is scanned.
pub(crate) async fn null_prefix_row_ids(
    reader: ReadHandle<'_>,
    index_id: u64,
    prefix: &CanonicalKey,
    order: ScanOrder,
) -> Result<Vec<u64>> {
    let mut scan_prefix = index_multi_value_prefix(index_id, prefix);
    // Dropping the exact-value terminator turns the bytes into a true
    // leading prefix.
    scan_prefix.pop();

    let mut iter = reader
        .scan_prefix_ordered(scan_prefix, .., ScanShape::Probe, order)
        .await
        .map_err(Error::from)?;

    collect_multi_row_ids(&mut iter, "null-prefix scan").await
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use object_store::memory::InMemory;
    use slatedb::IsolationLevel;

    use super::*;
    use crate::store::open::StoreBuilder;

    /// A unique put whose key the caller established absent is staged with
    /// its row id and no read.
    #[tokio::test]
    async fn a_known_absent_entry_stages_without_a_probe() {
        let (db, _) = StoreBuilder::new("known-absent", Arc::new(InMemory::new()))
            .open_writer()
            .await
            .unwrap();
        let tx = db.begin(IsolationLevel::Snapshot).await.unwrap();
        let key = Bytes::from_static(b"absent value");
        let entries = stream::once(async {
            Ok(StagedIndexEntry {
                index_id: 1,
                unique: true,
                key: Bytes::from_static(b"absent value"),
                row_id: 7,
                delete: false,
                building: true,
                known_absent: true,
                known_held: false,
            })
        });

        let staged = stage_index_entry_stream(&tx, stream::empty(), entries, 0)
            .await
            .unwrap();

        assert_eq!(staged.metrics.unique_probes, 0);
        assert_eq!(staged.metrics.known_absent, 1);
        assert_eq!(
            tx.get(&key).await.unwrap().as_deref(),
            Some(&7u64.to_be_bytes()[..])
        );
        tx.rollback();
        db.close().await.unwrap();
    }

    /// Two known-absent claims of one value in a batch still collide.
    #[tokio::test]
    async fn known_absent_claims_still_collide_in_batch() {
        let (db, _) = StoreBuilder::new("known-absent-collision", Arc::new(InMemory::new()))
            .open_writer()
            .await
            .unwrap();
        let tx = db.begin(IsolationLevel::Snapshot).await.unwrap();
        let entries = stream::iter([3u64, 4].map(|row_id| {
            Ok(StagedIndexEntry {
                index_id: 9,
                unique: true,
                key: Bytes::from_static(b"shared value"),
                row_id,
                delete: false,
                building: true,
                known_absent: true,
                known_held: false,
            })
        }));

        let staged = stage_index_entry_stream(&tx, stream::empty(), entries, 0)
            .await
            .unwrap();

        assert_eq!(staged.poisoned, [9]);
        assert_eq!(staged.metrics.unique_probes, 0);
        tx.rollback();
        db.close().await.unwrap();
    }

    /// A normal fact flush carries 250 source keys plus a handful of newly
    /// discovered reference keys. They must all enter the first probe window:
    /// leaving one or two for a second window charges a full remote-read tail
    /// for almost no work.
    #[test]
    fn normal_flush_unique_probes_share_one_window() {
        const FACTS_AND_REFERENCES: usize = 258;
        let concurrency = std::hint::black_box(UNIQUENESS_PROBE_CONCURRENCY);

        assert!(
            concurrency >= FACTS_AND_REFERENCES,
            "{FACTS_AND_REFERENCES} probes exceed the concurrency window of \
             {concurrency}"
        );
    }

    /// A stream larger than the former 1,024-entry group stays one rolling
    /// probe window rather than draining at an artificial boundary.
    #[tokio::test]
    async fn unique_probe_stream_crosses_the_old_group_boundary() {
        const ENTRIES: usize = 1_025;
        let (db, _) = StoreBuilder::new("continuous-probes", Arc::new(InMemory::new()))
            .open_writer()
            .await
            .unwrap();
        let tx = db.begin(IsolationLevel::Snapshot).await.unwrap();
        let entries = stream::iter((0..ENTRIES).map(|position| {
            let row_id = u64::try_from(position).unwrap();
            Ok(StagedIndexEntry {
                index_id: 1,
                unique: true,
                key: Bytes::copy_from_slice(&row_id.to_be_bytes()),
                row_id,
                delete: false,
                building: false,
                known_absent: false,
                known_held: false,
            })
        }));

        let staged = stage_index_entry_stream(&tx, stream::empty(), entries, 0)
            .await
            .unwrap();

        assert_eq!(staged.bytes, u64::try_from(ENTRIES * 16).unwrap());
        tx.rollback();
        db.close().await.unwrap();
    }

    /// A unique deletion removes the entry of the row it names, and only
    /// that row's.
    ///
    /// The store key is the value alone, so a deletion derived for a row
    /// that no longer holds the value would otherwise take the entry from
    /// whichever row does.
    #[tokio::test]
    async fn a_unique_deletion_removes_only_the_named_row_s_entry() {
        let (db, _) = StoreBuilder::new("guarded-deletions", Arc::new(InMemory::new()))
            .open_writer()
            .await
            .unwrap();
        let key = Bytes::from_static(b"one value");

        let tx = db.begin(IsolationLevel::Snapshot).await.unwrap();
        tx.put(key.clone(), 3u64.to_be_bytes()).unwrap();
        tx.commit().await.unwrap();

        let deletion = |row_id: u64, key: Bytes| {
            stream::once(async move {
                Ok(StagedIndexEntry {
                    index_id: 1,
                    unique: true,
                    key,
                    row_id,
                    delete: true,
                    building: false,
                    known_absent: false,
                    known_held: false,
                })
            })
        };

        let tx = db.begin(IsolationLevel::Snapshot).await.unwrap();
        let staged = stage_index_entry_stream(&tx, deletion(0, key.clone()), stream::empty(), 0)
            .await
            .unwrap();
        assert_eq!(staged.metrics.deletions, 0);
        assert_eq!(staged.metrics.guard_reads, 1);
        tx.commit().await.unwrap();

        let tx = db.begin(IsolationLevel::Snapshot).await.unwrap();
        assert!(
            ReadHandle::Tx(&tx)
                .get(key.clone())
                .await
                .unwrap()
                .is_some(),
            "row 3's entry survives a deletion derived for row 0"
        );
        let staged = stage_index_entry_stream(&tx, deletion(3, key.clone()), stream::empty(), 0)
            .await
            .unwrap();
        assert_eq!(staged.metrics.deletions, 1);
        tx.commit().await.unwrap();

        let tx = db.begin(IsolationLevel::Snapshot).await.unwrap();
        assert!(ReadHandle::Tx(&tx).get(key).await.unwrap().is_none());
        tx.rollback();

        db.close().await.unwrap();
    }

    /// A unique deletion the caller established names its own row is staged
    /// without reading the entry.
    #[tokio::test]
    async fn a_known_held_deletion_stages_without_a_guard_read() {
        let (db, _) = StoreBuilder::new("known-held", Arc::new(InMemory::new()))
            .open_writer()
            .await
            .unwrap();
        let key = Bytes::from_static(b"held value");

        let tx = db.begin(IsolationLevel::Snapshot).await.unwrap();
        tx.put(key.clone(), 3u64.to_be_bytes()).unwrap();
        tx.commit().await.unwrap();

        let tx = db.begin(IsolationLevel::Snapshot).await.unwrap();
        let deletion_key = key.clone();
        let deletion = stream::once(async move {
            Ok(StagedIndexEntry {
                index_id: 1,
                unique: true,
                key: deletion_key,
                row_id: 3,
                delete: true,
                building: false,
                known_absent: false,
                known_held: true,
            })
        });
        let staged = stage_index_entry_stream(&tx, deletion, stream::empty(), 0)
            .await
            .unwrap();
        assert_eq!(staged.metrics.deletions, 1);
        assert_eq!(staged.metrics.guard_reads, 0);
        assert!(tx.get(&key).await.unwrap().is_none());
        tx.rollback();

        db.close().await.unwrap();
    }

    /// Addition sources may fill the bounded prefetch while a deletion is
    /// pending, but the transaction must still stage the delete first.
    #[tokio::test]
    async fn addition_prefetch_preserves_the_deletion_phase_boundary() {
        let (db, _) = StoreBuilder::new("addition-prefetch", Arc::new(InMemory::new()))
            .open_writer()
            .await
            .unwrap();
        let key = Bytes::from_static(b"delete-before-add");
        // The entry the deletion removes: a unique deletion is staged only
        // for the row the committed entry names.
        let seed = db.begin(IsolationLevel::Snapshot).await.unwrap();
        seed.put(key.clone(), 7_u64.to_be_bytes()).unwrap();
        seed.commit().await.unwrap();

        let tx = db.begin(IsolationLevel::Snapshot).await.unwrap();
        let (release_deletion, deletion_released) = tokio::sync::oneshot::channel();
        let deletion_key = key.clone();
        let deletions = stream::once(async move {
            deletion_released.await.unwrap();
            Ok(StagedIndexEntry {
                index_id: 1,
                unique: true,
                key: deletion_key,
                row_id: 7,
                delete: true,
                building: false,
                known_absent: false,
                known_held: false,
            })
        });
        let (addition_polled, observed_addition) = tokio::sync::oneshot::channel();
        let addition_key = key.clone();
        let additions = stream::once(async move {
            addition_polled.send(()).unwrap();
            Ok(StagedIndexEntry {
                index_id: 1,
                unique: true,
                key: addition_key,
                row_id: 7,
                delete: false,
                building: false,
                known_absent: false,
                known_held: false,
            })
        });

        {
            let staging = stage_index_entry_stream(&tx, deletions, additions, 0);
            let mut staging = std::pin::pin!(staging);
            tokio::select! {
                _ = &mut staging => panic!("staging finished before deletion release"),
                observed = observed_addition => observed.unwrap(),
            }
            tokio::time::sleep(std::time::Duration::from_millis(25)).await;
            release_deletion.send(()).unwrap();
            let staged = staging.as_mut().await.unwrap();
            assert_eq!(staged.metrics.unique_probes, 1);
            assert_eq!(staged.metrics.probes_completed_during_deletions, 1);
            assert_eq!(staged.metrics.additions, 1);
            assert_eq!(staged.metrics.deletions, 1);
        }

        assert_eq!(
            tx.get(&key).await.unwrap(),
            Some(Bytes::copy_from_slice(&7_u64.to_be_bytes()))
        );
        tx.rollback();
        db.close().await.unwrap();
    }

    /// The oversized-commit refusal must be terminal: DuckLake re-runs a
    /// commit whose error text carries any of four substrings, and
    /// re-running this one only reaches the same answer more slowly.
    #[test]
    fn oversized_commit_refusal_avoids_ducklake_retry_substrings() {
        for text in [
            oversized_commit(1_000_001).to_string(),
            oversized_staged_bytes(64 * 1024 * 1024 + 1).to_string(),
            unique_violation(24, 900, 17, Collided::Committed).to_string(),
            unique_violation(24, 900, 901, Collided::InBatch).to_string(),
        ] {
            for substring in ["conflict", "concurrent", "unique", "primary key"] {
                assert!(
                    !text.contains(substring),
                    "{text:?} contains DuckLake's retry substring {substring:?}"
                );
            }
        }
    }

    /// A rejection names both rows and the state it collided with, so an
    /// entry naming a row that is no longer live is one query away.
    #[test]
    fn unique_violation_names_both_rows_and_the_state_it_collided_with() {
        let committed = unique_violation(24, 900, 17, Collided::Committed).to_string();
        assert!(
            committed.contains("index 24"),
            "names the index: {committed}"
        );
        assert!(
            committed.contains("row 900"),
            "names the claimant: {committed}"
        );
        assert!(
            committed.contains("row 17"),
            "names the holder: {committed}"
        );
        assert!(
            committed.contains("committed"),
            "names the state: {committed}"
        );

        let in_batch = unique_violation(24, 900, 901, Collided::InBatch).to_string();
        assert!(in_batch.contains("900"), "names the claimant: {in_batch}");
        assert!(in_batch.contains("901"), "names the holder: {in_batch}");
        assert!(
            in_batch.contains("one commit"),
            "names the state: {in_batch}"
        );
    }

    /// It names the count, the limit, the memory it would have needed, and
    /// the one thing the caller can do about it.
    #[test]
    fn oversized_commit_refusal_is_actionable() {
        let text = oversized_commit(1_000_001).to_string();
        assert!(text.contains("1000001"), "names the count: {text}");
        assert!(text.contains("1000000"), "names the limit: {text}");
        assert!(text.contains("smaller commits"), "names the remedy: {text}");
    }

    /// The entry ceiling bounds ordinary narrow keys, while the byte
    /// ceiling stops a smaller number of unusually wide keys.
    #[test]
    fn index_commit_budget_enforces_both_dimensions() {
        let mut budget = IndexCommitBudget::default();
        assert!(budget.observe_entry(MAX_INDEX_ENTRIES_PER_COMMIT).is_ok());
        assert!(budget.observe_entry(1).is_err());

        let mut budget = IndexCommitBudget::default();
        assert!(
            budget
                .stage_bytes(MAX_INDEX_STAGED_BYTES_PER_COMMIT)
                .is_ok()
        );
        let error = budget.stage_bytes(1).unwrap_err().to_string();
        assert!(error.contains("67108864"), "names the byte limit: {error}");
        assert!(
            error.contains("smaller commits"),
            "names the remedy: {error}"
        );
    }
}