luau-vm 0.732.0

Pure-Rust Luau virtual machine, garbage collector, and standard libraries
Documentation
use core::mem::MaybeUninit;
use core::ptr;

use crate::Table;
use crate::VmResult;
use crate::debug::DebugRuntime;
use crate::handle::RawHandle;
use crate::native::{NativeCallContext, NativeCallResult};
use crate::state::ThreadState;
use crate::thread::Thread;
use crate::types::{LUA_TFUNCTION, LUA_TTABLE};
use crate::value::{RawTValue, TValue};
use crate::vm::VmOperations;

/// `sort_func`
unsafe fn table_sort_function(thread: &Thread, left: TValue, right: TValue) -> VmResult<i32> {
    unsafe {
        debug_assert!(thread.get_top() == 2);

        let top = thread.stack_top();
        top.value_unchecked()
            .set_obj(thread.stack_base().add(1).value_unchecked());
        top.add(1).value_unchecked().set_obj(left);
        top.add(2).value_unchecked().set_obj(right);
        thread.set_stack_top(top.add(3));
        thread.call(2, 1)?;

        let result = thread.stack_top().sub(1);
        thread.set_stack_top(result);
        Ok(i32::from(!result.value_unchecked().is_false()))
    }
}

/// `sort_swap`
unsafe fn table_sort_swap(_thread: &Thread, table: Table, i: i32, j: i32) {
    unsafe {
        let table_ref = table.as_ptr().as_ref().unwrap_unchecked();
        let size_array = table_ref.size_array;
        debug_assert!((i as u32) < size_array as u32 && (j as u32) < size_array as u32);

        let left = table.array_slot(i as usize).as_ptr();
        let right = table.array_slot(j as usize).as_ptr();
        let mut temp = MaybeUninit::<RawTValue>::uninit();
        ptr::copy_nonoverlapping(left, temp.as_mut_ptr(), 1);
        ptr::copy(right, left, 1);
        ptr::copy_nonoverlapping(temp.as_ptr(), right, 1);
    }
}

/// `sort_less`
unsafe fn table_sort_less<const HAS_CUSTOM_PREDICATE: bool>(
    thread: &Thread,
    table: Table,
    i: i32,
    j: i32,
) -> VmResult<i32> {
    unsafe {
        let table_ref = table.as_ptr().as_ref().unwrap_unchecked();
        let size_array = table_ref.size_array;

        debug_assert!((i as u32) < size_array as u32 && (j as u32) < size_array as u32);

        let result = if HAS_CUSTOM_PREDICATE {
            table_sort_function(
                thread,
                table.array_slot(i as usize),
                table.array_slot(j as usize),
            )?
        } else {
            thread.less_than_internal(table.array_slot(i as usize), table.array_slot(j as usize))?
        };

        if table.as_ptr().as_ref().unwrap_unchecked().size_array != size_array {
            return crate::error!(thread, "table modified during sorting").map_err(Into::into);
        }

        Ok(result)
    }
}

/// `sort_siftheap`
unsafe fn table_sort_sift_heap<const HAS_CUSTOM_PREDICATE: bool>(
    thread: &Thread,
    table: Table,
    l: i32,
    u: i32,
    mut root: i32,
) -> VmResult {
    debug_assert!(l <= u);
    let count = u - l + 1;

    unsafe {
        while root * 2 + 2 < count {
            let left = root * 2 + 1;
            let right = root * 2 + 2;
            let mut next = root;
            next = if table_sort_less::<HAS_CUSTOM_PREDICATE>(thread, table, l + next, l + left)?
                != 0
            {
                left
            } else {
                next
            };
            next = if table_sort_less::<HAS_CUSTOM_PREDICATE>(thread, table, l + next, l + right)?
                != 0
            {
                right
            } else {
                next
            };

            if next == root {
                break;
            }

            table_sort_swap(thread, table, l + root, l + next);
            root = next;
        }

        let last_left = root * 2 + 1;
        if last_left == count - 1
            && table_sort_less::<HAS_CUSTOM_PREDICATE>(thread, table, l + root, l + last_left)? != 0
        {
            table_sort_swap(thread, table, l + root, l + last_left);
        }
    }
    Ok(())
}

/// `sort_heap`
unsafe fn table_sort_heap<const HAS_CUSTOM_PREDICATE: bool>(
    thread: &Thread,
    table: Table,
    l: i32,
    u: i32,
) -> VmResult {
    debug_assert!(l <= u);
    let count = u - l + 1;

    unsafe {
        for root in (0..=(count / 2 - 1)).rev() {
            table_sort_sift_heap::<HAS_CUSTOM_PREDICATE>(thread, table, l, u, root)?;
        }

        for index in (1..count).rev() {
            table_sort_swap(thread, table, l, l + index);
            table_sort_sift_heap::<HAS_CUSTOM_PREDICATE>(thread, table, l, l + index - 1, 0)?;
        }
    }
    Ok(())
}

/// `sort_rec`
unsafe fn table_sort_rec<const HAS_CUSTOM_PREDICATE: bool>(
    thread: &Thread,
    table: Table,
    mut l: i32,
    mut u: i32,
    mut limit: i32,
) -> VmResult {
    unsafe {
        while l < u {
            if limit == 0 {
                table_sort_heap::<HAS_CUSTOM_PREDICATE>(thread, table, l, u)?;
                return Ok(());
            }

            if table_sort_less::<HAS_CUSTOM_PREDICATE>(thread, table, u, l)? != 0 {
                table_sort_swap(thread, table, u, l);
            }
            if u - l == 1 {
                break;
            }

            let m = l + ((u - l) >> 1);
            if table_sort_less::<HAS_CUSTOM_PREDICATE>(thread, table, m, l)? != 0 {
                table_sort_swap(thread, table, m, l);
            } else if table_sort_less::<HAS_CUSTOM_PREDICATE>(thread, table, u, m)? != 0 {
                table_sort_swap(thread, table, m, u);
            }
            if u - l == 2 {
                break;
            }

            let p = u - 1;
            table_sort_swap(thread, table, m, p);

            let mut i = l;
            let mut j = u - 1;
            loop {
                loop {
                    i += 1;
                    if table_sort_less::<HAS_CUSTOM_PREDICATE>(thread, table, i, p)? == 0 {
                        break;
                    }
                    if i >= u {
                        return crate::error!(thread, "invalid order function for sorting")
                            .map_err(Into::into);
                    }
                }

                loop {
                    j -= 1;
                    if table_sort_less::<HAS_CUSTOM_PREDICATE>(thread, table, p, j)? == 0 {
                        break;
                    }
                    if j <= l {
                        return crate::error!(thread, "invalid order function for sorting")
                            .map_err(Into::into);
                    }
                }

                if j < i {
                    break;
                }

                table_sort_swap(thread, table, i, j);
            }

            table_sort_swap(thread, table, p, i);
            limit = (limit >> 1) + (limit >> 2);

            if i - l < u - i {
                table_sort_rec::<HAS_CUSTOM_PREDICATE>(thread, table, l, i - 1, limit)?;
                l = i + 1;
            } else {
                table_sort_rec::<HAS_CUSTOM_PREDICATE>(thread, table, i + 1, u, limit)?;
                u = i - 1;
            }
        }
    }
    Ok(())
}

/// `tsort`
pub(super) fn table_sort(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(1, LUA_TTABLE)?;

        let table = thread.stack_base().value_unchecked().table_value();
        let n = table.getn();
        if table.as_ptr().as_ref().unwrap_unchecked().readonly != 0 {
            return thread.readonly_error().map_err(Into::into);
        }

        let has_custom_predicate = thread.is_none_or_nil(2) == 0;
        if has_custom_predicate {
            thread.check_type(2, LUA_TFUNCTION)?;
        }

        thread.set_top(2)?;
        if n > 0 {
            if has_custom_predicate {
                table_sort_rec::<true>(thread, table, 0, n - 1, n)?;
            } else {
                table_sort_rec::<false>(thread, table, 0, n - 1, n)?;
            }
        }
    }
    Ok(0)
}