datalogic-rs 5.9.0

//! Shared infrastructure for the arithmetic operators: NaN handling,
//! coercion-pair helpers, the integer-checked-with-float-fallback pattern,
//! the variadic fold spec, and the small helpers around `DataValue::Number`
//! allocation.
//!
//! Coercion-pair helpers ([`coerce_pair_int`], [`coerce_pair_f64`]) are thin
//! wrappers over the engine-config-aware coercion in
//! [`crate::arena::coerce_to_number_cfg`] /
//! [`crate::arena::try_coerce_to_integer_cfg`]. See the module doc on
//! `src/arena/value/coercion.rs` for the full coercion-policy map across the
//! crate.

use crate::Engine;
use crate::Result;
use crate::arena::{ContextStack, DataValue, coerce_to_number_cfg, try_coerce_to_integer_cfg};
use crate::config::NanHandling;
use crate::operators::{NanForm, nan_error};
use bumpalo::Bump;
use datavalue::NumberValue;

/// Result of NaN handling check: what the caller should do with a non-numeric value.
pub(super) enum NanAction {
    /// Skip/ignore this value (`IgnoreValue`).
    Skip,
    /// Use `0` in its place (`CoerceToZero`).
    Zero,
    /// Return null immediately.
    ReturnNull,
}

/// Check the engine's NaN handling config and return the appropriate action.
/// Returns `Err` for `ThrowError`, `Ok(NanAction)` otherwise. Takes the
/// context so the `ThrowError` arm can use the deferred thrown-payload fast
/// lane (see [`nan_error`]) when a `try` is guaranteed to catch it.
#[inline]
pub(super) fn handle_nan(ctx: &mut ContextStack<'_>, engine: &Engine) -> Result<NanAction> {
    match engine.config().arithmetic_nan_handling {
        NanHandling::ThrowError => Err(nan_error(NanForm::Thrown, ctx)),
        NanHandling::IgnoreValue => Ok(NanAction::Skip),
        NanHandling::CoerceToZero => Ok(NanAction::Zero),
        NanHandling::ReturnNull => Ok(NanAction::ReturnNull),
    }
}

/// Wrap a [`NumberValue`] in an arena-resident [`DataValue::Number`].
///
/// Note: we do *not* route through `singleton_small_int` here. Most arithmetic
/// results are not small non-negative integers, so probing the singleton table
/// on every alloc_number call costs more than it saves. Sites where the result
/// is *typically* small (length, var-index, reduce-int accumulator) call
/// `singleton_small_int` directly themselves.
#[inline]
pub(super) fn alloc_number<'a>(arena: &'a Bump, n: NumberValue) -> &'a DataValue<'a> {
    arena.alloc(DataValue::Number(n))
}

/// Try the checked-integer op; on overflow promote both operands to `f64` and
/// apply the float fallback. Collapses the recurring
/// `match a.checked_op(b) { Some(r) => from_i64(r), None => from_f64(...) }`
/// pattern that appears across `+`, `-`, and `*`.
///
/// The overflow arm is `from_f64`, which collapses a whole,
/// exactly-i64-representable result back to `Integer`, unlike
/// `NumberValue`'s own `add`/`sub`/`mul` which leave it `Float`. The
/// `map` / `reduce` fast paths make this exact decision by calling this
/// same function (through `operators::array::fused::combine`);
/// `tests/fast_arith_test.rs` checks them against general dispatch.
///
/// Generic over the two operations so each caller's loop gets them
/// inlined; `inline(always)` because the fast paths call it per element.
#[inline(always)]
pub(crate) fn try_int_op(
    a: i64,
    b: i64,
    int_op: impl Fn(i64, i64) -> Option<i64>,
    float_op: impl Fn(f64, f64) -> f64,
) -> NumberValue {
    match int_op(a, b) {
        Some(r) => NumberValue::from_i64(r),
        None => NumberValue::from_f64(float_op(a as f64, b as f64)),
    }
}

/// Coerce both operands to `i64` using the engine's config-aware coercion.
/// Returns `None` if either operand can't be coerced.
#[inline]
pub(super) fn coerce_pair_int(
    a: &DataValue<'_>,
    b: &DataValue<'_>,
    engine: &Engine,
) -> Option<(i64, i64)> {
    Some((
        try_coerce_to_integer_cfg(a, engine)?,
        try_coerce_to_integer_cfg(b, engine)?,
    ))
}

/// Coerce both operands to `f64` using the engine's config-aware coercion.
/// Returns `None` if either operand can't be coerced.
#[inline]
pub(super) fn coerce_pair_f64(
    a: &DataValue<'_>,
    b: &DataValue<'_>,
    engine: &Engine,
) -> Option<(f64, f64)> {
    Some((
        coerce_to_number_cfg(a, engine)?,
        coerce_to_number_cfg(b, engine)?,
    ))
}

/// Spec for an integer-fast-path / float-fallback variadic fold:
/// inits, the integer combine (with overflow signaling via `None`), and
/// the float combine.
#[derive(Clone, Copy)]
pub(super) struct VariadicFoldSpec {
    pub(super) int_init: i64,
    pub(super) float_init: f64,
    pub(super) i_combine: fn(i64, i64) -> Option<i64>,
    pub(super) f_combine: fn(f64, f64) -> f64,
}

/// Running int-fast-path / f64-fallback accumulator state.
///
/// The three arithmetic loops (`variadic_fold`, `subtract_variadic`,
/// `one_arg_array_fold`) used to inline the same `int_acc` /
/// `float_acc` / `all_int` advancement. They diverge on *coercion*
/// (whether to use `as_i64()` strict, `try_coerce_to_integer_cfg`
/// permissive, or both) but agree on what to do once coerced — that
/// shared portion lives here. Each caller picks its coercion strategy
/// and feeds [`Self::step`] pre-coerced `(int_opt, float_opt)`.
pub(super) struct FoldState {
    pub(super) int_acc: i64,
    pub(super) float_acc: f64,
    pub(super) all_int: bool,
}

/// Outcome of a single fold step.
pub(super) enum FoldStepOutcome {
    /// Accumulator advanced; continue the loop.
    Continue,
    /// Operand is non-numeric AND `arithmetic_nan_handling` is `ReturnNull`.
    /// Caller should bail out and return a Null arena value.
    ReturnNull,
}

impl FoldState {
    #[inline]
    pub(super) fn new(int_init: i64, float_init: f64) -> Self {
        Self {
            int_acc: int_init,
            float_acc: float_init,
            all_int: true,
        }
    }

    /// Advance the accumulator by one operand. `int_opt` and `float_opt`
    /// are the caller-coerced views of the operand (caller picks
    /// `as_i64()` / `try_coerce_to_integer_cfg` / `as_f64()` / etc.).
    /// `int_opt.is_some()` short-circuits the f64 path on the
    /// dominant int-only sequence; the caller pre-decides if int is the
    /// right interpretation. When both are `None`, NaN handling kicks
    /// in per the engine config.
    #[inline]
    pub(super) fn step<I, F>(
        &mut self,
        int_opt: Option<i64>,
        float_opt: Option<f64>,
        i_combine: I,
        f_combine: F,
        ctx: &mut ContextStack<'_>,
        engine: &Engine,
    ) -> Result<FoldStepOutcome>
    where
        I: Fn(i64, i64) -> Option<i64>,
        F: Fn(f64, f64) -> f64,
    {
        if self.all_int
            && let Some(i) = int_opt
        {
            match i_combine(self.int_acc, i) {
                Some(r) => self.int_acc = r,
                None => {
                    self.all_int = false;
                    self.float_acc = f_combine(self.int_acc as f64, i as f64);
                }
            }
            return Ok(FoldStepOutcome::Continue);
        }
        // `all_int` already flipped to `false` (a previous arg was non-int)
        // or `int_opt` is `None`. In the former case an integer arg still
        // needs to flow through the float accumulator.
        let float_arg = float_opt.or_else(|| int_opt.map(|i| i as f64));
        if let Some(f) = float_arg {
            if self.all_int {
                self.all_int = false;
                self.float_acc = f_combine(self.int_acc as f64, f);
            } else {
                self.float_acc = f_combine(self.float_acc, f);
            }
            return Ok(FoldStepOutcome::Continue);
        }
        match handle_nan(ctx, engine)? {
            NanAction::Skip => Ok(FoldStepOutcome::Continue),
            // The operand is `0`: an integer, so an all-integer fold stays
            // on the integer track.
            NanAction::Zero => self.step(Some(0), None, i_combine, f_combine, ctx, engine),
            NanAction::ReturnNull => Ok(FoldStepOutcome::ReturnNull),
        }
    }

    /// Materialize the final accumulator as an arena `DataValue::Number`.
    #[inline]
    pub(super) fn finalize<'a>(self, arena: &'a Bump) -> &'a DataValue<'a> {
        if self.all_int {
            alloc_number(arena, NumberValue::from_i64(self.int_acc))
        } else {
            alloc_number(arena, NumberValue::from_f64(self.float_acc))
        }
    }
}

/// True when a compiled node is a literal array — either a structural
/// `Array` node or a `Value` wrapping an `OwnedDataValue::Array`. `+`/`*`
/// and `min`/`max` reject a single literal-array argument with this check.
#[inline]
pub(super) fn is_literal_array(node: &crate::CompiledNode) -> bool {
    matches!(node, crate::CompiledNode::Array { .. })
        || matches!(
            node,
            crate::CompiledNode::Value {
                value: datavalue::OwnedDataValue::Array(_),
                ..
            }
        )
}

/// Variadic fold over arena-evaluated args with integer-fast-path and
/// overflow promotion to `f64`. Used by `+` and `*` for the 3+ arg form.
/// Non-numeric args trigger NaN handling per engine config.
///
/// Coercion strategy: the same as the two-argument form and
/// `subtract_variadic` — native `as_i64()`, then the config-aware integer
/// coercion (so a numeric string stays exact on the int track), then the
/// config-aware `f64` coercion. Native integers, the dominant case, stop
/// at the first test.
#[inline]
pub(super) fn variadic_fold<'a>(
    args: &'a [crate::CompiledNode],
    ctx: &mut crate::arena::ContextStack<'a>,
    engine: &Engine,
    arena: &'a Bump,
    spec: VariadicFoldSpec,
) -> Result<&'a DataValue<'a>> {
    let mut state = FoldState::new(spec.int_init, spec.float_init);
    for arg in args {
        let av = engine.dispatch_node(arg, ctx, arena)?;
        if let FoldStepOutcome::ReturnNull = step_value(&mut state, av, ctx, engine, spec)? {
            return Ok(crate::arena::singletons::singleton_null());
        }
    }
    Ok(state.finalize(arena))
}

/// One operand of [`variadic_fold`], coerced as it coerces.
#[inline(always)]
fn step_value(
    state: &mut FoldState,
    av: &DataValue<'_>,
    ctx: &mut ContextStack<'_>,
    engine: &Engine,
    spec: VariadicFoldSpec,
) -> Result<FoldStepOutcome> {
    let int_opt = av
        .as_i64()
        .or_else(|| try_coerce_to_integer_cfg(av, engine));
    let float_opt = if int_opt.is_some() {
        None
    } else {
        coerce_to_number_cfg(av, engine)
    };
    state.step(
        int_opt,
        float_opt,
        spec.i_combine,
        spec.f_combine,
        ctx,
        engine,
    )
}

/// [`variadic_fold`] resumed from `state` with `first` already evaluated,
/// then the remaining `rest` operands.
#[cfg(feature = "datetime")]
pub(super) fn fold_values<'a>(
    mut state: FoldState,
    first: &'a DataValue<'a>,
    rest: &'a [crate::CompiledNode],
    ctx: &mut ContextStack<'a>,
    engine: &Engine,
    arena: &'a Bump,
    spec: VariadicFoldSpec,
) -> Result<&'a DataValue<'a>> {
    if let FoldStepOutcome::ReturnNull = step_value(&mut state, first, ctx, engine, spec)? {
        return Ok(crate::arena::singletons::singleton_null());
    }
    for arg in rest {
        let av = engine.dispatch_node(arg, ctx, arena)?;
        if let FoldStepOutcome::ReturnNull = step_value(&mut state, av, ctx, engine, spec)? {
            return Ok(crate::arena::singletons::singleton_null());
        }
    }
    Ok(state.finalize(arena))
}