tjs2dec\decompile/
cfg.rs

1use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
2
3use anyhow::{Result, bail};
4
5use crate::model::Tjs2Object;
6use crate::vmcodes::vm;
7
8use super::decode::{Insn, decode_object};
9
10#[derive(Debug, Clone)]
11pub struct BasicBlock {
12    pub id: usize,
13    pub start_pc: usize,
14    pub insns: Vec<Insn>,
15    pub succ: Vec<usize>, // successor block ids
16    pub pred: Vec<usize>, // predecessor block ids
17}
18
19#[derive(Debug, Clone)]
20pub struct Cfg {
21    pub obj_index: usize,
22    pub code_len: usize,
23    pub blocks: Vec<BasicBlock>,
24    pub pc_to_block: HashMap<usize, usize>,
25    pub entry_block: usize,
26    pub try_regions: Vec<TryRegion>,
27    pub catch_sites: HashMap<usize, i32>, // catch_pc -> ex_reg
28}
29
30#[derive(Debug, Clone)]
31pub struct TryRegion {
32    pub start_pc: usize,
33    pub end_pc: usize, // exclusive
34    pub catch_pc: usize,
35    pub ex_reg: i32,
36    pub depth: usize,
37}
38
39impl Cfg {
40    pub fn build(obj: &Tjs2Object) -> Result<Self> {
41        let insns = decode_object(obj)?;
42        let code_len = obj.code.len();
43        let (try_regions, catch_sites) = parse_try_regions(&insns, code_len)?;
44
45        let leaders = compute_leaders(&insns, code_len, &try_regions)?;
46        let (mut blocks, pc_to_block) = build_blocks(obj.index, code_len, &insns, &leaders)?;
47        connect_edges(&mut blocks, &pc_to_block, code_len, &try_regions)?;
48        let entry_block = *pc_to_block
49            .get(&0)
50            .ok_or_else(|| anyhow::anyhow!("missing entry block"))?;
51
52        Ok(Self {
53            obj_index: obj.index,
54            code_len,
55            blocks,
56            pc_to_block,
57            entry_block,
58            try_regions,
59            catch_sites,
60        })
61    }
62}
63
64fn parse_try_regions(
65    insns: &[Insn],
66    code_len: usize,
67) -> Result<(Vec<TryRegion>, HashMap<usize, i32>)> {
68    // Model nesting based on ENTRY/EXTRY.
69    // Range: from (entry_pc + 3) up to (extry_pc) exclusive.
70    let mut stack: Vec<(usize, usize, i32)> = Vec::new(); // (start_pc, catch_pc, ex_reg)
71    let mut regions: Vec<TryRegion> = Vec::new();
72    let mut catch_sites: HashMap<usize, i32> = HashMap::new();
73
74    for insn in insns {
75        match insn.op {
76            x if x == vm::VM_ENTRY => {
77                if insn.size != 3 {
78                    bail!("ENTRY size mismatch at {}", insn.pc);
79                }
80                let catch_rel = insn.words[1];
81                let catch_pc = (insn.pc as i32 + catch_rel) as isize;
82                if catch_pc < 0 || catch_pc as usize >= code_len {
83                    bail!(
84                        "ENTRY catch target out of range: pc={} rel={}",
85                        insn.pc,
86                        catch_rel
87                    );
88                }
89                let catch_pc = catch_pc as usize;
90                let ex_reg = insn.words[2];
91                let start_pc = insn.pc + insn.size;
92                stack.push((start_pc, catch_pc, ex_reg));
93                catch_sites.entry(catch_pc).or_insert(ex_reg);
94            }
95            x if x == vm::VM_EXTRY => {
96                // NOTE: In real TJS2 bytecode, the compiler may emit multiple EXTRY instructions
97                // for the same ENTRY (multi-exit try blocks). Linear matching therefore cannot
98                // assume a strict 1:1 nesting in the instruction stream.
99                if let Some((start_pc, catch_pc, ex_reg)) = stack.pop() {
100                    let end_pc = insn.pc; // exclusive
101                    let depth = stack.len();
102                    regions.push(TryRegion {
103                        start_pc,
104                        end_pc,
105                        catch_pc,
106                        ex_reg,
107                        depth,
108                    });
109                } else {
110                    // Orphan EXTRY: ignore for region construction (do not fail CFG/SSA build).
111                }
112            }
113            _ => {}
114        }
115    }
116    // If there are unmatched ENTRYs, we still keep them as open regions up to end.
117    while let Some((start_pc, catch_pc, ex_reg)) = stack.pop() {
118        let depth = stack.len();
119        regions.push(TryRegion {
120            start_pc,
121            end_pc: code_len,
122            catch_pc,
123            ex_reg,
124            depth,
125        });
126    }
127    regions.sort_by_key(|r| (r.start_pc, r.end_pc, r.depth));
128    Ok((regions, catch_sites))
129}
130
131fn compute_leaders(
132    insns: &[Insn],
133    code_len: usize,
134    try_regions: &[TryRegion],
135) -> Result<BTreeSet<usize>> {
136    let mut leaders: BTreeSet<usize> = BTreeSet::new();
137    leaders.insert(0);
138
139    for insn in insns {
140        let pc = insn.pc;
141        let next = pc + insn.size;
142
143        match insn.op {
144            x if x == vm::VM_JMP => {
145                let t = (pc as i32 + insn.words[1]) as isize;
146                if t < 0 || t as usize >= code_len {
147                    bail!("JMP target out of range at {}", pc);
148                }
149                leaders.insert(t as usize);
150                if next < code_len {
151                    leaders.insert(next);
152                }
153            }
154            x if x == vm::VM_JF || x == vm::VM_JNF => {
155                let t = (pc as i32 + insn.words[1]) as isize;
156                if t < 0 || t as usize >= code_len {
157                    bail!("JF/JNF target out of range at {}", pc);
158                }
159                leaders.insert(t as usize);
160                if next < code_len {
161                    leaders.insert(next);
162                }
163            }
164            x if x == vm::VM_ENTRY => {
165                // ENTRY's catch target is a leader.
166                let t = (pc as i32 + insn.words[1]) as isize;
167                if t >= 0 && (t as usize) < code_len {
168                    leaders.insert(t as usize);
169                }
170                if next < code_len {
171                    leaders.insert(next);
172                }
173            }
174            x if x == vm::VM_RET || x == vm::VM_THROW => {
175                if next < code_len {
176                    leaders.insert(next);
177                }
178            }
179            _ => {}
180        }
181    }
182
183    // Also mark catch entries as leaders (redundant but explicit).
184    for r in try_regions {
185        leaders.insert(r.catch_pc);
186    }
187    Ok(leaders)
188}
189
190fn build_blocks(
191    obj_index: usize,
192    code_len: usize,
193    insns: &[Insn],
194    leaders: &BTreeSet<usize>,
195) -> Result<(Vec<BasicBlock>, HashMap<usize, usize>)> {
196    // Map leader pc -> sequential id.
197    let mut leader_list: Vec<usize> = leaders.iter().copied().collect();
198    leader_list.sort();
199    if leader_list.first().copied().unwrap_or(usize::MAX) != 0 {
200        bail!("missing leader 0 in object {}", obj_index);
201    }
202
203    let mut pc_to_block: HashMap<usize, usize> = HashMap::new();
204    for (id, &pc) in leader_list.iter().enumerate() {
205        pc_to_block.insert(pc, id);
206    }
207
208    // Build pc->Insn lookup.
209    let mut insn_at: BTreeMap<usize, Insn> = BTreeMap::new();
210    for insn in insns {
211        insn_at.insert(insn.pc, insn.clone());
212    }
213
214    let mut blocks: Vec<BasicBlock> = Vec::new();
215    for (id, &start_pc) in leader_list.iter().enumerate() {
216        let end_pc = leader_list.get(id + 1).copied().unwrap_or(code_len);
217        let mut cur_pc = start_pc;
218        let mut blk_insns = Vec::new();
219        while cur_pc < end_pc {
220            let insn = insn_at
221                .get(&cur_pc)
222                .cloned()
223                .ok_or_else(|| anyhow::anyhow!("no instruction at pc {}", cur_pc))?;
224            cur_pc += insn.size;
225            blk_insns.push(insn);
226        }
227        blocks.push(BasicBlock {
228            id,
229            start_pc,
230            insns: blk_insns,
231            succ: Vec::new(),
232            pred: Vec::new(),
233        });
234    }
235    Ok((blocks, pc_to_block))
236}
237
238fn connect_edges(
239    blocks: &mut [BasicBlock],
240    pc_to_block: &HashMap<usize, usize>,
241    code_len: usize,
242    try_regions: &[TryRegion],
243) -> Result<()> {
244    // Precompute catch edges per block (conservative): if a block is inside a try region,
245    // it may have an exceptional successor to the innermost handler.
246    let mut exc_succ: HashMap<usize, usize> = HashMap::new(); // block_id -> catch_block_id
247    for b in blocks.iter() {
248        if let Some((catch_pc, _ex_reg)) = innermost_handler_for_pc(b.start_pc, try_regions) {
249            if let Some(&cid) = pc_to_block.get(&catch_pc) {
250                exc_succ.insert(b.id, cid);
251            }
252        }
253    }
254
255    for b in blocks.iter_mut() {
256        if b.insns.is_empty() {
257            b.succ = Vec::new();
258            continue;
259        }
260        let last = b.insns.last().unwrap();
261        let pc = last.pc;
262        let next = pc + last.size;
263
264        let mut succ: Vec<usize> = Vec::new();
265        match last.op {
266            x if x == vm::VM_JMP => {
267                let t = (pc as i32 + last.words[1]) as usize;
268                succ.push(
269                    *pc_to_block
270                        .get(&t)
271                        .ok_or_else(|| anyhow::anyhow!("missing block for JMP target {}", t))?,
272                );
273            }
274            x if x == vm::VM_JF || x == vm::VM_JNF => {
275                let t = (pc as i32 + last.words[1]) as usize;
276                succ.push(
277                    *pc_to_block
278                        .get(&t)
279                        .ok_or_else(|| anyhow::anyhow!("missing block for JF/JNF target {}", t))?,
280                );
281                if next < code_len {
282                    succ.push(*pc_to_block.get(&next).ok_or_else(|| {
283                        anyhow::anyhow!("missing block for fallthrough {}", next)
284                    })?);
285                }
286            }
287            x if x == vm::VM_RET || x == vm::VM_THROW => {
288                // no normal successors
289            }
290            _ => {
291                if next < code_len {
292                    if let Some(&nid) = pc_to_block.get(&next) {
293                        succ.push(nid);
294                    }
295                }
296            }
297        }
298
299        // Add an exceptional successor for blocks in a try region, conservatively.
300        if let Some(&cid) = exc_succ.get(&b.id) {
301            if !succ.contains(&cid) {
302                succ.push(cid);
303            }
304        }
305
306        // Dedup while preserving order.
307        let mut seen: HashSet<usize> = HashSet::new();
308        succ.retain(|x| seen.insert(*x));
309        b.succ = succ;
310    }
311
312    // Fill predecessor lists.
313    let mut preds: Vec<Vec<usize>> = vec![Vec::new(); blocks.len()];
314    for b in blocks.iter() {
315        for &s in &b.succ {
316            preds[s].push(b.id);
317        }
318    }
319    for (id, p) in preds.into_iter().enumerate() {
320        blocks[id].pred = p;
321    }
322    Ok(())
323}
324
325fn innermost_handler_for_pc(pc: usize, try_regions: &[TryRegion]) -> Option<(usize, i32)> {
326    // Pick the deepest (most nested) region that contains pc.
327    let mut best: Option<&TryRegion> = None;
328    for r in try_regions {
329        if pc >= r.start_pc && pc < r.end_pc {
330            match best {
331                None => best = Some(r),
332                Some(b) => {
333                    if r.depth >= b.depth {
334                        best = Some(r);
335                    }
336                }
337            }
338        }
339    }
340    best.map(|r| (r.catch_pc, r.ex_reg))
341}