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