luau-vm 0.732.0

use luau_bytecode::builder::{BytecodeBuilder, BytecodeDumpFlags};
use luau_bytecode::opcodes::{FEEDBACK_TYPE_CALLTARGET, PROTO_FLAG_INLINABLE};
use luau_common::{BString, ByteSlice, ScopedFastBool, ScopedFastInt, flags};
use luau_compiler::{CompileOptions, Compiler};
use luau_syntax::allocator::AstArena;
use luau_syntax::ast_names::AstNameTable;
use luau_syntax::parser::ParseOptions;
use luau_vm::internal::RawHandle;
use luau_vm::internal::api::RawStackAccess;
use luau_vm::internal::function::FeedbackVectorSlot;
use luau_vm::internal::function::{Closure, Proto};
use luau_vm::lua::Lua;
use luau_vm::thread::Thread;
use std::ptr::NonNull;

const INLINE_HITS_THRESHOLD: i32 = 2;

#[repr(C)]
struct AssertInlinerData {
    proto: Proto,
    target: Proto,
    pc: u32,
    called: bool,
}

struct CompiledFeedbackChunk {
    bytecode: Vec<u8>,
    function_one_dump: BString,
}

struct FeedbackVectorFixture {
    _emit_call_feedback: ScopedFastBool<'static>,
    _call_feedback: ScopedFastBool<'static>,
    _inline_hits_threshold: ScopedFastInt<'static>,
    state: Lua,
    chunk: Option<CompiledFeedbackChunk>,
}

impl FeedbackVectorFixture {
    fn new() -> Self {
        let emit_call_feedback = flags::LuauEmitCallFeedback.scoped(true);
        let call_feedback = flags::LuauCallFeedback.scoped(true);
        let inline_hits_threshold = flags::LuauInlineHitsThreshold.scoped(INLINE_HITS_THRESHOLD);

        Self {
            _emit_call_feedback: emit_call_feedback,
            _call_feedback: call_feedback,
            _inline_hits_threshold: inline_hits_threshold,
            state: Lua::new().expect("Lua::new should succeed"),
            chunk: None,
        }
    }

    fn thread(&self) -> &Thread {
        self.state.main_thread()
    }

    fn compile(&mut self, source: &str) {
        self.chunk = Some(compile_feedback_chunk(source));
    }

    fn function_one_dump(&self) -> &[u8] {
        &self
            .chunk
            .as_ref()
            .expect("fixture should be compiled before dump assertions")
            .function_one_dump
    }

    fn load(&self) -> Proto {
        let chunk = self
            .chunk
            .as_ref()
            .expect("fixture should be compiled before load");
        let thread = self.thread();
        let status = unsafe { thread.load("=FeedbackVectorTest", &chunk.bytecode, 0) };
        assert_eq!(
            status,
            Ok(()),
            "load failed: {}",
            thread_message(thread, -1),
        );

        let closure = unsafe {
            thread
                .to_object(-1)
                .expect("load should leave a closure on the stack")
                .closure_value()
        };
        unsafe { closure.proto().unwrap_unchecked() }
    }

    fn run(&self, callback: unsafe fn(&Thread, Closure, Closure, u32) -> Option<Proto>) {
        let thread = self.thread();
        unsafe {
            (*thread.global().execution_callbacks()).inline_function = Some(callback);
        }

        let status = unsafe { thread.resume(None, 0) };
        assert_eq!(
            status,
            Ok(()),
            "resume failed: {}",
            thread_message(thread, -1),
        );
    }
}

fn compile_feedback_chunk(source: &str) -> CompiledFeedbackChunk {
    let options = CompileOptions {
        optimization_level: 0,
        ..CompileOptions::default()
    };

    let arena = AstArena::new();
    let mut names = AstNameTable::new(&arena);
    let parse_options = ParseOptions::default();
    let parse_result =
        luau_syntax::parser::parse_bytes(source.as_bytes(), &arena, &mut names, parse_options)
            .unwrap_or_else(|errors| {
                let error = errors.first();
                panic!(
                    "feedback vector source should parse: line {}: {}",
                    error.location.begin.line + 1,
                    error.message
                );
            });

    if let Some(error) = parse_result.metadata.errors.first() {
        panic!(
            "feedback vector source should be error-free: {}",
            error.message
        );
    }

    let mut builder = BytecodeBuilder::new();
    let mut dump_flags = BytecodeDumpFlags::default();
    dump_flags.set_code(true);
    builder.set_dump_flags(dump_flags);
    Compiler::new(options)
        .compile_into(&parse_result, &mut names, &mut builder)
        .unwrap_or_else(|error| panic!("feedback vector source should compile: {error}"));

    CompiledFeedbackChunk {
        bytecode: builder.get_bytecode().to_vec(),
        function_one_dump: builder.dump_function(1),
    }
}

fn thread_message(thread: &Thread, index: i32) -> String {
    unsafe { thread.to_string(index) }
        .expect("thread message conversion should not fail")
        .map(|bytes| String::from_utf8_lossy(bytes.as_bytes()).into_owned())
        .unwrap_or_else(|| "<non-string error>".to_owned())
}

fn dump_with_leading_newline(dump: &[u8]) -> String {
    let mut result = String::from("\n");
    result.push_str(&String::from_utf8_lossy(dump));
    result
}

fn child_proto(proto: Proto, index: usize) -> Proto {
    unsafe { proto.child_proto(index) }.expect("child proto should be present")
}

fn proto_flags(proto: Proto) -> u8 {
    unsafe { proto.as_ptr().as_ref().unwrap_unchecked().flags }
}

fn proto_feedback_vec_size(proto: Proto) -> u32 {
    unsafe { proto.as_ptr().as_ref().unwrap_unchecked().feedback_vec_size }
}

fn proto_fun_id(proto: Proto) -> u32 {
    unsafe { proto.as_ptr().as_ref().unwrap_unchecked().fun_id }
}

fn feedback_slot(proto: Proto, index: usize) -> NonNull<FeedbackVectorSlot> {
    unsafe { proto.feedback_slot(index) }.expect("feedback slot should be materialized")
}

fn feedback_slot_kind(slot: NonNull<FeedbackVectorSlot>) -> i32 {
    unsafe { slot.as_ref().kind }
}

fn feedback_slot_pc(slot: NonNull<FeedbackVectorSlot>) -> u32 {
    unsafe { slot.as_ref().data.call_target.pc }
}

fn feedback_slot_proto(slot: NonNull<FeedbackVectorSlot>) -> u32 {
    unsafe { slot.as_ref().data.call_target.proto }
}

fn feedback_slot_hits(slot: NonNull<FeedbackVectorSlot>) -> u32 {
    unsafe { slot.as_ref().data.call_target.hits }
}

fn feedback_slot_aux(proto: Proto, slot: NonNull<FeedbackVectorSlot>) -> u32 {
    unsafe { proto.code_at(feedback_slot_pc(slot) as usize + 1) }
}

fn seal_feedback_slot(proto: Proto, slot: NonNull<FeedbackVectorSlot>) {
    unsafe { proto.set_code_at(feedback_slot_pc(slot) as usize + 1, u32::MAX) }
}

fn inline_data(thread: &Thread) -> *mut AssertInlinerData {
    unsafe {
        let global = thread.global();
        global
            .as_ptr()
            .as_mut()
            .unwrap_unchecked()
            .ecb_data
            .bytes
            .as_mut_ptr()
            .cast::<AssertInlinerData>()
    }
}

fn open_base_library(thread: &Thread) {
    assert_eq!(
        unsafe { thread.open_base() }.expect("base library should open"),
        1
    );
    unsafe { thread.pop(1) };
}

fn lua_closure_proto(closure: Closure) -> Proto {
    unsafe { closure.proto().unwrap_unchecked() }
}

unsafe fn id_inliner_with_assert(
    thread: &Thread,
    caller: Closure,
    target: Closure,
    pc: u32,
) -> Option<Proto> {
    unsafe {
        let data = inline_data(thread);

        assert!((*data).proto == lua_closure_proto(caller));
        assert!((*data).target == lua_closure_proto(target));
        assert_eq!((*data).pc, pc);
        (*data).called = true;
        Some(lua_closure_proto(caller))
    }
}

unsafe fn id_inliner(_: &Thread, caller: Closure, _: Closure, _: u32) -> Option<Proto> {
    let proto = unsafe { caller.proto().unwrap_unchecked() };
    Some(proto)
}

unsafe fn sealing_inliner(_: &Thread, _: Closure, _: Closure, _: u32) -> Option<Proto> {
    None
}

// FeedbackVector.test.cpp: simple_call
#[test]
fn simple_call() {
    let mut fixture = FeedbackVectorFixture::new();
    fixture.compile(
        r#"
        local function g() return 1 end
        local function f() return g() + 1 end
        f()
        f()
    "#,
    );

    assert_eq!(
        dump_with_leading_newline(fixture.function_one_dump()),
        r#"
GETUPVAL R1 0
CALLFB R1 0 1 [0]
LOADK R2 K0 [1]
ADD R0 R1 R2
RETURN R0 1
"#,
    );

    let top = fixture.load();
    let g = child_proto(top, 0);
    assert_ne!(proto_flags(g) & PROTO_FLAG_INLINABLE, 0);
    let f = child_proto(top, 1);

    assert_eq!(proto_feedback_vec_size(f), 1);

    let slot = feedback_slot(f, 0);
    assert_eq!(
        feedback_slot_kind(slot),
        i32::from(FEEDBACK_TYPE_CALLTARGET)
    );
    assert_eq!(feedback_slot_pc(slot), 1);
    assert_eq!(feedback_slot_proto(slot), 0);
    assert_eq!(feedback_slot_hits(slot), 0);
    assert_eq!(feedback_slot_aux(f, slot), 0);

    unsafe {
        inline_data(fixture.thread()).write(AssertInlinerData {
            proto: f,
            target: g,
            pc: feedback_slot_pc(slot),
            called: false,
        });
    }

    fixture.run(id_inliner_with_assert);

    assert_eq!(feedback_slot_pc(slot), 1);
    assert_eq!(feedback_slot_proto(slot), proto_fun_id(g));
    assert_eq!(feedback_slot_hits(slot), 2);
    assert!(unsafe { (*inline_data(fixture.thread())).called });
}

// FeedbackVector.test.cpp: simple_call_sealed
#[test]
fn simple_call_sealed() {
    let mut fixture = FeedbackVectorFixture::new();
    fixture.compile(
        r#"
        local function g() return 1 end
        local function f() return g() + 1 end
        f()
        f()
    "#,
    );

    let top = fixture.load();
    let f = child_proto(top, 1);
    let slot = feedback_slot(f, 0);

    assert_eq!(feedback_slot_aux(f, slot), 0);
    seal_feedback_slot(f, slot);

    fixture.run(id_inliner);

    assert_eq!(feedback_slot_proto(slot), 0);
    assert_eq!(feedback_slot_hits(slot), 0);
}

// FeedbackVector.test.cpp: simple_call_sealed_on_inline
#[test]
fn simple_call_sealed_on_inline() {
    let mut fixture = FeedbackVectorFixture::new();
    fixture.compile(
        r#"
        local function g() return 1 end
        local function f() return g() + 1 end
        f()
        f()
    "#,
    );

    let top = fixture.load();
    let f = child_proto(top, 1);
    let slot = feedback_slot(f, 0);

    assert_eq!(feedback_slot_aux(f, slot), 0);

    fixture.run(sealing_inliner);

    assert_eq!(feedback_slot_aux(f, slot), u32::MAX);
}

// FeedbackVector.test.cpp: high_order_call
#[test]
fn high_order_call() {
    let mut fixture = FeedbackVectorFixture::new();
    fixture.compile(
        r#"
        local function g() return 1 end
        local function f(h) return h() + 1 end
        f(g)
        f(g)
    "#,
    );

    assert_eq!(
        dump_with_leading_newline(fixture.function_one_dump()),
        r#"
MOVE R2 R0
CALLFB R2 0 1 [0]
LOADK R3 K0 [1]
ADD R1 R2 R3
RETURN R1 1
"#,
    );

    let top = fixture.load();
    let g = child_proto(top, 0);
    assert_ne!(proto_flags(g) & PROTO_FLAG_INLINABLE, 0);
    let f = child_proto(top, 1);

    assert_eq!(proto_feedback_vec_size(f), 1);

    let slot = feedback_slot(f, 0);
    assert_eq!(
        feedback_slot_kind(slot),
        i32::from(FEEDBACK_TYPE_CALLTARGET)
    );
    assert_eq!(feedback_slot_pc(slot), 1);
    assert_eq!(feedback_slot_proto(slot), 0);
    assert_eq!(feedback_slot_hits(slot), 0);
    assert_eq!(feedback_slot_aux(f, slot), 0);

    unsafe {
        inline_data(fixture.thread()).write(AssertInlinerData {
            proto: f,
            target: g,
            pc: feedback_slot_pc(slot),
            called: false,
        });
    }

    fixture.run(id_inliner_with_assert);

    assert_eq!(feedback_slot_pc(slot), 1);
    assert_eq!(feedback_slot_proto(slot), proto_fun_id(g));
    assert_eq!(feedback_slot_hits(slot), 2);
    assert!(unsafe { (*inline_data(fixture.thread())).called });
}

// FeedbackVector.test.cpp: polymorphic_call_sealed
#[test]
fn polymorphic_call_sealed() {
    let mut fixture = FeedbackVectorFixture::new();
    fixture.compile(
        r#"
        local function g() return 1 end
        local function y() return 2 end
        local function f(h) return h() + 1 end
        f(g)
        f(y)
    "#,
    );

    let top = fixture.load();
    let f = child_proto(top, 2);
    let slot = feedback_slot(f, 0);

    assert_eq!(feedback_slot_aux(f, slot), 0);

    fixture.run(id_inliner);

    assert_eq!(feedback_slot_aux(f, slot), u32::MAX);
}

// FeedbackVector.test.cpp: c_call_sealed
#[test]
fn c_call_sealed() {
    let mut fixture = FeedbackVectorFixture::new();
    fixture.compile(
        r#"
        local function f(h) return h(1) + 1 end
        f(tostring)
    "#,
    );

    let top = fixture.load();
    let f = child_proto(top, 0);
    let slot = feedback_slot(f, 0);

    assert_eq!(feedback_slot_aux(f, slot), 0);

    open_base_library(fixture.thread());
    fixture.run(id_inliner);

    assert_eq!(feedback_slot_aux(f, slot), u32::MAX);
}

// FeedbackVector.test.cpp: metamethod_call_sealed
#[test]
fn metamethod_call_sealed() {
    let mut fixture = FeedbackVectorFixture::new();
    fixture.compile(
        r#"
        local function f(h) return h(1) + 1 end

        local callableTable = {}

        setmetatable(callableTable, { __call = function(self, arg) return arg + 42 end })

        f(callableTable)
    "#,
    );

    let top = fixture.load();
    let f = child_proto(top, 0);
    let slot = feedback_slot(f, 0);

    assert_eq!(feedback_slot_aux(f, slot), 0);

    open_base_library(fixture.thread());
    fixture.run(id_inliner);

    assert_eq!(feedback_slot_aux(f, slot), u32::MAX);
}

// FeedbackVector.test.cpp: namecall
#[test]
fn namecall() {
    let mut fixture = FeedbackVectorFixture::new();
    fixture.compile(
        r#"
        local t = { x = 1 }
        function t.g(self) return self.x end
        local function f(t) return t:g() + 1 end
        f(t)
        f(t)
    "#,
    );

    assert_eq!(
        dump_with_leading_newline(fixture.function_one_dump()),
        r#"
NAMECALL R2 R0 K0 ['g']
CALLFB R2 1 1 [0]
LOADK R3 K1 [1]
ADD R1 R2 R3
RETURN R1 1
"#,
    );

    let top = fixture.load();
    let g = child_proto(top, 0);
    assert_ne!(proto_flags(g) & PROTO_FLAG_INLINABLE, 0);
    let f = child_proto(top, 1);

    assert_eq!(proto_feedback_vec_size(f), 1);

    let slot = feedback_slot(f, 0);
    assert_eq!(
        feedback_slot_kind(slot),
        i32::from(FEEDBACK_TYPE_CALLTARGET)
    );
    assert_eq!(feedback_slot_pc(slot), 2);
    assert_eq!(feedback_slot_proto(slot), 0);
    assert_eq!(feedback_slot_hits(slot), 0);
    assert_eq!(feedback_slot_aux(f, slot), 0);

    unsafe {
        inline_data(fixture.thread()).write(AssertInlinerData {
            proto: f,
            target: g,
            pc: feedback_slot_pc(slot),
            called: false,
        });
    }

    fixture.run(id_inliner_with_assert);

    assert_eq!(feedback_slot_proto(slot), proto_fun_id(g));
    assert_eq!(feedback_slot_hits(slot), 2);
    assert!(unsafe { (*inline_data(fixture.thread())).called });
}