core/type_inferencer: Fix location of unary/compare expressions

Codegen uses this location information to determine the CallId, and if
a function call is the operand of a unary expression or left-hand
operand of a compare expression, codegen will use the type of the
operator expression rather than the actual operand type.
This commit is contained in:
David Mak 2024-04-05 13:23:08 +08:00
parent 937a8b9698
commit e614dd4257
1 changed files with 9 additions and 17 deletions

View File

@ -553,7 +553,7 @@ impl<'a> Fold<()> for Inferencer<'a> {
Some(self.infer_unary_ops(expr.location, op, operand)?) Some(self.infer_unary_ops(expr.location, op, operand)?)
} }
ExprKind::Compare { left, ops, comparators } => { ExprKind::Compare { left, ops, comparators } => {
Some(self.infer_compare(left, ops, comparators)?) Some(self.infer_compare(expr.location, left, ops, comparators)?)
} }
ExprKind::Subscript { value, slice, ctx, .. } => { ExprKind::Subscript { value, slice, ctx, .. } => {
Some(self.infer_subscript(value.as_ref(), slice.as_ref(), ctx)?) Some(self.infer_subscript(value.as_ref(), slice.as_ref(), ctx)?)
@ -1269,11 +1269,12 @@ impl<'a> Inferencer<'a> {
operand.custom.unwrap(), operand.custom.unwrap(),
).map_err(|e| HashSet::from([format!("{e} (at {location})")]))?; ).map_err(|e| HashSet::from([format!("{e} (at {location})")]))?;
self.build_method_call(operand.location, method, operand.custom.unwrap(), vec![], ret) self.build_method_call(location, method, operand.custom.unwrap(), vec![], ret)
} }
fn infer_compare( fn infer_compare(
&mut self, &mut self,
location: Location,
left: &ast::Expr<Option<Type>>, left: &ast::Expr<Option<Type>>,
ops: &[ast::Cmpop], ops: &[ast::Cmpop],
comparators: &[ast::Expr<Option<Type>>], comparators: &[ast::Expr<Option<Type>>],
@ -1282,6 +1283,7 @@ impl<'a> Inferencer<'a> {
return Err(HashSet::from([String::from("Comparator chaining with ndarray types not supported")])) return Err(HashSet::from([String::from("Comparator chaining with ndarray types not supported")]))
} }
let mut res = None;
for (a, b, c) in izip!(once(left).chain(comparators), comparators, ops) { for (a, b, c) in izip!(once(left).chain(comparators), comparators, ops) {
let method = comparison_name(c) let method = comparison_name(c)
.ok_or_else(|| HashSet::from([ .ok_or_else(|| HashSet::from([
@ -1297,26 +1299,16 @@ impl<'a> Inferencer<'a> {
b.custom.unwrap(), b.custom.unwrap(),
).map_err(|e| HashSet::from([format!("{e} (at {})", b.location)]))?; ).map_err(|e| HashSet::from([format!("{e} (at {})", b.location)]))?;
self.build_method_call( res.replace(self.build_method_call(
a.location, location,
method, method,
a.custom.unwrap(), a.custom.unwrap(),
vec![b.custom.unwrap()], vec![b.custom.unwrap()],
ret, ret,
)?; )?);
} }
let res_lhs = comparators.iter().rev().nth(1).unwrap_or(left); Ok(res.unwrap())
let res_rhs = comparators.iter().rev().nth(0).unwrap();
let res_op = ops.iter().rev().nth(0).unwrap();
Ok(typeof_cmpop(
self.unifier,
self.primitives,
res_op,
res_lhs.custom.unwrap(),
res_rhs.custom.unwrap(),
).unwrap().unwrap())
} }
/// Infers the type of a subscript expression on an `ndarray`. /// Infers the type of a subscript expression on an `ndarray`.