datalogic-rs 5.8.0

//! `sort` — order an array, optionally with an extractor expression for keys.

use std::cmp::Ordering;

use crate::arena::{ContextStack, DataValue, IterGuard, bvec};
use crate::{CompiledNode, Engine, Result};
use bumpalo::Bump;
use datavalue::NumberValue;

use super::input::{IterArgKind, IterSrc, ResolvedInput, plain_var_segments, resolve_iter_input};

/// `sort`. Borrows input via `IterSrc` (no input clone), runs
/// `slice::sort_by` over indices, and emits `DataValue::Array` re-borrowing
/// the original arena items in their sorted order — avoids a deep-clone of
/// the input array, which dominates for object arrays.
///
/// Fast path (extractor is a root-scope `var`): keys come from
/// `traverse_segments` returning `&DataValue` directly, no key clones.
///
/// Scratch (index and key vectors) is bump-allocated in the eval arena,
/// which the session resets wholesale after each evaluation, so sorting
/// makes no per-call heap round-trips for its own buffers.
#[inline]
pub(crate) fn evaluate_sort<'a>(
    args: &'a [CompiledNode],
    iter_arg_kind: IterArgKind,
    ctx: &mut ContextStack<'a>,
    engine: &Engine,
    arena: &'a Bump,
) -> Result<&'a DataValue<'a>> {
    // Literal-null first arg is an error.
    if let CompiledNode::Value { value, .. } = &args[0]
        && value.is_null()
    {
        return Err(crate::Error::invalid_args());
    }

    let src = match resolve_iter_input(&args[0], iter_arg_kind, ctx, engine, arena)? {
        ResolvedInput::Iterable(s) => s,
        ResolvedInput::Empty => return Ok(crate::arena::singletons::singleton_null()),
        ResolvedInput::Bridge(av) => {
            // Bridge is never Array/Null (see ResolvedInput::Bridge), so a sort
            // input reaching here is a scalar or object: not sortable.
            debug_assert!(!matches!(av, DataValue::Array(_) | DataValue::Null));
            return Err(crate::Error::invalid_args());
        }
    };

    let len = src.len();
    if len == 0 {
        return Ok(crate::arena::singletons::singleton_empty_array());
    }
    // Comparisons, on top of the iterator's one per item: n * ceil(log2 n),
    // charged before sorting.
    ctx.charge(comparisons(len))?;

    let ascending = sort_direction(args, ctx, engine, arena)?;

    // No extractor — sort items directly by DataValue order.
    if args.len() <= 2 {
        return Ok(sort_no_extractor(&src, ascending, arena));
    }

    let extractor = &args[2];

    // Fast path: extractor is a root-scope `var` over non-empty segments —
    // keys come from `traverse_segments` directly. Not when traced, so the
    // trace records each item's key, as for every other iterator.
    if super::fast_paths::allowed(ctx, engine)
        && let Some(segments) = sort_key_field(extractor)
    {
        // The key expression's one operation per item, as the general path
        // charges it.
        ctx.charge(len as u64)?;
        return Ok(sort_by_field(&src, segments, ascending, arena));
    }

    // General extractor — push each item, evaluate, collect keys, sort indices.
    sort_general_extractor(&src, extractor, ascending, ctx, engine, arena)
}

/// Read the optional `args[1]` direction flag — defaults to ascending.
#[inline]
fn sort_direction<'a>(
    args: &'a [CompiledNode],
    ctx: &mut ContextStack<'a>,
    engine: &Engine,
    arena: &'a Bump,
) -> Result<bool> {
    if args.len() > 1 {
        let dir = engine.dispatch_node(&args[1], ctx, arena)?;
        Ok(match dir {
            DataValue::Bool(b) => *b,
            _ => true,
        })
    } else {
        Ok(true)
    }
}

#[inline]
fn sort_no_extractor<'a>(src: &IterSrc<'a>, ascending: bool, arena: &'a Bump) -> &'a DataValue<'a> {
    let len = src.len();

    // All-numeric fast path: sort scalar `(key, index)` pairs instead of
    // driving `compare_values` through an index indirection per comparison.
    // The index tiebreaker reproduces the stable sort's equal-key order, so
    // the output is identical to the general path's. The keys order
    // exactly only below 2^53 in magnitude (above it neighbouring integers
    // share a key), so a key that large (rare) sends the array to the
    // general path's exact comparison; the test rides along in the key
    // loop, on the key it already computed.
    if src.0.iter().all(|v| matches!(v, DataValue::Number(_))) {
        const EXACT: f64 = 9_007_199_254_740_992.0; // 2^53
        let mut keyed = bvec::<(f64, u32)>(arena, len);
        let mut wide = false;
        keyed.extend(src.0.iter().enumerate().map(|(i, v)| {
            let f = match v {
                DataValue::Number(n) => n.as_f64(),
                _ => unreachable!(),
            };
            wide |= f.abs() >= EXACT;
            // Collapse -0.0 to 0.0 so `total_cmp` can't order the two zero
            // representations, keeping zero-keyed ties on the index
            // tiebreaker exactly like the stable path.
            (if f == 0.0 { 0.0 } else { f }, i as u32)
        }));
        if !wide {
            keyed.sort_unstable_by(|(ka, ia), (kb, ib)| {
                // `total_cmp` keeps the comparator a total order even for
                // NaN keys (which JSON data can't contain anyway — pdqsort
                // panics on inconsistent comparators, so `partial_cmp` is
                // off the table). Ties break by input position, reproducing
                // the stable sort's equal-key order.
                let cmp = ka.total_cmp(kb);
                let cmp = if ascending { cmp } else { cmp.reverse() };
                cmp.then(ia.cmp(ib))
            });
            let slice =
                arena.alloc_slice_fill_iter(keyed.iter().map(|&(_, i)| *src.get(i as usize)));
            return arena.alloc(DataValue::Array(slice));
        }
    }

    let mut indices = bvec::<usize>(arena, len);
    indices.extend(0..len);
    indices.sort_by(|&a, &b| {
        let cmp = compare_values(src.get(a), src.get(b));
        if ascending { cmp } else { cmp.reverse() }
    });
    let slice = arena.alloc_slice_fill_iter(indices.iter().map(|&i| *src.get(i)));
    arena.alloc(DataValue::Array(slice))
}

/// The sort-by-field fast path's shape: a key expression that is a plain
/// element field (`{"var": "a.b"}`, non-empty path, no default). Listed in
/// [`super::fast_paths`].
#[inline]
pub(super) fn sort_key_field(extractor: &CompiledNode) -> Option<&[crate::node::PathSegment]> {
    plain_var_segments(extractor).filter(|segments| !segments.is_empty())
}

/// Extractor fast path: `{var: "field..."}` over non-empty segments at
/// scope 0 (see [`sort_key_field`]).
#[inline]
fn sort_by_field<'a>(
    src: &IterSrc<'a>,
    segments: &[crate::node::PathSegment],
    ascending: bool,
    arena: &'a Bump,
) -> &'a DataValue<'a> {
    let len = src.len();
    // A missing field reads as `null`, as the key expression would on the
    // general path, so it ties with a present `null` and keeps its place.
    let null = crate::arena::singletons::singleton_null();
    let mut keyed = bvec::<(usize, &'a DataValue<'a>)>(arena, len);
    keyed.extend((0..len).map(|i| {
        (
            i,
            crate::arena::value::traverse_segments(src.get(i), segments).unwrap_or(null),
        )
    }));
    keyed.sort_by(|(_, ka), (_, kb)| {
        let cmp = compare_values(ka, kb);
        if ascending { cmp } else { cmp.reverse() }
    });
    let slice = arena.alloc_slice_fill_iter(keyed.iter().map(|&(i, _)| *src.get(i)));
    arena.alloc(DataValue::Array(slice))
}

#[inline]
fn sort_general_extractor<'a>(
    src: &IterSrc<'a>,
    extractor: &'a CompiledNode,
    ascending: bool,
    ctx: &mut ContextStack<'a>,
    engine: &Engine,
    arena: &'a Bump,
) -> Result<&'a DataValue<'a>> {
    let len = src.len();
    let mut keys = bvec::<DataValue<'a>>(arena, len);
    let mut guard = IterGuard::new(ctx);
    for i in 0..len {
        let item = src.get(i);
        guard.step_indexed(item, i);
        let key_av = engine.dispatch_node(extractor, guard.stack(), arena)?;
        keys.push(*key_av);
    }
    drop(guard);

    let mut indices = bvec::<usize>(arena, len);
    indices.extend(0..len);
    indices.sort_by(|&a, &b| {
        let cmp = compare_values(&keys[a], &keys[b]);
        if ascending { cmp } else { cmp.reverse() }
    });
    let slice = arena.alloc_slice_fill_iter(indices.iter().map(|&i| *src.get(i)));
    Ok(arena.alloc(DataValue::Array(slice)))
}

/// Compare arena values for sorting.
/// Type order: null < bool < number < string < array < object.
#[inline]
fn compare_values(a: &DataValue<'_>, b: &DataValue<'_>) -> Ordering {
    #[inline]
    fn type_rank(v: &DataValue<'_>) -> u8 {
        match v {
            DataValue::Null => 0,
            DataValue::Bool(_) => 1,
            DataValue::Number(_) => 2,
            DataValue::String(_) => 3,
            DataValue::Array(_) => 4,
            DataValue::Object(_) => 5,
            #[cfg(feature = "datetime")]
            DataValue::DateTime(_) | DataValue::Duration(_) => 3,
            #[cfg(feature = "tensor")]
            DataValue::Tensor(_) => 6,
        }
    }

    match (a, b) {
        (DataValue::Null, DataValue::Null) => Ordering::Equal,
        (DataValue::Bool(a), DataValue::Bool(b)) => a.cmp(b),
        (DataValue::Number(a), DataValue::Number(b)) => cmp_numbers(a, b),
        (DataValue::String(a), DataValue::String(b)) => a.cmp(b),
        (DataValue::Array(_), DataValue::Array(_)) => Ordering::Equal,
        (DataValue::Object(_), DataValue::Object(_)) => Ordering::Equal,
        // Arrays and objects compare Equal (sort is stable, so they keep
        // input order). Tensors do better because they have a cheap total
        // order available: dtype, then shape, then the raw payload. It is
        // not numerically meaningful, but it is deterministic, which is
        // what a sort key needs.
        #[cfg(feature = "tensor")]
        (DataValue::Tensor(x), DataValue::Tensor(y)) => {
            (x.dtype().name(), x.shape(), x.data()).cmp(&(y.dtype().name(), y.shape(), y.data()))
        }
        _ => type_rank(a).cmp(&type_rank(b)),
    }
}

/// The numeric order `sort` uses: by value, exactly. Two integers compare
/// as integers (an `f64` comparison calls neighbours above 2^53 equal),
/// and an integer and a float by their exact values, so the order stays
/// total (a mix of exact and `f64` comparisons would not be, and a sort
/// needs a total order).
#[inline]
fn cmp_numbers(a: &NumberValue, b: &NumberValue) -> Ordering {
    match (*a, *b) {
        (NumberValue::Integer(x), NumberValue::Integer(y)) => x.cmp(&y),
        (NumberValue::Float(x), NumberValue::Float(y)) => cmp_floats(x, y),
        (NumberValue::Integer(x), NumberValue::Float(y)) => cmp_int_float(x, y),
        (NumberValue::Float(x), NumberValue::Integer(y)) => cmp_int_float(y, x).reverse(),
    }
}

/// Two floats; incomparable ones (NaN, which JSON cannot carry) tie.
#[inline]
fn cmp_floats(x: f64, y: f64) -> Ordering {
    x.partial_cmp(&y).unwrap_or(Ordering::Equal)
}

/// An integer against a float, exactly. `i as f64` rounds to the nearest
/// float, which keeps its order against any float, so only a tie needs a
/// second look: then `f` is a whole number, compared as an integer (or it
/// is 2^63, above every `i64`).
#[inline]
fn cmp_int_float(i: i64, f: f64) -> Ordering {
    match cmp_floats(i as f64, f) {
        Ordering::Equal if f.is_finite() => {
            if f >= 9_223_372_036_854_775_808.0 {
                Ordering::Less
            } else {
                i.cmp(&(f as i64))
            }
        }
        other => other,
    }
}

/// `n * ceil(log2 n)`: the comparison count a sort of `n` items is charged.
#[inline]
fn comparisons(n: usize) -> u64 {
    if n < 2 {
        return 0;
    }
    let log2 = u64::from(usize::BITS - (n - 1).leading_zeros());
    (n as u64).saturating_mul(log2)
}