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            } => {
44                if method == "as" && type_args.is_some() {
45                    let _receiver_type = self.check_expr(receiver)?;
46                    if !args.is_empty() {
47                        return Err(self.type_error(":as<T>() takes no arguments".to_string()));
48                    }
49
50                    let target_type = &type_args.as_ref().unwrap()[0];
51                    let actual_target_type = if let TypeKind::Named(name) = &target_type.kind {
52                        let resolved = self.resolve_type_key(name);
53                        if self.env.lookup_trait(&resolved).is_some() {
54                            Type::new(TypeKind::Trait(name.clone()), target_type.span)
55                        } else {
56                            target_type.clone()
57                        }
58                    } else {
59                        target_type.clone()
60                    };
61                    let span = Self::dummy_span();
62                    return Ok(Type::new(
63                        TypeKind::Option(Box::new(actual_target_type)),
64                        span,
65                    ));
66                }
67
68                let _ = type_args;
69                self.check_method_call(receiver, method, args)
70            }
71
72            ExprKind::FieldAccess { object, field } => {
73                self.check_field_access_with_hint(expr.span, object, field, expected_type)
74            }
75
76            ExprKind::Index { object, index } => self.check_index_expr(object, index),
77            ExprKind::Array(elements) => self.check_array_literal(elements, expected_type),
78            ExprKind::Map(entries) => self.check_map_literal(entries, expected_type),
79            ExprKind::StructLiteral { name, fields } => {
80                self.check_struct_literal(expr.span, name, fields)
81            }
82
83            ExprKind::Lambda {
84                params,
85                return_type,
86                body,
87            } => self.check_lambda(params, return_type.as_ref(), body),
88            ExprKind::Cast { expr, target_type } => {
89                let _expr_type = self.check_expr(expr)?;
90                Ok(target_type.clone())
91            }
92
93            ExprKind::TypeCheck {
94                expr,
95                check_type: _,
96            } => {
97                let _expr_type = self.check_expr(expr)?;
98                Ok(Type::new(TypeKind::Bool, Self::dummy_span()))
99            }
100
101            ExprKind::IsPattern { expr, pattern } => {
102                let scrutinee_type = self.check_expr(expr)?;
103                self.validate_is_pattern(pattern, &scrutinee_type)?;
104                Ok(Type::new(TypeKind::Bool, Self::dummy_span()))
105            }
106
107            ExprKind::If {
108                condition,
109                then_branch,
110                else_branch,
111            } => self.check_if_expr(condition, then_branch, else_branch),
112            ExprKind::Block(stmts) => {
113                self.env.push_scope();
114                let mut result_type = Type::new(TypeKind::Unit, Self::dummy_span());
115                for stmt in stmts {
116                    match &stmt.kind {
117                        StmtKind::Expr(expr) => {
118                            result_type = self.check_expr(expr)?;
119                        }
120
121                        StmtKind::Return(values) => {
122                            result_type = if values.is_empty() {
123                                Type::new(TypeKind::Unit, Self::dummy_span())
124                            } else if values.len() == 1 {
125                                let expected = self.current_function_return_type.clone();
126                                let mut raw =
127                                    self.check_expr_with_hint(&values[0], expected.as_ref())?;
128                                if raw.span.start_line == 0 && values[0].span.start_line > 0 {
129                                    raw.span = values[0].span;
130                                }
131
132                                self.canonicalize_type(&raw)
133                            } else {
134                                let mut el_types = Vec::new();
135                                for value in values {
136                                    let raw = self.check_expr(value)?;
137                                    el_types.push(self.canonicalize_type(&raw));
138                                }
139
140                                Type::new(TypeKind::Tuple(el_types), Self::dummy_span())
141                            };
142                            self.pending_generic_instances.take();
143                            self.check_stmt(stmt)?;
144                        }
145
146                        _ => {
147                            self.check_stmt(stmt)?;
148                        }
149                    }
150                }
151
152                self.env.pop_scope();
153                if result_type.span.start_line == 0 {
154                    result_type.span = expr.span;
155                }
156
157                Ok(result_type)
158            }
159
160            ExprKind::Range { .. } => Err(self.type_error_at(
161                "Range expressions are not supported; use numeric for-loops".to_string(),
162                expr.span,
163            )),
164            ExprKind::EnumConstructor {
165                enum_name,
166                variant,
167                args,
168            } => {
169                let enum_def = self
170                    .env
171                    .lookup_enum(enum_name)
172                    .ok_or_else(|| self.type_error(format!("Undefined enum '{}'", enum_name)))?
173                    .clone();
174                let variant_def = enum_def
175                    .variants
176                    .iter()
177                    .find(|v| v.name == *variant)
178                    .ok_or_else(|| {
179                        self.type_error(format!(
180                            "Enum '{}' has no variant '{}'",
181                            enum_name, variant
182                        ))
183                    })?;
184                if let Some(expected_fields) = &variant_def.fields {
185                    if args.len() != expected_fields.len() {
186                        return Err(self.type_error(format!(
187                            "Variant '{}::{}' expects {} arguments, got {}",
188                            enum_name,
189                            variant,
190                            expected_fields.len(),
191                            args.len()
192                        )));
193                    }
194
195                    let mut type_params = HashMap::new();
196                    for (arg, expected_type) in args.iter().zip(expected_fields.iter()) {
197                        let arg_type = self.check_expr(arg)?;
198                        if let TypeKind::Generic(type_param) = &expected_type.kind {
199                            type_params.insert(type_param.clone(), arg_type.clone());
200                        } else {
201                            self.unify(expected_type, &arg_type)?;
202                        }
203                    }
204
205                    if !type_params.is_empty() {
206                        self.pending_generic_instances = Some(type_params.clone());
207                    }
208
209                    if enum_name == "Option" {
210                        if let Some(inner_type) = type_params.get("T") {
211                            return Ok(Type::new(
212                                TypeKind::Option(Box::new(inner_type.clone())),
213                                Self::dummy_span(),
214                            ));
215                        }
216                    } else if enum_name == "Result" {
217                        if let (Some(ok_type), Some(err_type)) =
218                            (type_params.get("T"), type_params.get("E"))
219                        {
220                            return Ok(Type::new(
221                                TypeKind::Result(
222                                    Box::new(ok_type.clone()),
223                                    Box::new(err_type.clone()),
224                                ),
225                                Self::dummy_span(),
226                            ));
227                        }
228                    }
229                } else {
230                    if !args.is_empty() {
231                        return Err(self.type_error(format!(
232                            "Variant '{}::{}' is a unit variant and takes no arguments",
233                            enum_name, variant
234                        )));
235                    }
236                }
237
238                Ok(Type::new(
239                    TypeKind::Named(enum_name.clone()),
240                    Self::dummy_span(),
241                ))
242            }
243
244            ExprKind::Tuple(elements) => {
245                let expected_elements = expected_type.and_then(|ty| {
246                    if let TypeKind::Tuple(elems) = &ty.kind {
247                        Some(elems.clone())
248                    } else {
249                        None
250                    }
251                });
252                let mut element_types = Vec::new();
253                for (index, element) in elements.iter().enumerate() {
254                    let hint = expected_elements
255                        .as_ref()
256                        .and_then(|elems| elems.get(index));
257                    let mut raw_ty = if let Some(hint_ty) = hint {
258                        self.check_expr_with_hint(element, Some(hint_ty))?
259                    } else {
260                        self.check_expr(element)?
261                    };
262                    if raw_ty.span.start_line == 0 && element.span.start_line > 0 {
263                        raw_ty.span = element.span;
264                    }
265
266                    self.pending_generic_instances.take();
267                    element_types.push(self.canonicalize_type(&raw_ty));
268                }
269
270                Ok(Type::new(TypeKind::Tuple(element_types), expr.span))
271            }
272
273            ExprKind::Return(exprs) => {
274                let mut return_type = if exprs.is_empty() {
275                    Type::new(TypeKind::Unit, Self::dummy_span())
276                } else if exprs.len() == 1 {
277                    let expected = self.current_function_return_type.clone();
278                    let mut raw_ty = self.check_expr_with_hint(&exprs[0], expected.as_ref())?;
279                    if raw_ty.span.start_line == 0 && exprs[0].span.start_line > 0 {
280                        raw_ty.span = exprs[0].span;
281                    }
282
283                    self.pending_generic_instances.take();
284                    raw_ty
285                } else {
286                    let mut types = Vec::new();
287                    for value in exprs {
288                        let raw_ty = self.check_expr(value)?;
289                        let ty = self.canonicalize_type(&raw_ty);
290                        self.pending_generic_instances.take();
291                        types.push(ty);
292                    }
293
294                    Type::new(TypeKind::Tuple(types), Self::dummy_span())
295                };
296                if return_type.span.start_line == 0 {
297                    if let Some(first) = exprs.first() {
298                        return_type.span = first.span;
299                    } else {
300                        return_type.span = expr.span;
301                    }
302                }
303
304                if let Some(expected_return) = &self.current_function_return_type {
305                    self.unify(expected_return, &return_type)?;
306                } else {
307                    return Err(self.type_error("'return' outside of function".to_string()));
308                }
309
310                Ok(return_type)
311            }
312
313            ExprKind::Paren(inner) => self.check_expr(inner),
314        }
315    }
316}