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            match &expected.kind {
133                TypeKind::Map(key, value) => {
134                    expected_key_ty = Some(key.as_ref());
135                    expected_value_ty = Some(value.as_ref());
136                    allow_mixed_keys = matches!(key.kind, TypeKind::Unknown | TypeKind::Infer);
137                    allow_mixed_values = matches!(value.kind, TypeKind::Unknown | TypeKind::Infer);
138                }
139
140                TypeKind::Table => {
141                    allow_mixed_keys = true;
142                    allow_mixed_values = true;
143                }
144
145                _ => {}
146            }
147        }
148
149        if entries.is_empty() {
150            if let Some(expected) = expected_type {
151                match &expected.kind {
152                    TypeKind::Map(_, _) => {
153                        return Ok(self.canonicalize_type(expected));
154                    }
155
156                    TypeKind::Table => {
157                        let span = Self::dummy_span();
158                        return Ok(Type::new(
159                            TypeKind::Map(
160                                Box::new(Type::new(TypeKind::Unknown, span)),
161                                Box::new(Type::new(TypeKind::Unknown, span)),
162                            ),
163                            span,
164                        ));
165                    }
166
167                    _ => {}
168                }
169            }
170
171            let span = Self::dummy_span();
172            return Ok(Type::new(
173                TypeKind::Map(
174                    Box::new(Type::new(TypeKind::Unknown, span)),
175                    Box::new(Type::new(TypeKind::Unknown, span)),
176                ),
177                span,
178            ));
179        }
180
181        let key_hint = expected_key_ty.and_then(|ty| {
182            if matches!(ty.kind, TypeKind::Unknown | TypeKind::Infer) {
183                None
184            } else {
185                Some(ty)
186            }
187        });
188        let value_hint = expected_value_ty.and_then(|ty| {
189            if matches!(ty.kind, TypeKind::Unknown | TypeKind::Infer) {
190                None
191            } else {
192                Some(ty)
193            }
194        });
195
196        let mut inferred_key_type: Option<Type> = None;
197        let mut inferred_value_type: Option<Type> = None;
198        for (key_expr, value_expr) in entries {
199            let raw_key_type = if let Some(hint) = key_hint {
200                self.check_expr_with_hint(key_expr, Some(hint))?
201            } else {
202                self.check_expr(key_expr)?
203            };
204            if !self.env.type_implements_trait(&raw_key_type, "Hashable") {
205                return Err(self.type_error(format!(
206                    "Map key type '{}' must implement Hashable trait",
207                    raw_key_type
208                )));
209            }
210            let canonical_key = self.canonicalize_type(&raw_key_type);
211
212            let raw_value_type = if let Some(hint) = value_hint {
213                self.check_expr_with_hint(value_expr, Some(hint))?
214            } else {
215                self.check_expr(value_expr)?
216            };
217            let canonical_value = self.canonicalize_type(&raw_value_type);
218
219            if let Some(existing_key) = &inferred_key_type {
220                if !allow_mixed_keys {
221                    self.unify(existing_key, &canonical_key)?;
222                }
223            } else {
224                inferred_key_type = Some(canonical_key.clone());
225            }
226
227            if let Some(existing_value) = &inferred_value_type {
228                if !allow_mixed_values {
229                    self.unify(existing_value, &canonical_value)?;
230                }
231            } else {
232                inferred_value_type = Some(canonical_value.clone());
233            }
234        }
235
236        let span = Self::dummy_span();
237        let key_type = if allow_mixed_keys {
238            expected_key_ty
239                .and_then(|ty| Some(self.canonicalize_type(ty)))
240                .unwrap_or_else(|| Type::new(TypeKind::Unknown, span))
241        } else {
242            inferred_key_type.unwrap_or_else(|| Type::new(TypeKind::Unknown, span))
243        };
244        let value_type = if allow_mixed_values {
245            expected_value_ty
246                .and_then(|ty| Some(self.canonicalize_type(ty)))
247                .unwrap_or_else(|| Type::new(TypeKind::Unknown, span))
248        } else {
249            inferred_value_type.unwrap_or_else(|| Type::new(TypeKind::Unknown, span))
250        };
251
252        Ok(Type::new(
253            TypeKind::Map(Box::new(key_type), Box::new(value_type)),
254            Self::dummy_span(),
255        ))
256    }
257
258    pub fn check_struct_literal(
259        &mut self,
260        span: Span,
261        name: &str,
262        fields: &[StructLiteralField],
263    ) -> Result<Type> {
264        let key = self.resolve_type_key(name);
265        let struct_def = self
266            .env
267            .lookup_struct(&key)
268            .or_else(|| self.env.lookup_struct(name))
269            .ok_or_else(|| self.type_error_at(format!("Undefined struct '{}'", name), span))?
270            .clone();
271        if fields.len() != struct_def.fields.len() {
272            return Err(self.type_error_at(
273                format!(
274                    "Struct '{}' has {} fields, but {} were provided",
275                    name,
276                    struct_def.fields.len(),
277                    fields.len()
278                ),
279                span,
280            ));
281        }
282
283        for field in fields {
284            let expected_type = struct_def
285                .fields
286                .iter()
287                .find(|f| f.name == field.name)
288                .map(|f| &f.ty)
289                .ok_or_else(|| {
290                    self.type_error_at(
291                        format!("Struct '{}' has no field '{}'", name, field.name),
292                        field.span,
293                    )
294                })?;
295            let actual_type = self.check_expr(&field.value)?;
296            match &expected_type.kind {
297                TypeKind::Option(inner_expected) => {
298                    if self.unify(inner_expected, &actual_type).is_err() {
299                        self.unify(expected_type, &actual_type)?;
300                    }
301                }
302
303                _ => {
304                    self.unify(expected_type, &actual_type)?;
305                }
306            }
307        }
308
309        let ty_name = if self.env.lookup_struct(&key).is_some() {
310            key
311        } else {
312            name.to_string()
313        };
314        Ok(Type::new(TypeKind::Named(ty_name), Self::dummy_span()))
315    }
316
317    pub fn check_lambda(
318        &mut self,
319        params: &[(String, Option<Type>)],
320        return_type: Option<&Type>,
321        body: &Expr,
322    ) -> Result<Type> {
323        self.env.push_scope();
324        let expected_signature = self.expected_lambda_signature.take();
325        let mut param_types = Vec::new();
326        for (i, (param_name, param_type)) in params.iter().enumerate() {
327            let ty = if let Some(explicit_type) = param_type {
328                explicit_type.clone()
329            } else if let Some((ref expected_params, _)) = expected_signature {
330                if i < expected_params.len() {
331                    expected_params[i].clone()
332                } else {
333                    Type::new(TypeKind::Infer, Self::dummy_span())
334                }
335            } else {
336                Type::new(TypeKind::Infer, Self::dummy_span())
337            };
338            self.env.declare_variable(param_name.clone(), ty.clone())?;
339            param_types.push(ty);
340        }
341
342        let saved_return_type = self.current_function_return_type.clone();
343        let inferred_return_type = if let Some(explicit) = return_type {
344            Some(explicit.clone())
345        } else if let Some((_, expected_ret)) = expected_signature {
346            expected_ret.or_else(|| Some(Type::new(TypeKind::Infer, Self::dummy_span())))
347        } else {
348            Some(Type::new(TypeKind::Infer, Self::dummy_span()))
349        };
350        self.current_function_return_type = inferred_return_type.clone();
351        let body_type = self.check_expr(body)?;
352        self.current_function_return_type = saved_return_type;
353        let actual_return_type = if let Some(expected) = return_type {
354            expected.clone()
355        } else if let Some(inferred) = &inferred_return_type {
356            if !matches!(inferred.kind, TypeKind::Infer) {
357                inferred.clone()
358            } else {
359                body_type
360            }
361        } else {
362            body_type
363        };
364        self.env.pop_scope();
365        Ok(Type::new(
366            TypeKind::Function {
367                params: param_types,
368                return_type: Box::new(actual_return_type),
369            },
370            Self::dummy_span(),
371        ))
372    }
373
374    pub fn check_if_expr(
375        &mut self,
376        condition: &Expr,
377        then_branch: &Expr,
378        else_branch: &Option<Box<Expr>>,
379    ) -> Result<Type> {
380        let cond_type = self.check_expr(condition)?;
381        self.unify(&Type::new(TypeKind::Bool, Self::dummy_span()), &cond_type)?;
382        let then_type = self.check_expr(then_branch)?;
383        if let Some(else_expr) = else_branch {
384            let else_type = self.check_expr(else_expr)?;
385            self.unify(&then_type, &else_type)?;
386            Ok(then_type)
387        } else {
388            Ok(Type::new(TypeKind::Unit, Self::dummy_span()))
389        }
390    }
391}
lust/typechecker/expr_checker/collections.rs

lust/typechecker/expr_checker/
collections.rs
