luaur_bytecode/methods/
bytecode_graph_parser_rebuild_blocks.rs

1use crate::records::bytecode_graph_parser::BytecodeGraphParser;
2use crate::type_aliases::instruction::Instruction;
3use luaur_common::enums::luau_opcode::LuauOpcode;
4use luaur_common::functions::get_jump_target::get_jump_target;
5use luaur_common::functions::get_op_length::get_op_length;
6use luaur_common::functions::is_fallthrough::is_fallthrough;
7use luaur_common::functions::is_fast_call::is_fast_call;
8use luaur_common::functions::is_loop_jump::is_loop_jump;
9use luaur_common::macros::luau_assert::LUAU_ASSERT;
10use luaur_common::macros::luau_insn_op::LUAU_INSN_OP;
11
12impl<'a> BytecodeGraphParser<'a> {
13    pub fn rebuild_blocks(&mut self, code: *const Instruction, codesize: u32) -> usize {
14        let entry_block = self.make_block(0);
15        self.func.entry_block = entry_block;
16        let exit_block = self.make_block(0xFFFFFFFFu32);
17        self.func.exit_block = exit_block;
18        let mut i: u32 = 0;
19        let mut current_block = entry_block;
20        let mut instruction_count: usize = 0;
21
22        while i < codesize {
23            let insn = unsafe { *code.add(i as usize) };
24            let op: LuauOpcode = unsafe { std::mem::transmute((LUAU_INSN_OP(insn) & 0xff) as u8) };
25            let target = get_jump_target(insn, i);
26            if target >= 0 {
27                let target_insn = unsafe { *code.add(target as usize) };
28                let target_op: LuauOpcode =
29                    unsafe { std::mem::transmute((LUAU_INSN_OP(target_insn) & 0xff) as u8) };
30                if target_op == LuauOpcode::LOP_JUMPX {
31                    let target2 = get_jump_target(target_insn, target as u32);
32                    if target2 >= 0 {
33                        let target2_insn = unsafe { *code.add(target2 as usize) };
34                        let target2_op: LuauOpcode = unsafe {
35                            std::mem::transmute((LUAU_INSN_OP(target2_insn) & 0xff) as u8)
36                        };
37                        if target2_op == LuauOpcode::LOP_JUMPX {
38                            // Double jumpx is not expected, but handle it anyway
39                            let target3 = get_jump_target(target2_insn, target2 as u32);
40                            if target3 >= 0 {
41                                let target3_insn = unsafe { *code.add(target3 as usize) };
42                                let target3_op: LuauOpcode = unsafe {
43                                    std::mem::transmute((LUAU_INSN_OP(target3_insn) & 0xff) as u8)
44                                };
45                                if target3_op == LuauOpcode::LOP_JUMPX {
46                                    // Triple jumpx - stop here
47                                } else {
48                                    // Use target3 as final target
49                                }
50                            }
51                        } else {
52                            // Use target as final target
53                        }
54                    } else {
55                        // Use target as final target
56                    }
57                } else {
58                    // Use target as final target
59                }
60            }
61
62            let needs_block = target >= 0
63                && !is_fast_call(op)
64                && op != LuauOpcode::LOP_JUMPX
65                && !self.is_jump_trampoline(i, code, codesize);
66            if needs_block {
67                if !self.block_by_pc.contains_key(&(target as u32)) {
68                    let new_block_op = self.make_block(target as u32);
69                    if (target as u32) < i {
70                        // We are jumping back.
71                        // The new block was created in the middle of the existing one.
72                        // We need to maintain predecessor/successor relations.
73                        let mut block_start_pc = target as u32 - 1;
74                        while !self.block_by_pc.contains_key(&block_start_pc) && block_start_pc != 0
75                        {
76                            block_start_pc -= 1;
77                        }
78                        LUAU_ASSERT!(self.block_by_pc.contains_key(&block_start_pc));
79                        let prev_block_op = *self.block_by_pc.get(&block_start_pc).unwrap();
80                        // Steal successors of the previous block.
81                        let successors = self.func.blocks[prev_block_op.index as usize]
82                            .successors
83                            .clone();
84                        self.func.blocks[new_block_op.index as usize].successors =
85                            successors.clone();
86                        // Now it should only fallthrough to the new block.
87                        self.func.blocks[prev_block_op.index as usize]
88                            .successors
89                            .clear();
90                        self.add_successor(
91                            prev_block_op,
92                            new_block_op,
93                            crate::enums::bc_block_edge_kind::BcBlockEdgeKind::Fallthrough,
94                        );
95                        // Update all successors to have the new block as a predecessor instead of the old one.
96                        for edge in &successors {
97                            let target_block = self.func.block_op(edge.target);
98                            for back_edge in target_block.predecessors.iter_mut() {
99                                if back_edge.target == prev_block_op {
100                                    back_edge.target = new_block_op;
101                                }
102                            }
103                        }
104                    }
105                }
106                let edge_kind = if is_loop_jump(op) {
107                    crate::enums::bc_block_edge_kind::BcBlockEdgeKind::Loop
108                } else {
109                    crate::enums::bc_block_edge_kind::BcBlockEdgeKind::Branch
110                };
111                self.add_successor(
112                    current_block,
113                    *self.block_by_pc.get(&(target as u32)).unwrap(),
114                    edge_kind,
115                );
116            }
117            if op == LuauOpcode::LOP_RETURN {
118                self.add_successor(
119                    current_block,
120                    exit_block,
121                    crate::enums::bc_block_edge_kind::BcBlockEdgeKind::Fallthrough,
122                );
123            }
124            let op_len = get_op_length(op) as u32;
125            i += op_len;
126            if (needs_block || (op == LuauOpcode::LOP_RETURN && i < codesize))
127                && !self.block_by_pc.contains_key(&i)
128            {
129                self.make_block(i);
130            }
131
132            if self.block_by_pc.contains_key(&i) {
133                if is_fallthrough(op) {
134                    self.add_successor(
135                        current_block,
136                        *self.block_by_pc.get(&i).unwrap(),
137                        crate::enums::bc_block_edge_kind::BcBlockEdgeKind::Fallthrough,
138                    );
139                }
140                current_block = *self.block_by_pc.get(&i).unwrap();
141            }
142            instruction_count += 1;
143        }
144        instruction_count
145    }
146}
luaur_bytecode/methods/bytecode_graph_parser_rebuild_blocks.rs

luaur_bytecode/methods/
bytecode_graph_parser_rebuild_blocks.rs
