1use super::{Compiler, FnInferRet, ListElemState, Symbol};
2use anyhow::Result;
3use dynamic::{Dynamic, Type};
4use parser::{BinaryOp, Expr, ExprKind, Pattern, PatternKind, Span, Stmt, StmtKind, UnaryOp};
5use smol_str::SmolStr;
6
7#[derive(Clone)]
8struct ReturnInfo {
9 ty: Type,
10 shape: Option<Type>,
11}
12
13const MAX_INFER_DEPTH: usize = 64;
18
19impl Compiler {
20 fn current_infer_key(&self) -> Option<(u32, Vec<Type>, Vec<Type>)> {
21 self.infer_stack.last().cloned()
22 }
23
24 fn pending_return_seed(&self, id: u32, generic_args: &[Type], fn_tys: &[Type]) -> Option<Type> {
25 self.fns.get(&id).and_then(|fns| {
26 fns.iter().find_map(|item| {
27 if item.0 == generic_args
28 && item.1 == fn_tys
29 && let FnInferRet::Pending(seed) = &item.2
30 {
31 seed.clone()
32 } else {
33 None
34 }
35 })
36 })
37 }
38
39 fn update_pending_return_seed(&mut self, ty: &Type) {
40 if ty.is_any() {
41 return;
42 }
43 let Some((id, generic_args, fn_tys)) = self.current_infer_key() else {
44 return;
45 };
46 let Some(fns) = self.fns.get_mut(&id) else {
47 return;
48 };
49 if let Some(item) = fns.iter_mut().find(|item| item.0 == generic_args && item.1 == fn_tys)
50 && let FnInferRet::Pending(seed) = &mut item.2
51 {
52 let next = seed.take().map(|prev| prev + ty.clone()).unwrap_or_else(|| ty.clone());
53 *seed = Some(next);
54 }
55 }
56
57 fn try_find_base_return_ty(&self, body: &Stmt) -> Option<Type> {
59 match &body.kind {
60 StmtKind::Block(stmts) => stmts.iter().find_map(|s| self.try_find_base_return_ty(s)),
61 StmtKind::If { then_body, else_body, .. } => self.try_find_base_return_ty(then_body)
62 .or_else(|| else_body.as_ref().and_then(|b| self.try_find_base_return_ty(b))),
63 StmtKind::Return(Some(expr)) => Self::try_literal_type(expr),
64 StmtKind::Expr(expr, false) => Self::try_literal_type(expr),
65 _ => None,
66 }
67 }
68
69 fn try_find_base_return_ty_with_scope(&mut self, body: &Stmt, fn_id: u32, fn_name: &str, args: &[SmolStr], fn_tys: &[Type]) -> Option<Type> {
71 let saved_state = self.take_local_state();
72 self.frames.push(0);
73 for (arg, ty) in args.iter().zip(fn_tys.iter()) {
74 self.add_name(arg.clone());
75 self.add_ty(ty.clone());
76 }
77 let result = self.try_find_base_return_ty_with_scope_inner(body, fn_id, fn_name);
78 self.restore_local_state(saved_state);
79 result
80 }
81
82 fn try_find_base_return_ty_with_scope_inner(&mut self, body: &Stmt, fn_id: u32, fn_name: &str) -> Option<Type> {
83 match &body.kind {
84 StmtKind::Block(stmts) => stmts.iter().find_map(|s| self.try_find_base_return_ty_with_scope_inner(s, fn_id, fn_name)),
85 StmtKind::If { then_body, else_body, .. } => self.try_find_base_return_ty_with_scope_inner(then_body, fn_id, fn_name)
86 .or_else(|| else_body.as_ref().and_then(|b| self.try_find_base_return_ty_with_scope_inner(b, fn_id, fn_name))),
87 StmtKind::Return(Some(expr)) => {
88 if Self::expr_calls_fn(expr, fn_id, fn_name) { None }
89 else { self.infer_return_expr(expr).ok().map(|info| info.ty) }
90 }
91 StmtKind::Expr(expr, false) => {
92 if Self::expr_calls_fn(expr, fn_id, fn_name) { None }
93 else { self.infer_return_expr(expr).ok().map(|info| info.ty) }
94 }
95 _ => None,
96 }
97 }
98
99 fn expr_calls_fn(expr: &Expr, fn_id: u32, fn_name: &str) -> bool {
100 match &expr.kind {
101 ExprKind::Call { obj, params } => {
102 if let ExprKind::Id(id, _) = &obj.kind { return *id == fn_id; }
103 if let ExprKind::Ident(name) = &obj.kind {
104 if name.as_str() == fn_name || fn_name.ends_with(&format!("::{}", name)) { return true; }
105 }
106 params.iter().any(|p| Self::expr_calls_fn(p, fn_id, fn_name))
107 }
108 ExprKind::Binary { left, op: _, right } => Self::expr_calls_fn(left, fn_id, fn_name) || Self::expr_calls_fn(right, fn_id, fn_name),
109 ExprKind::Unary { op: _, value } => Self::expr_calls_fn(value, fn_id, fn_name),
110 ExprKind::Typed { value, ty: _ } => Self::expr_calls_fn(value, fn_id, fn_name),
111 _ => false,
112 }
113 }
114
115 fn try_literal_type(expr: &Expr) -> Option<Type> {
116 match &expr.kind {
117 ExprKind::Value(v) => Some(v.get_type()),
118 ExprKind::Unary { op: UnaryOp::Neg, value } => Self::try_literal_type(value),
119 _ => None,
120 }
121 }
122
123 fn add_pattern_bindings_for_infer(&mut self, pat: &Pattern, expr_ty: Type) -> Result<()> {
124 match &pat.kind {
125 PatternKind::Ident { name, ty } => {
126 let annotated_ty = self.symbols.get_type(ty)?;
127 self.add_name(name.clone());
128 self.add_ty(if annotated_ty.is_any() { expr_ty } else { annotated_ty });
129 }
130 PatternKind::Var { idx, .. } => self.set_ty(*idx, expr_ty),
131 PatternKind::Tuple(pats) => {
132 if let Type::Tuple(tys) = expr_ty {
133 for (pat, ty) in pats.iter().zip(tys) {
134 self.add_pattern_bindings_for_infer(pat, ty)?;
135 }
136 } else {
137 for pat in pats {
138 self.add_pattern_bindings_for_infer(pat, Type::Any)?;
139 }
140 }
141 }
142 PatternKind::List { elems, .. } => {
143 for pat in elems {
144 self.add_pattern_bindings_for_infer(pat, Type::Any)?;
145 }
146 }
147 PatternKind::Wildcard => {
148 self.add_name("".into());
149 self.add_ty(expr_ty);
150 }
151 PatternKind::Literal(_) | PatternKind::Member(_, _) | PatternKind::Idx(_, _) => {}
152 }
153 Ok(())
154 }
155
156 fn for_pattern_ty(&mut self, range: &Expr) -> Result<Type> {
157 if matches!(range.kind, ExprKind::Range { .. }) {
158 return self.infer_range_expr(range);
159 }
160 Ok(match self.infer_expr(range)? {
161 Type::Array(elem_ty, _) | Type::Vec(elem_ty, _) | Type::List(elem_ty) => elem_ty.as_ref().clone(),
162 _ => Type::Any,
163 })
164 }
165
166 fn infer_range_expr(&mut self, range: &Expr) -> Result<Type> {
167 let ExprKind::Range { start, stop, .. } = &range.kind else {
168 return self.infer_expr(range);
169 };
170 let start_ty = self.infer_expr(start)?;
171 let stop_ty = self.infer_expr(stop)?;
172 Ok(Self::merge_range_bound_types(start_ty, stop_ty))
173 }
174
175 fn merge_range_bound_types(start_ty: Type, stop_ty: Type) -> Type {
176 if start_ty.is_any() {
177 stop_ty
178 } else if stop_ty.is_any() {
179 start_ty
180 } else if matches!(start_ty, Type::I32 | Type::I64) && stop_ty.is_uint() {
183 stop_ty
184 } else if matches!(stop_ty, Type::I32 | Type::I64) && start_ty.is_uint() {
185 start_ty
186 } else {
187 start_ty + stop_ty
188 }
189 }
190
191 fn merge_return_type(span: Span, left: Option<Type>, right: Type) -> Result<Type> {
192 match left {
193 Some(left) if left == right => Ok(left),
194 Some(left) if left.is_void() || right.is_void() => Err(Self::semantic_error(span, format!("返回类型不一致: {:?} 和 {:?}", left, right))),
195 Some(left) if left.is_any() || right.is_any() => Ok(Type::Any),
196 Some(left) => Ok(left + right),
197 None => Ok(right),
198 }
199 }
200
201 fn return_shape(&self, expr: &Expr, ty: &Type) -> Option<Type> {
202 if !ty.is_any() {
203 return match ty {
204 Type::Struct { .. } => Some(ty.clone()),
205 Type::Map => Some(Type::Map),
206 Type::List(elem) | Type::Array(elem, _) => Some(Type::List(elem.clone())),
207 _ => None,
208 };
209 }
210 match &expr.kind {
211 ExprKind::List(_) | ExprKind::Tuple(_) => Some(Type::list_any()),
212 ExprKind::Dict(_) => Some(Type::Map),
213 ExprKind::Value(value) => Self::dynamic_return_shape(value.get_type()),
214 ExprKind::Const(idx) => self.consts.get(*idx).and_then(|value| Self::dynamic_return_shape(value.get_type())),
215 ExprKind::Typed { ty, .. } => Some(ty.clone()),
216 _ => None,
217 }
218 }
219
220 fn dynamic_return_shape(ty: Type) -> Option<Type> {
221 match ty {
222 Type::Map => Some(Type::Map),
223 Type::List(elem) => Some(Type::List(elem)),
224 Type::Array(elem, _) => Some(Type::List(elem)),
225 _ => None,
226 }
227 }
228
229 fn local_var_idx_for_expr(&self, expr: &Expr) -> Option<u32> {
230 match &expr.kind {
231 ExprKind::Var(idx) => Some(*idx),
232 ExprKind::Ident(name) => (self.top()..self.names.len()).rev().find(|idx| self.names[*idx].eq(name)).map(|idx| (idx - self.top()) as u32),
233 _ => None,
234 }
235 }
236
237 fn infer_list_method(&mut self, target: &Expr, elem_ty: &Type, method: &str, params: &[Expr]) -> Result<Option<Type>> {
238 match method {
239 "get_idx" | "pop" => Ok(Some(match self.local_var_idx_for_expr(target).and_then(|idx| self.list_elem_state(idx)) {
240 Some(ListElemState::Known(ty)) => ty,
241 Some(ListElemState::Unknown | ListElemState::Mixed) => Type::Any,
242 None => elem_ty.clone(),
243 })),
244 "push" => {
245 let pushed_ty = params
246 .first()
247 .map(|param| {
248 if let Some(value) = self.get_value(param)
249 && (value.is_str() || value.is_native())
250 {
251 Ok(value.get_type())
252 } else {
253 self.infer_expr(param)
254 }
255 })
256 .transpose()?
257 .unwrap_or(Type::Any);
258 if let Some(idx) = self.local_var_idx_for_expr(target) {
259 let state = self.list_elem_state(idx).unwrap_or_else(|| if elem_ty.is_any() { ListElemState::Unknown } else { ListElemState::Known(elem_ty.clone()) });
260 let next_state = match state {
261 ListElemState::Unknown if pushed_ty.is_any() => ListElemState::Mixed,
262 ListElemState::Unknown => ListElemState::Known(pushed_ty),
263 ListElemState::Known(_) if pushed_ty.is_any() => ListElemState::Mixed,
264 ListElemState::Known(prev) => {
265 let merged = if prev == pushed_ty {
266 prev
267 } else if (prev.is_int() || prev.is_uint() || prev.is_float()) && (pushed_ty.is_int() || pushed_ty.is_uint() || pushed_ty.is_float()) {
268 prev + pushed_ty
269 } else {
270 Type::Any
271 };
272 if merged.is_any() { ListElemState::Mixed } else { ListElemState::Known(merged) }
273 }
274 ListElemState::Mixed => ListElemState::Mixed,
275 };
276 let next_elem = if let ListElemState::Known(ty) = &next_state { ty.clone() } else { Type::Any };
277 self.set_ty(idx, Type::List(std::rc::Rc::new(next_elem)));
278 self.set_list_elem_state(idx, Some(next_state));
279 }
280 Ok(Some(Type::Void))
281 }
282 "len" => Ok(Some(Type::I32)),
283 "is_list" | "is_null" => Ok(Some(Type::Bool)),
284 _ => Ok(None),
285 }
286 }
287
288 fn infer_return_expr(&mut self, expr: &Expr) -> Result<ReturnInfo> {
289 let ty = self.infer_expr(expr)?;
290 let shape = self.return_shape(expr, &ty);
291 let ty = if matches!(shape, Some(Type::Map | Type::List(_))) { Type::Any } else { ty };
292 Ok(ReturnInfo { ty, shape })
293 }
294
295 fn merge_return_info(span: Span, left: Option<ReturnInfo>, right: ReturnInfo) -> Result<ReturnInfo> {
296 let Some(left) = left else {
297 return Ok(right);
298 };
299 if let (Some(left_shape), Some(right_shape)) = (&left.shape, &right.shape)
300 && left_shape != right_shape
301 {
302 return Err(Self::semantic_error(span, format!("返回类型不一致: {:?} 和 {:?}", left_shape, right_shape)));
303 }
304 if let Some(left_shape) = &left.shape
305 && left_shape.is_struct()
306 && right.ty.is_any()
307 && right.shape.is_none()
308 {
309 return Err(Self::semantic_error(span, format!("返回类型不一致: {:?} 和 {:?}", left_shape, Type::Any)));
310 }
311 if let Some(right_shape) = &right.shape
312 && right_shape.is_struct()
313 && left.ty.is_any()
314 && left.shape.is_none()
315 {
316 return Err(Self::semantic_error(span, format!("返回类型不一致: {:?} 和 {:?}", Type::Any, right_shape)));
317 }
318 let ty = Self::merge_return_type(span, Some(left.ty), right.ty)?;
319 Ok(ReturnInfo { ty, shape: left.shape.or(right.shape) })
320 }
321
322 fn infer_return_type(&mut self, stmt: &Stmt) -> Result<Option<Type>> {
323 self.infer_returns(stmt, true).map(|(info, _)| info.map(|info| info.ty))
324 }
325
326 pub(crate) fn check_return_type(&mut self, stmt: &Stmt) -> Result<()> {
327 self.infer_returns(stmt, true).map(|_| ())
328 }
329
330 fn infer_returns(&mut self, stmt: &Stmt, tail: bool) -> Result<(Option<ReturnInfo>, bool)> {
331 match &stmt.kind {
332 StmtKind::Return(Some(expr)) => Ok((Some(self.infer_return_expr(expr)?), true)),
333 StmtKind::Return(None) => Ok((Some(ReturnInfo { ty: Type::Void, shape: Some(Type::Void) }), true)),
334 StmtKind::Block(stmts) => {
335 let mut ret = None;
336 for (idx, stmt) in stmts.iter().enumerate() {
337 let (info, always_returns) = self.infer_returns(stmt, tail && idx == stmts.len().saturating_sub(1))?;
338 if let Some(info) = info {
339 self.update_pending_return_seed(&info.ty);
340 ret = Some(Self::merge_return_info(stmt.span, ret, info)?);
341 if let Some(ret) = &ret {
342 self.update_pending_return_seed(&ret.ty);
343 }
344 }
345 if always_returns {
346 return Ok((ret, true));
347 }
348 }
349 Ok((ret, false))
350 }
351 StmtKind::If { cond, then_body, else_body } => {
352 let cond_ty = self.infer_expr(cond)?;
353 if cond_ty != Type::Bool {
354 return Err(Self::semantic_error(cond.span, format!("条件表达式必须是布尔类型,实际是 {:?}", cond_ty)));
355 }
356 let (mut ret, then_returns) = self.infer_returns(then_body, tail)?;
357 if let Some(ret) = &ret {
358 self.update_pending_return_seed(&ret.ty);
359 }
360 let else_returns = if let Some(body) = else_body {
361 let (else_ty, else_returns) = self.infer_returns(body, tail)?;
362 if let Some(info) = else_ty {
363 self.update_pending_return_seed(&info.ty);
364 ret = Some(Self::merge_return_info(body.span, ret, info)?);
365 if let Some(ret) = &ret {
366 self.update_pending_return_seed(&ret.ty);
367 }
368 }
369 else_returns
370 } else {
371 false
372 };
373 Ok((ret, then_returns && else_returns))
374 }
375 StmtKind::While { cond, body } => {
376 let cond_ty = self.infer_expr(cond)?;
377 if cond_ty != Type::Bool {
378 return Err(Self::semantic_error(cond.span, format!("条件表达式必须是布尔类型,实际是 {:?}", cond_ty)));
379 }
380 self.infer_returns(body, false).map(|(ty, _)| (ty, false))
381 }
382 StmtKind::Loop(body) => self.infer_returns(body, false),
383 StmtKind::For { pat, range, body } => {
384 let ty = self.for_pattern_ty(range)?;
385 self.add_pattern_bindings_for_infer(pat, ty)?;
386 self.infer_returns(body, false).map(|(ty, _)| (ty, false))
387 }
388 StmtKind::Let { .. } => {
389 self.infer_stmt(stmt)?;
390 Ok((None, false))
391 }
392 StmtKind::Expr(expr, close) => {
393 let info = self.infer_return_expr(expr)?;
394 Ok(if *close || !tail { (None, false) } else { (Some(info), true) })
395 }
396 _ => {
397 self.infer_stmt(stmt)?;
398 Ok((None, false))
399 }
400 }
401 }
402
403 pub fn infer_expr(&mut self, expr: &Expr) -> Result<Type> {
404 match &expr.kind {
405 ExprKind::Value(Dynamic::Null) => Ok(Type::Any),
406 ExprKind::Value(v) if v.is_list() => Ok(v.get_type()),
407 ExprKind::Value(v) if v.is_map() => Ok(Type::Any),
408 ExprKind::Value(v) => Ok(v.get_type()),
409 ExprKind::Const(idx) => Ok(match self.consts.get(*idx) {
410 Some(value) if value.is_str() => Type::Str,
411 Some(value) if value.is_list() && value.len() == 0 => Type::list_any(),
412 _ => Type::Any,
413 }),
414 ExprKind::Var(idx) => {
415 let idx = self.top() + (*idx as usize);
416 if idx < self.tys.len() { self.symbols.get_type(&self.tys[idx]) } else { Ok(Type::Any) }
417 }
418 ExprKind::Ident(ident) => {
419 for idx in (self.top()..self.names.len()).rev() {
420 if self.names[idx].eq(ident) && idx < self.tys.len() {
421 return self.symbols.get_type(&self.tys[idx]);
422 }
423 }
424 let id = self.symbols.get_id(ident).map_err(|_| Self::semantic_error(expr.span, format!("未找到标识符 {}", ident)))?;
425 match self.symbols.get_symbol(id)?.1 {
426 Symbol::Const { ty, .. } => Ok(ty.clone()),
427 Symbol::Static { ty, .. } => Ok(ty.clone()),
428 Symbol::Struct(ty, _) => Ok(ty.clone()),
429 Symbol::Fn { .. } => Ok(Type::Symbol { id, params: Vec::new() }),
430 Symbol::Native(ty) => Ok(ty.clone()),
431 s => Err(Self::semantic_error(expr.span, format!("符号 {:?} 不是变量、常量、静态变量、结构体", s))),
432 }
433 }
434 ExprKind::Id(id, _) => match self.symbols.get_symbol(*id)?.1 {
435 Symbol::Const { ty, .. } => Ok(ty.clone()),
436 Symbol::Static { ty, .. } => Ok(ty.clone()),
437 Symbol::Struct(ty, _) => Ok(ty.clone()),
438 Symbol::Fn { .. } => Ok(Type::Symbol { id: *id, params: Vec::new() }),
439 Symbol::Native(ty) => Ok(ty.clone()),
440 s => Err(Self::semantic_error(expr.span, format!("符号 {:?} 不是变量、常量、静态变量、结构体", s))),
441 },
442 ExprKind::Generic { obj, params } => {
443 let params = params.iter().map(|param| self.symbols.get_type(param).unwrap_or_else(|_| param.clone())).collect();
444 match self.infer_expr(obj)? {
445 Type::Symbol { id, .. } => Ok(Type::Symbol { id, params }),
446 _ => Ok(Type::Any),
447 }
448 }
449 ExprKind::AssocId { id, params } => Ok(Type::Symbol { id: *id, params: params.clone() }),
450 ExprKind::Unary { op, value } => match op {
451 UnaryOp::Not => {
452 let ty = self.infer_expr(value.as_ref())?;
453 if ty.is_int() || ty.is_uint() { Ok(ty) } else { Ok(Type::Bool) }
454 }
455 UnaryOp::Neg => self.infer_expr(value.as_ref()),
456 UnaryOp::Unknow => Ok(Type::Any),
457 },
458 ExprKind::Binary { left, op, right } => {
459 if op == &BinaryOp::Assign
460 && let ExprKind::Tuple(left_items) | ExprKind::List(left_items) = &left.kind
461 {
462 if let ExprKind::Tuple(right_items) | ExprKind::List(right_items) = &right.kind {
463 if left_items.len() != right_items.len() {
464 return Err(Self::semantic_error(expr.span, format!("多重赋值数量不匹配: 左侧 {} 个,右侧 {} 个", left_items.len(), right_items.len())));
465 }
466 for item in right_items {
467 let _ = self.infer_expr(item)?;
468 }
469 } else {
470 let _ = self.infer_expr(right)?;
471 }
472 return Ok(Type::Void);
473 }
474 let assign_idx = if op.is_assign() { if let ExprKind::Var(idx) = &left.kind { Some(*idx) } else { None } } else { None };
475 let ty = if op.is_logic() {
476 Type::Bool
477 } else if op == &BinaryOp::Idx {
478 let left_ty = self.infer_expr(left)?;
479 if let Type::Array(elem_ty, _) = left_ty {
480 (*elem_ty).clone()
481 } else if let Type::Vec(elem_ty, _) = left_ty {
482 (*elem_ty).clone()
483 } else if let Type::List(elem_ty) = left_ty {
484 (*elem_ty).clone()
485 } else {
486 let left_ty = self.symbols.get_type(&left_ty)?;
487 let right_ty = if right.is_value() || right.is_const() {
488 let right_value = if let ExprKind::Const(c) = &right.kind { self.consts[*c].clone() } else { right.clone().value()? };
489 if right_value.is_str() {
490 if left_ty.is_any() {
491 return Ok(Type::Any);
492 }
493 if let Ok(field) = self.symbols.get_field(&left_ty, right_value.as_str()) {
494 return if let Type::Fn { ret, .. } = field.1 { Ok(ret.as_ref().clone()) } else { Ok(field.1.clone()) };
495 }
496 } else if let Type::Struct { fields, .. } = &left_ty
497 && let Some(idx) = right_value.as_int()
498 {
499 return fields.get(idx as usize).map(|(_, ty)| ty.clone()).ok_or_else(|| Self::semantic_error(right.span, format!("结构字段索引越界 {}", idx)));
500 }
501 right_value.get_type()
502 } else {
503 self.infer_expr(right)?
504 };
505 if right_ty.is_int() || right_ty.is_uint() {
506 if left_ty.is_any() {
507 return Ok(Type::Any);
508 }
509 let (_, s) = self.symbols.get_field(&left_ty, "get_idx")?;
510 let fn_ty = self.symbols.get_type(&s)?;
511 return if let Type::Fn { ret, .. } = &fn_ty { Ok(ret.as_ref().clone()) } else { Ok(fn_ty) };
512 }
513 if left_ty.is_any() {
514 return Ok(Type::Any);
515 }
516 Type::Any
517 }
518 } else {
519 let left_ty = self.infer_expr(left)?;
520 let right_ty = self.infer_expr(right)?;
521 if op == &BinaryOp::Assign {
522 if !left_ty.is_any() && right_ty.is_any() { left_ty } else { right_ty }
523 } else if op.is_assign() && !left_ty.is_any() && right_ty.is_any() {
524 left_ty
525 } else {
526 left_ty + right_ty
527 }
528 };
529 assign_idx.map(|idx| self.set_ty(idx, ty.clone()));
530 Ok(ty)
531 }
532 ExprKind::Call { obj, params } => {
533 if let ExprKind::Assoc { ty, name } = &obj.kind {
534 let base_name = match ty {
535 Type::Ident { name, .. } => name.clone(),
536 Type::Symbol { id, .. } => self.symbols.get_symbol(*id)?.0.clone(),
537 _ => return Ok(Type::Any),
538 };
539 let id = self.symbols.get_id(&format!("{}::{}", base_name, name))?;
540 let generic_args = match ty {
541 Type::Ident { params, .. } | Type::Symbol { params, .. } => params.iter().map(|param| self.symbols.get_type(param).unwrap_or_else(|_| param.clone())).collect::<Vec<_>>(),
542 _ => Vec::new(),
543 };
544 let mut args = Vec::new();
545 for p in params {
546 args.push(self.infer_expr(p)?);
547 }
548 self.infer_fn_with_params(id, &args, &generic_args)
549 } else if let ExprKind::AssocId { id, params: generic_args } = &obj.kind {
550 let mut args = Vec::new();
551 for p in params {
552 args.push(self.infer_expr(p)?);
553 }
554 self.infer_fn_with_params(*id, &args, generic_args)
555 } else if let ExprKind::Generic { obj, params: generic_args } = &obj.kind {
556 let Type::Symbol { id, .. } = self.infer_expr(obj)? else {
557 return Ok(Type::Any);
558 };
559 let generic_args = generic_args.iter().map(|param| self.symbols.get_type(param).unwrap_or_else(|_| param.clone())).collect::<Vec<_>>();
560 let mut args = Vec::new();
561 for p in params {
562 args.push(self.infer_expr(p)?);
563 }
564 self.infer_fn_with_params(id, &args, &generic_args)
565 } else if let ExprKind::TypedMethod { obj: target, ty, name } = &obj.kind {
566 let base_name = match ty {
567 Type::Ident { name, .. } => name.clone(),
568 Type::Symbol { id, .. } => self.symbols.get_symbol(*id)?.0.clone(),
569 _ => return Ok(Type::Any),
570 };
571 let id = self.symbols.get_id(&format!("{}::{}", base_name, name))?;
572 let mut args = vec![self.infer_expr(target)?];
573 for p in params {
574 args.push(self.infer_expr(p)?);
575 }
576 self.infer_fn(id, &args)
577 } else if let ExprKind::Id(id, obj_expr) = &obj.kind {
578 let method = self.symbols.get_symbol(*id).ok().and_then(|(name, _)| name.rsplit_once("::").map(|(_, method)| method.to_string()));
579 if let Some(target) = obj_expr
580 && let Some(method) = method
581 {
582 let target_ty = self.infer_expr(target)?;
583 if let Type::List(elem_ty) | Type::Array(elem_ty, _) = &target_ty
584 && let Some(ret_ty) = self.infer_list_method(target, elem_ty, method.as_str(), params)?
585 {
586 return Ok(ret_ty);
587 }
588 }
589 let mut args: Vec<Type> = if let Some(obj) = obj_expr { vec![self.infer_expr(obj)?] } else { Vec::new() };
590 for p in params {
591 args.push(self.infer_expr(p)?);
592 }
593 self.infer_fn(*id, &args)
594 } else if let ExprKind::Ident(name) = &obj.kind {
595 for idx in (self.top()..self.names.len()).rev() {
596 if self.names[idx].eq(name) && idx < self.tys.len() {
597 return if let Type::Symbol { id, .. } = &self.tys[idx] {
598 let id = *id;
599 let mut args = Vec::new();
600 for p in params {
601 args.push(self.infer_expr(p)?);
602 }
603 self.infer_fn(id, &args)
604 } else {
605 Ok(Type::Any)
606 };
607 }
608 }
609 let Ok(id) = self.symbols.get_id(name) else {
610 return Ok(Type::Any);
611 };
612 if !self.symbols.get_symbol(id)?.1.is_fn() {
613 return Err(Self::semantic_error(obj.span, format!("符号 {} 不是函数", name)));
614 }
615 let mut args = Vec::new();
616 for p in params {
617 args.push(self.infer_expr(p)?);
618 }
619 self.infer_fn(id, &args)
620 } else if obj.is_idx() {
621 let (target, _, method) = obj.clone().binary().unwrap();
622 let ty = self.infer_expr(&target)?;
623 if let Some(method) = self.get_value(&method) {
624 let method = method.as_str();
625 if let Type::List(elem_ty) | Type::Array(elem_ty, _) = &ty
626 && let Some(ret_ty) = self.infer_list_method(&target, elem_ty, method, params)?
627 {
628 return Ok(ret_ty);
629 }
630 let fn_ty = match self.get_field(&ty, method) {
631 Ok((_, fn_ty)) => fn_ty,
632 Err(_) => {
633 let id = self.symbols.get_id(method)?;
634 if self.symbols.get_symbol(id)?.1.is_fn() {
635 Type::Symbol { id, params: Vec::new() }
636 } else {
637 return Err(Self::semantic_error(obj.span, format!("符号 {method} 不是函数")));
638 }
639 }
640 };
641 if let Type::Symbol { id, .. } = fn_ty {
642 let mut args = vec![ty];
643 for p in params {
644 args.push(self.infer_expr(p)?);
645 }
646 self.infer_fn(id, &args)
647 } else {
648 Ok(fn_ty)
649 }
650 } else {
651 Ok(Type::Any)
652 }
653 } else if let ExprKind::Var(idx) = &obj.kind {
654 let idx = self.top() + (*idx as usize);
655 if idx < self.tys.len()
656 && let Type::Symbol { id, .. } = self.tys[idx]
657 {
658 let mut args = Vec::new();
659 for p in params {
660 args.push(self.infer_expr(p)?);
661 }
662 self.infer_fn(id, &args)
663 } else {
664 Ok(Type::Any)
665 }
666 } else if obj.is_value() {
667 Ok(Type::Void)
668 } else {
669 Ok(Type::Any)
670 }
671 }
672 ExprKind::Typed { ty, .. } => self.symbols.get_type(ty),
673 ExprKind::Stmt(stmt) => self.infer_stmt(stmt),
674 ExprKind::Repeat { value, len } => {
675 let value_ty = self.infer_expr(value)?;
676 let len = self.symbols.get_type(len).unwrap_or_else(|_| len.clone());
677 if let Type::ConstInt(len) = len {
678 let len = u32::try_from(len).map_err(|_| Self::semantic_error(expr.span, "重复数组长度必须是非负 u32"))?;
679 Ok(Type::Array(std::rc::Rc::new(value_ty), len))
680 } else {
681 Ok(Type::ArrayParam(std::rc::Rc::new(value_ty), std::rc::Rc::new(len)))
682 }
683 }
684 ExprKind::List(items) => {
685 if items.is_empty() {
686 return Ok(Type::list_any());
687 }
688 let mut elem_ty = Type::Any;
689 for item in items {
690 let item_ty = self.infer_expr(item)?;
691 elem_ty = if elem_ty.is_any() { item_ty } else { elem_ty + item_ty };
692 }
693 Ok(Type::Array(std::rc::Rc::new(elem_ty), items.len() as u32))
694 }
695 ExprKind::Range { start, stop, .. } => {
696 let start_ty = self.infer_expr(start)?;
697 let stop_ty = self.infer_expr(stop)?;
698 Ok(Self::merge_range_bound_types(start_ty, stop_ty))
699 }
700 _ => Ok(Type::Any),
701 }
702 }
703
704 fn get_fn_tys(&mut self, tys: &[Type], arg_tys: &[Type]) -> Result<Vec<Type>> {
705 let mut fn_tys = Vec::new();
706 for (i, ty) in tys.iter().enumerate() {
707 if !ty.is_any() {
708 fn_tys.push(ty.clone());
709 } else if let Some(arg_ty) = arg_tys.get(i) {
710 fn_tys.push(self.symbols.get_type(arg_ty)?);
711 } else {
712 fn_tys.push(Type::Any);
713 }
714 }
715 Ok(fn_tys)
716 }
717
718 fn is_optimizable_local_ty(ty: &Type) -> bool {
719 ty.is_bool() || ty.is_native()
720 }
721
722 fn is_optimizable_list_elem_ty(ty: &Type) -> bool {
723 matches!(ty, Type::Bool | Type::U8 | Type::I8 | Type::U16 | Type::I16 | Type::U32 | Type::I32 | Type::F32 | Type::U64 | Type::I64 | Type::F64 | Type::Str)
724 }
725
726 fn local_type_hint_at(&self, pos: usize) -> Option<Type> {
727 let ty = self.tys.get(pos)?;
728 match ty {
729 Type::List(_) => self.list_elem_states.get(pos).cloned().flatten().and_then(|state| {
730 if let ListElemState::Known(elem_ty) = state
731 && Self::is_optimizable_list_elem_ty(&elem_ty)
732 {
733 Some(Type::List(std::rc::Rc::new(elem_ty)))
734 } else {
735 None
736 }
737 }),
738 ty if Self::is_optimizable_local_ty(ty) => Some(ty.clone()),
739 _ => None,
740 }
741 }
742
743 fn collect_local_type_hints(&self) -> Vec<Option<Type>> {
744 (self.top()..self.tys.len()).map(|pos| self.local_type_hint_at(pos)).collect()
745 }
746
747 fn set_local_type_hints(&mut self, id: u32, generic_args: &[Type], fn_tys: &[Type], hints: Vec<Option<Type>>) {
748 let items = self.local_type_hints.entry(id).or_default();
749 if let Some(item) = items.iter_mut().find(|item| item.0 == generic_args && item.1 == fn_tys) {
750 item.2 = hints;
751 } else {
752 items.push((generic_args.to_vec(), fn_tys.to_vec(), hints));
753 }
754 }
755
756 pub fn inferred_local_type_hints(&self, id: u32, generic_args: &[Type], fn_tys: &[Type]) -> Vec<Option<Type>> {
757 self.local_type_hints.get(&id).and_then(|items| items.iter().find(|item| item.0 == generic_args && item.1 == fn_tys)).map(|item| item.2.clone()).unwrap_or_default()
758 }
759
760 pub fn infer_fn(&mut self, id: u32, arg_tys: &[Type]) -> Result<Type> {
761 self.infer_fn_with_params(id, arg_tys, &[])
762 }
763
764 pub fn infer_fn_with_params(&mut self, id: u32, arg_tys: &[Type], generic_args: &[Type]) -> Result<Type> {
765 if self.infer_stack.len() > MAX_INFER_DEPTH {
768 return Ok(Type::Any);
769 }
770 let (name, s) = self.symbols.get_symbol(id).map(|(n, s)| (n.clone(), s.clone()))?;
771 if let Symbol::Fn { ty, args, generic_params, cap, body, .. } = s {
772 if let Type::Fn { tys, ret: _ } = ty {
773 let resolved_generic_args = crate::resolve_generic_args_from_types(&generic_params, &tys, arg_tys, generic_args)?;
774 let generic_args = resolved_generic_args.as_slice();
775 let tys = if generic_params.is_empty() { tys } else { tys.iter().map(|ty| crate::substitute_type(ty, &generic_params, generic_args)).collect() };
776 let body = if generic_params.is_empty() { body.as_ref().clone() } else { crate::substitute_stmt(body.as_ref(), &generic_params, generic_args) };
777 let fn_tys = self.get_fn_tys(&tys, arg_tys)?;
778 let body = if generic_params.is_empty() {
779 body
780 } else {
781 let mut compile_tys = tys.clone();
782 let mut compile_cap = cap.clone();
783 let saved_state = self.take_local_state();
784 if let Some((module, _)) = name.split_once("::") {
785 self.symbols.push_module_scope(module.into());
786 }
787 let compiled = self.compile_fn(&args, &mut compile_tys, body, &mut compile_cap);
788 if name.contains("::") {
789 self.symbols.pop_module_scope();
790 }
791 self.restore_local_state(saved_state);
792 Stmt::new(StmtKind::Block(compiled?), Span::default())
793 };
794 if let Some(fns) = self.fns.get_mut(&id) {
795 for f in fns.iter() {
796 if f.0 == generic_args && f.1 == fn_tys {
797 return match &f.2 {
798 FnInferRet::Done(ret_ty) => self.symbols.get_type(ret_ty),
799 FnInferRet::Pending(seed) => seed.as_ref().map(|ty| self.symbols.get_type(ty)).unwrap_or_else(|| {
800 if self.infer_stack.iter().any(|(sid, sargs, _)| *sid == id && sargs == generic_args) {
802 if let Some(base_ty) = self.try_find_base_return_ty(&body) {
803 return self.symbols.get_type(&base_ty);
804 }
805 }
806 Ok(Type::Any)
807 }),
808 };
809 }
810 }
811 fns.push((generic_args.to_vec(), fn_tys.clone(), FnInferRet::Pending(None)));
812 } else {
813 self.fns.insert(id, vec![(generic_args.to_vec(), fn_tys.clone(), FnInferRet::Pending(None))]);
814 }
815 if self.pending_return_seed(id, generic_args, &fn_tys).is_none() {
817 if let Some(base_ty) = self.try_find_base_return_ty_with_scope(&body, id, &name, &args, &fn_tys) {
818 if let Some(fns) = self.fns.get_mut(&id) {
819 if let Some(item) = fns.iter_mut().find(|item| item.0 == generic_args && item.1 == fn_tys)
820 && let FnInferRet::Pending(seed) = &mut item.2
821 && seed.is_none()
822 {
823 *seed = Some(base_ty);
824 }
825 }
826 }
827 }
828 let mut ret_ty = None;
829 let mut local_type_hints = Vec::new();
830 for _ in 0..4 {
831 let before_seed = self.pending_return_seed(id, generic_args, &fn_tys);
832 let saved_state = self.take_local_state();
833 self.frames.push(0);
834 for (arg, ty) in args.iter().zip(fn_tys.iter()) {
835 self.add_name(arg.clone());
836 self.add_ty(ty.clone());
837 }
838 for c in cap.vars.iter() {
839 if let Some((name, ty)) = cap.names.get(*c) {
840 self.add_name(name.clone());
841 self.add_ty(ty.clone());
842 } else {
843 self.add_name("".into());
844 self.add_ty(Type::Any);
845 }
846 }
847 self.infer_stack.push((id, generic_args.to_vec(), fn_tys.clone()));
848 let pass_ret_ty = self.infer_return_type(&body).map(|ty| ty.unwrap_or(Type::Void));
849 self.infer_stack.pop();
850 let pass_local_type_hints = self.collect_local_type_hints();
851 self.restore_local_state(saved_state);
852 let pass_ret_ty = match pass_ret_ty {
853 Ok(pass_ret_ty) => self.symbols.get_type(&pass_ret_ty).unwrap_or(pass_ret_ty),
854 Err(err) => {
855 log::error!("infer_fn {} failed: {:?}", name, err);
856 let should_remove = self
857 .fns
858 .get_mut(&id)
859 .map(|fns| {
860 fns.retain(|item| item.0 != generic_args || item.1 != fn_tys || !matches!(item.2, FnInferRet::Pending(_)));
861 fns.is_empty()
862 })
863 .unwrap_or(false);
864 if should_remove {
865 self.fns.remove(&id);
866 }
867 return Err(err);
868 }
869 };
870 if !pass_ret_ty.is_any() {
871 self.update_pending_return_seed(&pass_ret_ty);
872 ret_ty = Some(pass_ret_ty.clone());
873 } else if ret_ty.is_none() {
874 ret_ty = Some(pass_ret_ty);
875 }
876 local_type_hints = pass_local_type_hints;
877 let after_seed = self.pending_return_seed(id, generic_args, &fn_tys);
878 if before_seed == after_seed {
879 break;
880 }
881 }
882 let ret_ty = ret_ty.unwrap_or(Type::Any);
883 self.fns.get_mut(&id).map(|f| {
884 f.iter_mut().find(|item| item.0 == generic_args && item.1 == fn_tys).map(|item| item.2 = FnInferRet::Done(ret_ty.clone()));
885 });
886 self.set_local_type_hints(id, generic_args, &fn_tys, local_type_hints);
887 if generic_args.is_empty()
888 && let Some((_, Symbol::Fn { ty: Type::Fn { ret, .. }, .. })) = self.symbols.get_symbol_mut(id)
889 && ret.is_any()
890 {
891 *ret = std::rc::Rc::new(ret_ty.clone());
892 }
893 Ok(ret_ty)
894 } else {
895 Ok(Type::Any)
896 }
897 } else if let Symbol::Native(f) = s {
898 if let Type::Fn { ret, .. } = f { Ok((*ret).clone()) } else { Ok(Type::Any) }
899 } else if matches!(s, Symbol::Null) {
900 Ok(Type::Any)
901 } else {
902 Err(Self::semantic_error(Span::default(), format!("符号 {:?} 不是函数", name)))
903 }
904 }
905
906 pub fn infer_stmt(&mut self, stmt: &Stmt) -> Result<Type> {
907 match &stmt.kind {
908 StmtKind::Expr(expr, close) => {
909 if !close {
910 self.infer_expr(expr)
911 } else {
912 self.infer_expr(expr)?;
913 Ok(Type::Void)
914 }
915 }
916 StmtKind::Return(expr) => {
917 if let Some(e) = expr {
918 self.infer_expr(e)
919 } else {
920 Ok(Type::Void)
921 }
922 }
923 StmtKind::Block(stmts) => {
924 for (idx, stmt) in stmts.iter().enumerate() {
925 let ty = self.infer_stmt(stmt)?;
926 if stmt.is_return() || idx == stmts.len() - 1 {
927 return Ok(ty);
928 }
929 }
930 Ok(Type::Void)
931 }
932 StmtKind::If { then_body, else_body, .. } => {
933 let then_ty = self.infer_stmt(then_body)?;
934 if let Some(e) = else_body {
935 let else_ty = self.infer_stmt(e)?;
936 if then_ty != else_ty {
937 log::debug!("then 和 else 有不同类型 {:?} {:?}", then_ty, else_ty);
938 return Self::merge_return_type(stmt.span, Some(then_ty), else_ty);
939 }
940 }
941 if else_body.is_none() {
942 return Ok(Type::Void);
943 }
944 Ok(then_ty)
945 }
946 StmtKind::While { cond, body } => {
947 let cond_ty = self.infer_expr(cond)?;
948 if cond_ty != Type::Bool {
949 return Err(Self::semantic_error(cond.span, format!("条件表达式必须是布尔类型,实际是 {:?}", cond_ty)));
950 }
951 self.infer_stmt(body)
952 }
953 StmtKind::For { pat, range, body } => {
954 let ty = self.for_pattern_ty(range)?;
955 self.add_pattern_bindings_for_infer(pat, ty)?;
956 self.infer_stmt(body)
957 }
958 StmtKind::Let { pat, value } => {
959 let expr_ty = if let StmtKind::Expr(expr, _) = &value.kind { self.infer_expr(expr)? } else { self.infer_stmt(value)? };
960 self.add_pattern_bindings_for_infer(pat, expr_ty)?;
961 Ok(Type::Void)
962 }
963 _ => Ok(Type::Void),
964 }
965 }
966}