Function eval_flat_elementwise

src/ir/flat.rs:75–109

Evaluate a compiled `FlatExpr` over input arguments, producing a new Buffer.

(flat: &FlatExpr, args: &[ArgValue])

/// Evaluate a compiled `FlatExpr` over input arguments, producing a new Buffer.
pub(crate) fn eval_flat_elementwise(flat: &FlatExpr, args: &[ArgValue]) -> ParsecResult<Buffer> {
    let len = infer_output_len(args)?;

    if len == 0 {
        return Ok(Buffer::empty_f64());
    }

    let arg_slices: Vec<Option<&[f64]>> = args
        .iter()
        .map(|a| match a {
            ArgValue::Scalar(_) => None,
            ArgValue::Buffer(b) => Some(b.as_f64_slice()),
        })
        .collect();

    const PAR_THRESHOLD: usize = 4096;

    let result: Vec<f64> = if len >= PAR_THRESHOLD {
        (0..len)
            .into_par_iter()
            .with_min_len(1024)
            .map_init(
                || Vec::with_capacity(flat.max_stack_depth),
                |scratch, i| eval_flat_at(flat, i, &arg_slices, args, scratch),
            )
            .collect()
    } else {
        let mut scratch = Vec::with_capacity(flat.max_stack_depth);
        (0..len)
            .map(|i| eval_flat_at(flat, i, &arg_slices, args, &mut scratch))
            .collect()
    };

    Ok(Buffer::from_f64_vec(result))
}

/// Evaluate the stack machine at element index `i`, reusing `scratch` as the operand stack.
#[inline]