arrow-tiberius 0.1.0

//! Row-major payload layout for direct raw TDS encoding.

use crate::{Diagnostic, DiagnosticCode, DiagnosticSet, Error, Result};

use super::payload::TDS_ROW_TOKEN;

const ROW_TOKEN_LEN: usize = 1;

/// Absolute byte position for one encoded cell inside an encoded rows payload.
///
/// A cell position belongs to one row and one column, but its `offset` is
/// measured from the start of the whole payload buffer, not from the row start.
/// The byte range for this cell is `offset..offset + len`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct CellPosition {
    /// Zero-based row number for this cell.
    row_index: usize,
    /// Zero-based column number inside the direct encoder mapping order.
    column_index: usize,
    /// Absolute start byte of this encoded cell inside the whole payload.
    ///
    /// This points at the first byte of the cell itself, not at the row token.
    /// For variable-width cells, this first byte is usually part of the length
    /// prefix or PLP marker.
    offset: usize,
    /// Encoded byte length of this cell.
    ///
    /// This includes any cell-local length prefix, null sentinel, PLP metadata,
    /// and value bytes, but does not include the row token.
    len: usize,
}

impl CellPosition {
    /// Creates a cell byte position.
    pub(crate) const fn new(
        row_index: usize,
        column_index: usize,
        offset: usize,
        len: usize,
    ) -> Self {
        Self {
            row_index,
            column_index,
            offset,
            len,
        }
    }

    /// Returns the row index.
    pub(crate) const fn row_index(&self) -> usize {
        self.row_index
    }

    /// Returns the column index.
    pub(crate) const fn column_index(&self) -> usize {
        self.column_index
    }

    /// Returns the absolute start byte of this cell in the whole payload.
    pub(crate) const fn offset(&self) -> usize {
        self.offset
    }

    /// Returns the encoded byte length of this cell.
    pub(crate) const fn len(&self) -> usize {
        self.len
    }

    /// Returns true when the position has no encoded bytes.
    pub(crate) const fn is_empty(&self) -> bool {
        self.len == 0
    }
}

/// Row-major layout metadata for one direct encoded rows payload.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct RowLayout {
    /// Absolute byte offsets where each TDS ROW token (`0xD1`) begins.
    ///
    /// One payload can contain zero, one, or many rows. For example,
    /// `[0, 12, 24]` means the payload has three rows and their row token bytes
    /// start at payload bytes 0, 12, and 24.
    row_token_offsets: Vec<usize>,
    /// Total byte length for each row, including the row token and all cells.
    ///
    /// The validator uses this with `row_token_offsets` to prove rows are
    /// contiguous and cover the whole payload.
    row_lengths: Vec<usize>,
    /// Absolute byte positions for every encoded cell.
    ///
    /// Positions are row-major: all cells for row 0, then all cells for row 1,
    /// and so on. Each position is measured from the start of the whole payload.
    cell_positions: Vec<CellPosition>,
    /// Total byte length of the complete payload buffer.
    payload_len: usize,
}

impl RowLayout {
    /// Creates row-major layout metadata for one complete payload buffer.
    pub(crate) fn new(
        row_token_offsets: Vec<usize>,
        row_lengths: Vec<usize>,
        cell_positions: Vec<CellPosition>,
        payload_len: usize,
    ) -> Result<Self> {
        validate_rows(&row_token_offsets, &row_lengths, payload_len)?;
        validate_cells(&cell_positions, row_token_offsets.len(), payload_len)?;

        Ok(Self {
            row_token_offsets,
            row_lengths,
            cell_positions,
            payload_len,
        })
    }

    /// Returns the number of rows described by this payload layout.
    pub(crate) fn row_count(&self) -> usize {
        self.row_token_offsets.len()
    }

    /// Returns absolute byte offsets where each `0xD1` row token begins.
    pub(crate) fn row_token_offsets(&self) -> &[usize] {
        &self.row_token_offsets
    }

    /// Returns row lengths in bytes, including each row's token byte.
    pub(crate) fn row_lengths(&self) -> &[usize] {
        &self.row_lengths
    }

    /// Returns absolute byte positions for encoded cells in row-major order.
    pub(crate) fn cell_positions(&self) -> &[CellPosition] {
        &self.cell_positions
    }

    /// Returns the total payload buffer length in bytes.
    pub(crate) const fn payload_len(&self) -> usize {
        self.payload_len
    }
}

pub(crate) fn build_fixed_width_row_layout(
    row_count: usize,
    column_count: usize,
    cell_lengths: &[usize],
) -> Result<RowLayout> {
    build_fixed_width_row_range_layout(0, row_count, column_count, cell_lengths)
}

pub(crate) fn build_fixed_width_row_range_layout(
    start_row: usize,
    row_count: usize,
    column_count: usize,
    cell_lengths: &[usize],
) -> Result<RowLayout> {
    let end_row = start_row
        .checked_add(row_count)
        .ok_or_else(|| invalid_layout("direct row range end overflowed usize"))?;
    let mut row_token_offsets = Vec::with_capacity(row_count);
    let mut row_lengths = Vec::with_capacity(row_count);
    let mut cell_positions = Vec::with_capacity(row_count * column_count);
    let mut offset = 0usize;

    for row_index in start_row..end_row {
        let row_offset = offset;
        row_token_offsets.push(row_offset);
        offset = checked_add(offset, ROW_TOKEN_LEN)?;

        for column_index in 0..column_count {
            let cell_len = cell_lengths[row_index * column_count + column_index];
            cell_positions.push(CellPosition::new(
                row_index - start_row,
                column_index,
                offset,
                cell_len,
            ));
            offset = checked_add(offset, cell_len)?;
        }

        // Row length is the byte span from this row's ROW token through the
        // last encoded cell. RowLayout uses it to prove rows are contiguous.
        row_lengths.push(offset - row_offset);
    }

    RowLayout::new(row_token_offsets, row_lengths, cell_positions, offset)
}

/// Allocates a complete payload buffer and writes every row token.
///
/// The measured layout already knows where each row starts inside the payload.
/// This function creates a zero-filled buffer of the final payload size and
/// writes `0xD1` at every absolute row start offset. Later fill steps write
/// encoded cell bytes into the remaining positions.
pub(crate) fn allocate_rows_payload_with_tokens(layout: &RowLayout) -> Vec<u8> {
    let mut bytes = vec![0; layout.payload_len()];

    // One payload can contain many rows. Each row must start with the TDS ROW
    // token byte, and row_token_offsets gives those absolute byte positions.
    for &row_offset in layout.row_token_offsets() {
        bytes[row_offset] = TDS_ROW_TOKEN;
    }

    bytes
}

fn checked_add(lhs: usize, rhs: usize) -> Result<usize> {
    lhs.checked_add(rhs)
        .ok_or_else(|| invalid_layout("direct primitive row layout length overflowed usize"))
}

fn validate_rows(
    row_token_offsets: &[usize],
    row_lengths: &[usize],
    payload_len: usize,
) -> Result<()> {
    if row_token_offsets.len() != row_lengths.len() {
        return Err(invalid_layout(format!(
            "row layout has {} row-token offset(s) but {} row length(s)",
            row_token_offsets.len(),
            row_lengths.len()
        )));
    }

    if row_token_offsets.is_empty() {
        if payload_len == 0 {
            return Ok(());
        }

        return Err(invalid_layout(format!(
            "empty row layout cannot describe non-empty payload length {payload_len}"
        )));
    }

    if row_token_offsets[0] != 0 {
        return Err(invalid_layout(format!(
            "first row token offset must be 0, got {}",
            row_token_offsets[0]
        )));
    }

    for (index, (&offset, &len)) in row_token_offsets.iter().zip(row_lengths).enumerate() {
        let end = offset.checked_add(len).ok_or_else(|| {
            invalid_layout(format!(
                "row {index} offset {offset} plus length {len} overflows usize"
            ))
        })?;

        if end > payload_len {
            return Err(invalid_layout(format!(
                "row {index} ends at {end}, outside payload length {payload_len}"
            )));
        }

        if let Some(&next_offset) = row_token_offsets.get(index + 1)
            && next_offset != end
        {
            return Err(invalid_layout(format!(
                "row {index} ends at {end}, but next row token offset is {next_offset}"
            )));
        }
    }

    let last_end = row_token_offsets[row_token_offsets.len() - 1]
        .checked_add(row_lengths[row_lengths.len() - 1])
        .ok_or_else(|| invalid_layout("last row end overflows usize"))?;

    if last_end != payload_len {
        return Err(invalid_layout(format!(
            "last row ends at {last_end}, but payload length is {payload_len}"
        )));
    }

    Ok(())
}

fn validate_cells(
    cell_positions: &[CellPosition],
    row_count: usize,
    payload_len: usize,
) -> Result<()> {
    for cell in cell_positions {
        if cell.row_index >= row_count {
            return Err(invalid_layout(format!(
                "cell row index {} is outside row count {row_count}",
                cell.row_index
            )));
        }

        let end = cell.offset.checked_add(cell.len).ok_or_else(|| {
            invalid_layout(format!(
                "cell offset {} plus length {} overflows usize",
                cell.offset, cell.len
            ))
        })?;

        if end > payload_len {
            return Err(invalid_layout(format!(
                "cell at row {} column {} ends at {end}, outside payload length {payload_len}",
                cell.row_index, cell.column_index
            )));
        }
    }

    Ok(())
}

fn invalid_layout(message: impl Into<String>) -> Error {
    Error::DirectEncoding {
        diagnostics: DiagnosticSet::from(vec![Diagnostic::error(
            DiagnosticCode::DirectEncodingInvalidPayload,
            message,
        )]),
    }
}

#[cfg(test)]
mod tests {
    use crate::{DiagnosticCode, Error};

    use super::{CellPosition, RowLayout};

    #[test]
    fn accepts_empty_layout() {
        let layout =
            RowLayout::new(Vec::new(), Vec::new(), Vec::new(), 0).expect("empty layout is valid");

        assert_eq!(layout.row_count(), 0);
        assert_eq!(layout.row_token_offsets(), []);
        assert_eq!(layout.row_lengths(), []);
        assert_eq!(layout.cell_positions(), []);
        assert_eq!(layout.payload_len(), 0);
    }

    #[test]
    fn accepts_contiguous_multi_row_layout_with_cells() {
        let cells = vec![CellPosition::new(0, 0, 1, 4), CellPosition::new(1, 0, 7, 4)];
        let layout = RowLayout::new(vec![0, 6], vec![6, 6], cells.clone(), 12)
            .expect("contiguous layout is valid");

        assert_eq!(layout.row_count(), 2);
        assert_eq!(layout.row_token_offsets(), [0, 6]);
        assert_eq!(layout.row_lengths(), [6, 6]);
        assert_eq!(layout.cell_positions(), cells);
        assert_eq!(layout.cell_positions()[0].row_index(), 0);
        assert_eq!(layout.cell_positions()[0].column_index(), 0);
        assert_eq!(layout.cell_positions()[0].offset(), 1);
        assert_eq!(layout.cell_positions()[0].len(), 4);
        assert!(!layout.cell_positions()[0].is_empty());
    }

    #[test]
    fn rejects_row_count_mismatch() {
        let err = RowLayout::new(vec![0], Vec::new(), Vec::new(), 1)
            .expect_err("row offsets and lengths must match");

        assert_invalid_layout(err);
    }

    #[test]
    fn rejects_non_zero_first_row_offset() {
        let err = RowLayout::new(vec![1], vec![1], Vec::new(), 2)
            .expect_err("first row must start at zero");

        assert_invalid_layout(err);
    }

    #[test]
    fn rejects_gaps_between_rows() {
        let err = RowLayout::new(vec![0, 3], vec![2, 1], Vec::new(), 4)
            .expect_err("rows must be contiguous");

        assert_invalid_layout(err);
    }

    #[test]
    fn rejects_layout_that_does_not_cover_payload() {
        let err = RowLayout::new(vec![0], vec![1], Vec::new(), 2)
            .expect_err("layout must cover payload exactly");

        assert_invalid_layout(err);
    }

    #[test]
    fn rejects_cell_outside_row_count() {
        let err = RowLayout::new(vec![0], vec![1], vec![CellPosition::new(1, 0, 0, 1)], 1)
            .expect_err("cell row must exist");

        assert_invalid_layout(err);
    }

    #[test]
    fn rejects_cell_outside_payload() {
        let err = RowLayout::new(vec![0], vec![1], vec![CellPosition::new(0, 0, 0, 2)], 1)
            .expect_err("cell must fit payload");

        assert_invalid_layout(err);
    }

    fn assert_invalid_layout(err: Error) {
        let Error::DirectEncoding { diagnostics } = err else {
            panic!("expected direct encoding error");
        };

        assert_eq!(diagnostics.len(), 1);
        assert_eq!(
            diagnostics.all()[0].code(),
            DiagnosticCode::DirectEncodingInvalidPayload
        );
    }
}