AstroLang/src/decompiler.rs

use std::collections::HashMap;

const ASXVERSION: [u8; 3] = [1,0,0];

#[derive(Debug, Clone)]
pub struct DecompiledFunction {
    pub body: Vec<DecompiledOperation>,
    pub variables: Vec<Variable>,
    pub strings: HashMap<u32, String>,
    pub functions: HashMap<u32, u32>,
    pub try_catch: u32,
}
#[derive(Debug, Clone)]
pub struct DecompiledOperation {
    pub opcode: u8,
    pub arg1: u8,
    pub arg2: f64,
    pub arg3: u8,
    pub pos: usize,
}

#[derive(Debug, Clone)]
pub struct Variable {
    pub name: String,
    pub id: u32,
    pub start: usize,
    pub end: usize
}

#[derive(Debug, Clone)]
pub struct DecompiledData {
    pub functions: Vec<DecompiledFunction>,
    pub func_count: usize,
    pub version: String,
}

pub fn operation_to_name(opcode: u8) -> String {
    match opcode {
        0 => "HLT".to_string(),
        1 => "LDS".to_string(),
        2 => "LDM".to_string(),
        3 => "LDI".to_string(),
        4 => "LDB".to_string(),
        5 => "LDF".to_string(),
        6 => "LDN".to_string(),
        7 => "ASS".to_string(),
        8 => "UNB".to_string(),
        9 => "MOV".to_string(),
        10 => "ADD".to_string(),
        11 => "SUB".to_string(),
        12 => "MUL".to_string(),
        13 => "DIV".to_string(),
        14 => "POW".to_string(),
        15 => "MOD".to_string(),
        16 => "AND".to_string(),
        17 => "OR".to_string(),
        18 => "EQ".to_string(),
        19 => "NEQ".to_string(),
        20 => "GRE".to_string(),
        21 => "GRQ".to_string(),
        22 => "LES".to_string(),
        23 => "LEQ".to_string(),
        24 => "NOT".to_string(),
        25 => "JMP".to_string(),
        26 => "CJP".to_string(),
        27 => "CAL".to_string(),
        28 => "PSH".to_string(),
        29 => "RET".to_string(),
        30 => "GET".to_string(),
        31 => "SET".to_string(),
        _ => panic!("Unknown operation code: {}", opcode),
    }
}

fn read_be_num(input: &[u8]) -> usize {
    let mut result: usize = 0;
    for i in 0..input.len() {
        result <<= 8;
        result |= input[i] as usize;
    }
    return result;
}

fn read_str(input: &[u8]) -> String {
    String::from_utf8(input.to_vec()).expect("Invalid UTF-8 sequence")
}

fn read_bin(input: &[u8]) -> String {
    let mut result = String::new();
    for &byte in input {
        result.push_str(&format!("{:08b}", byte));
    }
    return result;
}

fn bin_to_num(bin: String) -> usize {
    let mut result: usize = 0;
    for (i, c) in bin.chars().rev().enumerate() {
        if c == '1' {
            result += 1 << i;
        }
    }
    return result;
}

fn bin_to_snum(bin: String) -> f64 {
    f64::from_bits(bin_to_num(bin) as u64)
}

fn load_func(data: Vec<u8>, offset: &mut usize) -> DecompiledFunction {
    let var_id_len = read_be_num(&data[*offset..*offset + 4]);
    *offset += 4;
    let mut variables: Vec<Variable> = vec![];
    for _ in 0..var_id_len {
        let name_len = read_be_num(&data[*offset..*offset + 4]);
        *offset += 4;
        let name = read_str(&data[*offset..*offset + name_len]);
        *offset += name_len;
        let var_id = read_be_num(&data[*offset..*offset + 3]) as u32;
        *offset += 3;
        let start = read_be_num(&data[*offset..*offset + 4]);
        *offset += 4;
        let end = read_be_num(&data[*offset..*offset + 4]);
        *offset += 4;
        variables.push(Variable { name: name, id: var_id, start: start, end: end });
    }
    let mut strings: HashMap<u32, String> = HashMap::new();
    let string_count = read_be_num(&data[*offset..*offset + 4]);
    *offset += 4;
    for _ in 0..string_count {
        let str_id = read_be_num(&data[*offset..*offset + 3]) as u32;
        *offset += 3;
        let str_len = read_be_num(&data[*offset..*offset + 4]);
        *offset += 4;
        let str_value = read_str(&data[*offset..*offset + str_len]);
        *offset += str_len;
        strings.insert(str_id, str_value);
    }
    let mut functions: HashMap<u32, u32> = HashMap::new();
    let function_count = read_be_num(&data[*offset..*offset + 4]);
    *offset += 4;
    for _ in 0..function_count {
        let func_id = read_be_num(&data[*offset..*offset + 3]) as u32;
        *offset += 3;
        let func_pos = read_be_num(&data[*offset..*offset + 4]) as u32;
        *offset += 4;
        functions.insert(func_id, func_pos);
    }
    let try_catch = read_be_num(&data[*offset..*offset + 4]);
    *offset += 4;
    let mut body: Vec<DecompiledOperation> = Vec::new();
    let instr_len: usize = read_be_num(&data[*offset..*offset + 4]);
    *offset += 4;
    for _ in 0..instr_len {
        let op = read_bin(&data[*offset..*offset + 10]);
        *offset += 10;
        let op_code = bin_to_num(op[0..5].to_string());
        let arg1 = bin_to_num(op[5..9].to_string());
        let arg2 = bin_to_snum(op[9..73].to_string());
        let arg3 = bin_to_num(op[73..77].to_string());
        let pos = read_be_num(&data[*offset..*offset + 4]);
        *offset += 4;
        body.push(DecompiledOperation {
            opcode: op_code as u8,
            arg1: arg1 as u8,
            arg2: arg2,
            arg3: arg3 as u8,
            pos: pos,
        });
    }
    return DecompiledFunction {
        body: body,
        variables: variables,
        strings: strings,
        functions: functions,
        try_catch: try_catch as u32,
    };
}

pub fn process(data: &Vec<u8>) -> DecompiledData {
    if data[0..3] != *"ASX".as_bytes() {
        panic!("Invalid ASX file header");
    }
    let ver = &data[3..6].to_vec();
    let ver_str: String = ver.iter().map(|b| b.to_string()).collect::<Vec<String>>().join(".");
    if *ver != ASXVERSION {
        panic!("Unsupported ASX version ({})", ver_str);
    }
    let func_count = read_be_num(&data[6..10]);
    let mut offset = 10;
    let mut functions: Vec<DecompiledFunction> = Vec::new();
    while offset < data.len() {
        let func = load_func(data.clone(), &mut offset);
        functions.push(func);
    }
    return DecompiledData {
        functions: functions,
        func_count: func_count,
        version: ver_str,
    };
}