hcb2lua_decompiler/
cfg.rs

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

hcb2lua_decompiler/
cfg.rs
