runmat_runtime/builtins/array/sorting_sets/

ismembertol.rs

//! MATLAB-compatible `ismembertol` builtin for real numeric arrays.

use runmat_builtins::{
    BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
    BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
    CellArray, LiteralValue, LogicalArray, NumericDType, ResolveContext, Tensor, Type, Value,
};
use runmat_macros::runtime_builtin;

use super::type_resolvers::logical_output_type;
use crate::builtins::common::{
    gpu_helpers,
    spec::{
        BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
        ReductionNaN, ResidencyPolicy, ShapeRequirements,
    },
    tensor,
};
use crate::{build_runtime_error, BuiltinResult, RuntimeError};

const NAME: &str = "ismembertol";
const DEFAULT_TOL_DOUBLE: f64 = 1.0e-12;
const DEFAULT_TOL_SINGLE: f64 = 1.0e-6;

#[runmat_macros::register_gpu_spec(
    builtin_path = "crate::builtins::array::sorting_sets::ismembertol"
)]
pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
    name: NAME,
    op_kind: GpuOpKind::Custom("ismembertol"),
    supported_precisions: &[],
    broadcast: BroadcastSemantics::None,
    provider_hooks: &[],
    constant_strategy: ConstantStrategy::InlineLiteral,
    residency: ResidencyPolicy::GatherImmediately,
    nan_mode: ReductionNaN::Include,
    two_pass_threshold: None,
    workgroup_size: None,
    accepts_nan_mode: false,
    notes: "Tolerance membership is host materialised today; gpuArray inputs are gathered before comparison.",
};

#[runmat_macros::register_fusion_spec(
    builtin_path = "crate::builtins::array::sorting_sets::ismembertol"
)]
pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
    name: NAME,
    shape: ShapeRequirements::Any,
    constant_strategy: ConstantStrategy::InlineLiteral,
    elementwise: None,
    reduction: None,
    emits_nan: false,
    notes: "`ismembertol` materialises logical and index outputs and terminates fusion chains.",
};

const OUTPUT_MASK: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
    name: "LIA",
    ty: BuiltinParamType::LogicalArray,
    arity: BuiltinParamArity::Required,
    default: None,
    description: "Logical mask over A.",
}];

const OUTPUT_MASK_LOC: [BuiltinParamDescriptor; 2] = [
    BuiltinParamDescriptor {
        name: "LIA",
        ty: BuiltinParamType::LogicalArray,
        arity: BuiltinParamArity::Required,
        default: None,
        description: "Logical mask over A.",
    },
    BuiltinParamDescriptor {
        name: "LocB",
        ty: BuiltinParamType::Any,
        arity: BuiltinParamArity::Required,
        default: None,
        description: "First or all matching indices into B.",
    },
];

const INPUTS_AB: [BuiltinParamDescriptor; 2] = [
    BuiltinParamDescriptor {
        name: "A",
        ty: BuiltinParamType::NumericArray,
        arity: BuiltinParamArity::Required,
        default: None,
        description: "Query values or rows.",
    },
    BuiltinParamDescriptor {
        name: "B",
        ty: BuiltinParamType::NumericArray,
        arity: BuiltinParamArity::Required,
        default: None,
        description: "Reference values or rows.",
    },
];

const INPUTS_AB_OPTIONS: [BuiltinParamDescriptor; 4] = [
    BuiltinParamDescriptor {
        name: "A",
        ty: BuiltinParamType::NumericArray,
        arity: BuiltinParamArity::Required,
        default: None,
        description: "Query values or rows.",
    },
    BuiltinParamDescriptor {
        name: "B",
        ty: BuiltinParamType::NumericArray,
        arity: BuiltinParamArity::Required,
        default: None,
        description: "Reference values or rows.",
    },
    BuiltinParamDescriptor {
        name: "tol",
        ty: BuiltinParamType::NumericScalar,
        arity: BuiltinParamArity::Optional,
        default: Some("1e-12 for double, 1e-6 for single"),
        description: "Relative tolerance.",
    },
    BuiltinParamDescriptor {
        name: "Name,Value",
        ty: BuiltinParamType::Any,
        arity: BuiltinParamArity::Variadic,
        default: None,
        description: "Name-value options: ByRows, DataScale, OutputAllIndices.",
    },
];

const SIGNATURES: [BuiltinSignatureDescriptor; 4] = [
    BuiltinSignatureDescriptor {
        label: "LIA = ismembertol(A, B)",
        inputs: &INPUTS_AB,
        outputs: &OUTPUT_MASK,
    },
    BuiltinSignatureDescriptor {
        label: "LIA = ismembertol(A, B, tol, Name, Value)",
        inputs: &INPUTS_AB_OPTIONS,
        outputs: &OUTPUT_MASK,
    },
    BuiltinSignatureDescriptor {
        label: "[LIA, LocB] = ismembertol(A, B)",
        inputs: &INPUTS_AB,
        outputs: &OUTPUT_MASK_LOC,
    },
    BuiltinSignatureDescriptor {
        label: "[LIA, LocB] = ismembertol(A, B, tol, Name, Value)",
        inputs: &INPUTS_AB_OPTIONS,
        outputs: &OUTPUT_MASK_LOC,
    },
];

const ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
    code: "RM.ISMEMBERTOL.INVALID_INPUT",
    identifier: Some("RunMat:ismembertol:InvalidInput"),
    when: "A or B is not a supported real full numeric input.",
    message: "ismembertol: inputs must be real full numeric arrays",
};

const ERROR_INVALID_ARGUMENT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
    code: "RM.ISMEMBERTOL.INVALID_ARGUMENT",
    identifier: Some("RunMat:ismembertol:InvalidArgument"),
    when: "Tolerance or name-value arguments are malformed.",
    message: "ismembertol: invalid argument",
};

const ERROR_ROWS_COLUMN_MISMATCH: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
    code: "RM.ISMEMBERTOL.ROWS_COLUMN_MISMATCH",
    identifier: Some("RunMat:ismembertol:RowsColumnMismatch"),
    when: "ByRows is true and A/B column counts differ.",
    message: "ismembertol: inputs must have the same number of columns when ByRows is true",
};

const ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
    code: "RM.ISMEMBERTOL.INTERNAL",
    identifier: Some("RunMat:ismembertol:Internal"),
    when: "Internal conversion or allocation fails.",
    message: "ismembertol: internal error",
};

const ERRORS: [BuiltinErrorDescriptor; 4] = [
    ERROR_INVALID_INPUT,
    ERROR_INVALID_ARGUMENT,
    ERROR_ROWS_COLUMN_MISMATCH,
    ERROR_INTERNAL,
];

pub const ISMEMBERTOL_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
    signatures: &SIGNATURES,
    output_mode: BuiltinOutputMode::ByRequestedOutputCount,
    completion_policy: BuiltinCompletionPolicy::Public,
    errors: &ERRORS,
};

fn ismembertol_output_type(args: &[Type], ctx: &ResolveContext) -> Type {
    if literal_by_rows_enabled(ctx) {
        return match args.first().and_then(type_row_count) {
            Some(rows) => Type::Logical {
                shape: Some(vec![rows, Some(1)]),
            },
            None => Type::logical(),
        };
    }
    logical_output_type(args, ctx)
}

fn literal_by_rows_enabled(ctx: &ResolveContext) -> bool {
    let mut idx = 2usize;
    if !matches!(ctx.literal_args.get(idx), Some(LiteralValue::String(_))) {
        idx += 1;
    }
    while idx + 1 < ctx.literal_args.len() {
        if matches!(
            ctx.literal_args.get(idx),
            Some(LiteralValue::String(name)) if name.eq_ignore_ascii_case("ByRows")
        ) {
            return literal_truthy(ctx.literal_args.get(idx + 1));
        }
        idx += 2;
    }
    false
}

fn literal_truthy(value: Option<&LiteralValue>) -> bool {
    match value {
        Some(LiteralValue::Bool(flag)) => *flag,
        Some(LiteralValue::Number(num)) => *num == 1.0,
        _ => false,
    }
}

fn type_row_count(ty: &Type) -> Option<Option<usize>> {
    match ty {
        Type::Tensor { shape: Some(shape) } | Type::Logical { shape: Some(shape) } => {
            if shape.is_empty() {
                Some(Some(1))
            } else {
                Some(shape[0])
            }
        }
        Type::Tensor { shape: None } | Type::Logical { shape: None } => Some(None),
        Type::Num | Type::Int | Type::Bool => Some(Some(1)),
        _ => None,
    }
}

#[runtime_builtin(
    name = "ismembertol",
    category = "array/sorting_sets",
    summary = "Identify numeric array elements or rows that are within tolerance of another array.",
    keywords = "ismembertol,membership,tolerance,set,rows,indices",
    accel = "array_construct",
    sink = true,
    type_resolver(ismembertol_output_type),
    descriptor(crate::builtins::array::sorting_sets::ismembertol::ISMEMBERTOL_DESCRIPTOR),
    builtin_path = "crate::builtins::array::sorting_sets::ismembertol"
)]
async fn ismembertol_builtin(a: Value, b: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
    if let Some(out_count) = crate::output_count::current_output_count() {
        if out_count == 0 {
            return Ok(Value::OutputList(Vec::new()));
        }
        let opts = parse_options(&rest)?;
        if out_count == 1 {
            let eval = evaluate_with_options(a, b, opts, LocRequest::None).await?;
            return Ok(Value::OutputList(vec![eval.into_mask_value()]));
        }
        let loc_request = if opts.output_all_indices {
            LocRequest::All
        } else {
            LocRequest::First
        };
        let eval = evaluate_with_options(a, b, opts, loc_request).await?;
        let (mask, loc) = eval.into_pair();
        return Ok(crate::output_count::output_list_with_padding(
            out_count,
            vec![mask, loc],
        ));
    }
    let opts = parse_options(&rest)?;
    let eval = evaluate_with_options(a, b, opts, LocRequest::None).await?;
    Ok(eval.into_mask_value())
}

pub async fn evaluate(a: Value, b: Value, rest: &[Value]) -> BuiltinResult<IsMemberTolEvaluation> {
    let opts = parse_options(rest)?;
    let loc_request = if opts.output_all_indices {
        LocRequest::All
    } else {
        LocRequest::First
    };
    evaluate_with_options(a, b, opts, loc_request).await
}

async fn evaluate_with_options(
    a: Value,
    b: Value,
    opts: IsMemberTolOptions,
    loc_request: LocRequest,
) -> BuiltinResult<IsMemberTolEvaluation> {
    let a = gather_if_gpu(a).await?;
    let b = gather_if_gpu(b).await?;
    let tensor_a =
        tensor::value_into_tensor_for(NAME, a).map_err(|_| error(&ERROR_INVALID_INPUT))?;
    let tensor_b =
        tensor::value_into_tensor_for(NAME, b).map_err(|_| error(&ERROR_INVALID_INPUT))?;
    evaluate_tensors(tensor_a, tensor_b, opts, loc_request)
}

async fn gather_if_gpu(value: Value) -> BuiltinResult<Value> {
    match value {
        Value::GpuTensor(handle) => gpu_helpers::gather_value_async(&Value::GpuTensor(handle))
            .await
            .map_err(|e| error_with(&ERROR_INTERNAL, format!("{NAME}: {e}"))),
        other => Ok(other),
    }
}

fn evaluate_tensors(
    a: Tensor,
    b: Tensor,
    mut opts: IsMemberTolOptions,
    loc_request: LocRequest,
) -> BuiltinResult<IsMemberTolEvaluation> {
    if opts.tol.is_none() {
        opts.tol = Some(default_tolerance(a.dtype, b.dtype));
    }
    if opts.by_rows {
        evaluate_rows(a, b, opts, loc_request)
    } else {
        evaluate_elements(a, b, opts, loc_request)
    }
}

#[derive(Debug, Clone)]
struct IsMemberTolOptions {
    tol: Option<f64>,
    by_rows: bool,
    data_scale: Option<Vec<f64>>,
    output_all_indices: bool,
}

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum LocRequest {
    None,
    First,
    All,
}

fn parse_options(rest: &[Value]) -> BuiltinResult<IsMemberTolOptions> {
    let mut opts = IsMemberTolOptions {
        tol: None,
        by_rows: false,
        data_scale: None,
        output_all_indices: false,
    };

    let mut idx = 0usize;
    if let Some(first) = rest.first() {
        if option_name(first).is_none() {
            opts.tol = Some(parse_positive_scalar(first, "tolerance")?);
            idx = 1;
        }
    }

    let remaining = &rest[idx..];
    if !remaining.len().is_multiple_of(2) {
        return Err(error_with(
            &ERROR_INVALID_ARGUMENT,
            "ismembertol: name-value arguments must appear in pairs",
        ));
    }

    for pair in remaining.chunks_exact(2) {
        let name = option_name(&pair[0])
            .ok_or_else(|| {
                error_with(&ERROR_INVALID_ARGUMENT, "ismembertol: expected option name")
            })?
            .trim()
            .to_ascii_lowercase();
        match name.as_str() {
            "byrows" => opts.by_rows = parse_bool_option(&pair[1])?,
            "outputallindices" => opts.output_all_indices = parse_bool_option(&pair[1])?,
            "datascale" => opts.data_scale = Some(parse_data_scale(&pair[1])?),
            other => {
                return Err(error_with(
                    &ERROR_INVALID_ARGUMENT,
                    format!("ismembertol: unknown option '{other}'"),
                ))
            }
        }
    }

    Ok(opts)
}

fn option_name(value: &Value) -> Option<String> {
    match value {
        Value::String(_) | Value::StringArray(_) | Value::CharArray(_) => {
            tensor::value_to_string(value)
        }
        _ => None,
    }
}

fn evaluate_elements(
    a: Tensor,
    b: Tensor,
    opts: IsMemberTolOptions,
    loc_request: LocRequest,
) -> BuiltinResult<IsMemberTolEvaluation> {
    let tol = opts.tol.expect("defaulted tolerance");
    let scale = opts
        .data_scale
        .as_ref()
        .map(|values| {
            if values.len() == 1 {
                Ok(values[0])
            } else {
                Err(error_with(
                    &ERROR_INVALID_ARGUMENT,
                    "ismembertol: DataScale must be scalar unless ByRows is true",
                ))
            }
        })
        .transpose()?
        .unwrap_or_else(|| default_element_scale(&a, &b));
    let threshold = tol * scale.abs();

    let mut mask_data = Vec::with_capacity(a.data.len());
    let mut loc_data = match loc_request {
        LocRequest::First => Vec::with_capacity(a.data.len()),
        LocRequest::None | LocRequest::All => Vec::new(),
    };
    let mut all_cells = match loc_request {
        LocRequest::All => Vec::with_capacity(a.data.len()),
        LocRequest::None | LocRequest::First => Vec::new(),
    };

    for &value in &a.data {
        match loc_request {
            LocRequest::None => {
                mask_data.push(
                    if first_element_match(value, &b.data, threshold).is_some() {
                        1
                    } else {
                        0
                    },
                );
            }
            LocRequest::First => {
                let first = first_element_match(value, &b.data, threshold);
                mask_data.push(if first.is_some() { 1 } else { 0 });
                loc_data.push(first.unwrap_or(0) as f64);
            }
            LocRequest::All => {
                let matches = all_element_matches(value, &b.data, threshold);
                mask_data.push(if matches.is_empty() { 0 } else { 1 });
                all_cells.push(indices_cell(matches)?);
            }
        }
    }

    let mask = LogicalArray::new(mask_data, a.shape.clone())
        .map_err(|e| error_with(&ERROR_INTERNAL, format!("{NAME}: {e}")))?;
    let loc = match loc_request {
        LocRequest::None => None,
        LocRequest::First => Some(LocOutput::First(
            Tensor::new(loc_data, a.shape.clone())
                .map_err(|e| error_with(&ERROR_INTERNAL, format!("{NAME}: {e}")))?,
        )),
        LocRequest::All => Some(LocOutput::All(cell_array_from_linear_matches(
            all_cells,
            a.shape.clone(),
        )?)),
    };
    Ok(IsMemberTolEvaluation { mask, loc })
}

fn evaluate_rows(
    a: Tensor,
    b: Tensor,
    opts: IsMemberTolOptions,
    loc_request: LocRequest,
) -> BuiltinResult<IsMemberTolEvaluation> {
    let (rows_a, cols_a) = rows_cols(&a)?;
    let (rows_b, cols_b) = rows_cols(&b)?;
    if cols_a != cols_b {
        return Err(error(&ERROR_ROWS_COLUMN_MISMATCH));
    }
    let tol = opts.tol.expect("defaulted tolerance");
    let scales = row_scales(&a, rows_a, &b, rows_b, opts.data_scale.as_deref(), cols_a)?;
    let thresholds: Vec<f64> = scales.iter().map(|scale| tol * scale.abs()).collect();

    let mut mask_data = vec![0u8; rows_a];
    let mut loc_data = match loc_request {
        LocRequest::First => vec![0.0f64; rows_a],
        LocRequest::None | LocRequest::All => Vec::new(),
    };
    let mut all_cells = match loc_request {
        LocRequest::All => Vec::with_capacity(rows_a),
        LocRequest::None | LocRequest::First => Vec::new(),
    };

    for row_a in 0..rows_a {
        match loc_request {
            LocRequest::None => {
                mask_data[row_a] = if first_row_match(
                    &a,
                    row_a,
                    rows_a,
                    &b,
                    rows_b,
                    cols_a,
                    &thresholds,
                )
                .is_some()
                {
                    1
                } else {
                    0
                };
            }
            LocRequest::First => {
                let first = first_row_match(&a, row_a, rows_a, &b, rows_b, cols_a, &thresholds);
                mask_data[row_a] = if first.is_some() { 1 } else { 0 };
                loc_data[row_a] = first.unwrap_or(0) as f64;
            }
            LocRequest::All => {
                let matches = all_row_matches(&a, row_a, rows_a, &b, rows_b, cols_a, &thresholds);
                mask_data[row_a] = if matches.is_empty() { 0 } else { 1 };
                all_cells.push(indices_cell(matches)?);
            }
        }
    }

    let shape = vec![rows_a, 1];
    let mask = LogicalArray::new(mask_data, shape.clone())
        .map_err(|e| error_with(&ERROR_INTERNAL, format!("{NAME}: {e}")))?;
    let loc = match loc_request {
        LocRequest::None => None,
        LocRequest::First => Some(LocOutput::First(
            Tensor::new(loc_data, shape)
                .map_err(|e| error_with(&ERROR_INTERNAL, format!("{NAME}: {e}")))?,
        )),
        LocRequest::All => Some(LocOutput::All(cell_array_from_linear_matches(
            all_cells, shape,
        )?)),
    };
    Ok(IsMemberTolEvaluation { mask, loc })
}

fn first_element_match(value: f64, candidates: &[f64], threshold: f64) -> Option<usize> {
    candidates.iter().enumerate().find_map(|(idx, &candidate)| {
        within_tolerance(value, candidate, threshold).then_some(idx + 1)
    })
}

fn all_element_matches(value: f64, candidates: &[f64], threshold: f64) -> Vec<usize> {
    candidates
        .iter()
        .enumerate()
        .filter_map(|(idx, &candidate)| {
            within_tolerance(value, candidate, threshold).then_some(idx + 1)
        })
        .collect()
}

fn first_row_match(
    a: &Tensor,
    row_a: usize,
    rows_a: usize,
    b: &Tensor,
    rows_b: usize,
    cols: usize,
    thresholds: &[f64],
) -> Option<usize> {
    (0..rows_b).find_map(|row_b| {
        rows_match(a, row_a, rows_a, b, row_b, rows_b, cols, thresholds).then_some(row_b + 1)
    })
}

fn all_row_matches(
    a: &Tensor,
    row_a: usize,
    rows_a: usize,
    b: &Tensor,
    rows_b: usize,
    cols: usize,
    thresholds: &[f64],
) -> Vec<usize> {
    (0..rows_b)
        .filter_map(|row_b| {
            rows_match(a, row_a, rows_a, b, row_b, rows_b, cols, thresholds).then_some(row_b + 1)
        })
        .collect()
}

fn rows_match(
    a: &Tensor,
    row_a: usize,
    rows_a: usize,
    b: &Tensor,
    row_b: usize,
    rows_b: usize,
    cols: usize,
    thresholds: &[f64],
) -> bool {
    for col in 0..cols {
        let threshold = thresholds[col];
        if threshold.is_infinite() {
            continue;
        }
        let lhs = a.data[row_a + col * rows_a];
        let rhs = b.data[row_b + col * rows_b];
        if !within_tolerance(lhs, rhs, threshold) {
            return false;
        }
    }
    true
}

fn within_tolerance(lhs: f64, rhs: f64, threshold: f64) -> bool {
    if lhs.is_nan() || rhs.is_nan() {
        return false;
    }
    if lhs == rhs {
        return true;
    }
    (lhs - rhs).abs() <= threshold
}

fn default_tolerance(a: NumericDType, b: NumericDType) -> f64 {
    if matches!(a, NumericDType::F32) || matches!(b, NumericDType::F32) {
        DEFAULT_TOL_SINGLE
    } else {
        DEFAULT_TOL_DOUBLE
    }
}

fn default_element_scale(a: &Tensor, b: &Tensor) -> f64 {
    a.data
        .iter()
        .chain(b.data.iter())
        .filter_map(|value| (!value.is_nan()).then_some(value.abs()))
        .fold(0.0, f64::max)
}

fn row_scales(
    a: &Tensor,
    rows_a: usize,
    b: &Tensor,
    rows_b: usize,
    supplied: Option<&[f64]>,
    cols: usize,
) -> BuiltinResult<Vec<f64>> {
    if let Some(values) = supplied {
        if values.len() == 1 {
            return Ok(vec![values[0]; cols]);
        }
        if values.len() == cols {
            return Ok(values.to_vec());
        }
        return Err(error_with(
            &ERROR_INVALID_ARGUMENT,
            "ismembertol: DataScale vector length must match the number of columns",
        ));
    }

    let mut scales = vec![0.0f64; cols];
    for (col, scale) in scales.iter_mut().enumerate() {
        for row in 0..rows_a {
            let value = a.data[row + col * rows_a];
            if !value.is_nan() {
                *scale = (*scale).max(value.abs());
            }
        }
        for row in 0..rows_b {
            let value = b.data[row + col * rows_b];
            if !value.is_nan() {
                *scale = (*scale).max(value.abs());
            }
        }
    }
    Ok(scales)
}

fn rows_cols(tensor: &Tensor) -> BuiltinResult<(usize, usize)> {
    match tensor.shape.len() {
        0 => Ok((1, 1)),
        1 => Ok((tensor.shape[0], 1)),
        2 => Ok((tensor.shape[0], tensor.shape[1])),
        _ => Err(error_with(
            &ERROR_INVALID_ARGUMENT,
            "ismembertol: ByRows requires 2-D arrays",
        )),
    }
}

fn indices_cell(indices: Vec<usize>) -> BuiltinResult<Value> {
    if indices.is_empty() {
        return Ok(Value::Tensor(Tensor::new(vec![0.0], vec![1, 1]).map_err(
            |e| error_with(&ERROR_INTERNAL, format!("{NAME}: {e}")),
        )?));
    }
    let data = indices
        .into_iter()
        .map(|idx| idx as f64)
        .collect::<Vec<_>>();
    let len = data.len();
    Tensor::new(data, vec![len, 1])
        .map(Value::Tensor)
        .map_err(|e| error_with(&ERROR_INTERNAL, format!("{NAME}: {e}")))
}

fn cell_array_from_linear_matches(
    cells_column_major: Vec<Value>,
    shape: Vec<usize>,
) -> BuiltinResult<CellArray> {
    let cells = if shape.len() == 2 && shape[0] > 1 && shape[1] > 1 {
        let rows = shape[0];
        let cols = shape[1];
        let mut row_major = Vec::with_capacity(cells_column_major.len());
        for row in 0..rows {
            for col in 0..cols {
                row_major.push(cells_column_major[row + col * rows].clone());
            }
        }
        row_major
    } else {
        cells_column_major
    };
    CellArray::new_with_shape(cells, shape)
        .map_err(|e| error_with(&ERROR_INTERNAL, format!("{NAME}: {e}")))
}

fn parse_positive_scalar(value: &Value, label: &str) -> BuiltinResult<f64> {
    let scalar = numeric_scalar(value).ok_or_else(|| {
        error_with(
            &ERROR_INVALID_ARGUMENT,
            format!("ismembertol: {label} must be a positive real scalar"),
        )
    })?;
    if !scalar.is_finite() || scalar <= 0.0 {
        return Err(error_with(
            &ERROR_INVALID_ARGUMENT,
            format!("ismembertol: {label} must be a positive real scalar"),
        ));
    }
    Ok(scalar)
}

fn parse_bool_option(value: &Value) -> BuiltinResult<bool> {
    match numeric_scalar(value) {
        Some(0.0) => Ok(false),
        Some(1.0) => Ok(true),
        _ => Err(error_with(
            &ERROR_INVALID_ARGUMENT,
            "ismembertol: boolean options must be true/false or 0/1",
        )),
    }
}

fn parse_data_scale(value: &Value) -> BuiltinResult<Vec<f64>> {
    match value {
        Value::Num(n) => validate_data_scale(vec![*n]),
        Value::Int(i) => validate_data_scale(vec![i.to_i64() as f64]),
        Value::Bool(flag) => validate_data_scale(vec![if *flag { 1.0 } else { 0.0 }]),
        Value::Tensor(tensor) => validate_data_scale(tensor.data.clone()),
        Value::LogicalArray(logical) => validate_data_scale(
            logical
                .data
                .iter()
                .map(|&v| if v != 0 { 1.0 } else { 0.0 })
                .collect(),
        ),
        _ => Err(error_with(
            &ERROR_INVALID_ARGUMENT,
            "ismembertol: DataScale must be a numeric scalar or vector",
        )),
    }
}

fn validate_data_scale(values: Vec<f64>) -> BuiltinResult<Vec<f64>> {
    if values.is_empty() || values.iter().any(|value| value.is_nan() || *value < 0.0) {
        return Err(error_with(
            &ERROR_INVALID_ARGUMENT,
            "ismembertol: DataScale values must be non-negative or Inf",
        ));
    }
    Ok(values)
}

fn numeric_scalar(value: &Value) -> Option<f64> {
    match value {
        Value::Num(n) => Some(*n),
        Value::Int(i) => Some(i.to_i64() as f64),
        Value::Bool(flag) => Some(if *flag { 1.0 } else { 0.0 }),
        Value::Tensor(t) if t.data.len() == 1 => Some(t.data[0]),
        Value::LogicalArray(logical) if logical.data.len() == 1 => {
            Some(if logical.data[0] != 0 { 1.0 } else { 0.0 })
        }
        _ => None,
    }
}

#[derive(Debug, Clone)]
pub struct IsMemberTolEvaluation {
    mask: LogicalArray,
    loc: Option<LocOutput>,
}

#[derive(Debug, Clone)]
enum LocOutput {
    First(Tensor),
    All(CellArray),
}

impl IsMemberTolEvaluation {
    pub fn into_mask_value(self) -> Value {
        logical_array_into_value(self.mask)
    }

    pub fn mask_value(&self) -> Value {
        logical_array_into_value(self.mask.clone())
    }

    pub fn into_pair(self) -> (Value, Value) {
        let mask = logical_array_into_value(self.mask);
        let loc = match self.loc.expect("LocB was not requested") {
            LocOutput::First(tensor) => tensor::tensor_into_value(tensor),
            LocOutput::All(cell) => Value::Cell(cell),
        };
        (mask, loc)
    }
}

fn logical_array_into_value(logical: LogicalArray) -> Value {
    if logical.data.len() == 1 {
        Value::Bool(logical.data[0] != 0)
    } else {
        Value::LogicalArray(logical)
    }
}

fn error(error: &'static BuiltinErrorDescriptor) -> RuntimeError {
    error_with(error, error.message)
}

fn error_with(error: &'static BuiltinErrorDescriptor, message: impl Into<String>) -> RuntimeError {
    let mut builder = build_runtime_error(message).with_builtin(NAME);
    if let Some(identifier) = error.identifier {
        builder = builder.with_identifier(identifier);
    }
    builder.build()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::builtins::common::test_support;
    use futures::executor::block_on;
    use runmat_builtins::{LiteralValue, ResolveContext, Type};

    fn eval(a: Value, b: Value, rest: &[Value]) -> BuiltinResult<IsMemberTolEvaluation> {
        block_on(evaluate(a, b, rest))
    }

    #[test]
    fn type_resolver_returns_logical_shape() {
        assert_eq!(
            ismembertol_output_type(
                &[Type::Tensor {
                    shape: Some(vec![Some(1), Some(3)])
                }],
                &ResolveContext::new(Vec::new()),
            ),
            Type::Logical {
                shape: Some(vec![Some(1), Some(3)])
            }
        );
    }

    #[test]
    fn type_resolver_returns_row_shape_for_literal_by_rows() {
        assert_eq!(
            ismembertol_output_type(
                &[Type::Tensor {
                    shape: Some(vec![Some(4), Some(3)])
                }],
                &ResolveContext::new(vec![
                    LiteralValue::Unknown,
                    LiteralValue::Unknown,
                    LiteralValue::Number(0.01),
                    LiteralValue::String("ByRows".to_string()),
                    LiteralValue::Bool(true),
                ]),
            ),
            Type::Logical {
                shape: Some(vec![Some(4), Some(1)])
            }
        );
    }

    #[test]
    fn default_tolerance_uses_scaled_double_data() {
        let a = Tensor::new(vec![0.1, 1.0e10], vec![1, 2]).unwrap();
        let b = Tensor::new(vec![0.1 + 1.0e-14, 1.0e10 + 5.0e-3], vec![1, 2]).unwrap();
        let result = eval(Value::Tensor(a), Value::Tensor(b), &[]).unwrap();
        assert_eq!(result.mask.data, vec![1, 1]);
    }

    #[test]
    fn data_scale_one_uses_absolute_tolerance() {
        let a = Tensor::new(vec![1.0e10], vec![1, 1]).unwrap();
        let b = Tensor::new(vec![1.0e10 + 5.0e-3], vec![1, 1]).unwrap();
        let result = eval(
            Value::Tensor(a),
            Value::Tensor(b),
            &[
                Value::Num(1.0e-6),
                Value::from("DataScale"),
                Value::Num(1.0),
            ],
        )
        .unwrap();
        assert_eq!(result.mask_value(), Value::Bool(false));
    }

    #[test]
    fn loc_returns_first_matching_index() {
        let a = Tensor::new(vec![1.0, 2.0, 4.0], vec![1, 3]).unwrap();
        let b = Tensor::new(vec![2.01, 2.02, 4.2], vec![1, 3]).unwrap();
        let (_, loc) = eval(Value::Tensor(a), Value::Tensor(b), &[Value::Num(0.05)])
            .unwrap()
            .into_pair();
        match loc {
            Value::Tensor(tensor) => assert_eq!(tensor.data, vec![0.0, 1.0, 3.0]),
            other => panic!("expected tensor loc, got {other:?}"),
        }
    }

    #[test]
    fn mask_only_request_does_not_materialize_locations() {
        let a = Tensor::new(vec![2.0, 9.0], vec![1, 2]).unwrap();
        let b = Tensor::new(vec![1.99, 2.01, 9.5], vec![1, 3]).unwrap();
        let opts = parse_options(&[
            Value::Num(0.01),
            Value::from("OutputAllIndices"),
            Value::Bool(true),
        ])
        .unwrap();
        let result = block_on(evaluate_with_options(
            Value::Tensor(a),
            Value::Tensor(b),
            opts,
            LocRequest::None,
        ))
        .unwrap();
        assert_eq!(result.mask.data, vec![1, 0]);
        assert!(result.loc.is_none());
    }

    #[test]
    fn output_all_indices_returns_cell_array() {
        let a = Tensor::new(vec![2.0, 9.0], vec![1, 2]).unwrap();
        let b = Tensor::new(vec![1.99, 2.01, 9.5], vec![1, 3]).unwrap();
        let (_, loc) = eval(
            Value::Tensor(a),
            Value::Tensor(b),
            &[
                Value::Num(0.01),
                Value::from("OutputAllIndices"),
                Value::Bool(true),
            ],
        )
        .unwrap()
        .into_pair();
        let Value::Cell(cell) = loc else {
            panic!("expected cell loc");
        };
        assert_eq!(cell.shape, vec![1, 2]);
        match &cell.data[0] {
            Value::Tensor(indices) => assert_eq!(indices.data, vec![1.0, 2.0]),
            other => panic!("expected tensor indices, got {other:?}"),
        }
        match &cell.data[1] {
            Value::Tensor(indices) => assert_eq!(indices.data, vec![0.0]),
            other => panic!("expected tensor zero indices, got {other:?}"),
        }
    }

    #[test]
    fn output_all_indices_matrix_cells_follow_public_cell_indexing() {
        let a = Tensor::new(vec![1.0, 3.0, 2.0, 4.0], vec![2, 2]).unwrap();
        let b = Tensor::new(vec![10.0, 2.0, 3.0, 20.0], vec![1, 4]).unwrap();
        let (_, loc) = eval(
            Value::Tensor(a),
            Value::Tensor(b),
            &[
                Value::Num(1.0e-12),
                Value::from("OutputAllIndices"),
                Value::Bool(true),
            ],
        )
        .unwrap()
        .into_pair();
        let Value::Cell(cell) = loc else {
            panic!("expected cell loc");
        };
        assert_eq!(cell.shape, vec![2, 2]);
        let top_right = cell.get(0, 1).unwrap();
        match top_right {
            Value::Tensor(indices) => assert_eq!(indices.data, vec![2.0]),
            other => panic!("expected tensor indices, got {other:?}"),
        }
        let bottom_left = cell.get(1, 0).unwrap();
        match bottom_left {
            Value::Tensor(indices) => assert_eq!(indices.data, vec![3.0]),
            other => panic!("expected tensor indices, got {other:?}"),
        }
    }

    #[test]
    fn by_rows_uses_column_scales_and_inf_ignores_column() {
        let a = Tensor::new(vec![0.0, 10.0, 0.5, 20.0], vec![2, 2]).unwrap();
        let b = Tensor::new(vec![0.05, 10.2, 999.0, -999.0], vec![2, 2]).unwrap();
        let result = eval(
            Value::Tensor(a),
            Value::Tensor(b),
            &[
                Value::Num(0.1),
                Value::from("ByRows"),
                Value::Bool(true),
                Value::from("DataScale"),
                Value::Tensor(Tensor::new(vec![1.0, f64::INFINITY], vec![1, 2]).unwrap()),
            ],
        )
        .unwrap();
        assert_eq!(result.mask.data, vec![1, 0]);
        let (_, loc) = result.into_pair();
        match loc {
            Value::Tensor(tensor) => {
                assert_eq!(tensor.shape, vec![2, 1]);
                assert_eq!(tensor.data, vec![1.0, 0.0]);
            }
            other => panic!("expected tensor loc, got {other:?}"),
        }
    }

    #[test]
    fn by_rows_output_all_indices_returns_row_cells() {
        let a = Tensor::new(vec![0.0, 1.0], vec![1, 2]).unwrap();
        let b = Tensor::new(vec![0.01, -0.01, 5.0, 1.0, 1.0, 5.0], vec![3, 2]).unwrap();
        let (_, loc) = eval(
            Value::Tensor(a),
            Value::Tensor(b),
            &[
                Value::Num(0.02),
                Value::from("ByRows"),
                Value::Bool(true),
                Value::from("OutputAllIndices"),
                Value::Bool(true),
            ],
        )
        .unwrap()
        .into_pair();
        let Value::Cell(cell) = loc else {
            panic!("expected cell loc");
        };
        assert_eq!(cell.shape, vec![1, 1]);
        match &cell.data[0] {
            Value::Tensor(indices) => assert_eq!(indices.data, vec![1.0, 2.0]),
            other => panic!("expected tensor indices, got {other:?}"),
        }
    }

    #[test]
    fn rejects_bad_options_and_shape_mismatch() {
        let err = parse_options(&[Value::Num(-1.0)]).unwrap_err();
        assert_eq!(
            err.identifier.as_deref(),
            Some("RunMat:ismembertol:InvalidArgument")
        );

        let a = Tensor::new(vec![1.0, 2.0], vec![1, 2]).unwrap();
        let b = Tensor::new(vec![1.0, 2.0], vec![2, 1]).unwrap();
        let err = eval(
            Value::Tensor(a),
            Value::Tensor(b),
            &[Value::from("ByRows"), Value::Bool(true)],
        )
        .unwrap_err();
        assert_eq!(
            err.identifier.as_deref(),
            Some("RunMat:ismembertol:RowsColumnMismatch")
        );
    }

    #[test]
    fn gpu_inputs_gather_to_host() {
        test_support::with_test_provider(|provider| {
            let a = Tensor::new(vec![1.0, 2.0], vec![2, 1]).unwrap();
            let b = Tensor::new(vec![1.0 + 1e-13, 3.0], vec![2, 1]).unwrap();
            let handle_a = provider
                .upload(&runmat_accelerate_api::HostTensorView {
                    data: &a.data,
                    shape: &a.shape,
                })
                .expect("upload a");
            let handle_b = provider
                .upload(&runmat_accelerate_api::HostTensorView {
                    data: &b.data,
                    shape: &b.shape,
                })
                .expect("upload b");
            let result = eval(Value::GpuTensor(handle_a), Value::GpuTensor(handle_b), &[]).unwrap();
            assert_eq!(result.mask.data, vec![1, 0]);
        });
    }
}