lua-decompile 0.1.0

#![allow(dead_code)]

use std::slice::IterMut;

use itertools::{Itertools, MultiPeek};
use lua_bytecode::{
    Bytecode, Proto,
    constant::Constant,
    opcode::{Instruction, Opcode},
};

#[cfg(feature = "lua51")]
use lua_bytecode::lua51::LuaBytecode;
#[cfg(feature = "lua51")]
use lua_bytecode::opcode::{LuaInstruction, LuaOpcode};

const VARIABLE_NAME: &str = "r";
const CONSTANT_INDEX: u16 = 256;

#[derive(Clone, Debug)]
enum InstructionHistory {
    LuaInstruction,
}

#[derive(Clone, Default, Debug)]
struct DecompilerContext {
    indentation: u32,
}

#[derive(Clone, Default, Debug)]
pub struct Decompiler {
    output: String,
    blocks: Vec<Block>,
    history: Vec<InstructionHistory>,
    context: DecompilerContext,
    counter: u64,
}

#[derive(Clone, Default, Debug)]
enum LuaValue {
    #[default]
    Nil,
    Bool(bool),
    Number(f64),
    String(Vec<u8>),
}

impl LuaValue {
    fn from(constant: Constant) -> Self {
        match constant {
            Constant::Nil => LuaValue::Nil,
            Constant::Bool(value) => LuaValue::Bool(value),
            Constant::Number(value) => LuaValue::Number(value),
            Constant::String(value) => LuaValue::String(value),
        }
    }
}

impl std::fmt::Display for LuaValue {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            LuaValue::Nil => write!(f, "nil"),
            LuaValue::Number(value) => write!(f, "{value}"),
            LuaValue::Bool(value) => write!(f, "{value}"),
            LuaValue::String(value) => {
                write!(f, "\"{}\"", String::from_utf8_lossy(value).to_owned())
            }
        }
    }
}

#[derive(Clone, Default, Debug)]
struct Assignment {
    name: String,
}

#[derive(Clone, Default, Debug)]
struct IfStatement {}

#[derive(Clone, Debug)]
enum Statement {
    If(IfStatement),
    Assignment(Assignment),
}

// TODO: in the future we should be using an ast to re-create our lua code
#[derive(Clone, Default, Debug)]
struct Block {
    statements: Vec<Statement>,
}

enum VariableOrConstant {
    Constant(LuaValue),
    Variable(String),
}

impl std::fmt::Display for VariableOrConstant {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            VariableOrConstant::Constant(value) => write!(f, "{}", value),
            VariableOrConstant::Variable(string) => write!(f, "{}", string),
        }
    }
}

impl VariableOrConstant {
    fn from(proto: &Proto, index: u32) -> Self {
        if index >= CONSTANT_INDEX as u32 {
            VariableOrConstant::Constant(LuaValue::from(
                proto.constants[(index - CONSTANT_INDEX as u32) as usize].clone(),
            ))
        } else {
            VariableOrConstant::Variable(format!("{VARIABLE_NAME}{index}"))
        }
    }
}

fn opcode_operator(opcode: &Opcode) -> &str {
    match opcode {
        Opcode::LuaOpcode(op) => match op {
            LuaOpcode::Add => "+",
            LuaOpcode::Sub => "-",
            LuaOpcode::Mul => "*",
            LuaOpcode::Div => "/",
            LuaOpcode::Mod => "%",
            LuaOpcode::Pow => "^",

            LuaOpcode::Unm => "-",
            LuaOpcode::Not => "not",
            LuaOpcode::Len => "#",
            LuaOpcode::Concat => "..",

            LuaOpcode::Eq => "==",
            LuaOpcode::Lt => "<",
            LuaOpcode::Le => "<=",

            _ => "?",
        },

        _ => todo!(),
    }
}

fn is_opcode_condtional(opcode: &Opcode) -> bool {
    match opcode {
        Opcode::LuaOpcode(op) => match op {
            LuaOpcode::Test | LuaOpcode::Eq | LuaOpcode::Lt | LuaOpcode::Le => true,
            _ => false,
        },
    }
}

fn math_operation(output: &mut String, proto: &Proto, instruction: &Instruction) {
    let a = instruction.a();
    let b = VariableOrConstant::from(proto, instruction.b());
    let c = VariableOrConstant::from(proto, instruction.c());

    let binding = instruction.opcode();
    let sign = opcode_operator(&binding);
    output.push_str(&format!("{VARIABLE_NAME}{a} = {b} {sign} {c}\n"));
}

impl Decompiler {
    pub fn decompile(data: &[u8]) -> Result<String, String> {
        let bytecode = <Bytecode as LuaBytecode>::from(data)?;
        let mut decompiler = Decompiler::default();

        let main_proto = bytecode
            .protos
            .get(bytecode.main_proto_id as usize)
            .unwrap();

        decompiler.decompile_proto(main_proto)?;
        Ok(decompiler.output)
    }

    pub fn decompile_proto(self: &mut Self, proto: &Proto) -> Result<(), String> {
        let mut binding = proto.instructions.clone();
        let mut instructions = binding.iter_mut().multipeek();

        self.decompile_block(proto, &mut instructions, None)?;
        Ok(())
    }

    pub fn decompile_block(
        self: &mut Self,
        proto: &Proto,
        instructions: &mut MultiPeek<IterMut<Instruction>>,
        block_size: Option<u64>,
    ) -> Result<(), String> {
        let count = self.counter;

        loop {
            if let Some(i) = block_size {
                if self.counter - count == i {
                    break;
                }
            }

            let instruction = match instructions.next() {
                None => break,
                Some(instruction) => instruction,
            };

            self.output
                .push_str(&"\t".repeat(self.context.indentation as usize));

            self.counter += 1;
            match instruction.opcode() {
                Opcode::LuaOpcode(op) => match op {
                    LuaOpcode::Move => {
                        let a = instruction.a();
                        let b = instruction.b();
                        self.output
                            .push_str(&format!("{VARIABLE_NAME}{a} = {VARIABLE_NAME}{b}\n"));
                    }

                    LuaOpcode::LoadK => {
                        let a = instruction.a();
                        let index = instruction.bx();
                        let constant = LuaValue::from(proto.constants[index as usize].clone());

                        self.output
                            .push_str(&format!("{VARIABLE_NAME}{a} = {constant}\n"));
                    }

                    LuaOpcode::LoadBool => {
                        let a = instruction.a();
                        let value = instruction.b() > 0;
                        self.output
                            .push_str(&format!("{VARIABLE_NAME}{a} = {value}\n"));
                    }

                    LuaOpcode::LoadNil => todo!(),
                    LuaOpcode::GetUpval => todo!(),

                    LuaOpcode::GetGlobal => {
                        let a = instruction.a();
                        let index = instruction.bx();
                        let constant = LuaValue::from(proto.constants[index as usize].clone());

                        let global = match constant {
                            LuaValue::String(string) => {
                                String::from_utf8_lossy(&string).to_string()
                            }
                            _ => unreachable!(),
                        };

                        self.output
                            .push_str(&format!("{VARIABLE_NAME}{a} = {global}\n"));
                    }

                    LuaOpcode::GetTable => todo!(),
                    LuaOpcode::SetGlobal => todo!(),
                    LuaOpcode::SetUpval => todo!(),
                    LuaOpcode::SetTable => todo!(),
                    LuaOpcode::NewTable => todo!(),
                    LuaOpcode::Self_ => todo!(),

                    LuaOpcode::Add
                    | LuaOpcode::Sub
                    | LuaOpcode::Mul
                    | LuaOpcode::Div
                    | LuaOpcode::Mod
                    | LuaOpcode::Pow => math_operation(&mut self.output, proto, instruction),

                    LuaOpcode::Unm | LuaOpcode::Not | LuaOpcode::Len => {
                        let sign = instruction.opcode();
                        let sign = opcode_operator(&sign);

                        let a = instruction.a();
                        let b = instruction.b();
                        self.output
                            .push_str(&format!("{VARIABLE_NAME}{a} = {sign}{VARIABLE_NAME}{b}\n"));
                    }

                    LuaOpcode::Concat => {
                        let a = instruction.a();
                        self.output.push_str(&format!("{VARIABLE_NAME}{a} = "));

                        let b = instruction.b();
                        let c = instruction.c();

                        for i in b..c + 1 {
                            self.output.push_str(&format!("{VARIABLE_NAME}{i}"));
                            if i < c {
                                self.output.push_str(" .. ");
                            }
                        }

                        self.output.push_str("\n");
                    }

                    LuaOpcode::Jmp => {
                        todo!();
                        // positive jmps by themselves should skip the next instructions
                        //let sbx = instruction.sbx();

                        //for _ in 0..sbx {
                        //    instructions.next();
                        //}
                    }

                    LuaOpcode::Test | LuaOpcode::Eq | LuaOpcode::Lt | LuaOpcode::Le => {
                        self.if_condition(proto, instruction, instructions, false)?;
                    }

                    LuaOpcode::TestSet => todo!(),

                    LuaOpcode::Call => {
                        self.function_call(instruction);
                    }

                    LuaOpcode::TailCall => todo!(),

                    LuaOpcode::Return => {
                        let a = instruction.a();
                        let b = instruction.b();

                        if b > 1 {
                            self.output.push_str("return ");
                        }

                        for i in 0..b - 1 {
                            let x = a + i;
                            self.output.push_str(&format!("{VARIABLE_NAME}{x}"));
                            if i < b - 2 {
                                self.output.push_str(", ");
                            }
                        }
                    }

                    LuaOpcode::ForLoop => todo!(),
                    LuaOpcode::ForPrep => {
                        self.for_loop(proto, instruction, instructions)?;
                    }

                    LuaOpcode::TForLoop => todo!(),
                    LuaOpcode::SetList => todo!(),
                    LuaOpcode::Close => todo!(),
                    LuaOpcode::Closure => todo!(),
                    LuaOpcode::Vararg => todo!(),
                },

                _ => todo!(),
            }
        }

        Ok(())
    }

    fn function_call(&mut self, instruction: &mut Instruction) {
        let callee = instruction.a();
        let argument_count = instruction.b();

        let r = instruction.c();
        for i in (0..r - 1).rev() {
            let x = r - i;
            self.output.push_str(&format!("{VARIABLE_NAME}{x}"));
            if i > 0 {
                self.output.push_str(", ");
            }
        }

        if r > 1 {
            self.output.push_str(" = ");
        }

        self.output.push_str(&format!("{VARIABLE_NAME}{callee}("));

        let x = callee + argument_count;
        for i in callee + 1..x {
            self.output.push_str(&format!("{VARIABLE_NAME}{i}"));
            if i != x - 1 {
                self.output.push_str(", ");
            }
        }

        self.output.push_str(")\n");
    }

    fn if_condition(
        &mut self,
        proto: &Proto,
        instruction: &mut Instruction,
        instructions: &mut MultiPeek<IterMut<Instruction>>,
        else_if_continuation: bool,
    ) -> Result<(), String> {
        let opcode = instruction.opcode();

        let jump = instructions.next().unwrap();
        assert_eq!(jump.opcode(), Opcode::LuaOpcode(LuaOpcode::Jmp));

        self.counter += 1;
        let mut jump_count = jump.sbx();

        let mut condition = String::new();
        let mut condition_buffer = vec![(instruction, jump)];

        // if there are multiple TEST + JMP instructions in a row
        // it's a inlined conditional, i.e
        // ```lua
        // if ([x.. and y..]) then end
        // ```

        loop {
            match instructions.peek() {
                Some(instruction) => {
                    if !is_opcode_condtional(&instruction.opcode()) {
                        break;
                    }
                }
                None => break,
            }

            match instructions.peek() {
                Some(instruction) => match instruction.opcode() {
                    Opcode::LuaOpcode(op) => match op {
                        LuaOpcode::Jmp => {}
                        _ => break,
                    },
                },
                None => break,
            }

            let test = instructions.next().unwrap();
            let jump = instructions.next().unwrap();
            assert_eq!(jump.opcode(), Opcode::LuaOpcode(LuaOpcode::Jmp));

            jump_count = jump.sbx();
            self.counter += 2;

            condition_buffer.push((test, jump));
        }

        let count = condition_buffer.len();
        for (index, (test, jump)) in condition_buffer.iter().enumerate() {
            let index = index + 1;
            let r = test.a();

            if matches!(test.opcode(), Opcode::LuaOpcode(LuaOpcode::Test))
                && (jump.sbx() / 2) as usize == count - index
            {
                condition.push_str(&format!("{VARIABLE_NAME}{r} or "));
            } else {
                match test.opcode() {
                    Opcode::LuaOpcode(op) => match op {
                        LuaOpcode::Test => {
                            let c = test.c() != 0;

                            let not = if c { "not " } else { "" };
                            let and = if index != count { " and " } else { "" };
                            condition.push_str(&format!("{}{}{}{}", not, VARIABLE_NAME, r, and));
                        }

                        LuaOpcode::Eq | LuaOpcode::Lt | LuaOpcode::Le => {
                            let sign = opcode_operator(&opcode);

                            let a = test.a();
                            let b = VariableOrConstant::from(proto, test.b());
                            let c = VariableOrConstant::from(proto, test.c());

                            let is_even = jump.sbx() % 2 == 0;
                            let post_operator = if index != count {
                                if a == 0 && is_even { " and " } else { " or " }
                            } else {
                                ""
                            };

                            condition.push_str(&format!("{b} {sign} {c}{post_operator}"));
                        }

                        _ => unreachable!(),
                    },
                }
            }
        }

        let keyword = if !else_if_continuation {
            "if"
        } else {
            "elseif"
        };

        self.output
            .push_str(&format!("{keyword} {condition} then\n"));

        instructions.reset_peek();
        for _ in 0..jump_count - 1 {
            instructions.peek();
        }

        let has_else = if let Some(peak) = instructions.peek() {
            matches!(peak.opcode(), Opcode::LuaOpcode(LuaOpcode::Jmp))
        } else {
            false
        };

        if has_else {
            jump_count -= 1;
        }

        self.context.indentation += 1;
        self.decompile_block(proto, instructions, Some(jump_count as u64))?;

        let end_statement = if has_else {
            let jump = instructions.next().unwrap();
            self.counter += 1;

            let has_else_if = if let Some(peek) = instructions.peek() {
                is_opcode_condtional(&peek.opcode())
            } else {
                false
            };

            if has_else_if {
                self.context.indentation -= 1;
                self.if_condition(proto, instructions.next().unwrap(), instructions, true)?;
            } else {
                self.output.push_str("else\n");
                self.decompile_block(proto, instructions, Some(jump.sbx() as u64))?;
            }

            !has_else_if
        } else {
            true
        };

        if end_statement {
            self.output.push_str("end\n");
            self.context.indentation -= 1;
        }

        Ok(())
    }

    fn for_loop(
        &mut self,
        proto: &Proto,
        instruction: &mut Instruction,
        instructions: &mut MultiPeek<IterMut<'_, Instruction>>,
    ) -> Result<(), String> {
        let r = instruction.a();
        let r0 = format!("{VARIABLE_NAME}{}", r + 0);
        let r1 = format!("{VARIABLE_NAME}{}", r + 1);
        let r2 = format!("{VARIABLE_NAME}{}", r + 2);

        self.output
            .push_str(&format!("for i = {r0}, {r1}, {r2} do\n"));

        self.context.indentation += 1;
        self.decompile_block(proto, instructions, Some(instruction.sbx() as u64))?;
        self.context.indentation -= 1;

        let close = instructions.next().unwrap();
        assert_eq!(close.opcode(), Opcode::LuaOpcode(LuaOpcode::ForLoop));

        self.output.push_str("end\n");

        Ok(())
    }
}