psdisasm 0.2.0

mod arguments;
mod instruction;

use arguments::ArgumentTypes;
pub use instruction::Instruction;

const REGISTERS: [&str; 32] = [
    "0",   // Zero register
    "$at", // Assembler temporary register
    "$v0", "$v1", // Return value registers
    "$a0", "$a1", "$a2", "$a3", // Argument registers
    "$t0", "$t1", "$t2", "$t3", "$t4", "$t5", "$t6", "$t7", // Temporary registers
    "$s0", "$s1", "$s2", "$s3", "$s4", "$s5", "$s6", "$s7", // Saved registers
    "$t8", "$t9", // More temporary registers
    "$k0", "$k1", // Kernel registers
    "$gp", "$sp", // Global pointer, stack pointer
    "$fp", "$ra", // Frame pointer, return address
];

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Disassembler {
    /// Produce easier to read disassembly output.  
    /// Example transformations:
    /// - sll $zero, $zero, 0 -> nop
    /// - lui $t1, 1 -> move $t1, 0x10000
    /// - add $t1, $t1, 2 -> add $t1, 2
    pub pretty: bool,
}

impl Default for Disassembler {
    fn default() -> Self {
        Disassembler { pretty: true }
    }
}

impl Disassembler {
    /// Disassemble an instruction
    pub fn disassemble(&self, ins: u32, pc: u32) -> String {
        let ins = Instruction::new(ins);

        match ins.opcode() {
            0x00 => match ins.funct() {
                0x00 => self.format_shamt(ins, "sll"),
                0x02 => self.format_shamt(ins, "srl"),
                0x03 => self.format_shamt(ins, "sra"),
                0x04 => self.format_shift_r(ins, "sllv"),
                0x06 => self.format_shift_r(ins, "srlv"),
                0x07 => self.format_shift_r(ins, "srav"),
                0x08 => self.format_instruction(ins, "jr", ArgumentTypes::S, pc),
                0x09 => self.format_instruction(ins, "jalr", ArgumentTypes::S_D, pc),
                0x0c => self.format_instruction(ins, "syscall", ArgumentTypes::None, 0),
                0x0d => self.format_instruction(ins, "break", ArgumentTypes::None, 0),
                0x10 => self.format_instruction(ins, "mfhi", ArgumentTypes::D, 0),
                0x11 => self.format_instruction(ins, "mthi", ArgumentTypes::S, 0),
                0x12 => self.format_instruction(ins, "mflo", ArgumentTypes::D, 0),
                0x13 => self.format_instruction(ins, "mtlo", ArgumentTypes::S, 0),
                0x18 => self.format_instruction(ins, "mult", ArgumentTypes::S_T, 0),
                0x19 => self.format_instruction(ins, "multu", ArgumentTypes::S_T, 0),
                0x1a => self.format_instruction(ins, "div", ArgumentTypes::S_T, 0),
                0x1b => self.format_instruction(ins, "divu", ArgumentTypes::S_T, 0),
                0x20 => self.format_arith_r(ins, "add"),
                0x21 => self.format_arith_r(ins, "addu"),
                0x22 => self.format_arith_r(ins, "sub"),
                0x23 => self.format_arith_r(ins, "subu"),
                0x24 => self.format_arith_r(ins, "and"),
                0x25 => self.format_arith_r(ins, "or"),
                0x26 => self.format_arith_r(ins, "xor"),
                0x27 => self.format_instruction(ins, "nor", ArgumentTypes::D_S_T, 0),
                0x2a => self.format_instruction(ins, "slt", ArgumentTypes::D_S_T, 0),
                0x2b => self.format_instruction(ins, "sltu", ArgumentTypes::D_S_T, 0),
                _ => format!("Invalid opcode 0x00 with funct {:x}", ins.funct()),
            },
            0x01 => {
                // This format abuses the `rt` field for a sub-opcode
                let link = ins.rt() & 0x10 == 0x10;
                match (ins.rt() & 1, link) {
                    (0, false) => self.format_instruction(ins, "bltz", ArgumentTypes::S_Jump, pc),
                    (0, true) => self.format_instruction(ins, "bltzal", ArgumentTypes::S_Jump, pc),
                    (1, false) => self.format_instruction(ins, "bgez", ArgumentTypes::S_Jump, pc),
                    (1, true) => self.format_instruction(ins, "bgezal", ArgumentTypes::S_Jump, pc),
                    _ => unreachable!(),
                }
            }
            0x02 => self.format_instruction(ins, "j", ArgumentTypes::Jump, pc),
            0x03 => self.format_instruction(ins, "jal", ArgumentTypes::Jump, pc),
            0x04 => self.format_instruction(ins, "beq", ArgumentTypes::S_T_Jump, pc),
            0x05 => self.format_instruction(ins, "bne", ArgumentTypes::S_T_Jump, pc),
            0x06 => self.format_instruction(ins, "blez", ArgumentTypes::S_Jump, pc),
            0x07 => self.format_instruction(ins, "bgtz", ArgumentTypes::S_Jump, pc),
            0x08 => self.format_arith_imm_signed(ins, "addi"),
            0x09 => self.format_arith_imm_signed(ins, "addiu"),
            0x0a => self.format_instruction(ins, "slti", ArgumentTypes::T_S_SImm, 0),
            0x0b => self.format_instruction(ins, "sltiu", ArgumentTypes::T_S_Imm, 0),
            0x0c => self.format_arith_imm_unsigned(ins, "andi"),
            0x0d => self.format_arith_imm_unsigned(ins, "ori"),
            0x0e => self.format_arith_imm_unsigned(ins, "xori"),
            0x0f => self.format_lui(ins),
            0x10..=0x13 => self.format_coprocessor_instruction(ins),
            0x20 => self.format_memory(ins, "lb"),
            0x21 => self.format_memory(ins, "lh"),
            0x22 => self.format_memory(ins, "lwl"),
            0x23 => self.format_memory(ins, "lw"),
            0x24 => self.format_memory(ins, "lbu"),
            0x25 => self.format_memory(ins, "lhu"),
            0x26 => self.format_memory(ins, "lwr"),
            0x28 => self.format_memory(ins, "sb"),
            0x29 => self.format_memory(ins, "sh"),
            0x2a => self.format_memory(ins, "swl"),
            0x2b => self.format_memory(ins, "sw"),
            0x2e => self.format_memory(ins, "swr"),
            0x32 => self.format_memory_cop2(ins, "lwc2"),
            0x3a => self.format_memory_cop2(ins, "swc2"),
            _ => format!("Unknown opcode: {:#x}", ins.opcode()),
        }
    }

    /// Disassemble an instruction with additional context.
    pub fn disassemble_with_context(&self, ins: u32, pc: u32, registers: &[u32; 32]) -> String {
        let disassembled = self.disassemble(ins, pc);
        let ins = Instruction::new(ins);

        let mut kind = "";
        let mut branch_taken = false;
        let mut address = 0;

        match ins.opcode() {
            0x01 => {
                // This format abuses the `rt` field for a sub-opcode
                match ins.rt() & 1 {
                    0 => {
                        kind = "branch";
                        branch_taken = (registers[ins.rs()] as i32) < 0
                    }
                    1 => {
                        kind = "branch";
                        branch_taken = (registers[ins.rs()] as i32) >= 0
                    }
                    _ => {}
                }
            }
            0x04 => {
                kind = "branch";
                branch_taken = registers[ins.rs()] == registers[ins.rt()]
            }
            0x05 => {
                kind = "branch";
                branch_taken = registers[ins.rs()] != registers[ins.rt()]
            }
            0x06 => {
                kind = "branch";
                branch_taken = registers[ins.rs()] <= 0
            }
            0x07 => {
                kind = "branch";
                branch_taken = registers[ins.rs()] > 0
            }
            0x20..=0x26 | 0x28..=0x2e => {
                // Memory access instructions
                kind = "memory";
                address = (registers[ins.rs()] as i32 + ins.simm16()) as u32;
            }
            0x32 | 0x3a => {
                // Coprocessor memory access instructions
                kind = "memory";
                address = (registers[ins.rs()] as i32 + ins.simm16()) as u32;
            }
            _ => {}
        }

        match kind {
            "branch" => {
                if branch_taken {
                    format!("{disassembled}\t(taken)")
                } else {
                    format!("{disassembled}\t(not taken)")
                }
            }
            "memory" => {
                format!("{disassembled}\t({address:x})")
            }
            _ => disassembled,
        }
    }

    fn format_shamt(&self, ins: Instruction, name: &str) -> String {
        if self.pretty {
            if ins.rd() == 0 {
                return "nop".to_string();
            }

            if ins.shamt() == 0 {
                return format!("move {}, {}", REGISTERS[ins.rd()], REGISTERS[ins.rt()]);
            }
        }

        self.format_instruction(ins, name, ArgumentTypes::D_T_Shift, 0)
    }

    fn format_shift_r(&self, ins: Instruction, name: &str) -> String {
        if self.pretty {
            if ins.rd() == 0 && ins.rt() == 0 {
                return "nop".to_string();
            }

            if ins.rd() == ins.rt() {
                return format!("{name} {}, {}", REGISTERS[ins.rd()], REGISTERS[ins.rs()]);
            }
        }

        self.format_instruction(ins, name, ArgumentTypes::D_T_S, 0)
    }

    fn format_arith_r(&self, ins: Instruction, name: &str) -> String {
        if self.pretty {
            if ins.rd() == 0 {
                return "nop".to_string();
            }

            if ins.rd() == ins.rs() {
                return format!("{name} {}, {}", REGISTERS[ins.rd()], REGISTERS[ins.rt()]);
            }

            if ins.rs() == 0 && (name == "add" || name == "addu" || name == "or" || name == "xor") {
                return format!("move {}, {}", REGISTERS[ins.rd()], REGISTERS[ins.rt()]);
            }

            if ins.rt() == 0 && (name == "add" || name == "addu" || name == "or" || name == "xor") {
                return format!("move {}, {}", REGISTERS[ins.rd()], REGISTERS[ins.rs()]);
            }
        }

        self.format_instruction(ins, name, ArgumentTypes::D_S_T, 0)
    }

    fn format_arith_imm_signed(&self, ins: Instruction, name: &str) -> String {
        let mut immediate = ins.simm16();
        let sign = if immediate < 0 { "-" } else { "" };
        immediate = immediate.abs();

        let imm_str = if immediate < 10 {
            format!("{sign}{immediate}")
        } else {
            format!("{sign}{immediate:#x}")
        };

        if self.pretty {
            if ins.rt() == 0 {
                return "nop".to_string();
            }

            if ins.rt() == ins.rs() {
                return format!("{name} {}, {imm_str}", REGISTERS[ins.rt()]);
            }

            if ins.rs() == 0 {
                return format!("move {}, {imm_str}", REGISTERS[ins.rt()]);
            }
        }

        self.format_instruction(ins, name, ArgumentTypes::T_S_SImm, 0)
    }

    fn format_arith_imm_unsigned(&self, ins: Instruction, name: &str) -> String {
        let immediate = ins.imm16();

        let imm_str = if immediate < 10 {
            format!("{immediate}")
        } else {
            format!("{immediate:#x}")
        };

        if self.pretty {
            if ins.rt() == 0 {
                return "nop".to_string();
            }

            if ins.rt() == ins.rs() {
                return format!("{name} {}, {imm_str}", REGISTERS[ins.rt()]);
            }

            if ins.rs() == 0 && name != "andi" {
                return format!("move {}, {imm_str}", REGISTERS[ins.rt()]);
            }
        }

        self.format_instruction(ins, name, ArgumentTypes::T_S_SImm, 0)
    }

    fn format_lui(&self, ins: Instruction) -> String {
        format!(
            "{} {}, {:#x}",
            if self.pretty { "move" } else { "lui" },
            REGISTERS[ins.rt()],
            if self.pretty {
                ins.imm16() << 16
            } else {
                ins.imm16()
            }
        )
    }

    fn format_memory(&self, ins: Instruction, name: &str) -> String {
        let immediate = ins.simm16();
        let sign = if immediate < 0 { "-" } else { "" };
        let immediate = immediate.abs();

        let imm_str = if immediate == 0 {
            "".to_string()
        } else if immediate < 64 {
            format!("{sign}{immediate}")
        } else {
            format!("{sign}{immediate:#x}")
        };

        if self.pretty {
            return format!(
                "{name} {}, {imm_str}({})",
                REGISTERS[ins.rt()],
                REGISTERS[ins.rs()]
            );
        }

        self.format_instruction(ins, name, ArgumentTypes::T_Mem, 0)
    }

    fn format_memory_cop2(&self, ins: Instruction, name: &str) -> String {
        let immediate = ins.simm16();
        let sign = if immediate < 0 { "-" } else { "" };
        let immediate = immediate.abs();

        let imm_str = if immediate == 0 {
            "".to_string()
        } else if immediate < 64 {
            format!("{sign}{immediate}")
        } else {
            format!("{sign}{immediate:#x}")
        };

        if self.pretty {
            return format!(
                "{name} gte_r{}, {imm_str}({})",
                ins.rt(),
                REGISTERS[ins.rs()]
            );
        }

        self.format_instruction(ins, name, ArgumentTypes::T_Mem, 0)
    }

    /// Format the instruction with its name and arguments
    fn format_instruction(
        &self,
        ins: Instruction,
        name: &str,
        arg_types: ArgumentTypes,
        pc: u32,
    ) -> String {
        let args = match arg_types {
            ArgumentTypes::None => String::new(),
            ArgumentTypes::D_S_T => format!(
                "{}, {}, {}",
                REGISTERS[ins.rd()],
                REGISTERS[ins.rs()],
                REGISTERS[ins.rt()]
            ),
            ArgumentTypes::D_T_S => format!(
                "{}, {}, {}",
                REGISTERS[ins.rd()],
                REGISTERS[ins.rt()],
                REGISTERS[ins.rs()]
            ),
            ArgumentTypes::D_T_Shift => format!(
                "{}, {}, {}",
                REGISTERS[ins.rd()],
                REGISTERS[ins.rt()],
                ins.shamt()
            ),
            ArgumentTypes::T_S_SImm => format!(
                "{}, {}, {:#x}",
                REGISTERS[ins.rt()],
                REGISTERS[ins.rs()],
                ins.simm16()
            ),
            ArgumentTypes::T_Imm => format!("{}, {:#x}", REGISTERS[ins.rt()], ins.imm16()),
            ArgumentTypes::T_S_Imm => format!(
                "{}, {}, {:#x}",
                REGISTERS[ins.rt()],
                REGISTERS[ins.rs()],
                ins.imm16()
            ),
            ArgumentTypes::T_Mem => format!(
                "{}, {}({})",
                REGISTERS[ins.rt()],
                ins.simm16(),
                REGISTERS[ins.rs()]
            ),
            ArgumentTypes::S => format!("{}", REGISTERS[ins.rs()]),
            ArgumentTypes::D => format!("{}", REGISTERS[ins.rd()]),
            ArgumentTypes::S_D => format!("{}, {}", REGISTERS[ins.rs()], REGISTERS[ins.rd()]),
            ArgumentTypes::S_T => format!("{}, {}", REGISTERS[ins.rs()], REGISTERS[ins.rt()]),
            ArgumentTypes::S_T_Jump => {
                let target = pc.wrapping_add((ins.simm16() << 2) as u32);

                format!(
                    "{}, {}, {:#x}",
                    REGISTERS[ins.rs()],
                    REGISTERS[ins.rt()],
                    target
                )
            }
            ArgumentTypes::S_Jump => {
                let target = pc.wrapping_add((ins.simm16() << 2) as u32);
                format!("{}, {:#x}", REGISTERS[ins.rs()], target)
            }
            ArgumentTypes::Jump => {
                let target = (pc & 0xf000_0000) | (ins.jump_target() << 2);
                format!("{:#x}", target)
            }
        };

        format!("{name} {args}")
    }

    /// Coprocessor instruction have unique formatting rules. This function formats
    /// them based on the opcode and the specific coprocessor instruction.
    fn format_coprocessor_instruction(&self, ins: Instruction) -> String {
        // Get the coprocessor number from the opcode
        let cop = ins.opcode() & 3;

        if ins.cop_execute() {
            // Coprocessor specific opcode
            match cop {
                0 => self.format_cop0_instruction(ins),
                1 => format!("cop{} execute: {:x}", cop, ins.cop_instruction()),
                2 => self.format_gte_instruction(ins),
                3 => format!("cop{} execute: {:x}", cop, ins.cop_instruction()),
                _ => unreachable!(),
            }
        } else {
            let gpr = REGISTERS[ins.rt()];
            let cop_reg = ins.rd();

            match ins.cop_funct() {
                0 => format!("mfc{cop} {gpr}, cop{cop_reg}"),
                2 => format!("cfc{cop} {gpr}, cop{}", cop_reg + 32),
                4 => format!("mtc{cop} cop{cop_reg}, {gpr}"),
                6 => format!("ctc{cop} cop{}, {gpr}", cop_reg + 32),
                _ => format!("unknown cop{} funct: {:x}", cop, ins.cop_funct()),
            }
        }
    }

    fn format_cop0_instruction(&self, ins: Instruction) -> String {
        if ins.cop_instruction() == 0x10 {
            "rfe".to_string()
        } else {
            format!("cop0 execute: {:x}", ins.cop_instruction())
        }
    }

    fn format_gte_instruction(&self, ins: Instruction) -> String {
        let unsigned = if ins.cop_instruction() & (1 << 10) != 0 {
            "U"
        } else {
            "u"
        };
        let fixed_point = if ins.cop_instruction() & (1 << 19) != 0 {
            "F"
        } else {
            "f"
        };

        let mvmva_matrix = (ins.cop_instruction() >> 17) & 3;
        let mvmva_vector = (ins.cop_instruction() >> 15) & 3;
        let mvmva_tx_vector = (ins.cop_instruction() >> 13) & 3;

        let base = match ins.cop_instruction() & 0x3f {
            0x01 => "gte.rtps".to_string(),
            0x06 => "gte.nclip".to_string(),
            0x0c => "gte.op".to_string(),
            0x10 => "gte.dpcs".to_string(),
            0x11 => "gte.intpl".to_string(),
            0x12 => "gte.mvmva".to_string(),
            0x13 => "gte.ncds".to_string(),
            0x14 => "gte.cdp".to_string(),
            0x16 => "gte.ncdt".to_string(),
            0x1b => "gte.nccs".to_string(),
            0x1c => "gte.cc".to_string(),
            0x1e => "gte.ncs".to_string(),
            0x20 => "gte.nct".to_string(),
            0x28 => "gte.sqr".to_string(),
            0x29 => "gte.dcpl".to_string(),
            0x2a => "gte.dpct".to_string(),
            0x2d => "gte.avsz3".to_string(),
            0x2e => "gte.avsz4".to_string(),
            0x30 => "gte.rtpt".to_string(),
            0x3d => "gte.gpf".to_string(),
            0x3e => "gte.gpl".to_string(),
            0x3f => "gte.ncct".to_string(),
            _ => format!("gte execute: {:x}", ins.cop_instruction()),
        };

        let flags = unsigned.to_string() + fixed_point;
        if base == "gte.mvmva" {
            format!(
                "{base} {flags}, midx: {}, vidx: {}, tvidx: {}",
                mvmva_matrix, mvmva_vector, mvmva_tx_vector
            )
        } else {
            format!("{base} {flags}")
        }
    }
}