lust/bytecode/compiler/
mod.rs

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