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

hcb2lua_decompiler/
lua_opt.rs
