runmat-runtime 0.5.5

//! Matrix indexing and slicing operations
//!
//! Implements language-style tensor indexing and access patterns.

use crate::builtins::common::shape::normalize_scalar_shape;
use crate::{build_runtime_error, RuntimeError};
use runmat_builtins::{Tensor, Value};

fn indexing_error(message: impl Into<String>) -> RuntimeError {
    build_runtime_error(message).build()
}

fn indexing_error_with_identifier(message: impl Into<String>, identifier: &str) -> RuntimeError {
    build_runtime_error(message)
        .with_identifier(identifier)
        .build()
}

fn positive_integer_cell_index(value: f64, identifier: &str) -> Result<usize, RuntimeError> {
    if !value.is_finite() || value < 1.0 || value.fract() != 0.0 {
        return Err(indexing_error_with_identifier(
            format!("Cell index {value} must be a positive integer"),
            identifier,
        ));
    }
    Ok(value as usize)
}

fn cell_row_major_pos_from_linear(
    ca: &runmat_builtins::CellArray,
    idx: usize,
) -> Result<usize, RuntimeError> {
    if idx == 0 || idx > ca.data.len() {
        return Err(indexing_error_with_identifier(
            format!("Cell index {} out of bounds (1 to {})", idx, ca.data.len()),
            "RunMat:CellIndexOutOfBounds",
        ));
    }
    if ca.rows <= 1 || ca.cols <= 1 {
        return Ok(idx - 1);
    }
    let zero = idx - 1;
    let row = zero % ca.rows;
    let col = zero / ca.rows;
    Ok(row * ca.cols + col)
}

/// Get a single element from a tensor (1-based indexing like language)
pub fn matrix_get_element(tensor: &Tensor, row: usize, col: usize) -> Result<f64, RuntimeError> {
    if row == 0 || col == 0 {
        return Err(indexing_error_with_identifier(
            "MATLAB uses 1-based indexing",
            "RunMat:IndexOutOfBounds",
        ));
    }
    tensor
        .get2(row - 1, col - 1)
        .map_err(|err| indexing_error_with_identifier(err, "RunMat:IndexOutOfBounds"))
}

/// Set a single element in a tensor (1-based indexing like language)
pub fn matrix_set_element(
    tensor: &mut Tensor,
    row: usize,
    col: usize,
    value: f64,
) -> Result<(), RuntimeError> {
    if row == 0 || col == 0 {
        return Err(indexing_error_with_identifier(
            "The MATLAB language uses 1-based indexing",
            "RunMat:IndexOutOfBounds",
        ));
    }
    tensor
        .set2(row - 1, col - 1, value)
        .map_err(|err| indexing_error_with_identifier(err, "RunMat:IndexOutOfBounds"))
}

/// Get a row from a tensor
pub fn matrix_get_row(tensor: &Tensor, row: usize) -> Result<Tensor, RuntimeError> {
    if row == 0 || row > tensor.rows() {
        return Err(indexing_error_with_identifier(
            format!(
                "Row index {} out of bounds for {}x{} tensor",
                row,
                tensor.rows(),
                tensor.cols()
            ),
            "RunMat:IndexOutOfBounds",
        ));
    }

    // Column-major: row slice picks every element spaced by rows across columns
    let mut row_data = Vec::with_capacity(tensor.cols());
    for c in 0..tensor.cols() {
        row_data.push(tensor.data[(row - 1) + c * tensor.rows()]);
    }
    Tensor::new_2d(row_data, 1, tensor.cols()).map_err(|err| indexing_error(err))
}

/// Get a column from a tensor
pub fn matrix_get_col(tensor: &Tensor, col: usize) -> Result<Tensor, RuntimeError> {
    if col == 0 || col > tensor.cols() {
        return Err(indexing_error_with_identifier(
            format!(
                "Column index {} out of bounds for {}x{} tensor",
                col,
                tensor.rows(),
                tensor.cols()
            ),
            "RunMat:IndexOutOfBounds",
        ));
    }

    let mut col_data = Vec::with_capacity(tensor.rows());
    for row in 0..tensor.rows() {
        col_data.push(tensor.data[row + (col - 1) * tensor.rows()]);
    }
    Tensor::new_2d(col_data, tensor.rows(), 1).map_err(|err| indexing_error(err))
}

/// Array indexing operation (used by all interpreters/compilers)
/// In MATLAB, indexing is 1-based and supports:
/// - Single element: A(i) for vectors, A(i,j) for tensors
/// - Multiple indices: A(i1, i2, ..., iN)
pub async fn perform_indexing(base: &Value, indices: &[f64]) -> Result<Value, RuntimeError> {
    match base {
        Value::GpuTensor(h) => {
            let provider = runmat_accelerate_api::provider().ok_or_else(|| {
                indexing_error("Cannot index value of type GpuTensor without a provider")
            })?;
            if indices.is_empty() {
                return Err(indexing_error("At least one index is required"));
            }
            // Support scalar indexing cases mirroring Tensor branch
            if indices.len() == 1 {
                let idx = indices[0] as usize;
                let total = h.shape.iter().product();
                if idx < 1 || idx > total {
                    return Err(indexing_error_with_identifier(
                        format!("Index {} out of bounds (1 to {})", idx, total),
                        "RunMat:IndexOutOfBounds",
                    ));
                }
                let lin0 = idx - 1; // 0-based
                let val = gpu_index_scalar(provider, h, lin0).await?;
                return Ok(Value::Num(val));
            } else if indices.len() == 2 {
                let row = indices[0] as usize;
                let col = indices[1] as usize;
                let rows = h.shape.first().copied().unwrap_or(1);
                let cols = h.shape.get(1).copied().unwrap_or(1);
                if row < 1 || row > rows || col < 1 || col > cols {
                    return Err(indexing_error_with_identifier(
                        format!("Index ({row}, {col}) out of bounds for {rows}x{cols} tensor"),
                        "RunMat:IndexOutOfBounds",
                    ));
                }
                let lin0 = (row - 1) + (col - 1) * rows;
                let val = gpu_index_scalar(provider, h, lin0).await?;
                return Ok(Value::Num(val));
            }
            Err(indexing_error_with_identifier(
                format!("Cannot index value of type {base:?}"),
                "RunMat:SliceNonTensor",
            ))
        }
        Value::Tensor(tensor) => {
            if indices.is_empty() {
                return Err(indexing_error("At least one index is required"));
            }

            if indices.len() == 1 {
                // Linear indexing (1-based)
                let idx = indices[0] as usize;
                if idx < 1 || idx > tensor.data.len() {
                    return Err(indexing_error_with_identifier(
                        format!("Index {} out of bounds (1 to {})", idx, tensor.data.len()),
                        "RunMat:IndexOutOfBounds",
                    ));
                }
                Ok(Value::Num(tensor.data[idx - 1])) // Convert to 0-based
            } else if indices.len() == 2 {
                // Row-column indexing (1-based)
                let row = indices[0] as usize;
                let col = indices[1] as usize;
                let shape = normalize_scalar_shape(&tensor.shape);
                let rows = shape.first().copied().unwrap_or(1);
                let cols = shape.get(1).copied().unwrap_or(1);

                if row < 1 || row > rows {
                    return Err(indexing_error_with_identifier(
                        format!("Row index {} out of bounds (1 to {})", row, rows),
                        "RunMat:IndexOutOfBounds",
                    ));
                }
                if col < 1 || col > cols {
                    return Err(indexing_error_with_identifier(
                        format!("Column index {} out of bounds (1 to {})", col, cols),
                        "RunMat:IndexOutOfBounds",
                    ));
                }

                let linear_idx = (row - 1) + (col - 1) * rows; // Convert to 0-based, column-major
                Ok(Value::Num(tensor.data[linear_idx]))
            } else {
                Err(indexing_error(format!(
                    "Tensors support 1 or 2 indices, got {}",
                    indices.len()
                )))
            }
        }
        Value::ComplexTensor(tensor) => {
            if indices.is_empty() {
                return Err(indexing_error("At least one index is required"));
            }

            if indices.len() == 1 {
                let idx = indices[0] as usize;
                if idx < 1 || idx > tensor.data.len() {
                    return Err(indexing_error_with_identifier(
                        format!("Index {} out of bounds (1 to {})", idx, tensor.data.len()),
                        "RunMat:IndexOutOfBounds",
                    ));
                }
                let (re, im) = tensor.data[idx - 1];
                Ok(Value::Complex(re, im))
            } else if indices.len() == 2 {
                let row = indices[0] as usize;
                let col = indices[1] as usize;
                let shape = normalize_scalar_shape(&tensor.shape);
                let rows = shape.first().copied().unwrap_or(1);
                let cols = shape.get(1).copied().unwrap_or(1);

                if row < 1 || row > rows {
                    return Err(indexing_error_with_identifier(
                        format!("Row index {} out of bounds (1 to {})", row, rows),
                        "RunMat:IndexOutOfBounds",
                    ));
                }
                if col < 1 || col > cols {
                    return Err(indexing_error_with_identifier(
                        format!("Column index {} out of bounds (1 to {})", col, cols),
                        "RunMat:IndexOutOfBounds",
                    ));
                }

                let linear_idx = (row - 1) + (col - 1) * rows;
                let (re, im) = tensor.data[linear_idx];
                Ok(Value::Complex(re, im))
            } else {
                Err(indexing_error(format!(
                    "Complex tensors support 1 or 2 indices, got {}",
                    indices.len()
                )))
            }
        }
        Value::StringArray(sa) => {
            if indices.is_empty() {
                return Err(indexing_error("At least one index is required"));
            }
            if indices.len() == 1 {
                let idx = indices[0] as usize;
                let total = sa.data.len();
                if idx < 1 || idx > total {
                    return Err(indexing_error_with_identifier(
                        format!("Index {idx} out of bounds (1 to {total})"),
                        "RunMat:IndexOutOfBounds",
                    ));
                }
                Ok(Value::String(sa.data[idx - 1].clone()))
            } else if indices.len() == 2 {
                let row = indices[0] as usize;
                let col = indices[1] as usize;
                let shape = normalize_scalar_shape(&sa.shape);
                let rows = shape.first().copied().unwrap_or(1);
                let cols = shape.get(1).copied().unwrap_or(1);
                if row < 1 || row > rows || col < 1 || col > cols {
                    return Err(indexing_error_with_identifier(
                        "StringArray subscript out of bounds",
                        "RunMat:IndexOutOfBounds",
                    ));
                }
                let idx = (row - 1) + (col - 1) * rows;
                Ok(Value::String(sa.data[idx].clone()))
            } else {
                Err(indexing_error(format!(
                    "StringArray supports 1 or 2 indices, got {}",
                    indices.len()
                )))
            }
        }
        Value::Num(_) | Value::Int(_) => {
            if indices.len() == 1 && indices[0] == 1.0 {
                // Scalar indexing with A(1) returns the scalar itself
                Ok(base.clone())
            } else {
                Err(indexing_error_with_identifier(
                    "Slicing only supported on tensors",
                    "RunMat:SliceNonTensor",
                ))
            }
        }
        Value::Cell(ca) => {
            if indices.is_empty() {
                return Err(indexing_error("At least one index is required"));
            }
            if indices.len() == 1 {
                let idx = positive_integer_cell_index(indices[0], "RunMat:CellIndexOutOfBounds")?;
                if idx < 1 || idx > ca.data.len() {
                    return Err(indexing_error_with_identifier(
                        format!("Cell index {} out of bounds (1 to {})", idx, ca.data.len()),
                        "RunMat:CellIndexOutOfBounds",
                    ));
                }
                let pos = cell_row_major_pos_from_linear(ca, idx)?;
                Ok(ca.data[pos].clone())
            } else if indices.len() == 2 {
                let row =
                    positive_integer_cell_index(indices[0], "RunMat:CellSubscriptOutOfBounds")?;
                let col =
                    positive_integer_cell_index(indices[1], "RunMat:CellSubscriptOutOfBounds")?;
                if row < 1 || row > ca.rows || col < 1 || col > ca.cols {
                    return Err(indexing_error_with_identifier(
                        "Cell subscript out of bounds",
                        "RunMat:CellSubscriptOutOfBounds",
                    ));
                }
                Ok(ca.data[(row - 1) * ca.cols + (col - 1)].clone())
            } else {
                Err(indexing_error(format!(
                    "Cell arrays support 1 or 2 indices, got {}",
                    indices.len()
                )))
            }
        }
        Value::Struct(_) => {
            if matches!(indices, [1.0] | [1.0, 1.0]) {
                Ok(base.clone())
            } else {
                Err(indexing_error_with_identifier(
                    "Struct subscript out of bounds",
                    "RunMat:IndexOutOfBounds",
                ))
            }
        }
        _ => Err(indexing_error_with_identifier(
            format!("Cannot index value of type {base:?}"),
            "RunMat:SliceNonTensor",
        )),
    }
}

async fn gpu_index_scalar(
    provider: &dyn runmat_accelerate_api::AccelProvider,
    handle: &runmat_accelerate_api::GpuTensorHandle,
    lin0: usize,
) -> Result<f64, RuntimeError> {
    #[cfg(target_arch = "wasm32")]
    {
        let host = provider
            .download(handle)
            .await
            .map_err(|e| indexing_error(format!("gpu index: {e}")))?;
        if lin0 >= host.data.len() {
            return Err(indexing_error(format!(
                "gpu index: index {} out of bounds (len {})",
                lin0 + 1,
                host.data.len()
            )));
        }
        Ok(host.data[lin0])
    }
    #[cfg(not(target_arch = "wasm32"))]
    {
        provider
            .read_scalar(handle, lin0)
            .map_err(|e| indexing_error(format!("gpu index: {e}")))
    }
}

#[cfg(test)]
mod tests {
    use super::perform_indexing;
    use futures::executor::block_on;
    use runmat_builtins::{CellArray, Value};

    #[test]
    fn cell_index_rejects_fractional_before_cast() {
        let cell = CellArray::new(
            vec![
                Value::Num(1.0),
                Value::Num(2.0),
                Value::Num(3.0),
                Value::Num(4.0),
            ],
            1,
            4,
        )
        .expect("cell");
        let err = block_on(perform_indexing(&Value::Cell(cell), &[3.7]))
            .expect_err("fractional cell index should fail");
        assert_eq!(err.identifier(), Some("RunMat:CellIndexOutOfBounds"));
    }

    #[test]
    fn cell_subscript_rejects_nan_before_cast() {
        let cell = CellArray::new(vec![Value::Num(1.0)], 1, 1).expect("cell");
        let err = block_on(perform_indexing(&Value::Cell(cell), &[f64::NAN, 1.0]))
            .expect_err("NaN cell subscript should fail");
        assert_eq!(err.identifier(), Some("RunMat:CellSubscriptOutOfBounds"));
    }

    #[test]
    fn cell_linear_index_uses_column_major_semantics_for_2d_cells() {
        let cell = CellArray::new(
            vec![
                Value::String("r1c1".to_string()),
                Value::String("r1c2".to_string()),
                Value::String("r2c1".to_string()),
                Value::String("r2c2".to_string()),
            ],
            2,
            2,
        )
        .expect("cell");
        let value = block_on(perform_indexing(&Value::Cell(cell), &[2.0])).expect("cell read");
        assert_eq!(value, Value::String("r2c1".to_string()));
    }
}