luaur_code_gen/functions/
generate_vm_exit_blocks.rs

1use crate::enums::ir_block_kind::IrBlockKind;
2use crate::enums::ir_cmd::IrCmd;
3use crate::enums::ir_op_kind::IrOpKind;
4use crate::functions::add_use::add_use;
5use crate::functions::has_side_effects::has_side_effects;
6use crate::functions::is_unsafe_to_sink::is_unsafe_to_sink;
7use crate::functions::remove_use::remove_use;
8use crate::functions::replace_ir_utils::replace_ir_function_ir_op_ir_op;
9use crate::functions::visit_arguments::visit_arguments;
10use crate::macros::codegen_assert::CODEGEN_ASSERT;
11use crate::records::ir_builder::IrBuilder;
12use crate::records::ir_function::IrFunction;
13use crate::records::ir_inst::IrInst;
14use crate::records::ir_op::IrOp;
15use crate::records::vm_exit_sync_info::VmExitSyncInfo;
16use alloc::vec::Vec;
17use luaur_common::records::dense_hash_map::DenseHashMap;
18
19fn collect_inputs(inst: &mut IrInst, inputs: &mut Vec<(IrOp, u32)>) {
20    visit_arguments(inst, |op| {
21        if op.kind() == IrOpKind::Inst {
22            if let Some(slot) = inputs.iter_mut().find(|el| el.0 == op) {
23                slot.1 += 1;
24            } else {
25                inputs.push((op, 1));
26            }
27        }
28    });
29}
30
31fn redirect(op: &mut IrOp, inst_redir: &DenseHashMap<u32, u32>, inputs: &[(IrOp, u32)]) {
32    if op.kind() == IrOpKind::Inst {
33        if let Some(new_index) = inst_redir.find(&op.index()) {
34            *op = IrOp::ir_op_ir_op_kind_u32(IrOpKind::Inst, *new_index);
35        } else if !inputs.iter().any(|el| el.0 == *op) {
36            // Values can only be used if they are defined in the same block or be an input
37            CODEGEN_ASSERT!(false);
38        }
39    }
40}
41
42pub fn generate_vm_exit_blocks(build: &mut IrBuilder, recorded_vm_exit_syncs: &Vec<u32>) {
43    let function: *mut IrFunction = &mut build.function;
44
45    for &vm_exit_sync_location in recorded_vm_exit_syncs {
46        let sync_info: *mut VmExitSyncInfo = unsafe {
47            (*function)
48                .vm_exit_info
49                .get_or_insert(vm_exit_sync_location)
50        };
51
52        if unsafe { (*sync_info).reg_stores.is_empty() } {
53            continue;
54        }
55
56        // We will be collecting instructions we want to move into the VM exit in reverse order
57        let mut store_instructions: Vec<IrInst> = Vec::new();
58        let mut arg_instructions: Vec<u32> = Vec::new();
59        let mut inputs: Vec<(IrOp, u32)> = Vec::new();
60
61        // Start with the store instruction we got
62        let n_reg = unsafe { (*sync_info).reg_stores.len() };
63        for ri in 0..n_reg {
64            let n_st = unsafe { (&(*sync_info).reg_stores)[ri].stores.size() as usize };
65            for si in 0..n_st {
66                let mut backup = unsafe {
67                    (&(*sync_info).reg_stores)[ri].stores.as_slice()[si]
68                        .backup
69                        .clone()
70                };
71                collect_inputs(&mut backup, &mut inputs);
72                store_instructions.push(backup);
73            }
74        }
75
76        // For each input we got, see if we are the only user of it and if we are (and it has no side effects), schedule a move inside
77        let mut i = 0usize;
78        while i < inputs.len() {
79            let input_op = inputs[i].0;
80            let input_count = inputs[i].1;
81
82            let inst_use_count = unsafe { (*function).inst_op(input_op).use_count };
83            let inst_cmd = unsafe { (*function).inst_op(input_op).cmd };
84
85            if (inst_use_count as u32) == input_count
86                && !has_side_effects(inst_cmd)
87                && !is_unsafe_to_sink(inst_cmd)
88            {
89                let inst_idx = input_op.index();
90                arg_instructions.push(inst_idx);
91
92                inputs.remove(i); // Delete this input
93
94                let inst_ptr: *mut IrInst =
95                    unsafe { &mut (&mut (*function).instructions)[inst_idx as usize] };
96                collect_inputs(unsafe { &mut *inst_ptr }, &mut inputs);
97            } else {
98                i += 1;
99            }
100        }
101
102        for input in &inputs {
103            unsafe {
104                (*sync_info).arg_ops.push_back(input.0);
105            }
106        }
107
108        // We now should have an extracted instruction chain with no side effects in reverse order
109        let block_op = build.block(IrBlockKind::ExitSync);
110        unsafe {
111            (*sync_info).block = block_op;
112            *(*function)
113                .block_to_vm_exit_map
114                .get_or_insert(block_op.index()) = vm_exit_sync_location;
115        }
116        build.begin_block(block_op);
117
118        let mut inst_redir: DenseHashMap<u32, u32> = DenseHashMap::new(!0u32);
119
120        for k in (0..arg_instructions.len()).rev() {
121            let inst_idx = arg_instructions[k];
122
123            CODEGEN_ASSERT!((inst_idx as usize) < unsafe { (*function).instructions.len() });
124            let mut clone = unsafe { (&(*function).instructions)[inst_idx as usize].clone() };
125
126            let cn = clone.ops.size();
127            for oi in 0..cn {
128                redirect(
129                    &mut clone.ops.as_mut_slice()[oi as usize],
130                    &inst_redir,
131                    &inputs,
132                );
133            }
134
135            for oi in 0..cn {
136                let op = clone.ops.as_slice()[oi as usize];
137                add_use(unsafe { &mut *function }, op);
138            }
139
140            // Instructions that referenced the original will have to be adjusted to use the clone
141            unsafe {
142                *inst_redir.get_or_insert(inst_idx) = (*function).instructions.len() as u32;
143            }
144
145            // Reconstruct the fresh clone
146            build.inst_ir_cmd_ir_ops(clone.cmd, &clone.ops);
147        }
148
149        for si in 0..store_instructions.len() {
150            let mut clone = store_instructions[si].clone();
151
152            let cn = clone.ops.size();
153            for oi in 0..cn {
154                redirect(
155                    &mut clone.ops.as_mut_slice()[oi as usize],
156                    &inst_redir,
157                    &inputs,
158                );
159            }
160
161            for oi in 0..cn {
162                let op = clone.ops.as_slice()[oi as usize];
163                add_use(unsafe { &mut *function }, op);
164            }
165
166            // Reconstruct the fresh clone
167            build.inst_ir_cmd_ir_ops(clone.cmd, &clone.ops);
168
169            let mut orig = store_instructions[si].clone();
170            visit_arguments(&mut orig, |op| {
171                remove_use(unsafe { &mut *function }, op);
172            });
173        }
174
175        let vm_exit = unsafe { (*sync_info).vm_exit };
176        build.inst_ir_cmd_ir_op(IrCmd::JUMP, vm_exit);
177
178        // Replace guard VM exit with an exit sync block
179        let block_for_replace = unsafe { (*sync_info).block };
180        let guard_op_count = unsafe {
181            (&(*function).instructions)[vm_exit_sync_location as usize]
182                .ops
183                .size()
184        };
185        for oi in 0..guard_op_count {
186            let op = unsafe {
187                (&(*function).instructions)[vm_exit_sync_location as usize]
188                    .ops
189                    .as_slice()[oi as usize]
190            };
191            if op.kind() == IrOpKind::VmExit && op == vm_exit {
192                let op_ptr: *mut IrOp = unsafe {
193                    &mut (&mut (*function).instructions)[vm_exit_sync_location as usize]
194                        .ops
195                        .as_mut_slice()[oi as usize]
196                };
197                replace_ir_function_ir_op_ir_op(
198                    unsafe { &mut *function },
199                    unsafe { &mut *op_ptr },
200                    block_for_replace,
201                );
202                break;
203            }
204        }
205    }
206}
luaur_code_gen/functions/generate_vm_exit_blocks.rs

luaur_code_gen/functions/
generate_vm_exit_blocks.rs
