Skip to main content

luaur_compiler/methods/
compiler_compile_stat_for_in.rs

1use crate::records::compiler::Compiler;
2use crate::records::reg_scope::RegScope;
3use luaur_ast::records::ast_stat_for_in::AstStatForIn;
4use luaur_common::enums::luau_opcode::LuauOpcode;
5use luaur_common::macros::luau_assert::LUAU_ASSERT;
6
7impl Compiler {
8    pub fn compile_stat_for_in(&mut self, stat: *mut AstStatForIn) {
9        unsafe {
10            let stat_ref = &*stat;
11            let mut rs = self.reg_scope_compiler();
12
13            let old_locals = self.local_stack.len();
14            let old_jumps = self.loop_jumps.len();
15
16            self.loops.push(crate::records::r#loop::Loop {
17                local_offset: old_locals,
18                local_offset_continue: old_locals,
19                continue_used: core::ptr::null_mut(),
20            });
21            self.has_loops = true;
22
23            let regs = self.alloc_reg(stat as *mut _, 3);
24
25            self.compile_expr_list_temp(&stat_ref.values, regs, 3, true);
26
27            let vars = self.alloc_reg(stat as *mut _, core::cmp::max(stat_ref.vars.size as u32, 2));
28            LUAU_ASSERT!(vars == regs + 3);
29            let vars_alloc_pc = (*self.bytecode).get_debug_pc();
30
31            let mut skip_op = LuauOpcode::LOP_FORGPREP;
32
33            if self.options.optimization_level >= 1 && stat_ref.vars.size <= 2 {
34                if stat_ref.values.size == 1
35                    && luaur_ast::rtti::ast_node_is::<luaur_ast::records::ast_expr_call::AstExprCall>(
36                        *stat_ref.values.data as *mut luaur_ast::records::ast_node::AstNode,
37                    )
38                {
39                    // C++ passes the CALL's `func` (the `ipairs`/`pairs` reference),
40                    // NOT the call expression itself, to getBuiltin. The model passed
41                    // the whole `ipairs(t)` call, so it never matched the builtin and
42                    // the loop fell back to generic FORGPREP.
43                    let call = luaur_ast::rtti::ast_node_as::<
44                        luaur_ast::records::ast_expr_call::AstExprCall,
45                    >(
46                        *stat_ref.values.data as *mut luaur_ast::records::ast_node::AstNode,
47                    );
48                    let builtin = crate::functions::get_builtin::get_builtin(
49                        (*call).func,
50                        &self.globals,
51                        &self.variables,
52                    );
53
54                    if builtin.is_global("ipairs") {
55                        skip_op = LuauOpcode::LOP_FORGPREP_INEXT;
56                    } else if builtin.is_global("pairs") {
57                        skip_op = LuauOpcode::LOP_FORGPREP_NEXT;
58                    }
59                } else if stat_ref.values.size == 2 {
60                    let builtin = crate::functions::get_builtin::get_builtin(
61                        *stat_ref.values.data,
62                        &self.globals,
63                        &self.variables,
64                    );
65
66                    if builtin.is_global("next") {
67                        skip_op = LuauOpcode::LOP_FORGPREP_NEXT;
68                    }
69                }
70            }
71
72            let skip_label = (*self.bytecode).emit_label();
73
74            (*self.bytecode).emit_ad(skip_op, regs, 0);
75
76            let loop_label = (*self.bytecode).emit_label();
77
78            for i in 0..stat_ref.vars.size {
79                self.push_local(
80                    *stat_ref.vars.data.add(i as usize),
81                    (vars + i as u8) as u8,
82                    vars_alloc_pc,
83                );
84            }
85
86            self.compile_stat(stat_ref.body as *mut _);
87
88            self.close_locals(old_locals);
89            self.pop_locals(old_locals);
90
91            self.set_debug_line_ast_node(stat as *mut _);
92
93            let cont_label = (*self.bytecode).emit_label();
94            let back_label = (*self.bytecode).emit_label();
95
96            (*self.bytecode).emit_ad(LuauOpcode::LOP_FORGLOOP, regs, 0);
97            (*self.bytecode).emit_aux(
98                if skip_op == LuauOpcode::LOP_FORGPREP_INEXT {
99                    0x80000000
100                } else {
101                    0
102                } | stat_ref.vars.size as u32,
103            );
104
105            let end_label = (*self.bytecode).emit_label();
106
107            self.patch_jump(stat as *mut _, skip_label, back_label);
108            self.patch_jump(stat as *mut _, back_label, loop_label);
109
110            self.patch_loop_jumps(stat as *mut _, old_jumps, end_label, cont_label);
111            self.loop_jumps.resize(
112                old_jumps,
113                crate::records::loop_jump::LoopJump {
114                    r#type: crate::enums::type_compiler::Type::Break,
115                    label: 0,
116                },
117            );
118
119            self.loops.pop();
120        }
121    }
122}