hcb2lua_decompiler/
cfg.rs

1use std::collections::{BTreeMap, BTreeSet, VecDeque};
2
3use crate::decode::{Function, Instruction, Op};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6pub enum BlockTerm {
7    Fallthrough,
8    Jmp,
9    Jz,
10    Ret,
11    RetV,
12}
13
14#[derive(Debug, Clone)]
15pub struct BasicBlock {
16    pub id: usize,
17    pub start: u32,
18    pub end: u32, // exclusive (next leader or function end)
19    pub inst_indices: std::ops::Range<usize>,
20    pub preds: Vec<usize>,
21    pub succs: Vec<usize>,
22    pub term: BlockTerm,
23    pub in_depth: usize,
24    pub out_depth: usize,
25    pub is_loop_header: bool,
26}
27
28#[derive(Debug, Clone)]
29pub struct FunctionCfg {
30    pub blocks: Vec<BasicBlock>,
31    pub addr_to_block: BTreeMap<u32, usize>,
32    pub max_depth: usize,
33    pub stack_consistent: bool,
34}
35
36fn inst_stack_delta(inst: &Instruction, callee_args: &BTreeMap<u32, u8>) -> i32 {
37    match &inst.op {
38        Op::Nop | Op::InitStack { .. } | Op::Jmp { .. } | Op::Ret => 0,
39        Op::Jz { .. } => -1,
40        Op::RetV => -1,
41        Op::PushNil
42        | Op::PushTrue
43        | Op::PushI8(_)
44        | Op::PushI16(_)
45        | Op::PushI32(_)
46        | Op::PushF32(_)
47        | Op::PushString(_)
48        | Op::PushGlobal(_)
49        | Op::PushStack(_) => 1,
50        Op::PushGlobalTable(_) | Op::PushLocalTable(_) => 0,
51        Op::PushTop | Op::PushReturn => 1,
52        Op::PopGlobal(_) | Op::PopStack(_) => -1,
53        Op::PopGlobalTable(_) | Op::PopLocalTable(_) => -2,
54        Op::Neg => 0,
55        Op::Add
56        | Op::Sub
57        | Op::Mul
58        | Op::Div
59        | Op::Mod
60        | Op::BitTest
61        | Op::And
62        | Op::Or
63        | Op::SetE
64        | Op::SetNE
65        | Op::SetG
66        | Op::SetLE
67        | Op::SetL
68        | Op::SetGE => -1,
69        Op::Call { target } => -(callee_args.get(target).copied().unwrap_or(0) as i32),
70        Op::Syscall { args, .. } => -(*args as i32),
71        Op::Unknown(_) => 0,
72    }
73}
74
75fn block_term(last: Option<&Instruction>) -> BlockTerm {
76    match last.map(|i| &i.op) {
77        Some(Op::Jmp { .. }) => BlockTerm::Jmp,
78        Some(Op::Jz { .. }) => BlockTerm::Jz,
79        Some(Op::Ret) => BlockTerm::Ret,
80        Some(Op::RetV) => BlockTerm::RetV,
81        _ => BlockTerm::Fallthrough,
82    }
83}
84
85pub fn build_cfg(func: &Function, callee_args: &BTreeMap<u32, u8>) -> FunctionCfg {
86    // Leaders.
87    let mut leaders: BTreeSet<u32> = BTreeSet::new();
88    if let Some(first) = func.insts.first() {
89        leaders.insert(first.addr);
90    }
91
92    // Map addr -> inst index.
93    let mut addr_to_idx: BTreeMap<u32, usize> = BTreeMap::new();
94    for (i, inst) in func.insts.iter().enumerate() {
95        addr_to_idx.insert(inst.addr, i);
96    }
97
98    for (i, inst) in func.insts.iter().enumerate() {
99        match &inst.op {
100            Op::Jmp { target } | Op::Jz { target } => {
101                leaders.insert(*target);
102                if let Some(next) = func.insts.get(i + 1) {
103                    leaders.insert(next.addr);
104                }
105            }
106            Op::Ret | Op::RetV => {
107                if let Some(next) = func.insts.get(i + 1) {
108                    leaders.insert(next.addr);
109                }
110            }
111            _ => {}
112        }
113    }
114
115    let func_end = func
116        .insts
117        .last()
118        .map(|i| i.addr + 1)
119        .unwrap_or(func.start_addr);
120
121    // Build blocks by leader ranges.
122    let mut leader_vec: Vec<u32> = leaders.into_iter().collect();
123    leader_vec.sort();
124
125    let mut blocks: Vec<BasicBlock> = Vec::new();
126    let mut addr_to_block: BTreeMap<u32, usize> = BTreeMap::new();
127
128    for (bid, &start) in leader_vec.iter().enumerate() {
129        let end = leader_vec.get(bid + 1).copied().unwrap_or(func_end);
130        let si = addr_to_idx.get(&start).copied().unwrap_or(0);
131        let ei = addr_to_idx
132            .range(end..)
133            .next()
134            .map(|(_, &idx)| idx)
135            .unwrap_or(func.insts.len());
136
137        for (&a, _) in addr_to_idx.range(start..end) {
138            addr_to_block.insert(a, bid);
139        }
140
141        let last = if ei > 0 { func.insts.get(ei - 1) } else { None };
142        blocks.push(BasicBlock {
143            id: bid,
144            start,
145            end,
146            inst_indices: si..ei,
147            preds: Vec::new(),
148            succs: Vec::new(),
149            term: block_term(last),
150            in_depth: 0,
151            out_depth: 0,
152            is_loop_header: false,
153        });
154    }
155
156    // Fill succs.
157    for b in &mut blocks {
158        let last_idx = b.inst_indices.clone().last();
159        let last = last_idx.and_then(|i| func.insts.get(i));
160        match last.map(|i| &i.op) {
161            Some(Op::Jmp { target }) => {
162                if let Some(&tid) = addr_to_block.get(target) {
163                    b.succs.push(tid);
164                }
165            }
166            Some(Op::Jz { target }) => {
167                if let Some(&tid) = addr_to_block.get(target) {
168                    b.succs.push(tid);
169                }
170                // fallthrough
171                let ft = func.insts.get(b.inst_indices.end).map(|i| i.addr);
172                if let Some(ft) = ft {
173                    if let Some(&fid) = addr_to_block.get(&ft) {
174                        b.succs.push(fid);
175                    }
176                }
177            }
178            Some(Op::Ret) | Some(Op::RetV) => {}
179            _ => {
180                // Fallthrough to next block by address.
181                let ft = func.insts.get(b.inst_indices.end).map(|i| i.addr);
182                if let Some(ft) = ft {
183                    if let Some(&fid) = addr_to_block.get(&ft) {
184                        b.succs.push(fid);
185                    }
186                }
187            }
188        }
189    }
190
191    // Fill preds.
192    for b in 0..blocks.len() {
193        let succs = blocks[b].succs.clone();
194        for s in succs {
195            if let Some(sb) = blocks.get_mut(s) {
196                sb.preds.push(b);
197            }
198        }
199    }
200
201    // Compute in/out depths using a simple worklist.
202    let mut in_depth: Vec<Option<usize>> = vec![None; blocks.len()];
203    if !blocks.is_empty() {
204        in_depth[0] = Some(0);
205    }
206    let mut q: VecDeque<usize> = VecDeque::new();
207    if !blocks.is_empty() {
208        q.push_back(0);
209    }
210    let mut max_depth = 0usize;
211    let mut stack_consistent = true;
212
213    while let Some(bid) = q.pop_front() {
214        let mut d = in_depth[bid].unwrap_or(0) as i32;
215        let b = &blocks[bid];
216
217        for i in b.inst_indices.clone() {
218            let inst = &func.insts[i];
219            d += inst_stack_delta(inst, callee_args);
220            if d < 0 {
221                d = 0;
222            }
223            max_depth = max_depth.max(d as usize);
224        }
225
226        let out = d as usize;
227        for &sid in &blocks[bid].succs {
228            match in_depth[sid] {
229                None => {
230                    in_depth[sid] = Some(out);
231                    q.push_back(sid);
232                }
233                Some(existing) => {
234                    // If mismatched, keep the max to avoid panics; we will still emit
235                    // readable code, but the stack-model is approximate.
236                    if existing != out {
237                        stack_consistent = false;
238                        let newd = existing.max(out);
239                        if newd != existing {
240                            in_depth[sid] = Some(newd);
241                            q.push_back(sid);
242                        }
243                    }
244                }
245            }
246        }
247    }
248
249    for b in &mut blocks {
250        b.in_depth = in_depth[b.id].unwrap_or(0);
251        // Simulate to compute out depth.
252        let mut d = b.in_depth as i32;
253        for i in b.inst_indices.clone() {
254            let inst = &func.insts[i];
255            d += inst_stack_delta(inst, callee_args);
256            if d < 0 {
257                d = 0;
258            }
259        }
260        b.out_depth = d as usize;
261    }
262
263    // Loop header classification: any block that is target of a back-edge.
264    let mut loop_headers: BTreeSet<usize> = BTreeSet::new();
265    for b in &blocks {
266        for &s in &b.succs {
267            if let (Some(src), Some(dst)) = (blocks.get(b.id), blocks.get(s)) {
268                if dst.start < src.start {
269                    loop_headers.insert(s);
270                }
271            }
272        }
273    }
274    for b in &mut blocks {
275        b.is_loop_header = loop_headers.contains(&b.id);
276    }
277
278    FunctionCfg {
279        blocks,
280        addr_to_block,
281        max_depth,
282        stack_consistent,
283    }
284}
hcb2lua_decompiler/cfg.rs

hcb2lua_decompiler/
cfg.rs
