Skip to main content

luaur_compiler/functions/
fold_binary.rs

1use alloc::vec::Vec;
2use core::ffi::c_char;
3use luaur_ast::records::ast_expr_binary::AstExprBinaryOp;
4use luaur_ast::records::ast_name_table::AstNameTable;
5use luaur_common::macros::luau_assert::LUAU_ASSERT;
6
7use crate::enums::type_constant_folding::Type;
8use crate::functions::constants_equal::constants_equal;
9use crate::records::constant::Constant;
10
11pub fn fold_binary(
12    result: &mut Constant,
13    op: AstExprBinaryOp,
14    la: &Constant,
15    ra: &Constant,
16    string_table: &mut AstNameTable,
17) {
18    const K_CONSTANT_FOLD_STRING_LIMIT: u32 = 4096;
19
20    match op {
21        AstExprBinaryOp::Add => {
22            if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Number {
23                result.r#type = Type::Type_Number;
24                unsafe {
25                    result.data.value_number = la.data.value_number + ra.data.value_number;
26                }
27            } else if la.r#type == Type::Type_Vector && ra.r#type == Type::Type_Vector {
28                result.r#type = Type::Type_Vector;
29                unsafe {
30                    for i in 0..4 {
31                        result.data.value_vector[i] =
32                            la.data.value_vector[i] + ra.data.value_vector[i];
33                    }
34                }
35            }
36        }
37        AstExprBinaryOp::Sub => {
38            if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Number {
39                result.r#type = Type::Type_Number;
40                unsafe {
41                    result.data.value_number = la.data.value_number - ra.data.value_number;
42                }
43            } else if la.r#type == Type::Type_Vector && ra.r#type == Type::Type_Vector {
44                result.r#type = Type::Type_Vector;
45                unsafe {
46                    for i in 0..4 {
47                        result.data.value_vector[i] =
48                            la.data.value_vector[i] - ra.data.value_vector[i];
49                    }
50                }
51            }
52        }
53        AstExprBinaryOp::Mul => {
54            if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Number {
55                result.r#type = Type::Type_Number;
56                unsafe {
57                    result.data.value_number = la.data.value_number * ra.data.value_number;
58                }
59            } else if la.r#type == Type::Type_Vector && ra.r#type == Type::Type_Vector {
60                unsafe {
61                    let had_w = la.data.value_vector[3] != 0.0 || ra.data.value_vector[3] != 0.0;
62                    let result_w = la.data.value_vector[3] * ra.data.value_vector[3];
63                    if result_w == 0.0 || had_w {
64                        result.r#type = Type::Type_Vector;
65                        for i in 0..3 {
66                            result.data.value_vector[i] =
67                                la.data.value_vector[i] * ra.data.value_vector[i];
68                        }
69                        result.data.value_vector[3] = result_w;
70                    }
71                }
72            } else if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Vector {
73                unsafe {
74                    let had_w = ra.data.value_vector[3] != 0.0;
75                    let result_w = (la.data.value_number as f32) * ra.data.value_vector[3];
76                    if result_w == 0.0 || had_w {
77                        result.r#type = Type::Type_Vector;
78                        for i in 0..3 {
79                            result.data.value_vector[i] =
80                                (la.data.value_number as f32) * ra.data.value_vector[i];
81                        }
82                        result.data.value_vector[3] = result_w;
83                    }
84                }
85            } else if la.r#type == Type::Type_Vector && ra.r#type == Type::Type_Number {
86                unsafe {
87                    let had_w = la.data.value_vector[3] != 0.0;
88                    let result_w = la.data.value_vector[3] * (ra.data.value_number as f32);
89                    if result_w == 0.0 || had_w {
90                        result.r#type = Type::Type_Vector;
91                        for i in 0..3 {
92                            result.data.value_vector[i] =
93                                la.data.value_vector[i] * (ra.data.value_number as f32);
94                        }
95                        result.data.value_vector[3] = result_w;
96                    }
97                }
98            }
99        }
100        AstExprBinaryOp::Div => {
101            if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Number {
102                result.r#type = Type::Type_Number;
103                unsafe {
104                    result.data.value_number = la.data.value_number / ra.data.value_number;
105                }
106            } else if la.r#type == Type::Type_Vector && ra.r#type == Type::Type_Vector {
107                unsafe {
108                    let had_w = la.data.value_vector[3] != 0.0 || ra.data.value_vector[3] != 0.0;
109                    let result_w = la.data.value_vector[3] / ra.data.value_vector[3];
110                    if result_w == 0.0 || had_w {
111                        result.r#type = Type::Type_Vector;
112                        for i in 0..3 {
113                            result.data.value_vector[i] =
114                                la.data.value_vector[i] / ra.data.value_vector[i];
115                        }
116                        result.data.value_vector[3] = result_w;
117                    }
118                }
119            } else if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Vector {
120                unsafe {
121                    let had_w = ra.data.value_vector[3] != 0.0;
122                    let result_w = (la.data.value_number as f32) / ra.data.value_vector[3];
123                    if result_w == 0.0 || had_w {
124                        result.r#type = Type::Type_Vector;
125                        for i in 0..3 {
126                            result.data.value_vector[i] =
127                                (la.data.value_number as f32) / ra.data.value_vector[i];
128                        }
129                        result.data.value_vector[3] = result_w;
130                    }
131                }
132            } else if la.r#type == Type::Type_Vector && ra.r#type == Type::Type_Number {
133                unsafe {
134                    let had_w = la.data.value_vector[3] != 0.0;
135                    let result_w = la.data.value_vector[3] / (ra.data.value_number as f32);
136                    if result_w == 0.0 || had_w {
137                        result.r#type = Type::Type_Vector;
138                        for i in 0..3 {
139                            result.data.value_vector[i] =
140                                la.data.value_vector[i] / (ra.data.value_number as f32);
141                        }
142                        result.data.value_vector[3] = result_w;
143                    }
144                }
145            }
146        }
147        AstExprBinaryOp::FloorDiv => {
148            if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Number {
149                result.r#type = Type::Type_Number;
150                unsafe {
151                    result.data.value_number =
152                        (la.data.value_number / ra.data.value_number).floor();
153                }
154            } else if la.r#type == Type::Type_Vector && ra.r#type == Type::Type_Vector {
155                unsafe {
156                    let had_w = la.data.value_vector[3] != 0.0 || ra.data.value_vector[3] != 0.0;
157                    let result_w = (la.data.value_vector[3] / ra.data.value_vector[3]).floor();
158                    if result_w == 0.0 || had_w {
159                        result.r#type = Type::Type_Vector;
160                        for i in 0..3 {
161                            result.data.value_vector[i] =
162                                (la.data.value_vector[i] / ra.data.value_vector[i]).floor();
163                        }
164                        result.data.value_vector[3] = result_w;
165                    }
166                }
167            } else if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Vector {
168                unsafe {
169                    let had_w = ra.data.value_vector[3] != 0.0;
170                    let result_w =
171                        ((la.data.value_number as f32) / ra.data.value_vector[3]).floor();
172                    if result_w == 0.0 || had_w {
173                        result.r#type = Type::Type_Vector;
174                        for i in 0..3 {
175                            result.data.value_vector[i] =
176                                ((la.data.value_number as f32) / ra.data.value_vector[i]).floor();
177                        }
178                        result.data.value_vector[3] = result_w;
179                    }
180                }
181            } else if la.r#type == Type::Type_Vector && ra.r#type == Type::Type_Number {
182                unsafe {
183                    let had_w = la.data.value_vector[3] != 0.0;
184                    let result_w =
185                        (la.data.value_vector[3] / (ra.data.value_number as f32)).floor();
186                    if result_w == 0.0 || had_w {
187                        result.r#type = Type::Type_Vector;
188                        for i in 0..3 {
189                            result.data.value_vector[i] =
190                                (la.data.value_vector[i] / (ra.data.value_number as f32)).floor();
191                        }
192                        result.data.value_vector[3] = result_w;
193                    }
194                }
195            }
196        }
197        AstExprBinaryOp::Mod => {
198            if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Number {
199                result.r#type = Type::Type_Number;
200                unsafe {
201                    result.data.value_number = la.data.value_number
202                        - (la.data.value_number / ra.data.value_number).floor()
203                            * ra.data.value_number;
204                }
205            }
206        }
207        AstExprBinaryOp::Pow => {
208            if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Number {
209                result.r#type = Type::Type_Number;
210                unsafe {
211                    result.data.value_number = la.data.value_number.powf(ra.data.value_number);
212                }
213            }
214        }
215        AstExprBinaryOp::Concat => {
216            if la.r#type == Type::Type_String
217                && ra.r#type == Type::Type_String
218                && (la.string_length + ra.string_length) <= K_CONSTANT_FOLD_STRING_LIMIT
219            {
220                result.r#type = Type::Type_String;
221                result.string_length = la.string_length + ra.string_length;
222                if la.string_length == 0 {
223                    unsafe {
224                        result.data.value_string = ra.data.value_string;
225                    }
226                } else if ra.string_length == 0 {
227                    unsafe {
228                        result.data.value_string = la.data.value_string;
229                    }
230                } else {
231                    let mut tmp = Vec::with_capacity(result.string_length as usize);
232                    let la_slice = la.get_string();
233                    let ra_slice = ra.get_string();
234                    tmp.extend_from_slice(unsafe {
235                        core::slice::from_raw_parts(
236                            la_slice.as_slice().as_ptr() as *const u8,
237                            la_slice.len(),
238                        )
239                    });
240                    tmp.extend_from_slice(unsafe {
241                        core::slice::from_raw_parts(
242                            ra_slice.as_slice().as_ptr() as *const u8,
243                            ra_slice.len(),
244                        )
245                    });
246                    let name = string_table.get_or_add(tmp.as_ptr() as *const c_char, tmp.len());
247                    unsafe {
248                        result.data.value_string = name.value;
249                    }
250                }
251            }
252        }
253        AstExprBinaryOp::CompareNe => {
254            if la.r#type != Type::Type_Unknown && ra.r#type != Type::Type_Unknown {
255                result.r#type = Type::Type_Boolean;
256                unsafe {
257                    result.data.value_boolean = !constants_equal(la, ra);
258                }
259            }
260        }
261        AstExprBinaryOp::CompareEq => {
262            if la.r#type != Type::Type_Unknown && ra.r#type != Type::Type_Unknown {
263                result.r#type = Type::Type_Boolean;
264                unsafe {
265                    result.data.value_boolean = constants_equal(la, ra);
266                }
267            }
268        }
269        AstExprBinaryOp::CompareLt => {
270            if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Number {
271                result.r#type = Type::Type_Boolean;
272                unsafe {
273                    result.data.value_boolean = la.data.value_number < ra.data.value_number;
274                }
275            }
276        }
277        AstExprBinaryOp::CompareLe => {
278            if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Number {
279                result.r#type = Type::Type_Boolean;
280                unsafe {
281                    result.data.value_boolean = la.data.value_number <= ra.data.value_number;
282                }
283            }
284        }
285        AstExprBinaryOp::CompareGt => {
286            if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Number {
287                result.r#type = Type::Type_Boolean;
288                unsafe {
289                    result.data.value_boolean = la.data.value_number > ra.data.value_number;
290                }
291            }
292        }
293        AstExprBinaryOp::CompareGe => {
294            if la.r#type == Type::Type_Number && ra.r#type == Type::Type_Number {
295                result.r#type = Type::Type_Boolean;
296                unsafe {
297                    result.data.value_boolean = la.data.value_number >= ra.data.value_number;
298                }
299            }
300        }
301        AstExprBinaryOp::And => {
302            if la.r#type != Type::Type_Unknown {
303                *result = if la.is_truthful() { *ra } else { *la };
304            }
305        }
306        AstExprBinaryOp::Or => {
307            if la.r#type != Type::Type_Unknown {
308                *result = if la.is_truthful() { *la } else { *ra };
309            }
310        }
311        _ => {
312            LUAU_ASSERT!(false);
313        }
314    }
315}