luau-vm 0.732.0

use crate::support::{ConformanceOptions, run_source_chunk};
use core::alloc::Layout;
use core::cell::Cell;
use core::ptr::NonNull;
use luau_common::ByteSlice;
use luau_compiler::CompileOptions;
use luau_vm::internal::gc::GcRuntime;
use luau_vm::lua::Lua;
use luau_vm::state::{LuaAllocator, SystemLuaAllocator};
use luau_vm::{VmError, VmExit};

fn make_huge_function_source() -> String {
    let mut source = String::new();

    source += "if ... then\n";
    source += "local _ = {\n";

    for index in 0..40_000 {
        source += "0.";
        source += &index.to_string();
        source += ",";
    }

    source += "}\n";
    source += "end\n";
    source += "return bit32.lshift('84', -1)";

    source
}

fn default_large_compile_options() -> CompileOptions {
    CompileOptions::default()
}

// Conformance.test.cpp: HugeFunction
#[test]
fn huge_function() {
    let options = ConformanceOptions {
        compile: default_large_compile_options(),
        ..ConformanceOptions::default()
    };
    let source = make_huge_function_source();
    let state = run_source_chunk("=HugeFunction", source.as_bytes(), &options);
    let thread = state.main_thread();

    assert_eq!(unsafe { thread.to_number(-1) }, Some(42.0));
}

struct HugeFunctionLoadFailureState {
    large_allocation_to_fail: usize,
    large_allocation_count: Cell<usize>,
}

unsafe impl LuaAllocator for HugeFunctionLoadFailureState {
    unsafe fn allocate(&self, layout: Layout) -> Option<NonNull<u8>> {
        if layout.size() > 32_768 {
            let large_allocation_count = self.large_allocation_count.get();
            if large_allocation_count == self.large_allocation_to_fail {
                return None;
            }

            self.large_allocation_count.set(large_allocation_count + 1);
        }

        unsafe { SystemLuaAllocator.allocate(layout) }
    }

    unsafe fn reallocate(
        &self,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Option<NonNull<u8>> {
        if new_layout.size() > old_layout.size() && new_layout.size() > 32_768 {
            let large_allocation_count = self.large_allocation_count.get();
            if large_allocation_count == self.large_allocation_to_fail {
                return None;
            }

            self.large_allocation_count.set(large_allocation_count + 1);
        }

        unsafe { SystemLuaAllocator.reallocate(ptr, old_layout, new_layout) }
    }

    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
        unsafe { SystemLuaAllocator.deallocate(ptr, layout) };
    }
}

// Conformance.test.cpp: HugeFunctionLoadFailure
#[test]
fn huge_function_load_failure() {
    let source = make_huge_function_source();
    let bytecode = luau_compiler::compile_bytes(source, default_large_compile_options());

    const EXPECTED_TOTAL_LARGE_ALLOCATIONS: usize = 2;
    for large_allocation_to_fail in 0..EXPECTED_TOTAL_LARGE_ALLOCATIONS {
        let failure_state = HugeFunctionLoadFailureState {
            large_allocation_to_fail,
            large_allocation_count: Cell::new(0),
        };

        let state =
            Lua::new_with_allocator(failure_state).expect("Lua::new_with_allocator should succeed");
        let thread = state.main_thread();

        unsafe {
            thread.open_libs().expect("libraries should open");
            thread.sandbox().expect("state should sandbox");
            thread.sandbox_thread().expect("thread should sandbox");
        }

        let status = unsafe { thread.load("=HugeFunction", &bytecode, 0) };
        assert_eq!(status, Err(VmExit::Error(VmError::Memory)));
        assert_eq!(
            unsafe { thread.to_string(-1) }
                .expect("load failure should produce a stringable error")
                .unwrap()
                .as_bytes(),
            b"not enough memory"
        );

        unsafe { thread.full_gc() };
    }
}

// Conformance.test.cpp: HugeConstantTable
#[test]
fn huge_constant_table() {
    let mut source = String::from("function foo(...)\n");

    source += "    local args = ...\n";
    source += "    local t = args and {\n";

    for i in 0..400 {
        for k in 0..100 {
            source += "call(";
            source += &(i * 100 + k).to_string();
            source += ".125), ";
        }

        source += "\n        ";
    }

    source += "    }\n";
    source += "    return { a = 1, b = 2, c = 3 }\n";
    source += "end\n";
    source += "return foo().a + foo().b\n";

    let options = ConformanceOptions {
        compile: default_large_compile_options(),
        ..ConformanceOptions::default()
    };
    let state = run_source_chunk("=HugeConstantTable", source.as_bytes(), &options);
    let thread = state.main_thread();

    assert_eq!(unsafe { thread.to_number(-1) }, Some(3.0));
}

// Conformance.test.cpp: LargeNestedClosure
#[test]
fn large_nested_closure() {
    const COUNT: i32 = 2048;

    let mut source = String::from("local function test()\n");
    source += "local x = 0\n";

    for index in 0..COUNT {
        let name = (index + 1).to_string();
        source += "    function f";
        source += &name;
        source += "() x = x + 1; return ";
        source += &name;
        source += " end\n";
    }

    source += "    return f";
    source += &COUNT.to_string();
    source += "\nend\n";
    source += "return test()()\n";

    let options = ConformanceOptions {
        compile: default_large_compile_options(),
        ..ConformanceOptions::default()
    };
    let state = run_source_chunk("=LargeNestedClosure", source.as_bytes(), &options);
    let thread = state.main_thread();

    assert_eq!(unsafe { thread.to_number(-1) }, Some(COUNT as f64));
}