runmat-runtime 0.5.0

//! MATLAB-compatible `cumprod` builtin with GPU-aware semantics for RunMat.

use runmat_accelerate_api::{GpuTensorHandle, ProviderNanMode, ProviderScanDirection};
use runmat_builtins::{
    BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
    BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
    ComplexTensor, ResolveContext, Tensor, Type, Value,
};
use runmat_macros::runtime_builtin;

use crate::{build_runtime_error, BuiltinResult, RuntimeError};

const NAME: &str = "cumprod";

fn cumprod_type(args: &[Type], ctx: &ResolveContext) -> Type {
    cumulative_numeric_type(args, ctx)
}

const CUMPROD_OUTPUT: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
    name: "B",
    ty: BuiltinParamType::NumericArray,
    arity: BuiltinParamArity::Required,
    default: None,
    description: "Cumulative product result.",
}];

const CUMPROD_PARAM_A: BuiltinParamDescriptor = BuiltinParamDescriptor {
    name: "A",
    ty: BuiltinParamType::Any,
    arity: BuiltinParamArity::Required,
    default: None,
    description: "Input scalar or array.",
};

const CUMPROD_PARAM_DIM: BuiltinParamDescriptor = BuiltinParamDescriptor {
    name: "dim",
    ty: BuiltinParamType::Any,
    arity: BuiltinParamArity::Optional,
    default: Some("[]"),
    description: "Dimension selector (placeholder [] keeps default dimension).",
};

const CUMPROD_PARAM_DIRECTION: BuiltinParamDescriptor = BuiltinParamDescriptor {
    name: "direction",
    ty: BuiltinParamType::StringScalar,
    arity: BuiltinParamArity::Optional,
    default: Some("\"forward\""),
    description: "Scan direction: \"forward\" or \"reverse\".",
};

const CUMPROD_PARAM_NANFLAG: BuiltinParamDescriptor = BuiltinParamDescriptor {
    name: "nanflag",
    ty: BuiltinParamType::StringScalar,
    arity: BuiltinParamArity::Optional,
    default: Some("\"includenan\""),
    description:
        "Missing-value mode: \"includenan\"/\"includemissing\" or \"omitnan\"/\"omitmissing\".",
};

const CUMPROD_INPUTS_CORE: [BuiltinParamDescriptor; 1] = [CUMPROD_PARAM_A];

const CUMPROD_INPUTS_DIM: [BuiltinParamDescriptor; 2] = [CUMPROD_PARAM_A, CUMPROD_PARAM_DIM];

const CUMPROD_INPUTS_DIRECTION: [BuiltinParamDescriptor; 2] =
    [CUMPROD_PARAM_A, CUMPROD_PARAM_DIRECTION];

const CUMPROD_INPUTS_NANFLAG: [BuiltinParamDescriptor; 2] =
    [CUMPROD_PARAM_A, CUMPROD_PARAM_NANFLAG];

const CUMPROD_INPUTS_DIM_DIRECTION: [BuiltinParamDescriptor; 3] =
    [CUMPROD_PARAM_A, CUMPROD_PARAM_DIM, CUMPROD_PARAM_DIRECTION];

const CUMPROD_INPUTS_DIRECTION_DIM: [BuiltinParamDescriptor; 3] =
    [CUMPROD_PARAM_A, CUMPROD_PARAM_DIRECTION, CUMPROD_PARAM_DIM];

const CUMPROD_INPUTS_DIM_NANFLAG: [BuiltinParamDescriptor; 3] =
    [CUMPROD_PARAM_A, CUMPROD_PARAM_DIM, CUMPROD_PARAM_NANFLAG];

const CUMPROD_INPUTS_NANFLAG_DIM: [BuiltinParamDescriptor; 3] =
    [CUMPROD_PARAM_A, CUMPROD_PARAM_NANFLAG, CUMPROD_PARAM_DIM];

const CUMPROD_INPUTS_DIRECTION_NANFLAG: [BuiltinParamDescriptor; 3] = [
    CUMPROD_PARAM_A,
    CUMPROD_PARAM_DIRECTION,
    CUMPROD_PARAM_NANFLAG,
];

const CUMPROD_INPUTS_NANFLAG_DIRECTION: [BuiltinParamDescriptor; 3] = [
    CUMPROD_PARAM_A,
    CUMPROD_PARAM_NANFLAG,
    CUMPROD_PARAM_DIRECTION,
];

const CUMPROD_INPUTS_DIM_DIRECTION_NANFLAG: [BuiltinParamDescriptor; 4] = [
    CUMPROD_PARAM_A,
    CUMPROD_PARAM_DIM,
    CUMPROD_PARAM_DIRECTION,
    CUMPROD_PARAM_NANFLAG,
];

const CUMPROD_INPUTS_DIM_NANFLAG_DIRECTION: [BuiltinParamDescriptor; 4] = [
    CUMPROD_PARAM_A,
    CUMPROD_PARAM_DIM,
    CUMPROD_PARAM_NANFLAG,
    CUMPROD_PARAM_DIRECTION,
];

const CUMPROD_INPUTS_DIRECTION_DIM_NANFLAG: [BuiltinParamDescriptor; 4] = [
    CUMPROD_PARAM_A,
    CUMPROD_PARAM_DIRECTION,
    CUMPROD_PARAM_DIM,
    CUMPROD_PARAM_NANFLAG,
];

const CUMPROD_INPUTS_DIRECTION_NANFLAG_DIM: [BuiltinParamDescriptor; 4] = [
    CUMPROD_PARAM_A,
    CUMPROD_PARAM_DIRECTION,
    CUMPROD_PARAM_NANFLAG,
    CUMPROD_PARAM_DIM,
];

const CUMPROD_INPUTS_NANFLAG_DIM_DIRECTION: [BuiltinParamDescriptor; 4] = [
    CUMPROD_PARAM_A,
    CUMPROD_PARAM_NANFLAG,
    CUMPROD_PARAM_DIM,
    CUMPROD_PARAM_DIRECTION,
];

const CUMPROD_INPUTS_NANFLAG_DIRECTION_DIM: [BuiltinParamDescriptor; 4] = [
    CUMPROD_PARAM_A,
    CUMPROD_PARAM_NANFLAG,
    CUMPROD_PARAM_DIRECTION,
    CUMPROD_PARAM_DIM,
];

const CUMPROD_SIGNATURES: [BuiltinSignatureDescriptor; 16] = [
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A)",
        inputs: &CUMPROD_INPUTS_CORE,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, dim)",
        inputs: &CUMPROD_INPUTS_DIM,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, direction)",
        inputs: &CUMPROD_INPUTS_DIRECTION,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, nanflag)",
        inputs: &CUMPROD_INPUTS_NANFLAG,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, dim, direction)",
        inputs: &CUMPROD_INPUTS_DIM_DIRECTION,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, direction, dim)",
        inputs: &CUMPROD_INPUTS_DIRECTION_DIM,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, dim, nanflag)",
        inputs: &CUMPROD_INPUTS_DIM_NANFLAG,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, nanflag, dim)",
        inputs: &CUMPROD_INPUTS_NANFLAG_DIM,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, direction, nanflag)",
        inputs: &CUMPROD_INPUTS_DIRECTION_NANFLAG,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, nanflag, direction)",
        inputs: &CUMPROD_INPUTS_NANFLAG_DIRECTION,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, dim, direction, nanflag)",
        inputs: &CUMPROD_INPUTS_DIM_DIRECTION_NANFLAG,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, dim, nanflag, direction)",
        inputs: &CUMPROD_INPUTS_DIM_NANFLAG_DIRECTION,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, direction, dim, nanflag)",
        inputs: &CUMPROD_INPUTS_DIRECTION_DIM_NANFLAG,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, direction, nanflag, dim)",
        inputs: &CUMPROD_INPUTS_DIRECTION_NANFLAG_DIM,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, nanflag, dim, direction)",
        inputs: &CUMPROD_INPUTS_NANFLAG_DIM_DIRECTION,
        outputs: &CUMPROD_OUTPUT,
    },
    BuiltinSignatureDescriptor {
        label: "B = cumprod(A, nanflag, direction, dim)",
        inputs: &CUMPROD_INPUTS_NANFLAG_DIRECTION_DIM,
        outputs: &CUMPROD_OUTPUT,
    },
];

const CUMPROD_ERROR_INVALID_ARGUMENT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
    code: "RM.CUMPROD.INVALID_ARGUMENT",
    identifier: Some("RunMat:cumprod:InvalidArgument"),
    when: "Dimension, direction, or missing-value argument grammar is invalid.",
    message: "cumprod: invalid argument",
};

const CUMPROD_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
    code: "RM.CUMPROD.INVALID_INPUT",
    identifier: Some("RunMat:cumprod:InvalidInput"),
    when: "Input value type is unsupported for cumulative product reduction.",
    message: "cumprod: invalid input",
};

const CUMPROD_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
    code: "RM.CUMPROD.INTERNAL",
    identifier: Some("RunMat:cumprod:Internal"),
    when: "Reduction execution fails due to conversion, provider, or allocation operations.",
    message: "cumprod: internal reduction failure",
};

const CUMPROD_ERRORS: [BuiltinErrorDescriptor; 3] = [
    CUMPROD_ERROR_INVALID_ARGUMENT,
    CUMPROD_ERROR_INVALID_INPUT,
    CUMPROD_ERROR_INTERNAL,
];

pub const CUMPROD_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
    signatures: &CUMPROD_SIGNATURES,
    output_mode: BuiltinOutputMode::Fixed,
    completion_policy: BuiltinCompletionPolicy::Public,
    errors: &CUMPROD_ERRORS,
};

use crate::builtins::common::spec::{
    BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
    ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
};
use crate::builtins::common::{gpu_helpers, tensor};
use crate::builtins::math::reduction::type_resolvers::cumulative_numeric_type;

#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::math::reduction::cumprod")]
pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
    name: "cumprod",
    op_kind: GpuOpKind::Custom("scan"),
    supported_precisions: &[ScalarType::F32, ScalarType::F64],
    broadcast: BroadcastSemantics::Matlab,
    provider_hooks: &[ProviderHook::Custom("cumprod_scan")],
    constant_strategy: ConstantStrategy::InlineLiteral,
    residency: ResidencyPolicy::NewHandle,
    nan_mode: ReductionNaN::Include,
    two_pass_threshold: None,
    workgroup_size: None,
    accepts_nan_mode: false,
    notes: "Providers may expose device prefix-product kernels; the runtime gathers to host when hooks are absent or options are unsupported.",
};

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

fn cumprod_error_with_detail(
    error: &'static BuiltinErrorDescriptor,
    detail: impl AsRef<str>,
) -> RuntimeError {
    cumprod_error_with_message(format!("{}: {}", error.message, detail.as_ref()), error)
}

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

fn cumprod_internal_error(detail: impl AsRef<str>) -> RuntimeError {
    cumprod_error_with_detail(&CUMPROD_ERROR_INTERNAL, detail)
}

#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::math::reduction::cumprod")]
pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
    name: "cumprod",
    shape: ShapeRequirements::BroadcastCompatible,
    constant_strategy: ConstantStrategy::InlineLiteral,
    elementwise: None,
    reduction: None,
    emits_nan: false,
    notes: "Fusion planner currently lowers cumprod to the runtime implementation; providers can substitute specialised scan kernels.",
};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum CumprodDirection {
    Forward,
    Reverse,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum CumprodNanMode {
    Include,
    Omit,
}

#[runtime_builtin(
    name = "cumprod",
    category = "math/reduction",
    summary = "Compute cumulative products.",
    keywords = "cumprod,cumulative product,running product,reverse,omitnan,gpu",
    accel = "reduction",
    type_resolver(cumprod_type),
    descriptor(crate::builtins::math::reduction::cumprod::CUMPROD_DESCRIPTOR),
    builtin_path = "crate::builtins::math::reduction::cumprod"
)]
async fn cumprod_builtin(value: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
    let (dim, direction, nan_mode) = parse_arguments(&rest)?;
    match value {
        Value::GpuTensor(handle) => cumprod_gpu(handle, dim, direction, nan_mode).await,
        Value::Complex(re, im) => {
            let tensor = ComplexTensor::new(vec![(re, im)], vec![1, 1])
                .map_err(|e| cumprod_internal_error(&e))?;
            let target_dim = dim.unwrap_or(1);
            let result = cumprod_complex_tensor(&tensor, target_dim, direction, nan_mode)?;
            Ok(complex_tensor_into_value(result))
        }
        Value::ComplexTensor(ct) => {
            let target_dim = dim.unwrap_or_else(|| default_dimension_from_shape(&ct.shape));
            let result = cumprod_complex_tensor(&ct, target_dim, direction, nan_mode)?;
            Ok(complex_tensor_into_value(result))
        }
        other => cumprod_host(other, dim, direction, nan_mode),
    }
}

fn parse_arguments(
    args: &[Value],
) -> BuiltinResult<(Option<usize>, CumprodDirection, CumprodNanMode)> {
    if args.len() > 3 {
        return Err(cumprod_error(&CUMPROD_ERROR_INVALID_ARGUMENT));
    }

    let mut dim: Option<usize> = None;
    let mut direction = CumprodDirection::Forward;
    let mut direction_set = false;
    let mut nan_mode = CumprodNanMode::Include;
    let mut nan_set = false;

    for value in args {
        match value {
            Value::Int(_) | Value::Num(_) => {
                if dim.is_some() {
                    return Err(cumprod_error_with_detail(
                        &CUMPROD_ERROR_INVALID_ARGUMENT,
                        "dimension specified more than once",
                    ));
                }
                dim = Some(tensor::parse_dimension(value, "cumprod").map_err(|err| {
                    cumprod_error_with_detail(&CUMPROD_ERROR_INVALID_ARGUMENT, err)
                })?);
            }
            Value::Tensor(t) if t.data.is_empty() => {
                // MATLAB allows [] as a placeholder for the default dimension; ignore it.
            }
            Value::LogicalArray(la) if la.data.is_empty() => {
                // Treat logical [] the same way.
            }
            _ => {
                if let Some(text) = tensor::value_to_string(value) {
                    let keyword = text.trim().to_ascii_lowercase();
                    match keyword.as_str() {
                        "forward" => {
                            if direction_set {
                                return Err(cumprod_error_with_detail(
                                    &CUMPROD_ERROR_INVALID_ARGUMENT,
                                    "direction specified more than once",
                                ));
                            }
                            direction = CumprodDirection::Forward;
                            direction_set = true;
                        }
                        "reverse" => {
                            if direction_set {
                                return Err(cumprod_error_with_detail(
                                    &CUMPROD_ERROR_INVALID_ARGUMENT,
                                    "direction specified more than once",
                                ));
                            }
                            direction = CumprodDirection::Reverse;
                            direction_set = true;
                        }
                        "omitnan" | "omitmissing" => {
                            if nan_set {
                                return Err(cumprod_error_with_detail(
                                    &CUMPROD_ERROR_INVALID_ARGUMENT,
                                    "missing-value handling specified more than once",
                                ));
                            }
                            nan_mode = CumprodNanMode::Omit;
                            nan_set = true;
                        }
                        "includenan" | "includemissing" => {
                            if nan_set {
                                return Err(cumprod_error_with_detail(
                                    &CUMPROD_ERROR_INVALID_ARGUMENT,
                                    "missing-value handling specified more than once",
                                ));
                            }
                            nan_mode = CumprodNanMode::Include;
                            nan_set = true;
                        }
                        "" => {
                            return Err(cumprod_error_with_detail(
                                &CUMPROD_ERROR_INVALID_ARGUMENT,
                                "empty string option is not supported",
                            ));
                        }
                        other => {
                            return Err(cumprod_error_with_detail(
                                &CUMPROD_ERROR_INVALID_ARGUMENT,
                                format!("unrecognised option '{other}'"),
                            ));
                        }
                    }
                } else {
                    return Err(cumprod_error_with_detail(
                        &CUMPROD_ERROR_INVALID_ARGUMENT,
                        format!("unsupported argument type {value:?}"),
                    ));
                }
            }
        }
    }

    Ok((dim, direction, nan_mode))
}

fn cumprod_host(
    value: Value,
    dim: Option<usize>,
    direction: CumprodDirection,
    nan_mode: CumprodNanMode,
) -> BuiltinResult<Value> {
    let tensor = tensor::value_into_tensor_for("cumprod", value)
        .map_err(|err| cumprod_error_with_detail(&CUMPROD_ERROR_INVALID_INPUT, err))?;
    let target_dim = dim.unwrap_or_else(|| default_dimension(&tensor));
    let result = cumprod_tensor(&tensor, target_dim, direction, nan_mode)?;
    Ok(tensor::tensor_into_value(result))
}

async fn cumprod_gpu(
    handle: GpuTensorHandle,
    dim: Option<usize>,
    direction: CumprodDirection,
    nan_mode: CumprodNanMode,
) -> BuiltinResult<Value> {
    #[cfg(all(test, feature = "wgpu"))]
    {
        if handle.device_id != 0 {
            let _ = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
                runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
            );
        }
    }
    if matches!(direction, CumprodDirection::Reverse) && matches!(nan_mode, CumprodNanMode::Omit) {
        let tensor = gpu_helpers::gather_tensor_async(&handle)
            .await
            .map_err(|err| cumprod_internal_error(err.message()))?;
        let fallback_dim = dim.unwrap_or_else(|| default_dimension_from_shape(&tensor.shape));
        let result = cumprod_tensor(&tensor, fallback_dim, direction, nan_mode)?;
        return Ok(tensor::tensor_into_value(result));
    }

    if let Some(target) = dim {
        if target == 0 {
            return Err(cumprod_error_with_detail(
                &CUMPROD_ERROR_INVALID_ARGUMENT,
                "dimension must be >= 1",
            ));
        }
        if target > handle.shape.len() {
            return Ok(Value::GpuTensor(handle));
        }
    }

    let fallback_dim = dim.unwrap_or_else(|| default_dimension_from_shape(&handle.shape));
    if fallback_dim == 0 {
        return Err(cumprod_error_with_detail(
            &CUMPROD_ERROR_INVALID_ARGUMENT,
            "dimension must be >= 1",
        ));
    }

    if let Some(provider) = runmat_accelerate_api::provider() {
        let zero_based_dim = fallback_dim.saturating_sub(1);
        if zero_based_dim < handle.shape.len() {
            let provider_direction = match direction {
                CumprodDirection::Forward => ProviderScanDirection::Forward,
                CumprodDirection::Reverse => ProviderScanDirection::Reverse,
            };
            let provider_nan_mode = match nan_mode {
                CumprodNanMode::Include => ProviderNanMode::Include,
                CumprodNanMode::Omit => ProviderNanMode::Omit,
            };
            if let Ok(device_result) = provider.cumprod_scan(
                &handle,
                zero_based_dim,
                provider_direction,
                provider_nan_mode,
            ) {
                return Ok(Value::GpuTensor(device_result));
            }
        }
    }

    let tensor = gpu_helpers::gather_tensor_async(&handle)
        .await
        .map_err(|err| cumprod_internal_error(err.message()))?;
    let result = cumprod_tensor(&tensor, fallback_dim, direction, nan_mode)?;
    Ok(tensor::tensor_into_value(result))
}

fn cumprod_tensor(
    tensor: &Tensor,
    dim: usize,
    direction: CumprodDirection,
    nan_mode: CumprodNanMode,
) -> BuiltinResult<Tensor> {
    if dim == 0 {
        return Err(cumprod_error_with_detail(
            &CUMPROD_ERROR_INVALID_ARGUMENT,
            "dimension must be >= 1",
        ));
    }
    if tensor.data.is_empty() || dim > tensor.shape.len() {
        return Ok(tensor.clone());
    }

    let dim_index = dim - 1;
    let segment_len = tensor.shape[dim_index];
    if segment_len == 0 {
        return Ok(tensor.clone());
    }

    let stride_before = dim_product(&tensor.shape[..dim_index]);
    let stride_after = dim_product(&tensor.shape[dim..]);
    let block = stride_before * segment_len;
    let mut output = vec![0.0f64; tensor.data.len()];

    for after in 0..stride_after {
        let base = after * block;
        for before in 0..stride_before {
            match direction {
                CumprodDirection::Forward => {
                    let mut prod = 1.0f64;
                    let mut prod_is_nan = false;
                    for k in 0..segment_len {
                        let idx = base + before + k * stride_before;
                        let value = tensor.data[idx];
                        match nan_mode {
                            CumprodNanMode::Include => {
                                if prod_is_nan {
                                    output[idx] = f64::NAN;
                                    continue;
                                }
                                if value.is_nan() {
                                    prod_is_nan = true;
                                    output[idx] = f64::NAN;
                                } else {
                                    prod *= value;
                                    output[idx] = prod;
                                }
                            }
                            CumprodNanMode::Omit => {
                                if !value.is_nan() {
                                    prod *= value;
                                }
                                output[idx] = prod;
                            }
                        }
                    }
                }
                CumprodDirection::Reverse => {
                    let mut prod = 1.0f64;
                    let mut prod_is_nan = false;
                    for offset in (0..segment_len).rev() {
                        let idx = base + before + offset * stride_before;
                        let value = tensor.data[idx];
                        match nan_mode {
                            CumprodNanMode::Include => {
                                if prod_is_nan {
                                    output[idx] = f64::NAN;
                                    continue;
                                }
                                if value.is_nan() {
                                    prod_is_nan = true;
                                    output[idx] = f64::NAN;
                                } else {
                                    prod *= value;
                                    output[idx] = prod;
                                }
                            }
                            CumprodNanMode::Omit => {
                                if !value.is_nan() {
                                    prod *= value;
                                }
                                output[idx] = prod;
                            }
                        }
                    }
                }
            }
        }
    }

    Tensor::new(output, tensor.shape.clone()).map_err(|e| cumprod_internal_error(&e))
}

fn cumprod_complex_tensor(
    tensor: &ComplexTensor,
    dim: usize,
    direction: CumprodDirection,
    nan_mode: CumprodNanMode,
) -> BuiltinResult<ComplexTensor> {
    if dim == 0 {
        return Err(cumprod_error_with_detail(
            &CUMPROD_ERROR_INVALID_ARGUMENT,
            "dimension must be >= 1",
        ));
    }
    if tensor.data.is_empty() || dim > tensor.shape.len() {
        return Ok(tensor.clone());
    }

    let dim_index = dim - 1;
    let segment_len = tensor.shape[dim_index];
    if segment_len == 0 {
        return Ok(tensor.clone());
    }

    let stride_before = dim_product(&tensor.shape[..dim_index]);
    let stride_after = dim_product(&tensor.shape[dim..]);
    let block = stride_before * segment_len;
    let mut output = vec![(0.0f64, 0.0f64); tensor.data.len()];

    for after in 0..stride_after {
        let base = after * block;
        for before in 0..stride_before {
            match direction {
                CumprodDirection::Forward => {
                    let mut prod = (1.0f64, 0.0f64);
                    let mut prod_is_nan = false;
                    for k in 0..segment_len {
                        let idx = base + before + k * stride_before;
                        let value = tensor.data[idx];
                        let value_is_nan = value.0.is_nan() || value.1.is_nan();
                        match nan_mode {
                            CumprodNanMode::Include => {
                                if prod_is_nan {
                                    output[idx] = (f64::NAN, f64::NAN);
                                    continue;
                                }
                                if value_is_nan {
                                    prod_is_nan = true;
                                    output[idx] = (f64::NAN, f64::NAN);
                                } else {
                                    prod = complex_mul(prod, value);
                                    output[idx] = prod;
                                }
                            }
                            CumprodNanMode::Omit => {
                                if !value_is_nan {
                                    prod = complex_mul(prod, value);
                                }
                                output[idx] = prod;
                            }
                        }
                    }
                }
                CumprodDirection::Reverse => {
                    let mut prod = (1.0f64, 0.0f64);
                    let mut prod_is_nan = false;
                    for offset in (0..segment_len).rev() {
                        let idx = base + before + offset * stride_before;
                        let value = tensor.data[idx];
                        let value_is_nan = value.0.is_nan() || value.1.is_nan();
                        match nan_mode {
                            CumprodNanMode::Include => {
                                if prod_is_nan {
                                    output[idx] = (f64::NAN, f64::NAN);
                                    continue;
                                }
                                if value_is_nan {
                                    prod_is_nan = true;
                                    output[idx] = (f64::NAN, f64::NAN);
                                } else {
                                    prod = complex_mul(prod, value);
                                    output[idx] = prod;
                                }
                            }
                            CumprodNanMode::Omit => {
                                if !value_is_nan {
                                    prod = complex_mul(prod, value);
                                }
                                output[idx] = prod;
                            }
                        }
                    }
                }
            }
        }
    }

    ComplexTensor::new(output, tensor.shape.clone()).map_err(|e| cumprod_internal_error(&e))
}

fn complex_tensor_into_value(tensor: ComplexTensor) -> Value {
    if tensor.data.len() == 1 {
        let (re, im) = tensor.data[0];
        Value::Complex(re, im)
    } else {
        Value::ComplexTensor(tensor)
    }
}

fn complex_mul(a: (f64, f64), b: (f64, f64)) -> (f64, f64) {
    (a.0 * b.0 - a.1 * b.1, a.0 * b.1 + a.1 * b.0)
}

fn default_dimension(tensor: &Tensor) -> usize {
    default_dimension_from_shape(&tensor.shape)
}

fn default_dimension_from_shape(shape: &[usize]) -> usize {
    if shape.is_empty() {
        return 1;
    }
    shape
        .iter()
        .position(|&extent| extent != 1)
        .map(|idx| idx + 1)
        .unwrap_or(1)
}

fn dim_product(dims: &[usize]) -> usize {
    dims.iter()
        .copied()
        .fold(1usize, |acc, value| acc.saturating_mul(value))
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use crate::builtins::common::test_support;
    use futures::executor::block_on;
    use runmat_builtins::{IntValue, Tensor as BuiltinsTensor};

    #[test]
    fn cumprod_type_keeps_shape() {
        let out = cumprod_type(
            &[Type::Tensor {
                shape: Some(vec![Some(2), Some(2)]),
            }],
            &ResolveContext::new(Vec::new()),
        );
        assert_eq!(
            out,
            Type::Tensor {
                shape: Some(vec![Some(2), Some(2)])
            }
        );
    }

    #[cfg(feature = "wgpu")]
    use runmat_accelerate::backend::wgpu::provider as wgpu_provider;
    #[cfg(feature = "wgpu")]
    use runmat_accelerate_api::HostTensorView;

    fn cumprod_builtin(value: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
        block_on(super::cumprod_builtin(value, rest))
    }

    fn error_identifier(error: &crate::RuntimeError) -> Option<&str> {
        error.identifier()
    }

    #[test]
    fn cumprod_descriptor_signatures_and_errors() {
        let labels: Vec<&str> = CUMPROD_DESCRIPTOR
            .signatures
            .iter()
            .map(|sig| sig.label)
            .collect();
        assert!(labels.contains(&"B = cumprod(A)"));
        assert!(labels.contains(&"B = cumprod(A, dim)"));
        assert!(labels.contains(&"B = cumprod(A, direction)"));
        assert!(labels.contains(&"B = cumprod(A, nanflag)"));
        assert!(labels.contains(&"B = cumprod(A, dim, direction)"));
        assert!(labels.contains(&"B = cumprod(A, direction, dim)"));
        assert!(labels.contains(&"B = cumprod(A, dim, nanflag)"));
        assert!(labels.contains(&"B = cumprod(A, nanflag, dim)"));
        assert!(labels.contains(&"B = cumprod(A, direction, nanflag)"));
        assert!(labels.contains(&"B = cumprod(A, nanflag, direction)"));
        assert!(labels.contains(&"B = cumprod(A, dim, direction, nanflag)"));
        assert!(labels.contains(&"B = cumprod(A, nanflag, direction, dim)"));
        assert!(CUMPROD_DESCRIPTOR
            .errors
            .iter()
            .any(|err| err.code == CUMPROD_ERROR_INVALID_ARGUMENT.code));
        assert!(CUMPROD_DESCRIPTOR
            .errors
            .iter()
            .any(|err| err.code == CUMPROD_ERROR_INVALID_INPUT.code));
        assert!(CUMPROD_DESCRIPTOR
            .errors
            .iter()
            .any(|err| err.code == CUMPROD_ERROR_INTERNAL.code));
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_scalar_num() {
        let result = cumprod_builtin(Value::Num(7.0), Vec::new()).expect("cumprod scalar");
        assert_eq!(result, Value::Num(7.0));
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_matrix_default_dimension() {
        let tensor = BuiltinsTensor::new(vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0], vec![2, 3]).unwrap();
        let result = cumprod_builtin(Value::Tensor(tensor), Vec::new()).expect("cumprod");
        match result {
            Value::Tensor(out) => {
                assert_eq!(out.shape, vec![2, 3]);
                assert_eq!(out.data, vec![1.0, 4.0, 2.0, 10.0, 3.0, 18.0]);
            }
            other => panic!("expected tensor result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_matrix_dimension_two() {
        let tensor = BuiltinsTensor::new(vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0], vec![2, 3]).unwrap();
        let args = vec![Value::Int(IntValue::I32(2))];
        let result = cumprod_builtin(Value::Tensor(tensor), args).expect("cumprod");
        match result {
            Value::Tensor(out) => {
                assert_eq!(out.shape, vec![2, 3]);
                assert_eq!(out.data, vec![1.0, 4.0, 2.0, 20.0, 6.0, 120.0]);
            }
            other => panic!("expected tensor result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_reverse_direction() {
        let tensor = BuiltinsTensor::new(vec![2.0, 3.0, 4.0, 5.0], vec![4, 1]).unwrap();
        let result =
            cumprod_builtin(Value::Tensor(tensor), vec![Value::from("reverse")]).expect("cumprod");
        match result {
            Value::Tensor(out) => {
                assert_eq!(out.data, vec![120.0, 60.0, 20.0, 5.0]);
            }
            other => panic!("expected tensor result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_omit_nan_forward() {
        let tensor =
            BuiltinsTensor::new(vec![f64::NAN, 2.0, f64::NAN, 4.0], vec![4, 1]).expect("tensor");
        let result =
            cumprod_builtin(Value::Tensor(tensor), vec![Value::from("omitnan")]).expect("cumprod");
        match result {
            Value::Tensor(out) => {
                assert_eq!(out.data, vec![1.0, 2.0, 2.0, 8.0]);
            }
            other => panic!("expected tensor result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_include_nan_propagates() {
        let tensor = BuiltinsTensor::new(vec![1.0, f64::NAN, 3.0], vec![3, 1]).unwrap();
        let result = cumprod_builtin(Value::Tensor(tensor), Vec::new()).expect("cumprod");
        match result {
            Value::Tensor(out) => {
                assert!((out.data[0] - 1.0).abs() < 1e-12);
                assert!(out.data[1].is_nan());
                assert!(out.data[2].is_nan());
            }
            other => panic!("expected tensor result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_dim_greater_than_ndims_returns_input() {
        let tensor = BuiltinsTensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
        let args = vec![Value::Int(IntValue::I32(5))];
        let result = cumprod_builtin(Value::Tensor(tensor.clone()), args).expect("cumprod");
        match result {
            Value::Tensor(out) => assert_eq!(out, tensor),
            other => panic!("expected tensor result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_complex_tensor() {
        let data = vec![(1.0, 2.0), (3.0, -1.0)];
        let tensor = ComplexTensor::new(data, vec![2, 1]).unwrap();
        let result = cumprod_builtin(Value::ComplexTensor(tensor), Vec::new()).expect("cumprod");
        match result {
            Value::ComplexTensor(out) => {
                assert_eq!(out.shape, vec![2, 1]);
                assert!((out.data[0].0 - 1.0).abs() < 1e-12);
                assert!((out.data[0].1 - 2.0).abs() < 1e-12);
                assert!((out.data[1].0 - 5.0).abs() < 1e-12);
                assert!((out.data[1].1 - 5.0).abs() < 1e-12);
            }
            other => panic!("expected complex tensor result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_gpu_provider_roundtrip() {
        test_support::with_test_provider(|provider| {
            let tensor = BuiltinsTensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![4, 1]).unwrap();
            let view = runmat_accelerate_api::HostTensorView {
                data: &tensor.data,
                shape: &tensor.shape,
            };
            let handle = provider.upload(&view).expect("upload");
            let result = cumprod_builtin(Value::GpuTensor(handle), Vec::new()).expect("cumprod");
            let gathered = test_support::gather(result).expect("gather");
            assert_eq!(gathered.shape, vec![4, 1]);
            assert_eq!(gathered.data, vec![1.0, 2.0, 6.0, 24.0]);
        });
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_reverse_with_omit_nan() {
        let tensor =
            BuiltinsTensor::new(vec![1.0, f64::NAN, 4.0, 2.0], vec![4, 1]).expect("tensor");
        let result = cumprod_builtin(
            Value::Tensor(tensor),
            vec![Value::from("omitnan"), Value::from("reverse")],
        )
        .expect("cumprod");
        match result {
            Value::Tensor(out) => {
                assert_eq!(out.data, vec![8.0, 8.0, 8.0, 2.0]);
            }
            other => panic!("expected tensor result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_accepts_omitmissing_synonym() {
        let tensor =
            BuiltinsTensor::new(vec![f64::NAN, 2.0, 3.0, f64::NAN], vec![4, 1]).expect("tensor");
        let result = cumprod_builtin(Value::Tensor(tensor), vec![Value::from("omitmissing")])
            .expect("cumprod");
        match result {
            Value::Tensor(out) => {
                assert_eq!(out.data, vec![1.0, 2.0, 6.0, 6.0]);
            }
            other => panic!("expected tensor result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_accepts_includemissing_synonym() {
        let tensor = BuiltinsTensor::new(vec![1.0, f64::NAN, 4.0], vec![3, 1]).unwrap();
        let result = cumprod_builtin(Value::Tensor(tensor), vec![Value::from("includemissing")])
            .expect("cumprod");
        match result {
            Value::Tensor(out) => {
                assert!((out.data[0] - 1.0).abs() < 1e-12);
                assert!(out.data[1].is_nan());
                assert!(out.data[2].is_nan());
            }
            other => panic!("expected tensor result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_dimension_placeholder_is_ignored() {
        let tensor = BuiltinsTensor::new(vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0], vec![2, 3]).unwrap();
        let placeholder = BuiltinsTensor::new(Vec::<f64>::new(), vec![0, 0]).unwrap();
        let args = vec![Value::Tensor(placeholder), Value::from("reverse")];
        let result = cumprod_builtin(Value::Tensor(tensor), args).expect("cumprod");
        match result {
            Value::Tensor(out) => {
                assert_eq!(out.shape, vec![2, 3]);
                assert_eq!(out.data, vec![4.0, 4.0, 10.0, 5.0, 18.0, 6.0]);
            }
            other => panic!("expected tensor result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn cumprod_dimension_zero_errors() {
        let tensor = BuiltinsTensor::new(vec![1.0, 2.0], vec![2, 1]).unwrap();
        let result = cumprod_builtin(Value::Tensor(tensor), vec![Value::Int(IntValue::I32(0))]);
        match result {
            Err(err) => {
                assert_eq!(
                    error_identifier(&err),
                    CUMPROD_ERROR_INVALID_ARGUMENT.identifier
                );
                assert!(
                    err.message()
                        .contains(CUMPROD_ERROR_INVALID_ARGUMENT.message),
                    "unexpected result: {err}"
                );
            }
            Ok(_) => panic!("expected error"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    #[cfg(feature = "wgpu")]
    fn cumprod_wgpu_forward_matches_cpu() {
        let _ =
            wgpu_provider::register_wgpu_provider(wgpu_provider::WgpuProviderOptions::default());
        let tensor = BuiltinsTensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![4, 1]).unwrap();

        let cpu = cumprod_host(
            Value::Tensor(tensor.clone()),
            Some(1),
            CumprodDirection::Forward,
            CumprodNanMode::Include,
        )
        .unwrap();

        let view = HostTensorView {
            data: &tensor.data,
            shape: &tensor.shape,
        };
        let provider = runmat_accelerate_api::provider().expect("wgpu provider");
        let handle = provider.upload(&view).expect("upload");
        let gpu_value = block_on(cumprod_gpu(
            handle,
            Some(1),
            CumprodDirection::Forward,
            CumprodNanMode::Include,
        ))
        .expect("cumprod gpu");
        let gathered = test_support::gather(gpu_value).expect("gather");

        match cpu {
            Value::Tensor(ct) => {
                assert_eq!(gathered.shape, ct.shape);
                assert_eq!(gathered.data, ct.data);
            }
            other => panic!("expected tensor result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    #[cfg(feature = "wgpu")]
    fn cumprod_wgpu_reverse_omitnan_matches_cpu() {
        let _ =
            wgpu_provider::register_wgpu_provider(wgpu_provider::WgpuProviderOptions::default());
        let tensor = BuiltinsTensor::new(vec![1.0, f64::NAN, 4.0, 2.0], vec![4, 1]).unwrap();

        let cpu = cumprod_host(
            Value::Tensor(tensor.clone()),
            Some(1),
            CumprodDirection::Reverse,
            CumprodNanMode::Omit,
        )
        .unwrap();

        let view = HostTensorView {
            data: &tensor.data,
            shape: &tensor.shape,
        };
        let provider = runmat_accelerate_api::provider().expect("wgpu provider");
        let handle = provider.upload(&view).expect("upload");
        let gpu_value = block_on(cumprod_gpu(
            handle,
            Some(1),
            CumprodDirection::Reverse,
            CumprodNanMode::Omit,
        ))
        .expect("cumprod gpu");
        let gathered = test_support::gather(gpu_value).expect("gather");

        match cpu {
            Value::Tensor(ct) => {
                assert_eq!(gathered.shape, ct.shape);
                let tol = match provider.precision() {
                    runmat_accelerate_api::ProviderPrecision::F64 => 1e-12,
                    runmat_accelerate_api::ProviderPrecision::F32 => 1e-5,
                };
                for (a, b) in gathered.data.iter().zip(ct.data.iter()) {
                    if b.is_nan() {
                        assert!(a.is_nan());
                    } else {
                        assert!((a - b).abs() < tol);
                    }
                }
            }
            other => panic!("expected tensor result, got {other:?}"),
        }
    }
}