mintrt 0.1.0

//     <<*>> Revenge of snowflake.
//     -^v-  SASWE
//     @ This source code is licensed under a dual license model:
//       1. MPL-2.0 for non-commercial use only
//       2. Commercial License for commercial use
//       See LICENSE and LICENSE_COMMERCIAL files for details.

//! Lua 字节码数据结构定义
//! 
//! 本模块包含字节码解析所需的所有数据结构

use std::io::{Read, BufReader};
use std::fs::File;
use std::path::Path;

/// Lua 5.5 字节码头部结构
#[derive(Debug)]
pub struct LuaHeader {
    pub signature: [u8; 4],      // Lua 字节码签名：0x1B 0x4C 0x75 0x61 (\033Lua)
    pub version: u8,              // Lua 版本 (0x53 = 5.3)
    pub format_version: u8,       // 格式版本 (0 = 官方)
}

/// 函数原型信息
#[derive(Debug, Clone)]
pub struct FunctionPrototype {
    pub source_name: String,      // 源文件名
    pub line_defined: i32,        // 起始行号
    pub last_line_defined: i32,   // 结束行号
    pub num_params: u8,           // 参数个数
    pub max_stack_size: i32,      // 最大栈大小
    pub instructions: Vec<u32>,   // 指令数组
    pub upvalues_count: usize,    // 上值数量
    pub constants_count: usize,   // 常量数量
    pub functions_count: usize,   // 内部函数数量
}

/// CALL 指令信息
#[derive(Debug)]
pub struct CallInfo {
    pub instruction_index: usize,  // 指令位置
    pub opcode: OpCode,            // 操作码 (CALL 或 TAILCALL)
    pub base_register: u32,        // 基址寄存器 A
    pub num_args: u32,             // 参数个数 B
    pub num_returns: u32,          // 返回值个数 C
    pub line_number: Option<i32>,  // 源代码行号 (如果有调试信息)
}

/// 字节码分析结果
#[derive(Debug)]
pub struct BytecodeAnalysis {
    pub has_main: bool,                    // 是否有 main 函数
    pub main_calls: Vec<CallInfo>,         // main 函数中的 CALL 指令
    pub main_tailcalls: Vec<CallInfo>,     // main 函数中的 TAILCALL 指令
    pub total_functions: usize,            // 总函数数量
    pub main_function: Option<FunctionPrototype>, // main 函数原型
}

/// 操作码枚举 (Lua 5.5)
#[derive(Debug, Clone, Copy, PartialEq)]
#[repr(u8)]
pub enum OpCode {
    MOVE = 0,
    LOADI = 1,
    LOADF = 2,
    LOADK = 3,
    LOADKX = 4,
    LOADFALSE = 5,
    LFALSESKIP = 6,
    LOADTRUE = 7,
    LOADNIL = 8,
    GETUPVAL = 9,
    SETUPVAL = 10,
    GETTABUP = 11,
    GETTABLE = 12,
    GETI = 13,
    GETFIELD = 14,
    SETTABUP = 15,
    SETTABLE = 16,
    SETI = 17,
    SETFIELD = 18,
    NEWTABLE = 19,
    SELF = 20,
    ADDI = 21,
    ADDK = 22,
    SUBK = 23,
    MULK = 24,
    MODK = 25,
    POWK = 26,
    DIVK = 27,
    IDIVK = 28,
    BANDK = 29,
    BORK = 30,
    BXORK = 31,
    SHLI = 32,
    SHRI = 33,
    ADD = 34,
    SUB = 35,
    MUL = 36,
    MOD = 37,
    POW = 38,
    DIV = 39,
    IDIV = 40,
    BAND = 41,
    BOR = 42,
    BXOR = 43,
    SHL = 44,
    SHR = 45,
    MMBIN = 46,
    MMBINI = 47,
    MMBINK = 48,
    UNM = 49,
    BNOT = 50,
    NOT = 51,
    LEN = 52,
    CONCAT = 53,
    CLOSE = 54,
    TBC = 55,
    JMP = 56,
    EQ = 57,
    LT = 58,
    LE = 59,
    EQK = 60,
    EQI = 61,
    LTI = 62,
    LEI = 63,
    GTI = 64,
    GEI = 65,
    TEST = 66,
    TESTSET = 67,
    CALL = 68,
    TAILCALL = 69,
    RETURN = 70,
    RETURN0 = 71,
    RETURN1 = 72,
    FORLOOP = 73,
    FORPREP = 74,
    TFORPREP = 75,
    TFORCALL = 76,
    TFORLOOP = 77,
    SETLIST = 78,
    CLOSURE = 79,
    VARARG = 80,
    GETVARG = 81,
    ERRNNIL = 82,
    VARARGPREP = 83,
    EXTRAARG = 84,
}

impl OpCode {
    /// 从 u32 指令中提取操作码
    pub fn from_instruction(instruction: u32) -> OpCode {
        let opcode = (instruction & 0x7F) as u8;
        unsafe { std::mem::transmute(opcode) }//ai 牛逼的，还能这么写
    }
    
    /// 获取操作码名称
    pub fn name(&self) -> &'static str {
        match self {
            OpCode::MOVE => "MOVE",
            OpCode::LOADI => "LOADI",
            OpCode::LOADF => "LOADF",
            OpCode::LOADK => "LOADK",
            OpCode::LOADKX => "LOADKX",
            OpCode::LOADFALSE => "LOADFALSE",
            OpCode::LFALSESKIP => "LFALSESKIP",
            OpCode::LOADTRUE => "LOADTRUE",
            OpCode::LOADNIL => "LOADNIL",
            OpCode::GETUPVAL => "GETUPVAL",
            OpCode::SETUPVAL => "SETUPVAL",
            OpCode::GETTABUP => "GETTABUP",
            OpCode::GETTABLE => "GETTABLE",
            OpCode::GETI => "GETI",
            OpCode::GETFIELD => "GETFIELD",
            OpCode::SETTABUP => "SETTABUP",
            OpCode::SETTABLE => "SETTABLE",
            OpCode::SETI => "SETI",
            OpCode::SETFIELD => "SETFIELD",
            OpCode::NEWTABLE => "NEWTABLE",
            OpCode::SELF => "SELF",
            OpCode::ADDI => "ADDI",
            OpCode::ADDK => "ADDK",
            OpCode::SUBK => "SUBK",
            OpCode::MULK => "MULK",
            OpCode::MODK => "MODK",
            OpCode::POWK => "POWK",
            OpCode::DIVK => "DIVK",
            OpCode::IDIVK => "IDIVK",
            OpCode::BANDK => "BANDK",
            OpCode::BORK => "BORK",
            OpCode::BXORK => "BXORK",
            OpCode::SHLI => "SHLI",
            OpCode::SHRI => "SHRI",
            OpCode::ADD => "ADD",
            OpCode::SUB => "SUB",
            OpCode::MUL => "MUL",
            OpCode::MOD => "MOD",
            OpCode::POW => "POW",
            OpCode::DIV => "DIV",
            OpCode::IDIV => "IDIV",
            OpCode::BAND => "BAND",
            OpCode::BOR => "BOR",
            OpCode::BXOR => "BXOR",
            OpCode::SHL => "SHL",
            OpCode::SHR => "SHR",
            OpCode::MMBIN => "MMBIN",
            OpCode::MMBINI => "MMBINI",
            OpCode::MMBINK => "MMBINK",
            OpCode::UNM => "UNM",
            OpCode::BNOT => "BNOT",
            OpCode::NOT => "NOT",
            OpCode::LEN => "LEN",
            OpCode::CONCAT => "CONCAT",
            OpCode::CLOSE => "CLOSE",
            OpCode::TBC => "TBC",
            OpCode::JMP => "JMP",
            OpCode::EQ => "EQ",
            OpCode::LT => "LT",
            OpCode::LE => "LE",
            OpCode::EQK => "EQK",
            OpCode::EQI => "EQI",
            OpCode::LTI => "LTI",
            OpCode::LEI => "LEI",
            OpCode::GTI => "GTI",
            OpCode::GEI => "GEI",
            OpCode::TEST => "TEST",
            OpCode::TESTSET => "TESTSET",
            OpCode::CALL => "CALL",
            OpCode::TAILCALL => "TAILCALL",
            OpCode::RETURN => "RETURN",
            OpCode::RETURN0 => "RETURN0",
            OpCode::RETURN1 => "RETURN1",
            OpCode::FORLOOP => "FORLOOP",
            OpCode::FORPREP => "FORPREP",
            OpCode::TFORPREP => "TFORPREP",
            OpCode::TFORCALL => "TFORCALL",
            OpCode::TFORLOOP => "TFORLOOP",
            OpCode::SETLIST => "SETLIST",
            OpCode::CLOSURE => "CLOSURE",
            OpCode::VARARG => "VARARG",
            OpCode::GETVARG => "GETVARG",
            OpCode::ERRNNIL => "ERRNNIL",
            OpCode::VARARGPREP => "VARARGPREP",
            OpCode::EXTRAARG => "EXTRAARG",
        }
    }
    
    /// 检查是否为 CALL 或 TAILCALL 指令
    pub fn is_call(&self) -> bool {
        matches!(self, OpCode::CALL | OpCode::TAILCALL)
    }
}

// ============================================================================
// 以下是字节码解析的辅助函数
// ============================================================================

/// 读取 u32 (小端序)
pub fn read_u32_le(reader: &mut impl Read) -> std::io::Result<u32> {
    let mut buf = [0u8; 4];
    reader.read_exact(&mut buf)?;
    Ok(u32::from_le_bytes(buf))
}

/// 读取 i32 (小端序)
pub fn read_i32_le(reader: &mut impl Read) -> std::io::Result<i32> {
    let mut buf = [0u8; 4];
    reader.read_exact(&mut buf)?;
    Ok(i32::from_le_bytes(buf))
}

/// 读取 u8
pub fn read_u8(reader: &mut impl Read) -> std::io::Result<u8> {
    let mut buf = [0u8; 1];
    reader.read_exact(&mut buf)?;
    Ok(buf[0])
}

/// 读取字符串 (带长度前缀)
pub fn read_string(reader: &mut impl Read) -> std::io::Result<String> {
    let size = read_u8(reader)? as usize;
    if size == 0 {
        return Ok(String::new());
    }
    
    let mut buf = vec![0u8; size];
    reader.read_exact(&mut buf)?;
    
    // 移除末尾的 null 终止符
    if buf.last() == Some(&0) {
        buf.pop();
    }
    
    String::from_utf8(buf)
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
}

/// 解析 Lua 字节码头部
pub fn parse_header(reader: &mut impl Read) -> std::io::Result<LuaHeader> {
    let mut signature = [0u8; 4];
    reader.read_exact(&mut signature)?;
    
    // 验证签名
    if &signature != b"\x1B\x4C\x75\x61" {
        return Err(std::io::Error::new(
            std::io::ErrorKind::InvalidData,
            "无效的 Lua 字节码签名"
        ));
    }
    
    let version = read_u8(reader)?;
    let format_version = read_u8(reader)?;
    
    // 跳过剩余头部信息 (integer size, size_t size, instruction size 等)
    let mut skip_buf = [0u8; 6];
    reader.read_exact(&mut skip_buf)?;
    
    Ok(LuaHeader {
        signature,
        version,
        format_version,
    })
}

/// 解析函数原型
pub fn parse_function(reader: &mut impl Read, depth: usize) -> std::io::Result<FunctionPrototype> {
    // 读取源文件名
    let source_name = read_string(reader)?;
    
    // 读取行号信息
    let line_defined = read_i32_le(reader)?;
    let last_line_defined = read_i32_le(reader)?;
    
    // 读取函数属性
    let num_params = read_u8(reader)?;
    let max_stack_size = read_i32_le(reader)?;
    
    // 读取指令数组
    let code_size = read_u32_le(reader)? as usize;
    let mut instructions = Vec::with_capacity(code_size);
    for _ in 0..code_size {
        instructions.push(read_u32_le(reader)?);
    }
    
    // 读取常量数量 (这里简化处理，实际需要根据具体格式解析)
    let constants_count = read_u32_le(reader)? as usize;
    
    // 读取上值数量
    let upvalues_count = read_u32_le(reader)? as usize;
    
    // 读取内部函数数量
    let functions_count = read_u32_le(reader)? as usize;
    
    // 递归解析内部函数
    for _ in 0..functions_count {
        let _ = parse_function(reader, depth + 1)?;
    }
    
    // 跳过调试信息 (简化处理)
    // 实际实现需要解析行号表、局部变量信息等
    
    Ok(FunctionPrototype {
        source_name,
        line_defined,
        last_line_defined,
        num_params,
        max_stack_size,
        instructions,
        upvalues_count,
        constants_count,
        functions_count,
    })
}

/// 从指令中提取 CALL/TLACALL 的参数
pub fn extract_call_info(instruction: u32, index: usize) -> CallInfo {
    let opcode = OpCode::from_instruction(instruction);
    
    // Lua 5.5 指令格式 (iABC):
    // bits 0-6: opcode (7 bits)
    // bit 7: k flag
    // bits 8-15: A (8 bits)
    // bits 16-23: B (8 bits)
    // bits 24-31: C (8 bits)
    
    let a = ((instruction >> 8) & 0xFF) as u32;
    let b = ((instruction >> 16) & 0xFF) as u32;
    let c = ((instruction >> 24) & 0xFF) as u32;
    
    CallInfo {
        instruction_index: index,
        opcode,
        base_register: a,
        num_args: if b == 0 { 0 } else { b }, // B=0 表示可变参数
        num_returns: if c == 0 { 0 } else { c }, // C=0 表示可变返回值
        line_number: None, // 需要调试信息才能获取行号
    }
}

/// 分析函数中的 CALL 和 TAILCALL 指令
pub fn analyze_function(func: &FunctionPrototype) -> (Vec<CallInfo>, Vec<CallInfo>) {
    let mut calls = Vec::new();
    let mut tailcalls = Vec::new();
    
    for (index, &instruction) in func.instructions.iter().enumerate() {
        let opcode = OpCode::from_instruction(instruction);
        
        match opcode {
            OpCode::CALL => {
                let call_info = extract_call_info(instruction, index);
                calls.push(call_info);
            }
            OpCode::TAILCALL => {
                let call_info = extract_call_info(instruction, index);
                tailcalls.push(call_info);
            }
            _ => {}
        }
    }
    
    (calls, tailcalls)
}

/// 解析 Lua 字节码文件
pub fn parse_bytecode_file<P: AsRef<Path>>(file_path: P) -> std::io::Result<Vec<FunctionPrototype>> {
    let file = File::open(file_path.as_ref())?;
    let mut reader = BufReader::new(file);
    
    // 解析头部
    let _header = parse_header(&mut reader)?;
    
    // 解析主函数 (第一个函数)
    let mut functions = Vec::new();
    let main_func = parse_function(&mut reader, 0)?;
    functions.push(main_func);
    
    Ok(functions)
}

/// 分析字节码文件中的 CALL 和 TAILCALL 指令
pub fn analyze_bytecode<P: AsRef<Path>>(file_path: P) -> std::io::Result<BytecodeAnalysis> {
    let functions = parse_bytecode_file(file_path)?;
    
    if functions.is_empty() {
        return Ok(BytecodeAnalysis {
            has_main: false,
            main_calls: Vec::new(),
            main_tailcalls: Vec::new(),
            total_functions: 0,
            main_function: None,
        });
    }
    
    // 第一个函数是 main 函数
    let main_func = &functions[0];
    let (calls, tailcalls) = analyze_function(main_func);
    
    log::info!("字节码分析完成:");
    log::info!("  - 找到 {} 个函数", functions.len());
    log::info!("  - main 函数中有 {} 条 CALL 指令", calls.len());
    log::info!("  - main 函数中有 {} 条 TAILCALL 指令", tailcalls.len());
    
    // 打印详细信息
    for call in &calls {
        log::warn!("CALL 指令 @{}: 寄存器={}, 参数={}, 返回值={}", 
              call.instruction_index, call.base_register, 
              call.num_args, call.num_returns);
    }
    
    for tailcall in &tailcalls {
        log::warn!("TAILCALL 指令 @{}: 寄存器={}, 参数={}, 返回值={}", 
              tailcall.instruction_index, tailcall.base_register, 
              tailcall.num_args, tailcall.num_returns);
    }
    
    Ok(BytecodeAnalysis {
        has_main: true,
        main_calls: calls,
        main_tailcalls: tailcalls,
        total_functions: functions.len(),
        main_function: Some(functions[0].clone()),
    })
}