datalogic-rs 5.9.0

//! `slice` — array and string slicing with optional start/end/step.

use crate::arena::{ContextStack, DataValue, bvec};
use crate::operators::string::char_to_byte_offset;
use crate::operators::{NanForm, nan_error};
use crate::{CompiledNode, Engine, Result};
use bumpalo::Bump;

/// Native arena-mode `slice`. Returns array slices as views over arena items;
/// string slices are allocated in the arena.
#[inline]
pub(crate) fn evaluate_slice<'a>(
    args: &'a [CompiledNode],
    ctx: &mut ContextStack<'a>,
    engine: &Engine,
    arena: &'a Bump,
) -> Result<&'a DataValue<'a>> {
    let coll_av = engine.dispatch_node(&args[0], ctx, arena)?;

    // Null passthrough.
    if matches!(coll_av, DataValue::Null) {
        return Ok(crate::arena::singletons::singleton_null());
    }

    // Resolve start/end/step.
    let start = if args.len() > 1 {
        extract_opt_i64_arena(&args[1], ctx, engine, arena)?
    } else {
        None
    };
    let end = if args.len() > 2 {
        extract_opt_i64_arena(&args[2], ctx, engine, arena)?
    } else {
        None
    };
    let step = if args.len() > 3 {
        let s = extract_opt_i64_arena(&args[3], ctx, engine, arena)?.unwrap_or(1);
        if s == 0 {
            return Err(crate::Error::invalid_args());
        }
        s
    } else {
        1
    };

    if let DataValue::Array(items) = coll_av {
        return slice_array(items, start, end, step, ctx, arena);
    }
    if let DataValue::String(s) = coll_av {
        // Char-indexing walks the string, stepped or not.
        ctx.charge_bytes(s.len())?;
        return Ok(slice_string(s, start, end, step, arena));
    }
    Err(crate::Error::invalid_args())
}

/// Composite arena array — slice through the arena items.
///
/// A contiguous slice borrows the source and costs nothing beyond its node.
/// A stepped one copies, and is charged one operation per item it produces
/// before copying them.
#[inline]
fn slice_array<'a>(
    items: &'a [DataValue<'a>],
    start: Option<i64>,
    end: Option<i64>,
    step: i64,
    ctx: &mut ContextStack<'a>,
    arena: &'a Bump,
) -> Result<&'a DataValue<'a>> {
    let len = items.len() as i64;

    // Contiguous step==1 (the default, no explicit step): the selected
    // indices form the run [s, e), so return a borrowed arena sub-slice
    // directly instead of materializing an index Vec and copying elements.
    if step == 1 {
        let s = normalize_index(start.unwrap_or(0), len);
        let e = normalize_index(end.unwrap_or(len), len);
        if s >= e {
            return Ok(crate::arena::singletons::singleton_empty_array());
        }
        return Ok(arena.alloc(DataValue::Array(&items[s as usize..e as usize])));
    }

    let indices = slice_indices(len, start, end, step);
    let count = indices.clone().count();
    if count == 0 {
        return Ok(crate::arena::singletons::singleton_empty_array());
    }
    ctx.charge(count as u64)?;
    let mut out = bvec::<DataValue<'a>>(arena, count);
    for i in indices {
        out.push(items[i as usize]);
    }
    Ok(arena.alloc(DataValue::Array(out.into_bump_slice())))
}

/// String slice — char-indexed; the contiguous `step == 1` case borrows a
/// sub-slice of the source, other steps allocate the result in the arena.
#[inline]
fn slice_string<'a>(
    s: &'a str,
    start: Option<i64>,
    end: Option<i64>,
    step: i64,
    arena: &'a Bump,
) -> &'a DataValue<'a> {
    if s.is_ascii() {
        let len = s.len() as i64;
        if step == 1 {
            let a = normalize_index(start.unwrap_or(0), len) as usize;
            let b = normalize_index(end.unwrap_or(len), len) as usize;
            if a >= b {
                return crate::arena::singletons::singleton_empty_string();
            }
            return arena.alloc(DataValue::String(&s[a..b]));
        }
        let bytes = s.as_bytes();
        let indices = slice_indices(len, start, end, step);
        let mut out = bumpalo::collections::String::new_in(arena);
        for i in indices {
            out.push(bytes[i as usize] as char);
        }
        if out.is_empty() {
            return crate::arena::singletons::singleton_empty_string();
        }
        return arena.alloc(DataValue::String(out.into_bump_str()));
    }

    // Contiguous step==1: mirror the array arm. The selected chars form the
    // run [a, b), and `s` is already an arena `&'a str`, so return a borrowed
    // sub-slice at the char boundaries instead of collecting chars.
    if step == 1 {
        let char_len = s.chars().count() as i64;
        let a = normalize_index(start.unwrap_or(0), char_len);
        let b = normalize_index(end.unwrap_or(char_len), char_len);
        if a >= b {
            return crate::arena::singletons::singleton_empty_string();
        }
        let byte_a = char_to_byte_offset(s, a as usize);
        let byte_b = byte_a + char_to_byte_offset(&s[byte_a..], (b - a) as usize);
        return arena.alloc(DataValue::String(&s[byte_a..byte_b]));
    }

    // Random access by char index: the chars go into an arena buffer, and
    // the result is built straight into the arena (no heap round-trip).
    let mut chars = bvec::<char>(arena, s.len());
    chars.extend(s.chars());
    let mut out = bumpalo::collections::String::new_in(arena);
    for i in slice_indices(chars.len() as i64, start, end, step) {
        out.push(chars[i as usize]);
    }
    if out.is_empty() {
        return crate::arena::singletons::singleton_empty_string();
    }
    arena.alloc(DataValue::String(out.into_bump_str()))
}

#[inline]
fn extract_opt_i64_arena<'a>(
    node: &'a CompiledNode,
    ctx: &mut ContextStack<'a>,
    engine: &Engine,
    arena: &'a Bump,
) -> Result<Option<i64>> {
    // A literal skips the dispatch but reads exactly as the evaluated value
    // below: a fractional number is not an index, written or computed.
    if let CompiledNode::Value { value, .. } = node {
        return match value {
            datavalue::OwnedDataValue::Null => Ok(None),
            datavalue::OwnedDataValue::Number(n) => match n.as_i64() {
                Some(i) => Ok(Some(i)),
                None => Err(nan_error(NanForm::InvalidArguments, ctx)),
            },
            _ => Err(nan_error(NanForm::InvalidArguments, ctx)),
        };
    }
    let av = engine.dispatch_node(node, ctx, arena)?;
    match av {
        DataValue::Null => Ok(None),
        _ => match av.as_i64() {
            Some(i) => Ok(Some(i)),
            None => Err(nan_error(NanForm::InvalidArguments, ctx)),
        },
    }
}

/// The indices a slice selects, in order, given start/end/step. An iterator
/// rather than a list so a caller can count the selection (to charge for
/// it) before allocating anything.
#[inline]
fn slice_indices(len: i64, start: Option<i64>, end: Option<i64>, step: i64) -> SliceIndices {
    let (actual_start, actual_end) = if step > 0 {
        (
            normalize_index(start.unwrap_or(0), len),
            normalize_index(end.unwrap_or(len), len),
        )
    } else {
        let default_start = len.saturating_sub(1);
        let s = normalize_index(start.unwrap_or(default_start), len);
        let e = if let Some(e) = end {
            if e < -len {
                -1
            } else {
                normalize_index(e, len)
            }
        } else {
            -1
        };
        (s, e)
    };
    SliceIndices {
        next: Some(actual_start),
        end: actual_end,
        step,
        len,
    }
}

/// See [`slice_indices`]. `next` is `None` once a step would overflow.
#[derive(Clone)]
struct SliceIndices {
    next: Option<i64>,
    end: i64,
    step: i64,
    len: i64,
}

impl Iterator for SliceIndices {
    type Item = i64;

    #[inline]
    fn next(&mut self) -> Option<i64> {
        while let Some(i) = self.next {
            if !((self.step > 0 && i < self.end) || (self.step < 0 && i > self.end)) {
                self.next = None;
                return None;
            }
            self.next = i.checked_add(self.step);
            if i >= 0 && i < self.len {
                return Some(i);
            }
        }
        None
    }
}

/// Normalize slice indices with overflow protection.
#[inline]
fn normalize_index(index: i64, len: i64) -> i64 {
    if index < 0 {
        // Use saturating_add to prevent overflow when index is very negative
        let adjusted = len.saturating_add(index);
        adjusted.max(0)
    } else {
        index.min(len)
    }
}