1pub(super) use super::{Chunk, Function, Instruction, Register, Value};
2pub(super) use crate::ast::{
3 BinaryOp, ExprKind, ExternItem, Item, ItemKind, Literal, Span, Stmt, StmtKind, Type, TypeKind,
4 UnaryOp,
5};
6use crate::config::LustConfig;
7pub(super) use crate::number::LustInt;
8use crate::typechecker::FunctionSignature;
9pub(super) use crate::{Expr, LustError, Result};
10pub(super) use alloc::{
11 format,
12 string::{String, ToString},
13 vec,
14 vec::Vec,
15};
16pub(super) use hashbrown::{HashMap, HashSet};
17mod closures;
18mod expressions;
19mod methods;
20mod module;
21mod patterns;
22mod registers;
23mod statements;
24pub struct Compiler {
25 pub(super) functions: Vec<Function>,
26 pub(super) function_table: HashMap<String, usize>,
27 pub(super) trait_impls: Vec<(String, String)>,
28 pub(super) trait_names: HashSet<String>,
29 pub(super) current_function: usize,
30 pub(super) scopes: Vec<Scope>,
31 pub(super) loop_contexts: Vec<LoopContext>,
32 pub(super) next_register: Register,
33 pub(super) max_register: Register,
34 pub(super) current_line: usize,
35 pub(super) imports_by_module: HashMap<String, crate::modules::ModuleImports>,
36 pub(super) current_module: Option<String>,
37 pub(super) entry_module: Option<String>,
38 pub(super) module_locals: HashMap<String, HashSet<String>>,
39 pub(super) current_function_name: Option<String>,
40 pub(super) extern_value_aliases: HashMap<String, String>,
41 pub(super) stdlib_symbols: HashSet<String>,
42 option_coercions: HashMap<String, HashSet<Span>>,
43 function_signatures: HashMap<String, FunctionSignature>,
44}
45
46#[derive(Debug, Clone)]
47pub(super) struct Scope {
48 pub(super) locals: HashMap<String, (Register, bool)>,
49 pub(super) depth: usize,
50}
51
52#[derive(Debug, Clone)]
53pub(super) struct LoopContext {
54 pub(super) continue_target: Option<usize>,
55 pub(super) continue_jumps: Vec<usize>,
56 pub(super) break_jumps: Vec<usize>,
57}
58
59impl Compiler {
60 pub fn new() -> Self {
61 let mut compiler = Self {
62 functions: Vec::new(),
63 function_table: HashMap::new(),
64 trait_impls: Vec::new(),
65 trait_names: HashSet::new(),
66 current_function: 0,
67 scopes: Vec::new(),
68 loop_contexts: Vec::new(),
69 next_register: 0,
70 max_register: 0,
71 current_line: 0,
72 imports_by_module: HashMap::new(),
73 current_module: None,
74 entry_module: None,
75 module_locals: HashMap::new(),
76 current_function_name: None,
77 extern_value_aliases: HashMap::new(),
78 stdlib_symbols: HashSet::new(),
79 option_coercions: HashMap::new(),
80 function_signatures: HashMap::new(),
81 };
82 compiler.configure_stdlib(&LustConfig::default());
83 compiler
84 }
85
86 pub fn set_imports_by_module(&mut self, map: HashMap<String, crate::modules::ModuleImports>) {
87 self.imports_by_module = map;
88 }
89
90 pub fn set_entry_module(&mut self, module: impl Into<String>) {
91 self.entry_module = Some(module.into());
92 }
93
94 pub fn get_trait_impls(&self) -> &[(String, String)] {
95 &self.trait_impls
96 }
97
98 pub fn configure_stdlib(&mut self, config: &LustConfig) {
99 self.stdlib_symbols.clear();
100 self.stdlib_symbols.extend(
101 [
102 "print",
103 "println",
104 "type",
105 "tostring",
106 "unpack",
107 "select",
108 "task",
109 "lua",
110 "error",
111 "assert",
112 "tonumber",
113 "pairs",
114 "ipairs",
115 "setmetatable",
116 ]
117 .into_iter()
118 .map(String::from),
119 );
120 for module in config.enabled_modules() {
121 match module {
122 "io" | "os" | "string" | "math" | "table" => {
123 self.stdlib_symbols.insert(module.to_string());
124 }
125
126 _ => {}
127 }
128 }
129 }
130
131 pub fn set_option_coercions(&mut self, map: HashMap<String, HashSet<Span>>) {
132 self.option_coercions = map;
133 }
134
135 pub fn set_function_signatures(&mut self, signatures: HashMap<String, FunctionSignature>) {
136 self.function_signatures = signatures;
137 }
138
139 pub fn take_function_signatures(&mut self) -> HashMap<String, FunctionSignature> {
140 core::mem::take(&mut self.function_signatures)
141 }
142
143 pub(super) fn is_stdlib_symbol(&self, name: &str) -> bool {
144 self.stdlib_symbols.contains(name)
145 }
146
147 pub(super) fn should_wrap_option(&self, span: Span) -> bool {
148 let module = self.current_module.as_deref().unwrap_or("");
149 self.option_coercions
150 .get(module)
151 .map_or(false, |set| set.contains(&span))
152 }
153
154 fn assign_signature_by_name(&mut self, func_idx: usize, name: &str) {
155 if let Some(signature) = self.function_signatures.get(name).cloned() {
156 self.functions[func_idx].set_signature(signature);
157 }
158 }
159
160 fn try_set_lambda_signature(
161 &mut self,
162 func_idx: usize,
163 params: &[(String, Option<Type>)],
164 return_type: &Option<Type>,
165 ) {
166 if let Some(signature) = Self::lambda_signature(params, return_type) {
167 self.functions[func_idx].set_signature(signature);
168 }
169 }
170
171 fn lambda_signature(
172 params: &[(String, Option<Type>)],
173 return_type: &Option<Type>,
174 ) -> Option<FunctionSignature> {
175 let mut param_types = Vec::with_capacity(params.len());
176 for (_, ty) in params {
177 if let Some(ty) = ty {
178 param_types.push(ty.clone());
179 } else {
180 return None;
181 }
182 }
183
184 let return_type = return_type
185 .clone()
186 .unwrap_or_else(|| Type::new(TypeKind::Unit, Span::dummy()));
187 Some(FunctionSignature {
188 params: param_types,
189 return_type,
190 is_method: false,
191 })
192 }
193
194 pub fn record_extern_value(&mut self, name: &str) {
195 let runtime_name = name.to_string();
196 self.extern_value_aliases
197 .entry(runtime_name.clone())
198 .or_insert(runtime_name.clone());
199 let module_name = self
200 .current_module
201 .clone()
202 .or_else(|| self.entry_module.clone());
203 if let Some(module) = module_name {
204 if !name.contains('.') {
205 let qualified = format!("{}.{}", module, name);
206 self.extern_value_aliases
207 .entry(qualified)
208 .or_insert(runtime_name);
209 }
210 }
211 }
212
213 pub(super) fn describe_expr_kind(kind: &ExprKind) -> &'static str {
214 match kind {
215 ExprKind::Literal(_) => "literal expression",
216 ExprKind::Identifier(_) => "identifier expression",
217 ExprKind::Binary { .. } => "binary expression",
218 ExprKind::Unary { .. } => "unary expression",
219 ExprKind::Call { .. } => "function call",
220 ExprKind::MethodCall { .. } => "method call",
221 ExprKind::FieldAccess { .. } => "field access",
222 ExprKind::Index { .. } => "index access",
223 ExprKind::Array(_) => "array literal",
224 ExprKind::Map(_) => "map literal",
225 ExprKind::Tuple(_) => "tuple literal",
226 ExprKind::StructLiteral { .. } => "struct literal",
227 ExprKind::EnumConstructor { .. } => "enum constructor",
228 ExprKind::Lambda { .. } => "lambda expression",
229 ExprKind::Paren(_) => "parenthesized expression",
230 ExprKind::Cast { .. } => "cast expression",
231 ExprKind::TypeCheck { .. } => "`is` type check",
232 ExprKind::IsPattern { .. } => "`is` pattern expression",
233 ExprKind::If { .. } => "`if` expression",
234 ExprKind::Block(_) => "block expression",
235 ExprKind::Return(_) => "return expression",
236 ExprKind::Range { .. } => "range expression",
237 }
238 }
239
240 pub(super) fn type_to_string(type_kind: &crate::ast::TypeKind) -> String {
241 use crate::ast::TypeKind;
242 match type_kind {
243 TypeKind::Int => "int".to_string(),
244 TypeKind::Float => "float".to_string(),
245 TypeKind::String => "string".to_string(),
246 TypeKind::Bool => "bool".to_string(),
247 TypeKind::Named(name) => name.clone(),
248 TypeKind::Array(inner) => format!("Array<{}>", Self::type_to_string(&inner.kind)),
249 TypeKind::Map(key, val) => format!(
250 "Map<{}, {}>",
251 Self::type_to_string(&key.kind),
252 Self::type_to_string(&val.kind)
253 ),
254 TypeKind::Option(inner) => format!("Option<{}>", Self::type_to_string(&inner.kind)),
255 TypeKind::Result(ok, err) => format!(
256 "Result<{}, {}>",
257 Self::type_to_string(&ok.kind),
258 Self::type_to_string(&err.kind)
259 ),
260 TypeKind::Function {
261 params,
262 return_type,
263 } => {
264 let param_strs: Vec<String> = params
265 .iter()
266 .map(|p| Self::type_to_string(&p.kind))
267 .collect();
268 format!(
269 "function({}) -> {}",
270 param_strs.join(", "),
271 Self::type_to_string(&return_type.kind)
272 )
273 }
274
275 TypeKind::Tuple(elements) => {
276 let element_strs: Vec<String> = elements
277 .iter()
278 .map(|t| Self::type_to_string(&t.kind))
279 .collect();
280 format!("Tuple<{}>", element_strs.join(", "))
281 }
282
283 TypeKind::Generic(name) => name.clone(),
284 TypeKind::GenericInstance { name, type_args } => {
285 let arg_strs: Vec<String> = type_args
286 .iter()
287 .map(|t| Self::type_to_string(&t.kind))
288 .collect();
289 format!("{}<{}>", name, arg_strs.join(", "))
290 }
291
292 TypeKind::Unknown => "unknown".to_string(),
293 TypeKind::Union(types) => {
294 let type_strs: Vec<String> = types
295 .iter()
296 .map(|t| Self::type_to_string(&t.kind))
297 .collect();
298 format!("{}", type_strs.join(" | "))
299 }
300
301 TypeKind::Unit => "()".to_string(),
302 TypeKind::Infer => "_".to_string(),
303 TypeKind::Ref(inner) => format!("&{}", Self::type_to_string(&inner.kind)),
304 TypeKind::MutRef(inner) => format!("&mut {}", Self::type_to_string(&inner.kind)),
305 TypeKind::Pointer { mutable, pointee } => {
306 if *mutable {
307 format!("*mut {}", Self::type_to_string(&pointee.kind))
308 } else {
309 format!("*{}", Self::type_to_string(&pointee.kind))
310 }
311 }
312
313 TypeKind::Trait(name) => name.clone(),
314 TypeKind::TraitBound(traits) => traits.join(" + "),
315 }
316 }
317
318 fn module_context_name(&self) -> Option<&str> {
319 self.current_module
320 .as_deref()
321 .or_else(|| self.entry_module.as_deref())
322 }
323
324 fn is_builtin_type_name(name: &str) -> bool {
325 matches!(
326 name,
327 "int"
328 | "float"
329 | "string"
330 | "bool"
331 | "unknown"
332 | "Array"
333 | "Map"
334 | "Option"
335 | "Result"
336 | "Iterator"
337 | "Task"
338 | "TaskStatus"
339 | "TaskInfo"
340 | "LuaValue"
341 | "LuaTable"
342 | "LuaFunction"
343 | "LuaUserdata"
344 | "LuaThread"
345 )
346 }
347
348 pub(super) fn resolve_type_name(&self, name: &str) -> String {
349 if let Some((head, tail)) = name.split_once('.') {
350 if let Some(module) = self.module_context_name() {
351 if let Some(imports) = self.imports_by_module.get(module) {
352 if let Some(real_module) = imports.module_aliases.get(head) {
353 if tail.is_empty() {
354 return real_module.clone();
355 } else {
356 return format!("{}.{}", real_module, tail);
357 }
358 }
359 }
360 }
361
362 return name.to_string();
363 }
364
365 if Self::is_builtin_type_name(name) {
366 return name.to_string();
367 }
368
369 if let Some(module) = self.module_context_name() {
370 if let Some(imports) = self.imports_by_module.get(module) {
371 if let Some(fq) = imports.type_aliases.get(name) {
372 return fq.clone();
373 }
374 }
375
376 return format!("{}.{}", module, name);
377 }
378
379 name.to_string()
380 }
381}
382
383impl Default for Compiler {
384 fn default() -> Self {
385 Self::new()
386 }
387}