luau-vm 0.732.0

use luau_common::{BStr, BString, ByteSlice};
use luau_compiler::CompileOptions;
use luau_vm::internal::gc::{GcHeapVisitor, GcRuntime};
use luau_vm::lua::Lua;
use luau_vm::types::{LUA_TFUNCTION, LUA_TPROTO, LUA_TSTRING, LUA_TTHREAD, LUA_TUSERDATA};
use luau_vm::{NativeCallContext, NativeCallResult};
use std::collections::HashMap;

struct HeapNode {
    tag: u8,
    name: BString,
}

struct HeapEnumContext {
    nodes: HashMap<usize, HeapNode>,
    edges: HashMap<usize, usize>,
    seen_target_string: bool,
}

impl GcHeapVisitor for HeapEnumContext {
    fn node(&mut self, ptr: *mut (), tt: u8, _memcat: u8, size: usize, name: Option<&BStr>) {
        if let Some(name) = name {
            let bytes = name.as_bytes();

            if tt == LUA_TUSERDATA as u8 {
                assert_eq!(bytes, b"u42");
            } else if tt == LUA_TPROTO as u8 {
                assert!(
                    bytes == b"proto unnamed:1 =GCDump"
                        || bytes == b"proto foo:7 =GCDump"
                        || bytes == b"proto f:4 =GCDump"
                );
            } else if tt == LUA_TFUNCTION as u8 {
                assert!(
                    bytes == b"test"
                        || bytes == b"unnamed:1 =GCDump"
                        || bytes == b"foo:7 =GCDump"
                        || bytes == b"f:4 =GCDump"
                );
            } else if tt == LUA_TTHREAD as u8 {
                assert_eq!(bytes, b"thread at unnamed:1 =GCDump");
            }
        } else if tt == LUA_TSTRING as u8 && size >= 100_000 {
            assert!(!self.seen_target_string);
            self.seen_target_string = true;
            assert!(size > 100_000);
        }

        self.nodes.insert(
            ptr as usize,
            HeapNode {
                tag: tt,
                name: name
                    .map(|bytes| BString::from(bytes.as_bytes()))
                    .unwrap_or_default(),
            },
        );
    }

    fn edge(&mut self, from: *mut (), to: *mut (), _: &BStr) {
        self.edges.insert(from as usize, to as usize);
    }
}

fn silence(_: NativeCallContext) -> NativeCallResult {
    Ok(0)
}

fn enum_heap_node(
    context: *mut (),
    ptr: *mut (),
    tt: u8,
    memcat: u8,
    size: usize,
    name: Option<&BStr>,
) {
    let context = unsafe { &mut *context.cast::<HeapEnumContext>() };
    context.node(ptr, tt, memcat, size, name);
}

fn enum_heap_edge(context: *mut (), from: *mut (), to: *mut (), name: &BStr) {
    let context = unsafe { &mut *context.cast::<HeapEnumContext>() };
    context.edge(from, to, name);
}

// Conformance.test.cpp: GCDump
#[test]
fn gc_dump() {
    let state = Lua::new().expect("Lua::new should succeed");
    let l = state.main_thread();

    unsafe {
        l.create_table(1, 2).expect("table should create");
        l.push_optional_string(Some("value"))
            .expect("value string should push");
        l.set_field(-2, "key").expect("key field should set");

        l.push_optional_string(Some("u42"))
            .expect("type string should push");
        l.set_field(-2, "__type").expect("__type field should set");

        l.push_integer(42).expect("integer should push");
        l.raw_seti(-2, 1000).expect("array field should set");

        l.push_integer(42).expect("integer should push");
        l.raw_seti(-2, 1).expect("array field should set");

        l.push_value(-1).expect("table should push");
        let _ = l.set_metatable(-2);

        l.new_userdata_tagged(42, 0)
            .expect("userdata should allocate");
        l.push_value(-2).expect("metatable should push");
        let _ = l.set_metatable(-2);

        l.push_integer(1).expect("integer upvalue should push");
        l.push_native_closure_k(silence, Some("test"), 1, None)
            .expect("test closure should push");

        l.new_buffer(100).expect("buffer should allocate");
    }

    let cl = unsafe { l.new_thread().expect("GCDump thread should be created") };
    let source = br#"
local x
x = {}
local function f()
    x[1] = math.abs(42)
end
function foo()
    x[2] = ''
    for i = 1, 10000 do x[2] ..= '1234567890' end
end
foo()
return f
"#;
    let bytecode = luau_compiler::compile_bytes(source, CompileOptions::default());
    assert_eq!(unsafe { cl.load("=GCDump", &bytecode, 0) }, Ok(()));
    assert_eq!(unsafe { cl.resume(None, 0) }, Ok(()));

    unsafe { l.full_gc() };
    let mut dump = BString::new(Vec::new());
    unsafe { l.dump((&raw mut dump).cast(), None) };

    let mut context = HeapEnumContext {
        nodes: HashMap::new(),
        edges: HashMap::new(),
        seen_target_string: false,
    };
    unsafe { l.enum_heap((&raw mut context).cast(), enum_heap_node, enum_heap_edge) };

    assert!(!context.nodes.is_empty());
    assert!(!context.edges.is_empty());
    assert!(context.seen_target_string);

    let has_named_userdata = context
        .nodes
        .values()
        .any(|node| node.tag == LUA_TUSERDATA as u8 && node.name.as_slice() == b"u42");
    assert!(has_named_userdata);
}