hcb2lua_decompiler/
lua_opt.rs

1use std::collections::BTreeMap;
2use std::collections::BTreeSet;
3use std::fmt::Write as _;
4
5use crate::cfg::BlockTerm;
6use crate::decode::{Instruction, Op};
7
8#[derive(Clone, Debug, PartialEq, Eq)]
9pub enum UnOp {
10    Neg,
11}
12
13#[derive(Clone, Debug, PartialEq, Eq)]
14pub enum BinOp {
15    Add,
16    Sub,
17    Mul,
18    Div,
19    Mod,
20    BitAnd,
21    And,
22    Or,
23    Eq,
24    Ne,
25    Gt,
26    Ge,
27    Lt,
28    Le,
29}
30
31#[derive(Clone, Debug, PartialEq, Eq)]
32pub enum Expr {
33    Nil,
34    Bool(bool),
35    Int(i64),
36    Float(String),
37    Str(String),
38    Var(String),
39    Unary(UnOp, Box<Expr>),
40    Binary(BinOp, Box<Expr>, Box<Expr>),
41    Index(Box<Expr>, Box<Expr>),
42}
43
44impl Expr {
45    pub fn var(name: impl Into<String>) -> Self {
46        Expr::Var(name.into())
47    }
48
49    pub fn stack_var(idx: usize) -> Self {
50        Expr::Var(format!("S{}", idx))
51    }
52
53    pub fn global_var() -> Self {
54        Expr::Var("G".to_string())
55    }
56
57    pub fn global_table_var() -> Self {
58        Expr::Var("GT".to_string())
59    }
60
61    pub fn local_table_var() -> Self {
62        Expr::Var("LT".to_string())
63    }
64
65    pub fn index(base: Expr, idx: Expr) -> Self {
66        Expr::Index(Box::new(base), Box::new(idx))
67    }
68
69    pub fn unary(op: UnOp, a: Expr) -> Self {
70        match (&op, &a) {
71            (UnOp::Neg, Expr::Int(v)) => Expr::Int(-v),
72            _ => Expr::Unary(op, Box::new(a)),
73        }
74    }
75
76    pub fn binary(op: BinOp, a: Expr, b: Expr) -> Self {
77        // Constant folding (safe subset).
78        match (&op, &a, &b) {
79            (BinOp::Add, Expr::Int(x), Expr::Int(y)) => return Expr::Int(x + y),
80            (BinOp::Sub, Expr::Int(x), Expr::Int(y)) => return Expr::Int(x - y),
81            (BinOp::Mul, Expr::Int(x), Expr::Int(y)) => return Expr::Int(x * y),
82            (BinOp::Eq, Expr::Int(x), Expr::Int(y)) => return Expr::Bool(x == y),
83            (BinOp::Ne, Expr::Int(x), Expr::Int(y)) => return Expr::Bool(x != y),
84            (BinOp::Gt, Expr::Int(x), Expr::Int(y)) => return Expr::Bool(x > y),
85            (BinOp::Ge, Expr::Int(x), Expr::Int(y)) => return Expr::Bool(x >= y),
86            (BinOp::Lt, Expr::Int(x), Expr::Int(y)) => return Expr::Bool(x < y),
87            (BinOp::Le, Expr::Int(x), Expr::Int(y)) => return Expr::Bool(x <= y),
88
89            // nil comparisons
90            (BinOp::Eq, Expr::Nil, Expr::Nil) => return Expr::Bool(true),
91            (BinOp::Ne, Expr::Nil, Expr::Nil) => return Expr::Bool(false),
92            (BinOp::Eq, Expr::Nil, _) | (BinOp::Eq, _, Expr::Nil) => {
93                // If the other side is a known non-nil literal.
94                if matches!(a, Expr::Nil) {
95                    if matches!(
96                        b,
97                        Expr::Bool(_) | Expr::Int(_) | Expr::Float(_) | Expr::Str(_)
98                    ) {
99                        return Expr::Bool(false);
100                    }
101                }
102                if matches!(b, Expr::Nil) {
103                    if matches!(
104                        a,
105                        Expr::Bool(_) | Expr::Int(_) | Expr::Float(_) | Expr::Str(_)
106                    ) {
107                        return Expr::Bool(false);
108                    }
109                }
110            }
111            (BinOp::Ne, Expr::Nil, _) | (BinOp::Ne, _, Expr::Nil) => {
112                if matches!(a, Expr::Nil) {
113                    if matches!(
114                        b,
115                        Expr::Bool(_) | Expr::Int(_) | Expr::Float(_) | Expr::Str(_)
116                    ) {
117                        return Expr::Bool(true);
118                    }
119                }
120                if matches!(b, Expr::Nil) {
121                    if matches!(
122                        a,
123                        Expr::Bool(_) | Expr::Int(_) | Expr::Float(_) | Expr::Str(_)
124                    ) {
125                        return Expr::Bool(true);
126                    }
127                }
128            }
129
130            // Simple algebraic identities that do not drop evaluation.
131            (BinOp::Add, x, Expr::Int(0)) => return x.clone(),
132            (BinOp::Add, Expr::Int(0), x) => return x.clone(),
133            (BinOp::Sub, x, Expr::Int(0)) => return x.clone(),
134            (BinOp::Mul, x, Expr::Int(1)) => return x.clone(),
135            (BinOp::Mul, Expr::Int(1), x) => return x.clone(),
136
137            _ => {}
138        }
139        Expr::Binary(op, Box::new(a), Box::new(b))
140    }
141
142    fn precedence(&self) -> u8 {
143        match self {
144            Expr::Nil
145            | Expr::Bool(_)
146            | Expr::Int(_)
147            | Expr::Float(_)
148            | Expr::Str(_)
149            | Expr::Var(_)
150            | Expr::Index(_, _) => 100,
151            Expr::Unary(_, _) => 90,
152            Expr::Binary(op, _, _) => match op {
153                BinOp::Mul | BinOp::Div | BinOp::Mod => 80,
154                BinOp::Add | BinOp::Sub => 70,
155                BinOp::BitAnd => 65,
156                BinOp::Eq | BinOp::Ne | BinOp::Gt | BinOp::Ge | BinOp::Lt | BinOp::Le => 50,
157                BinOp::And => 40,
158                BinOp::Or => 30,
159            },
160        }
161    }
162
163    fn mark_stack_var(name: &str, used_s: &mut BTreeSet<usize>) {
164        if let Some(rest) = name.strip_prefix('S') {
165            if !rest.is_empty() && rest.bytes().all(|c| c.is_ascii_digit()) {
166                if let Ok(v) = rest.parse::<usize>() {
167                    used_s.insert(v);
168                }
169            }
170        }
171    }
172
173    pub fn render(&self, used_s: &mut BTreeSet<usize>, parent_prec: u8) -> String {
174        let my_prec = self.precedence();
175        let mut s = match self {
176            Expr::Nil => "nil".to_string(),
177            Expr::Bool(v) => {
178                if *v {
179                    "true".to_string()
180                } else {
181                    "false".to_string()
182                }
183            }
184            Expr::Int(v) => v.to_string(),
185            Expr::Float(v) => v.clone(),
186            Expr::Str(v) => {
187                let lit = v.replace('\\', "\\\\").replace('"', "\\\"");
188                format!("\"{}\"", lit)
189            }
190            Expr::Var(v) => {
191                Self::mark_stack_var(v, used_s);
192                v.clone()
193            }
194            Expr::Index(base, idx) => {
195                let b = base.render(used_s, 100);
196                let i = idx.render(used_s, 0);
197                format!("{}[{}]", b, i)
198            }
199            Expr::Unary(UnOp::Neg, a) => {
200                let aa = a.render(used_s, my_prec);
201                format!("-{}", aa)
202            }
203            Expr::Binary(op, a, b) => {
204                let aa = a.render(used_s, my_prec);
205                let bb = b.render(used_s, my_prec + 1);
206                let op_s = match op {
207                    BinOp::Add => "+",
208                    BinOp::Sub => "-",
209                    BinOp::Mul => "*",
210                    BinOp::Div => "/",
211                    BinOp::Mod => "%",
212                    BinOp::BitAnd => "&",
213                    BinOp::And => "and",
214                    BinOp::Or => "or",
215                    BinOp::Eq => "==",
216                    BinOp::Ne => "~=",
217                    BinOp::Gt => ">",
218                    BinOp::Ge => ">=",
219                    BinOp::Lt => "<",
220                    BinOp::Le => "<=",
221                };
222                if matches!(op, BinOp::And | BinOp::Or) {
223                    format!("({}) {} ({})", aa, op_s, bb)
224                } else {
225                    format!("{} {} {}", aa, op_s, bb)
226                }
227            }
228        };
229
230        if my_prec < parent_prec {
231            s = format!("({})", s);
232        }
233        s
234    }
235
236    pub fn const_eq_zero(&self) -> Option<bool> {
237        match self {
238            Expr::Int(v) => Some(*v == 0),
239            // Do NOT treat booleans/nil as numeric; the generated code uses `cond == 0`.
240            Expr::Float(s) => s.parse::<f64>().ok().map(|v| v == 0.0),
241            _ => None,
242        }
243    }
244}
245
246fn global_name(idx: u16, non_volatile_count: u16, volatile_count: u16) -> String {
247    if idx < non_volatile_count {
248        return format!("g{}", idx);
249    }
250
251    let vbase = non_volatile_count;
252    let vlimit = non_volatile_count.saturating_add(volatile_count);
253    if idx >= vbase && idx < vlimit {
254        return format!("vg{}", idx - vbase);
255    }
256
257    format!("G[{}]", idx)
258}
259
260pub struct BlockEmitter<'a> {
261    indent: &'a str,
262    func_args: u8,
263    callee_args: &'a BTreeMap<u32, u8>,
264    used_s: &'a mut BTreeSet<usize>,
265    out: String,
266    stack: Vec<Expr>,
267    non_volatile_global_count: u16,
268    volatile_global_count: u16,
269}
270
271impl<'a> BlockEmitter<'a> {
272    pub fn new(
273        indent: &'a str,
274        func_args: u8,
275        callee_args: &'a BTreeMap<u32, u8>,
276        used_s: &'a mut BTreeSet<usize>,
277        non_volatile_global_count: u16,
278        volatile_global_count: u16,
279    ) -> Self {
280        BlockEmitter {
281            indent,
282            func_args,
283            callee_args,
284            used_s,
285            out: String::new(),
286            stack: Vec::new(),
287            non_volatile_global_count,
288            volatile_global_count,
289        }
290    }
291
292    pub fn init_stack(&mut self, depth: usize) {
293        self.stack.clear();
294        for i in 0..depth {
295            self.stack.push(Expr::stack_var(i));
296        }
297    }
298
299    pub fn take_output(self) -> String {
300        self.out
301    }
302
303    fn emit_line(&mut self, line: &str) {
304        let _ = writeln!(&mut self.out, "{}{}", self.indent, line);
305    }
306
307    fn pop(&mut self) -> Expr {
308        self.stack.pop().unwrap_or(Expr::Nil)
309    }
310
311    fn push(&mut self, e: Expr) {
312        self.stack.push(e);
313    }
314
315    fn stack_slot_get(&self, idx: i8) -> String {
316        // This is *frame* stack (args + locals), not the operand stack.
317        if idx < 0 {
318            let abs = (-idx) as u8 - 2;
319            if abs <= self.func_args {
320                let a = (self.func_args - abs) as usize;
321                return format!("a{}", a);
322            }
323            return format!("a_{}", idx);
324        }
325
326        let u = idx as u8;
327        if u < self.func_args {
328            format!("a{}", u as usize)
329        } else {
330            let l = (u - self.func_args) as usize;
331            format!("l{}", l)
332        }
333    }
334
335    fn materialize_all_stack_slots(&mut self) {
336        // Ensure S0..S{n-1} hold the stack values for successor blocks.
337        for i in 0..self.stack.len() {
338            let want = Expr::stack_var(i);
339            if self.stack[i] != want {
340                let rhs = self.stack[i].render(self.used_s, 0);
341                self.used_s.insert(i);
342                self.emit_line(&format!("S{} = {}", i, rhs));
343                self.stack[i] = Expr::stack_var(i);
344            } else {
345                // Mark it as used only if it is referenced later; render() will do that.
346            }
347        }
348    }
349
350    pub fn emit_inst(&mut self, inst: &Instruction) {
351        match &inst.op {
352            Op::Nop | Op::InitStack { .. } => {}
353
354            Op::PushNil => self.push(Expr::Nil),
355            Op::PushTrue => self.push(Expr::Bool(true)),
356            Op::PushI8(v) => self.push(Expr::Int(*v as i64)),
357            Op::PushI16(v) => self.push(Expr::Int(*v as i64)),
358            Op::PushI32(v) => self.push(Expr::Int(*v as i64)),
359            Op::PushF32(v) => self.push(Expr::Float(format!("{}", v))),
360            Op::PushString(s) => self.push(Expr::Str(s.clone())),
361
362            Op::PushTop => {
363                if let Some(top) = self.stack.last().cloned() {
364                    self.push(top);
365                } else {
366                    self.push(Expr::Nil);
367                }
368            }
369            Op::PushReturn => self.push(Expr::var("__ret")),
370
371            Op::PushGlobal(idx) => {
372                self.push(Expr::var(global_name(
373                    *idx,
374                    self.non_volatile_global_count,
375                    self.volatile_global_count,
376                )));
377            }
378            Op::PopGlobal(idx) => {
379                let v = self.pop();
380                let rhs = v.render(self.used_s, 0);
381                self.emit_line(&format!(
382                    "{} = {}",
383                    global_name(
384                        *idx,
385                        self.non_volatile_global_count,
386                        self.volatile_global_count,
387                    ),
388                    rhs
389                ));
390            }
391
392            Op::PushStack(idx) => {
393                let v = self.stack_slot_get(*idx);
394                self.push(Expr::var(v));
395            }
396            Op::PopStack(idx) => {
397                let rhs_expr = self.pop();
398                let rhs = rhs_expr.render(self.used_s, 0);
399                let lhs = self.stack_slot_get(*idx);
400                self.emit_line(&format!("{} = {}", lhs, rhs));
401            }
402
403            Op::PushGlobalTable(idx) => {
404                // top = GT[idx][top]
405                if let Some(last) = self.stack.pop() {
406                    let base = Expr::index(Expr::global_table_var(), Expr::Int(*idx as i64));
407                    let e = Expr::index(base, last);
408                    self.stack.push(e);
409                } else {
410                    self.stack.push(Expr::Nil);
411                }
412            }
413            Op::PopGlobalTable(idx) => {
414                // GT[idx][key] = value
415                let v = self.pop();
416                let k = self.pop();
417                let base = Expr::index(Expr::global_table_var(), Expr::Int(*idx as i64));
418                let lhs = Expr::index(base, k).render(self.used_s, 0);
419                let rhs = v.render(self.used_s, 0);
420                self.emit_line(&format!("{} = {}", lhs, rhs));
421            }
422
423            Op::PushLocalTable(idx) => {
424                if let Some(last) = self.stack.pop() {
425                    let base = Expr::index(Expr::local_table_var(), Expr::Int(*idx as i64));
426                    let e = Expr::index(base, last);
427                    self.stack.push(e);
428                } else {
429                    self.stack.push(Expr::Nil);
430                }
431            }
432            Op::PopLocalTable(idx) => {
433                let v = self.pop();
434                let k = self.pop();
435                let base = Expr::index(Expr::local_table_var(), Expr::Int(*idx as i64));
436                let lhs = Expr::index(base, k).render(self.used_s, 0);
437                let rhs = v.render(self.used_s, 0);
438                self.emit_line(&format!("{} = {}", lhs, rhs));
439            }
440
441            Op::Neg => {
442                let a = self.pop();
443                self.push(Expr::unary(UnOp::Neg, a));
444            }
445
446            Op::Add
447            | Op::Sub
448            | Op::Mul
449            | Op::Div
450            | Op::Mod
451            | Op::BitTest
452            | Op::And
453            | Op::Or
454            | Op::SetE
455            | Op::SetNE
456            | Op::SetG
457            | Op::SetGE
458            | Op::SetL
459            | Op::SetLE => {
460                let b = self.pop();
461                let a = self.pop();
462                let e = match &inst.op {
463                    Op::Add => Expr::binary(BinOp::Add, a, b),
464                    Op::Sub => Expr::binary(BinOp::Sub, a, b),
465                    Op::Mul => Expr::binary(BinOp::Mul, a, b),
466                    Op::Div => Expr::binary(BinOp::Div, a, b),
467                    Op::Mod => Expr::binary(BinOp::Mod, a, b),
468                    Op::BitTest => {
469                        let and = Expr::binary(BinOp::BitAnd, a, b);
470                        Expr::binary(BinOp::Ne, and, Expr::Int(0))
471                    }
472                    Op::And => {
473                        let aa = Expr::binary(BinOp::Ne, a, Expr::Nil);
474                        let bb = Expr::binary(BinOp::Ne, b, Expr::Nil);
475                        Expr::binary(BinOp::And, aa, bb)
476                    }
477                    Op::Or => {
478                        let aa = Expr::binary(BinOp::Ne, a, Expr::Nil);
479                        let bb = Expr::binary(BinOp::Ne, b, Expr::Nil);
480                        Expr::binary(BinOp::Or, aa, bb)
481                    }
482                    Op::SetE => Expr::binary(BinOp::Eq, a, b),
483                    Op::SetNE => Expr::binary(BinOp::Ne, a, b),
484                    Op::SetG => Expr::binary(BinOp::Gt, a, b),
485                    Op::SetGE => Expr::binary(BinOp::Ge, a, b),
486                    Op::SetL => Expr::binary(BinOp::Lt, a, b),
487                    Op::SetLE => Expr::binary(BinOp::Le, a, b),
488                    _ => Expr::Nil,
489                };
490                self.push(e);
491            }
492
493            Op::Call { target } => {
494                let argc = self.callee_args.get(target).copied().unwrap_or(0) as usize;
495                if self.stack.len() < argc {
496                    self.emit_line(&format!(
497                        "-- call f_{:08X} argc={} on short stack",
498                        target, argc
499                    ));
500                    self.emit_line(&format!("__ret = f_{:08X}()", target));
501                    self.stack.clear();
502                    return;
503                }
504
505                let base = self.stack.len() - argc;
506                let mut args_s = String::new();
507                for i in 0..argc {
508                    if i > 0 {
509                        args_s.push_str(", ");
510                    }
511                    let a = self.stack[base + i].clone();
512                    args_s.push_str(&a.render(self.used_s, 0));
513                }
514                self.emit_line(&format!("__ret = f_{:08X}({})", target, args_s));
515                self.stack.truncate(base);
516            }
517
518            Op::Syscall { name, args, id } => {
519                let argc = *args as usize;
520                if self.stack.len() < argc {
521                    self.emit_line(&format!(
522                        "-- syscall {} (id={}) argc={} on short stack",
523                        name, id, argc
524                    ));
525                    self.emit_line(&format!("__ret = {}()", name));
526                    self.stack.clear();
527                    return;
528                }
529
530                let base = self.stack.len() - argc;
531                let mut args_s = String::new();
532                for i in 0..argc {
533                    if i > 0 {
534                        args_s.push_str(", ");
535                    }
536                    let a = self.stack[base + i].clone();
537                    args_s.push_str(&a.render(self.used_s, 0));
538                }
539                self.emit_line(&format!("__ret = {}({})", name, args_s));
540                self.stack.truncate(base);
541            }
542
543            // Control-flow is handled at the basic-block terminator level.
544            Op::Jmp { .. } | Op::Jz { .. } | Op::Ret | Op::RetV => {}
545
546            Op::Unknown(opcode) => {
547                self.emit_line(&format!("-- unknown opcode 0x{:02X}", opcode));
548            }
549        }
550    }
551
552    pub fn emit_terminator(&mut self, term: BlockTerm, succs: &[usize]) {
553        match term {
554            BlockTerm::Ret => {
555                self.emit_line("return");
556            }
557            BlockTerm::RetV => {
558                let v = self.pop();
559                let rhs = v.render(self.used_s, 0);
560                self.emit_line(&format!("return {}", rhs));
561            }
562            BlockTerm::Jmp | BlockTerm::Fallthrough => {
563                self.materialize_all_stack_slots();
564                if let Some(&t) = succs.get(0) {
565                    self.emit_line(&format!("__pc = {}", t));
566                } else {
567                    self.emit_line("return");
568                }
569            }
570            BlockTerm::Jz => {
571                let cond = self.pop();
572                self.materialize_all_stack_slots();
573                let t = succs.get(0).copied();
574                let f = succs.get(1).copied();
575
576                if let Some(is_zero) = cond.const_eq_zero() {
577                    // Fold constant condition.
578                    let dst = if is_zero { t } else { f };
579                    match dst {
580                        Some(id) => self.emit_line(&format!("__pc = {}", id)),
581                        None => self.emit_line("return"),
582                    }
583                    return;
584                }
585
586                let c = cond.render(self.used_s, 0);
587                self.emit_line(&format!("if {} == 0 then", c));
588                match t {
589                    Some(tid) => self.emit_line(&format!("  __pc = {}", tid)),
590                    None => self.emit_line("  return"),
591                }
592                self.emit_line("else");
593                match f {
594                    Some(fid) => self.emit_line(&format!("  __pc = {}", fid)),
595                    None => self.emit_line("  return"),
596                }
597                self.emit_line("end");
598            }
599        }
600    }
601}
hcb2lua_decompiler/lua_opt.rs

hcb2lua_decompiler/
lua_opt.rs
