lust/typechecker/expr_checker/
collections.rs

1use super::*;
2use alloc::{
3    boxed::Box,
4    format,
5    string::{String, ToString},
6    vec::Vec,
7};
8impl TypeChecker {
9    pub fn check_array_literal(
10        &mut self,
11        elements: &[Expr],
12        expected_type: Option<&Type>,
13    ) -> Result<Type> {
14        if elements.is_empty() {
15            if let Some(expected) = expected_type {
16                return Ok(expected.clone());
17            }
18
19            let span = Self::dummy_span();
20            return Ok(Type::new(
21                TypeKind::Array(Box::new(Type::new(TypeKind::Unknown, span))),
22                span,
23            ));
24        }
25
26        let expected_elem_type = expected_type.and_then(|t| {
27            if let TypeKind::Array(elem_type) = &t.kind {
28                Some(elem_type.as_ref())
29            } else {
30                None
31            }
32        });
33        if let Some(expected_elem) = expected_elem_type {
34            if let TypeKind::Union(union_types) = &expected_elem.kind {
35                for elem in elements {
36                    let elem_type = self.check_expr(elem)?;
37                    let mut matches = false;
38                    for union_variant in union_types {
39                        if self.types_equal(&elem_type, union_variant) {
40                            matches = true;
41                            break;
42                        }
43                    }
44
45                    if !matches {
46                        let union_desc = union_types
47                            .iter()
48                            .map(|t| t.to_string())
49                            .collect::<Vec<_>>()
50                            .join(" | ");
51                        return Err(self.type_error(format!(
52                            "Array element type '{}' does not match any type in union [{}]",
53                            elem_type, union_desc
54                        )));
55                    }
56                }
57
58                return Ok(expected_type.unwrap().clone());
59            }
60        }
61
62        if let Some(expected_elem) = expected_elem_type {
63            if matches!(expected_elem.kind, TypeKind::Unknown) {
64                for elem in elements {
65                    self.check_expr(elem)?;
66                }
67
68                return Ok(expected_type.unwrap().clone());
69            }
70
71            if let TypeKind::Option(inner) = &expected_elem.kind {
72                if matches!(inner.kind, TypeKind::Unknown) {
73                    for elem in elements {
74                        let elem_type = self.check_expr(elem)?;
75                        let is_option = matches!(&elem_type.kind, TypeKind::Option(_))
76                            || matches!(&elem_type.kind, TypeKind::Named(name) if name == "Option");
77                        if !is_option {
78                            return Err(self.type_error(format!(
79                                "Expected Option type for Array<Option<unknown>>, got '{}'",
80                                elem_type
81                            )));
82                        }
83                    }
84
85                    return Ok(expected_type.unwrap().clone());
86                }
87            }
88
89            if let TypeKind::Result(ok_inner, err_inner) = &expected_elem.kind {
90                if matches!(ok_inner.kind, TypeKind::Unknown)
91                    || matches!(err_inner.kind, TypeKind::Unknown)
92                {
93                    for elem in elements {
94                        let elem_type = self.check_expr(elem)?;
95                        let is_result = matches!(&elem_type.kind, TypeKind::Result(_, _))
96                            || matches!(&elem_type.kind, TypeKind::Named(name) if name == "Result");
97                        if !is_result {
98                            return Err(self.type_error(format!(
99                                "Expected Result type for Array<Result<unknown, ...>>, got '{}'",
100                                elem_type
101                            )));
102                        }
103                    }
104
105                    return Ok(expected_type.unwrap().clone());
106                }
107            }
108        }
109
110        let first_type = self.check_expr(&elements[0])?;
111        for elem in &elements[1..] {
112            let elem_type = self.check_expr(elem)?;
113            self.unify(&first_type, &elem_type)?;
114        }
115
116        Ok(Type::new(
117            TypeKind::Array(Box::new(first_type)),
118            Self::dummy_span(),
119        ))
120    }
121
122    pub fn check_map_literal(
123        &mut self,
124        entries: &[(Expr, Expr)],
125        expected_type: Option<&Type>,
126    ) -> Result<Type> {
127        let mut expected_key_ty: Option<&Type> = None;
128        let mut expected_value_ty: Option<&Type> = None;
129        let mut allow_mixed_keys = false;
130        let mut allow_mixed_values = false;
131        if let Some(expected) = expected_type {
132            if let TypeKind::Map(key, value) = &expected.kind {
133                expected_key_ty = Some(key.as_ref());
134                expected_value_ty = Some(value.as_ref());
135                allow_mixed_keys = matches!(key.kind, TypeKind::Unknown | TypeKind::Infer);
136                allow_mixed_values = matches!(value.kind, TypeKind::Unknown | TypeKind::Infer);
137            }
138        }
139
140        if entries.is_empty() {
141            if let Some(expected) = expected_type {
142                if let TypeKind::Map(_, _) = &expected.kind {
143                    return Ok(self.canonicalize_type(expected));
144                }
145            }
146
147            let span = Self::dummy_span();
148            return Ok(Type::new(
149                TypeKind::Map(
150                    Box::new(Type::new(TypeKind::Unknown, span)),
151                    Box::new(Type::new(TypeKind::Unknown, span)),
152                ),
153                span,
154            ));
155        }
156
157        let key_hint = expected_key_ty.and_then(|ty| {
158            if matches!(ty.kind, TypeKind::Unknown | TypeKind::Infer) {
159                None
160            } else {
161                Some(ty)
162            }
163        });
164        let value_hint = expected_value_ty.and_then(|ty| {
165            if matches!(ty.kind, TypeKind::Unknown | TypeKind::Infer) {
166                None
167            } else {
168                Some(ty)
169            }
170        });
171
172        let mut inferred_key_type: Option<Type> = None;
173        let mut inferred_value_type: Option<Type> = None;
174        for (key_expr, value_expr) in entries {
175            let raw_key_type = if let Some(hint) = key_hint {
176                self.check_expr_with_hint(key_expr, Some(hint))?
177            } else {
178                self.check_expr(key_expr)?
179            };
180            let canonical_key = self.canonicalize_type(&raw_key_type);
181
182            let raw_value_type = if let Some(hint) = value_hint {
183                self.check_expr_with_hint(value_expr, Some(hint))?
184            } else {
185                self.check_expr(value_expr)?
186            };
187            let canonical_value = self.canonicalize_type(&raw_value_type);
188
189            if let Some(existing_key) = &mut inferred_key_type {
190                if !allow_mixed_keys && self.unify(existing_key, &canonical_key).is_err() {
191                    *existing_key = Type::new(TypeKind::Unknown, Self::dummy_span());
192                }
193            } else {
194                inferred_key_type = Some(canonical_key.clone());
195            }
196
197            if let Some(existing_value) = &mut inferred_value_type {
198                if !allow_mixed_values && self.unify(existing_value, &canonical_value).is_err() {
199                    *existing_value = Type::new(TypeKind::Unknown, Self::dummy_span());
200                }
201            } else {
202                inferred_value_type = Some(canonical_value.clone());
203            }
204        }
205
206        let span = Self::dummy_span();
207        let key_type = if allow_mixed_keys {
208            expected_key_ty
209                .and_then(|ty| Some(self.canonicalize_type(ty)))
210                .unwrap_or_else(|| Type::new(TypeKind::Unknown, span))
211        } else {
212            inferred_key_type.unwrap_or_else(|| Type::new(TypeKind::Unknown, span))
213        };
214        let value_type = if allow_mixed_values {
215            expected_value_ty
216                .and_then(|ty| Some(self.canonicalize_type(ty)))
217                .unwrap_or_else(|| Type::new(TypeKind::Unknown, span))
218        } else {
219            inferred_value_type.unwrap_or_else(|| Type::new(TypeKind::Unknown, span))
220        };
221
222        Ok(Type::new(
223            TypeKind::Map(Box::new(key_type), Box::new(value_type)),
224            Self::dummy_span(),
225        ))
226    }
227
228    pub fn check_struct_literal(
229        &mut self,
230        span: Span,
231        name: &str,
232        fields: &[StructLiteralField],
233    ) -> Result<Type> {
234        let key = self.resolve_type_key(name);
235        let struct_def = self
236            .env
237            .lookup_struct(&key)
238            .or_else(|| self.env.lookup_struct(name))
239            .ok_or_else(|| self.type_error_at(format!("Undefined struct '{}'", name), span))?
240            .clone();
241        if fields.len() != struct_def.fields.len() {
242            return Err(self.type_error_at(
243                format!(
244                    "Struct '{}' has {} fields, but {} were provided",
245                    name,
246                    struct_def.fields.len(),
247                    fields.len()
248                ),
249                span,
250            ));
251        }
252
253        for field in fields {
254            let expected_type = struct_def
255                .fields
256                .iter()
257                .find(|f| f.name == field.name)
258                .map(|f| &f.ty)
259                .ok_or_else(|| {
260                    self.type_error_at(
261                        format!("Struct '{}' has no field '{}'", name, field.name),
262                        field.span,
263                    )
264                })?;
265            let actual_type = self.check_expr(&field.value)?;
266            match &expected_type.kind {
267                TypeKind::Option(inner_expected) => {
268                    if self.unify(inner_expected, &actual_type).is_err() {
269                        self.unify(expected_type, &actual_type)?;
270                    }
271                }
272
273                _ => {
274                    self.unify(expected_type, &actual_type)?;
275                }
276            }
277        }
278
279        let ty_name = if self.env.lookup_struct(&key).is_some() {
280            key
281        } else {
282            name.to_string()
283        };
284        Ok(Type::new(TypeKind::Named(ty_name), Self::dummy_span()))
285    }
286
287    pub fn check_lambda(
288        &mut self,
289        params: &[(String, Option<Type>)],
290        return_type: Option<&Type>,
291        body: &Expr,
292    ) -> Result<Type> {
293        self.env.push_scope();
294        let expected_signature = self.expected_lambda_signature.take();
295        let mut param_types = Vec::new();
296        for (i, (param_name, param_type)) in params.iter().enumerate() {
297            let ty = if let Some(explicit_type) = param_type {
298                explicit_type.clone()
299            } else if let Some((ref expected_params, _)) = expected_signature {
300                if i < expected_params.len() {
301                    expected_params[i].clone()
302                } else {
303                    Type::new(TypeKind::Infer, Self::dummy_span())
304                }
305            } else {
306                Type::new(TypeKind::Infer, Self::dummy_span())
307            };
308            self.env.declare_variable(param_name.clone(), ty.clone())?;
309            param_types.push(ty);
310        }
311
312        let saved_return_type = self.current_function_return_type.clone();
313        let inferred_return_type = if let Some(explicit) = return_type {
314            Some(explicit.clone())
315        } else if let Some((_, expected_ret)) = expected_signature {
316            expected_ret.or_else(|| Some(Type::new(TypeKind::Infer, Self::dummy_span())))
317        } else {
318            Some(Type::new(TypeKind::Infer, Self::dummy_span()))
319        };
320        self.current_function_return_type = inferred_return_type.clone();
321        let body_type = self.check_expr(body)?;
322        self.current_function_return_type = saved_return_type;
323        let actual_return_type = if let Some(expected) = return_type {
324            expected.clone()
325        } else if let Some(inferred) = &inferred_return_type {
326            if !matches!(inferred.kind, TypeKind::Infer) {
327                inferred.clone()
328            } else {
329                body_type
330            }
331        } else {
332            body_type
333        };
334        self.env.pop_scope();
335        Ok(Type::new(
336            TypeKind::Function {
337                params: param_types,
338                return_type: Box::new(actual_return_type),
339            },
340            Self::dummy_span(),
341        ))
342    }
343
344    pub fn check_if_expr(
345        &mut self,
346        condition: &Expr,
347        then_branch: &Expr,
348        else_branch: &Option<Box<Expr>>,
349    ) -> Result<Type> {
350        let cond_type = self.check_expr(condition)?;
351        self.unify(&Type::new(TypeKind::Bool, Self::dummy_span()), &cond_type)?;
352        let then_type = self.check_expr(then_branch)?;
353        if let Some(else_expr) = else_branch {
354            let else_type = self.check_expr(else_expr)?;
355            self.unify(&then_type, &else_type)?;
356            Ok(then_type)
357        } else {
358            Ok(Type::new(TypeKind::Unit, Self::dummy_span()))
359        }
360    }
361}
lust/typechecker/expr_checker/collections.rs

lust/typechecker/expr_checker/
collections.rs
