aedb 0.3.1

//! Bounded primary-key-ordered scans: `ORDER BY pk ASC|DESC LIMIT n` (no
//! predicate) walks the hot PK map directly in key order — forward or reverse
//! — and stops at the page window, instead of materializing the whole table
//! and top-k sorting it. Cursor pages seek straight to the continuation
//! point. Shares its order-safety helpers with the streaming join fast path.

use crate::catalog::schema::TableSchema;
use crate::catalog::types::{ColumnType, Row, Value};
use crate::query::error::QueryError;
use crate::query::plan::{Order, Query};
use crate::storage::encoded_key::EncodedKey;
use crate::storage::keyspace::{KeyspaceSnapshot, TableData};
use std::ops::Bound;

use super::CursorToken;
use super::pagination::PageWindow;

/// Encoded primary keys sort identically to `value_compare` for these column
/// types, so walking the PK map in encoded order satisfies an `ORDER BY pk`.
/// Float PKs are excluded (NaN ordering differs), as are types whose value
/// comparison is not plain byte order.
pub(super) fn pk_encoded_order_matches_value_order(col_type: Option<&ColumnType>) -> bool {
    matches!(
        col_type,
        Some(
            ColumnType::Integer
                | ColumnType::U8
                | ColumnType::U64
                | ColumnType::Timestamp
                | ColumnType::Text
        )
    )
}

/// Whether a cursor-carried sort-key value is exactly of the base PK's column
/// type, making `EncodedKey` order and `Value::cmp` order interchangeable for
/// a cursor seek. Only types admitted by
/// [`pk_encoded_order_matches_value_order`] appear here.
pub(super) fn cursor_value_matches_pk_type(value: &Value, col_type: &ColumnType) -> bool {
    matches!(
        (value, col_type),
        (Value::Integer(_), ColumnType::Integer)
            | (Value::U8(_), ColumnType::U8)
            | (Value::U64(_), ColumnType::U64)
            | (Value::Timestamp(_), ColumnType::Timestamp)
            | (Value::Text(_), ColumnType::Text)
    )
}

pub(super) struct PkOrderedScanRequest<'a> {
    pub snapshot: &'a KeyspaceSnapshot,
    pub schema: &'a TableSchema,
    pub query: &'a Query,
    pub table: Option<&'a TableData>,
    pub cursor_state: &'a Option<CursorToken>,
    pub page_window: &'a PageWindow,
    pub max_scan_rows: usize,
}

pub(super) struct PkOrderedScan {
    pub rows: Vec<Row>,
}

/// Bounded PK-ordered row source for `ORDER BY pk ASC|DESC LIMIT n` with no
/// predicate. Returns `Ok(None)` — falling back to the full-scan + top-k sort
/// path — unless the shape guarantees the bounded walk equals full
/// evaluation:
///
/// - single-column primary key of an order-preserving type, named exactly by
///   the only ORDER BY term;
/// - an explicit LIMIT ≥ 1 (without one, `estimated_rows` must reflect the
///   true table size for the `ScanBoundExceeded` guard; LIMIT 0 keeps the
///   general path's cursor-emission quirk);
/// - no GROUP BY / aggregates / HAVING / DISTINCT (each needs every row);
/// - a fully-resident table (cold segments or tombstones make the hot map
///   incomplete);
/// - cursors only with OFFSET 0 and a single sort-key value of exactly the
///   PK's type: the seek then starts strictly past the last returned row, so
///   every walked row is admitted and the bounded window is exact.
///
/// OFFSET is applied in-walk (skipped rows are never materialized); the
/// caller must treat the source as offset-applied and order-satisfying.
pub(super) fn pk_ordered_scan_for_query(
    request: PkOrderedScanRequest<'_>,
) -> Result<Option<PkOrderedScan>, QueryError> {
    let query = request.query;
    if query.predicate.is_some()
        || query.distinct
        || !query.group_by.is_empty()
        || !query.aggregates.is_empty()
        || query.having.is_some()
        || query.limit.is_none_or(|limit| limit == 0)
    {
        return Ok(None);
    }
    let [(order_col, order)] = query.order_by.as_slice() else {
        return Ok(None);
    };
    if request.schema.primary_key.len() != 1 || request.schema.primary_key[0] != *order_col {
        return Ok(None);
    }
    let pk_type = request
        .schema
        .columns
        .iter()
        .find(|c| c.name == *order_col)
        .map(|c| &c.col_type);
    if !pk_encoded_order_matches_value_order(pk_type) {
        return Ok(None);
    }
    let Some(table) = request.table else {
        // No table data in the snapshot: the full scan would see no rows.
        return Ok(Some(PkOrderedScan { rows: Vec::new() }));
    };
    if !table.row_segments.is_empty() || !table.row_tombstones.is_empty() {
        return Ok(None);
    }

    let reverse = matches!(order, Order::Desc);
    let bounds: (Bound<EncodedKey>, Bound<EncodedKey>) = match request.cursor_state {
        None => (Bound::Unbounded, Bound::Unbounded),
        Some(cursor) => {
            // Cursor pages require full evaluation, whose ScanBoundExceeded
            // guard compares the whole table against max_scan_rows — a
            // cursor must not let pagination bypass the bound. Replicate the
            // general path's rejection before walking anything.
            if table.rows.len() > request.max_scan_rows {
                return Err(QueryError::ScanBoundExceeded {
                    estimated_rows: table.rows.len() as u64,
                    max_scan_rows: request.max_scan_rows as u64,
                });
            }
            if request.page_window.row_offset_count > 0
                || cursor.last_sort_key.len() != 1
                || !cursor_value_matches_pk_type(
                    &cursor.last_sort_key[0],
                    pk_type.expect("checked by order gate"),
                )
            {
                return Ok(None);
            }
            let key = EncodedKey::from_values(&cursor.last_sort_key);
            if reverse {
                (Bound::Unbounded, Bound::Excluded(key))
            } else {
                (Bound::Excluded(key), Bound::Unbounded)
            }
        }
    };
    // A cursor already positions the walk past everything returned, so the
    // page offset only applies to cursorless first pages.
    let offset = match request.cursor_state {
        Some(_) => 0,
        None => request.page_window.row_offset_count,
    };
    let limit = request.page_window.page_read_limit;

    let mut rows = Vec::with_capacity(limit.min(1024));
    let range = table.rows.range(bounds);
    let stored_rows: Box<dyn Iterator<Item = &crate::storage::keyspace::StoredRow>> = if reverse {
        Box::new(range.rev().map(|(_, stored)| stored))
    } else {
        Box::new(range.map(|(_, stored)| stored))
    };
    let mut skipped = 0usize;
    for stored in stored_rows {
        if rows.len() >= limit {
            break;
        }
        if skipped < offset {
            skipped += 1;
            continue;
        }
        rows.push(request.snapshot.materialize_row(stored)?.into_owned());
    }
    Ok(Some(PkOrderedScan { rows }))
}