lust/typechecker/expr_checker/
entry.rs

1use super::*;
2use alloc::{boxed::Box, format, string::ToString, vec::Vec};
3use hashbrown::HashMap;
4impl TypeChecker {
5    pub fn check_expr(&mut self, expr: &Expr) -> Result<Type> {
6        let mut ty = self.check_expr_with_hint(expr, None)?;
7        if ty.span.start_line == 0 && expr.span.start_line > 0 {
8            ty.span = expr.span;
9        }
10
11        if expr.span.start_line > 0 {
12            if let Some(module) = &self.current_module {
13                self.expr_types_by_module
14                    .entry(module.clone())
15                    .or_default()
16                    .insert(expr.span, ty.clone());
17            }
18        }
19
20        Ok(ty)
21    }
22
23    pub fn check_expr_with_hint(
24        &mut self,
25        expr: &Expr,
26        expected_type: Option<&Type>,
27    ) -> Result<Type> {
28        match &expr.kind {
29            ExprKind::Literal(lit) => self.check_literal(lit),
30            ExprKind::Identifier(name) => self.env.lookup_variable(name).ok_or_else(|| {
31                self.type_error_at(format!("Undefined variable '{}'", name), expr.span)
32            }),
33            ExprKind::Binary { left, op, right } => {
34                self.check_binary_expr(expr.span, left, op, right)
35            }
36            ExprKind::Unary { op, operand } => self.check_unary_expr(op, operand),
37            ExprKind::Call { callee, args } => self.check_call_expr(expr.span, callee, args),
38            ExprKind::MethodCall {
39                receiver,
40                method,
41                type_args: _,
42                args,
43            } => self.check_method_call(receiver, method, args),
44
45            ExprKind::FieldAccess { object, field } => {
46                self.check_field_access_with_hint(expr.span, object, field, expected_type)
47            }
48
49            ExprKind::Index { object, index } => self.check_index_expr(object, index),
50            ExprKind::Array(elements) => self.check_array_literal(elements, expected_type),
51            ExprKind::Map(entries) => self.check_map_literal(entries, expected_type),
52            ExprKind::StructLiteral { name, fields } => {
53                self.check_struct_literal(expr.span, name, fields)
54            }
55
56            ExprKind::Lambda {
57                params,
58                return_type,
59                body,
60            } => self.check_lambda(params, return_type.as_ref(), body),
61            ExprKind::Cast { expr, target_type } => {
62                let _expr_type = self.check_expr(expr)?;
63                Ok(target_type.clone())
64            }
65
66            ExprKind::TypeCheck {
67                expr,
68                check_type: _,
69            } => {
70                let _expr_type = self.check_expr(expr)?;
71                Ok(Type::new(TypeKind::Bool, Self::dummy_span()))
72            }
73
74            ExprKind::IsPattern { expr, pattern } => {
75                let scrutinee_type = self.check_expr(expr)?;
76                self.validate_is_pattern(pattern, &scrutinee_type)?;
77                Ok(Type::new(TypeKind::Bool, Self::dummy_span()))
78            }
79
80            ExprKind::If {
81                condition,
82                then_branch,
83                else_branch,
84            } => self.check_if_expr(condition, then_branch, else_branch),
85            ExprKind::Block(stmts) => {
86                self.env.push_scope();
87                let mut result_type = Type::new(TypeKind::Unit, Self::dummy_span());
88                for stmt in stmts {
89                    match &stmt.kind {
90                        StmtKind::Expr(expr) => {
91                            result_type = self.check_expr(expr)?;
92                        }
93
94                        StmtKind::Return(values) => {
95                            result_type = if values.is_empty() {
96                                Type::new(TypeKind::Unit, Self::dummy_span())
97                            } else if values.len() == 1 {
98                                let expected = self.current_function_return_type.clone();
99                                let mut raw =
100                                    self.check_expr_with_hint(&values[0], expected.as_ref())?;
101                                if raw.span.start_line == 0 && values[0].span.start_line > 0 {
102                                    raw.span = values[0].span;
103                                }
104
105                                self.canonicalize_type(&raw)
106                            } else {
107                                let mut el_types = Vec::new();
108                                for value in values {
109                                    let raw = self.check_expr(value)?;
110                                    el_types.push(self.canonicalize_type(&raw));
111                                }
112
113                                Type::new(TypeKind::Tuple(el_types), Self::dummy_span())
114                            };
115                            self.pending_generic_instances.take();
116                            self.check_stmt(stmt)?;
117                        }
118
119                        _ => {
120                            self.check_stmt(stmt)?;
121                        }
122                    }
123                }
124
125                self.env.pop_scope();
126                if result_type.span.start_line == 0 {
127                    result_type.span = expr.span;
128                }
129
130                Ok(result_type)
131            }
132
133            ExprKind::Range { .. } => Err(self.type_error_at(
134                "Range expressions are not supported; use numeric for-loops".to_string(),
135                expr.span,
136            )),
137            ExprKind::EnumConstructor {
138                enum_name,
139                variant,
140                args,
141            } => {
142                let enum_def = self
143                    .env
144                    .lookup_enum(enum_name)
145                    .ok_or_else(|| self.type_error(format!("Undefined enum '{}'", enum_name)))?
146                    .clone();
147                let variant_def = enum_def
148                    .variants
149                    .iter()
150                    .find(|v| v.name == *variant)
151                    .ok_or_else(|| {
152                        self.type_error(format!(
153                            "Enum '{}' has no variant '{}'",
154                            enum_name, variant
155                        ))
156                    })?;
157                if let Some(expected_fields) = &variant_def.fields {
158                    if args.len() != expected_fields.len() {
159                        return Err(self.type_error(format!(
160                            "Variant '{}::{}' expects {} arguments, got {}",
161                            enum_name,
162                            variant,
163                            expected_fields.len(),
164                            args.len()
165                        )));
166                    }
167
168                    let mut type_params = HashMap::new();
169                    for (arg, expected_type) in args.iter().zip(expected_fields.iter()) {
170                        let arg_type = self.check_expr(arg)?;
171                        if let TypeKind::Generic(type_param) = &expected_type.kind {
172                            type_params.insert(type_param.clone(), arg_type.clone());
173                        } else {
174                            self.unify(expected_type, &arg_type)?;
175                        }
176                    }
177
178                    if !type_params.is_empty() {
179                        self.pending_generic_instances = Some(type_params.clone());
180                    }
181
182                    if enum_name == "Option" {
183                        if let Some(inner_type) = type_params.get("T") {
184                            return Ok(Type::new(
185                                TypeKind::Option(Box::new(inner_type.clone())),
186                                Self::dummy_span(),
187                            ));
188                        }
189                    } else if enum_name == "Result" {
190                        if let (Some(ok_type), Some(err_type)) =
191                            (type_params.get("T"), type_params.get("E"))
192                        {
193                            return Ok(Type::new(
194                                TypeKind::Result(
195                                    Box::new(ok_type.clone()),
196                                    Box::new(err_type.clone()),
197                                ),
198                                Self::dummy_span(),
199                            ));
200                        }
201                    }
202                } else {
203                    if !args.is_empty() {
204                        return Err(self.type_error(format!(
205                            "Variant '{}::{}' is a unit variant and takes no arguments",
206                            enum_name, variant
207                        )));
208                    }
209                }
210
211                Ok(Type::new(
212                    TypeKind::Named(enum_name.clone()),
213                    Self::dummy_span(),
214                ))
215            }
216
217            ExprKind::Tuple(elements) => {
218                let expected_elements = expected_type.and_then(|ty| {
219                    if let TypeKind::Tuple(elems) = &ty.kind {
220                        Some(elems.clone())
221                    } else {
222                        None
223                    }
224                });
225                let mut element_types = Vec::new();
226                for (index, element) in elements.iter().enumerate() {
227                    let hint = expected_elements
228                        .as_ref()
229                        .and_then(|elems| elems.get(index));
230                    let mut raw_ty = if let Some(hint_ty) = hint {
231                        self.check_expr_with_hint(element, Some(hint_ty))?
232                    } else {
233                        self.check_expr(element)?
234                    };
235                    if raw_ty.span.start_line == 0 && element.span.start_line > 0 {
236                        raw_ty.span = element.span;
237                    }
238
239                    self.pending_generic_instances.take();
240                    element_types.push(self.canonicalize_type(&raw_ty));
241                }
242
243                Ok(Type::new(TypeKind::Tuple(element_types), expr.span))
244            }
245
246            ExprKind::Return(exprs) => {
247                let mut return_type = if exprs.is_empty() {
248                    Type::new(TypeKind::Unit, Self::dummy_span())
249                } else if exprs.len() == 1 {
250                    let expected = self.current_function_return_type.clone();
251                    let mut raw_ty = self.check_expr_with_hint(&exprs[0], expected.as_ref())?;
252                    if raw_ty.span.start_line == 0 && exprs[0].span.start_line > 0 {
253                        raw_ty.span = exprs[0].span;
254                    }
255
256                    self.pending_generic_instances.take();
257                    raw_ty
258                } else {
259                    let mut types = Vec::new();
260                    for value in exprs {
261                        let raw_ty = self.check_expr(value)?;
262                        let ty = self.canonicalize_type(&raw_ty);
263                        self.pending_generic_instances.take();
264                        types.push(ty);
265                    }
266
267                    Type::new(TypeKind::Tuple(types), Self::dummy_span())
268                };
269                if return_type.span.start_line == 0 {
270                    if let Some(first) = exprs.first() {
271                        return_type.span = first.span;
272                    } else {
273                        return_type.span = expr.span;
274                    }
275                }
276
277                if let Some(expected_return) = &self.current_function_return_type {
278                    self.unify(expected_return, &return_type)?;
279                } else {
280                    return Err(self.type_error("'return' outside of function".to_string()));
281                }
282
283                Ok(return_type)
284            }
285
286            ExprKind::Paren(inner) => self.check_expr(inner),
287        }
288    }
289}