Skip to main content

compiler/
infer.rs

1use super::{Compiler, Symbol};
2use anyhow::Result;
3use dynamic::{Dynamic, Type};
4use parser::{BinaryOp, Expr, ExprKind, PatternKind, Span, Stmt, StmtKind};
5
6impl Compiler {
7    fn merge_return_type(span: Span, left: Option<Type>, right: Type) -> Result<Type> {
8        match left {
9            Some(left) if left == right => Ok(left),
10            Some(left) if left.is_void() || right.is_void() => Err(Self::semantic_error(span, format!("返回类型不一致: {:?} 和 {:?}", left, right))),
11            Some(left) => Ok(left + right),
12            None => Ok(right),
13        }
14    }
15
16    fn infer_return_type(&mut self, stmt: &Stmt) -> Result<Option<Type>> {
17        self.infer_returns(stmt, true).map(|(ty, _)| ty)
18    }
19
20    fn infer_returns(&mut self, stmt: &Stmt, tail: bool) -> Result<(Option<Type>, bool)> {
21        match &stmt.kind {
22            StmtKind::Return(Some(expr)) => Ok((Some(self.infer_expr(expr)?), true)),
23            StmtKind::Return(None) => Ok((Some(Type::Void), true)),
24            StmtKind::Block(stmts) => {
25                let mut ret = None;
26                for (idx, stmt) in stmts.iter().enumerate() {
27                    let (ty, always_returns) = self.infer_returns(stmt, tail && idx == stmts.len().saturating_sub(1))?;
28                    if let Some(ty) = ty {
29                        ret = Some(Self::merge_return_type(stmt.span, ret, ty)?);
30                    }
31                    if always_returns {
32                        return Ok((ret, true));
33                    }
34                }
35                Ok((ret, false))
36            }
37            StmtKind::If { cond, then_body, else_body } => {
38                let cond_ty = self.infer_expr(cond)?;
39                if cond_ty != Type::Bool {
40                    return Err(Self::semantic_error(cond.span, format!("条件表达式必须是布尔类型,实际是 {:?}", cond_ty)));
41                }
42                let (mut ret, then_returns) = self.infer_returns(then_body, tail)?;
43                let else_returns = if let Some(body) = else_body {
44                    let (else_ty, else_returns) = self.infer_returns(body, tail)?;
45                    if let Some(ty) = else_ty {
46                        ret = Some(Self::merge_return_type(body.span, ret, ty)?);
47                    }
48                    else_returns
49                } else {
50                    false
51                };
52                Ok((ret, then_returns && else_returns))
53            }
54            StmtKind::While { cond, body } => {
55                let cond_ty = self.infer_expr(cond)?;
56                if cond_ty != Type::Bool {
57                    return Err(Self::semantic_error(cond.span, format!("条件表达式必须是布尔类型,实际是 {:?}", cond_ty)));
58                }
59                self.infer_returns(body, false).map(|(ty, _)| (ty, false))
60            }
61            StmtKind::Loop(body) => self.infer_returns(body, false),
62            StmtKind::For { pat, range, body } => {
63                if let PatternKind::Var { idx, .. } = &pat.kind {
64                    let ty = self.infer_expr(range)?;
65                    self.set_ty(*idx, ty);
66                } else if let PatternKind::Tuple(pats) = &pat.kind {
67                    let ty = self.infer_expr(range)?;
68                    assert!(ty.is_any());
69                    for pat in pats {
70                        if let Some(idx) = pat.var() {
71                            self.set_ty(idx, Type::Any);
72                        }
73                    }
74                }
75                self.infer_returns(body, false).map(|(ty, _)| (ty, false))
76            }
77            StmtKind::Let { .. } => {
78                self.infer_stmt(stmt)?;
79                Ok((None, false))
80            }
81            StmtKind::Expr(expr, close) => {
82                let ty = self.infer_expr(expr)?;
83                Ok(if *close || !tail { (None, false) } else { (Some(ty), true) })
84            }
85            _ => {
86                self.infer_stmt(stmt)?;
87                Ok((None, false))
88            }
89        }
90    }
91
92    pub fn infer_expr(&mut self, expr: &Expr) -> Result<Type> {
93        match &expr.kind {
94            ExprKind::Value(Dynamic::Null) => Ok(Type::Any),
95            ExprKind::Value(v) => Ok(v.get_type()),
96            ExprKind::Var(idx) => {
97                let idx = self.top() + (*idx as usize);
98                if idx < self.tys.len() { self.symbols.get_type(&self.tys[idx]) } else { Ok(Type::Any) }
99            }
100            ExprKind::Id(id, _) => match self.symbols.get_symbol(*id)?.1 {
101                Symbol::Const { ty, .. } => Ok(ty.clone()),
102                Symbol::Static { ty, .. } => Ok(ty.clone()),
103                Symbol::Struct(ty, _) => Ok(ty.clone()),
104                Symbol::Fn { .. } => Ok(Type::Symbol { id: *id, params: Vec::new() }),
105                Symbol::Native(ty) => Ok(ty.clone()),
106                s => Err(Self::semantic_error(expr.span, format!("符号 {:?} 不是变量、常量、静态变量、结构体", s))),
107            },
108            ExprKind::AssocId { id, params } => Ok(Type::Symbol { id: *id, params: params.clone() }),
109            ExprKind::Unary { value, .. } => self.infer_expr(value.as_ref()),
110            ExprKind::Binary { left, op, right } => {
111                let assign_idx = if op.is_assign() { if let ExprKind::Var(idx) = &left.kind { Some(*idx) } else { None } } else { None };
112                let ty = if op.is_logic() {
113                    let left_ty = self.infer_expr(left)?;
114                    if matches!(op, BinaryOp::And | BinaryOp::Or) && left_ty.is_any() { Type::Any } else { Type::Bool }
115                } else if op == &BinaryOp::Idx {
116                    let left_ty = self.infer_expr(left)?;
117                    if let Type::Array(elem_ty, _) = left_ty {
118                        (*elem_ty).clone()
119                    } else if let Type::Vec(elem_ty, _) = left_ty {
120                        (*elem_ty).clone()
121                    } else {
122                        let left_ty = self.symbols.get_type(&left_ty)?;
123                        let right_ty = if right.is_value() || right.is_const() {
124                            let right_value = if let ExprKind::Const(c) = &right.kind { self.consts[*c].clone() } else { right.clone().value()? };
125                            if right_value.is_str() {
126                                if left_ty.is_any() {
127                                    return Ok(Type::Any);
128                                }
129                                if let Ok(field) = self.symbols.get_field(&left_ty, right_value.as_str()) {
130                                    return if let Type::Fn { ret, .. } = field.1 { Ok(ret.as_ref().clone()) } else { Ok(field.1.clone()) };
131                                }
132                            } else if let Type::Struct { fields, .. } = &left_ty
133                                && let Some(idx) = right_value.as_int()
134                            {
135                                return fields.get(idx as usize).map(|(_, ty)| ty.clone()).ok_or_else(|| Self::semantic_error(right.span, format!("结构字段索引越界 {}", idx)));
136                            }
137                            right_value.get_type()
138                        } else {
139                            self.infer_expr(right)?
140                        };
141                        if right_ty.is_int() || right_ty.is_uint() {
142                            if left_ty.is_any() {
143                                return Ok(Type::Any);
144                            }
145                            let (_, s) = self.symbols.get_field(&left_ty, "get_idx")?;
146                            let fn_ty = self.symbols.get_type(&s)?;
147                            return if let Type::Fn { ret, .. } = &fn_ty { Ok(ret.as_ref().clone()) } else { Ok(fn_ty) };
148                        }
149                        if left_ty.is_any() {
150                            return Ok(Type::Any);
151                        }
152                        Type::Any
153                    }
154                } else {
155                    let right_ty = self.infer_expr(right)?;
156                    if op == &BinaryOp::Assign { right_ty } else { self.infer_expr(left)? + right_ty }
157                };
158                assign_idx.map(|idx| self.set_ty(idx, ty.clone()));
159                Ok(ty)
160            }
161            ExprKind::Call { obj, params } => {
162                if let ExprKind::AssocId { id, params: generic_args } = &obj.kind {
163                    let mut args = Vec::new();
164                    for p in params {
165                        args.push(self.infer_expr(p)?);
166                    }
167                    self.infer_fn_with_params(*id, &args, generic_args)
168                } else if let ExprKind::TypedMethod { obj: target, ty, name } = &obj.kind {
169                    let base_name = match ty {
170                        Type::Ident { name, .. } => name.clone(),
171                        Type::Symbol { id, .. } => self.symbols.get_symbol(*id)?.0.clone(),
172                        _ => return Ok(Type::Any),
173                    };
174                    let id = self.symbols.get_id(&format!("{}::{}", base_name, name))?;
175                    let mut args = vec![self.infer_expr(target)?];
176                    for p in params {
177                        args.push(self.infer_expr(p)?);
178                    }
179                    self.infer_fn(id, &args)
180                } else if let ExprKind::Id(id, obj_expr) = &obj.kind {
181                    let mut args: Vec<Type> = if let Some(obj) = obj_expr { vec![self.infer_expr(obj)?] } else { Vec::new() };
182                    for p in params {
183                        args.push(self.infer_expr(p)?);
184                    }
185                    self.infer_fn(*id, &args)
186                } else if obj.is_idx() {
187                    let (target, _, method) = obj.clone().binary().unwrap();
188                    let ty = self.infer_expr(&target)?;
189                    if let Some(method) = self.get_value(&method) {
190                        let method = method.as_str();
191                        let fn_ty = match self.get_field(&ty, method) {
192                            Ok((_, fn_ty)) => fn_ty,
193                            Err(_) => {
194                                let id = self.symbols.get_id(method)?;
195                                if self.symbols.get_symbol(id)?.1.is_fn() {
196                                    Type::Symbol { id, params: Vec::new() }
197                                } else {
198                                    return Err(Self::semantic_error(obj.span, format!("符号 {method} 不是函数")));
199                                }
200                            }
201                        };
202                        if let Type::Symbol { id, .. } = fn_ty {
203                            let mut args = vec![ty];
204                            for p in params {
205                                args.push(self.infer_expr(p)?);
206                            }
207                            self.infer_fn(id, &args)
208                        } else {
209                            Ok(fn_ty)
210                        }
211                    } else {
212                        Ok(Type::Any)
213                    }
214                } else if let ExprKind::Var(idx) = &obj.kind {
215                    let idx = self.top() + (*idx as usize);
216                    if idx < self.tys.len()
217                        && let Type::Symbol { id, .. } = self.tys[idx]
218                    {
219                        let mut args = Vec::new();
220                        for p in params {
221                            args.push(self.infer_expr(p)?);
222                        }
223                        self.infer_fn(id, &args)
224                    } else {
225                        Ok(Type::Any)
226                    }
227                } else if obj.is_value() {
228                    Ok(Type::Void)
229                } else {
230                    Ok(Type::Any)
231                }
232            }
233            ExprKind::Typed { ty, .. } => Ok(ty.clone()),
234            ExprKind::Stmt(stmt) => self.infer_stmt(stmt),
235            ExprKind::Range { start, stop, .. } => {
236                let start_ty = self.infer_expr(start)?;
237                let stop_ty = self.infer_expr(stop)?;
238                Ok(if start_ty.is_any() {
239                    stop_ty
240                } else if stop_ty.is_any() {
241                    start_ty
242                } else {
243                    stop_ty
244                })
245            }
246            _ => Ok(Type::Any),
247        }
248    }
249
250    fn get_fn_tys(&mut self, tys: &[Type], arg_tys: &[Type]) -> Result<Vec<Type>> {
251        let mut fn_tys = Vec::new();
252        for (i, ty) in tys.iter().enumerate() {
253            if !ty.is_any() {
254                fn_tys.push(ty.clone());
255            } else if let Some(arg_ty) = arg_tys.get(i) {
256                fn_tys.push(self.symbols.get_type(arg_ty)?);
257            } else {
258                fn_tys.push(Type::Any);
259            }
260        }
261        Ok(fn_tys)
262    }
263
264    pub fn infer_fn(&mut self, id: u32, arg_tys: &[Type]) -> Result<Type> {
265        self.infer_fn_with_params(id, arg_tys, &[])
266    }
267
268    pub fn infer_fn_with_params(&mut self, id: u32, arg_tys: &[Type], generic_args: &[Type]) -> Result<Type> {
269        let (name, s) = self.symbols.get_symbol(id).map(|(n, s)| (n.clone(), s.clone()))?;
270        if let Symbol::Fn { ty, args, generic_params, cap, body, .. } = s {
271            if let Type::Fn { tys, ret: _ } = ty {
272                let inferred_generic_args = if generic_args.is_empty() { crate::infer_generic_args_from_types(&generic_params, &tys, arg_tys) } else { generic_args.to_vec() };
273                let generic_args = if generic_params.is_empty() { &[] } else { inferred_generic_args.as_slice() };
274                let tys = if generic_params.is_empty() { tys } else { tys.iter().map(|ty| crate::substitute_type(ty, &generic_params, generic_args)).collect() };
275                let body = if generic_params.is_empty() { body.as_ref().clone() } else { crate::substitute_stmt(body.as_ref(), &generic_params, generic_args) };
276                let fn_tys = self.get_fn_tys(&tys, arg_tys)?;
277                let body = if generic_params.is_empty() {
278                    body
279                } else {
280                    let mut compile_tys = tys.clone();
281                    let mut compile_cap = cap.clone();
282                    let saved_state = self.take_local_state();
283                    let compiled = self.compile_fn(&args, &mut compile_tys, body, &mut compile_cap);
284                    self.restore_local_state(saved_state);
285                    Stmt::new(StmtKind::Block(compiled?), Span::default())
286                };
287                if let Some(fns) = self.fns.get_mut(&id) {
288                    for f in fns.iter() {
289                        if f.0 == generic_args && f.1 == fn_tys {
290                            return Ok(f.2.clone());
291                        }
292                    }
293                    fns.push((generic_args.to_vec(), fn_tys.clone(), Type::Any));
294                } else {
295                    self.fns.insert(id, vec![(generic_args.to_vec(), fn_tys.clone(), Type::Any)]);
296                }
297                let top = self.tys.len();
298                self.tys.append(&mut fn_tys.clone());
299                for c in cap.vars.iter() {
300                    self.tys.push(self.tys[self.top() + *c].clone());
301                }
302                self.frames.push(top);
303                let ret_ty = self.infer_return_type(&body).map(|ty| ty.unwrap_or(Type::Void));
304                if let Some(top) = self.frames.pop() {
305                    self.tys.truncate(top);
306                }
307                let ret_ty = match ret_ty {
308                    Ok(ret_ty) => ret_ty,
309                    Err(err) => {
310                        log::error!("infer_fn {} failed: {:?}", name, err);
311                        let should_remove = self
312                            .fns
313                            .get_mut(&id)
314                            .map(|fns| {
315                                fns.retain(|item| item.0 != generic_args || item.1 != fn_tys || item.2 != Type::Any);
316                                fns.is_empty()
317                            })
318                            .unwrap_or(false);
319                        if should_remove {
320                            self.fns.remove(&id);
321                        }
322                        return Err(err);
323                    }
324                };
325                self.fns.get_mut(&id).map(|f| {
326                    f.iter_mut().find(|item| item.0 == generic_args && item.1 == fn_tys).map(|item| item.2 = ret_ty.clone());
327                });
328                Ok(ret_ty)
329            } else {
330                Ok(Type::Any)
331            }
332        } else if let Symbol::Native(f) = s {
333            if let Type::Fn { ret, .. } = f { Ok((*ret).clone()) } else { Ok(Type::Any) }
334        } else if matches!(s, Symbol::Null) {
335            Ok(Type::Any)
336        } else {
337            Err(Self::semantic_error(Span::default(), format!("符号 {:?} 不是函数", name)))
338        }
339    }
340
341    pub fn infer_stmt(&mut self, stmt: &Stmt) -> Result<Type> {
342        match &stmt.kind {
343            StmtKind::Expr(expr, close) => {
344                if !close {
345                    self.infer_expr(expr)
346                } else {
347                    self.infer_expr(expr)?;
348                    Ok(Type::Void)
349                }
350            }
351            StmtKind::Return(expr) => {
352                if let Some(e) = expr {
353                    self.infer_expr(e)
354                } else {
355                    Ok(Type::Void)
356                }
357            }
358            StmtKind::Block(stmts) => {
359                for (idx, stmt) in stmts.iter().enumerate() {
360                    let ty = self.infer_stmt(stmt)?;
361                    if stmt.is_return() || idx == stmts.len() - 1 {
362                        return Ok(ty);
363                    }
364                }
365                Ok(Type::Void)
366            }
367            StmtKind::If { then_body, else_body, .. } => {
368                let then_ty = self.infer_stmt(then_body)?;
369                if let Some(e) = else_body {
370                    let else_ty = self.infer_stmt(e)?;
371                    if then_ty != else_ty {
372                        log::info!("then 和 else 有不同类型 {:?} {:?}", then_ty, else_ty);
373                        return Ok(if then_ty.is_any() { else_ty } else { then_ty });
374                    }
375                }
376                if else_body.is_none() {
377                    return Ok(Type::Void);
378                }
379                Ok(then_ty)
380            }
381            StmtKind::While { cond, body } => {
382                let cond_ty = self.infer_expr(cond)?;
383                if cond_ty != Type::Bool {
384                    return Err(Self::semantic_error(cond.span, format!("条件表达式必须是布尔类型,实际是 {:?}", cond_ty)));
385                }
386                self.infer_stmt(body)
387            }
388            StmtKind::For { pat, range, body } => {
389                if let PatternKind::Var { idx, .. } = &pat.kind {
390                    let ty = self.infer_expr(range)?;
391                    self.set_ty(*idx, ty);
392                } else if let PatternKind::Tuple(pats) = &pat.kind {
393                    let ty = self.infer_expr(range)?;
394                    assert!(ty.is_any());
395                    for pat in pats {
396                        if let Some(idx) = pat.var() {
397                            self.set_ty(idx, Type::Any);
398                        }
399                    }
400                }
401                self.infer_stmt(body)
402            }
403            StmtKind::Let { pat, value } => {
404                let expr_ty = if let StmtKind::Expr(expr, _) = &value.kind { self.infer_expr(expr)? } else { self.infer_stmt(value)? };
405                if let PatternKind::Ident { ty, .. } = &pat.kind {
406                    let annotated_ty = self.symbols.get_type(ty)?;
407                    if annotated_ty.is_any() {
408                        self.add_ty(expr_ty);
409                    } else {
410                        self.add_ty(annotated_ty);
411                    }
412                } else if let PatternKind::Var { idx, .. } = &pat.kind {
413                    self.set_ty(*idx, expr_ty);
414                } else if matches!(pat.kind, PatternKind::Wildcard) {
415                    self.add_ty(expr_ty);
416                }
417                Ok(Type::Void)
418            }
419            _ => Ok(Type::Void),
420        }
421    }
422}