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(&mut self, entries: &[(Expr, Expr)]) -> Result<Type> {
123        if entries.is_empty() {
124            let span = Self::dummy_span();
125            return Ok(Type::new(
126                TypeKind::Map(
127                    Box::new(Type::new(TypeKind::Unknown, span)),
128                    Box::new(Type::new(TypeKind::Unknown, span)),
129                ),
130                span,
131            ));
132        }
133
134        let (first_key, first_value) = &entries[0];
135        let key_type = self.check_expr(first_key)?;
136        let value_type = self.check_expr(first_value)?;
137        if !self.env.type_implements_trait(&key_type, "Hashable") {
138            return Err(self.type_error(format!(
139                "Map key type '{}' must implement Hashable trait",
140                key_type
141            )));
142        }
143
144        for (key, value) in &entries[1..] {
145            let k_type = self.check_expr(key)?;
146            let v_type = self.check_expr(value)?;
147            self.unify(&key_type, &k_type)?;
148            self.unify(&value_type, &v_type)?;
149        }
150
151        Ok(Type::new(
152            TypeKind::Map(Box::new(key_type), Box::new(value_type)),
153            Self::dummy_span(),
154        ))
155    }
156
157    pub fn check_struct_literal(
158        &mut self,
159        span: Span,
160        name: &str,
161        fields: &[StructLiteralField],
162    ) -> Result<Type> {
163        let key = self.resolve_type_key(name);
164        let struct_def = self
165            .env
166            .lookup_struct(&key)
167            .or_else(|| self.env.lookup_struct(name))
168            .ok_or_else(|| self.type_error_at(format!("Undefined struct '{}'", name), span))?
169            .clone();
170        if fields.len() != struct_def.fields.len() {
171            return Err(self.type_error_at(
172                format!(
173                    "Struct '{}' has {} fields, but {} were provided",
174                    name,
175                    struct_def.fields.len(),
176                    fields.len()
177                ),
178                span,
179            ));
180        }
181
182        for field in fields {
183            let expected_type = struct_def
184                .fields
185                .iter()
186                .find(|f| f.name == field.name)
187                .map(|f| &f.ty)
188                .ok_or_else(|| {
189                    self.type_error_at(
190                        format!("Struct '{}' has no field '{}'", name, field.name),
191                        field.span,
192                    )
193                })?;
194            let actual_type = self.check_expr(&field.value)?;
195            match &expected_type.kind {
196                TypeKind::Option(inner_expected) => {
197                    if self.unify(inner_expected, &actual_type).is_err() {
198                        self.unify(expected_type, &actual_type)?;
199                    }
200                }
201
202                _ => {
203                    self.unify(expected_type, &actual_type)?;
204                }
205            }
206        }
207
208        let ty_name = if self.env.lookup_struct(&key).is_some() {
209            key
210        } else {
211            name.to_string()
212        };
213        Ok(Type::new(TypeKind::Named(ty_name), Self::dummy_span()))
214    }
215
216    pub fn check_lambda(
217        &mut self,
218        params: &[(String, Option<Type>)],
219        return_type: Option<&Type>,
220        body: &Expr,
221    ) -> Result<Type> {
222        self.env.push_scope();
223        let expected_signature = self.expected_lambda_signature.take();
224        let mut param_types = Vec::new();
225        for (i, (param_name, param_type)) in params.iter().enumerate() {
226            let ty = if let Some(explicit_type) = param_type {
227                explicit_type.clone()
228            } else if let Some((ref expected_params, _)) = expected_signature {
229                if i < expected_params.len() {
230                    expected_params[i].clone()
231                } else {
232                    Type::new(TypeKind::Infer, Self::dummy_span())
233                }
234            } else {
235                Type::new(TypeKind::Infer, Self::dummy_span())
236            };
237            self.env.declare_variable(param_name.clone(), ty.clone())?;
238            param_types.push(ty);
239        }
240
241        let saved_return_type = self.current_function_return_type.clone();
242        let inferred_return_type = if let Some(explicit) = return_type {
243            Some(explicit.clone())
244        } else if let Some((_, expected_ret)) = expected_signature {
245            expected_ret.or_else(|| Some(Type::new(TypeKind::Infer, Self::dummy_span())))
246        } else {
247            Some(Type::new(TypeKind::Infer, Self::dummy_span()))
248        };
249        self.current_function_return_type = inferred_return_type.clone();
250        let body_type = self.check_expr(body)?;
251        self.current_function_return_type = saved_return_type;
252        let actual_return_type = if let Some(expected) = return_type {
253            expected.clone()
254        } else if let Some(inferred) = &inferred_return_type {
255            if !matches!(inferred.kind, TypeKind::Infer) {
256                inferred.clone()
257            } else {
258                body_type
259            }
260        } else {
261            body_type
262        };
263        self.env.pop_scope();
264        Ok(Type::new(
265            TypeKind::Function {
266                params: param_types,
267                return_type: Box::new(actual_return_type),
268            },
269            Self::dummy_span(),
270        ))
271    }
272
273    pub fn check_if_expr(
274        &mut self,
275        condition: &Expr,
276        then_branch: &Expr,
277        else_branch: &Option<Box<Expr>>,
278    ) -> Result<Type> {
279        let cond_type = self.check_expr(condition)?;
280        self.unify(&Type::new(TypeKind::Bool, Self::dummy_span()), &cond_type)?;
281        let then_type = self.check_expr(then_branch)?;
282        if let Some(else_expr) = else_branch {
283            let else_type = self.check_expr(else_expr)?;
284            self.unify(&then_type, &else_type)?;
285            Ok(then_type)
286        } else {
287            Ok(Type::new(TypeKind::Unit, Self::dummy_span()))
288        }
289    }
290}
lust/typechecker/expr_checker/collections.rs

lust/typechecker/expr_checker/
collections.rs
