Merge #9988

9988: fix: Refactor & improve handling of overloaded binary operators r=flodiebold a=flodiebold

Fixes #9971. Also records them as method resolutions, which we could use later.

Co-authored-by: Florian Diebold <flodiebold@gmail.com>

Author: bors[bot]

crates/hir_def/src/data.rs

@@ -199,6 +199,13 @@ impl TraitData {
             _ => None,
         })
     }
+
+    pub fn method_by_name(&self, name: &Name) -> Option<FunctionId> {
+        self.items.iter().find_map(|(item_name, item)| match item {
+            AssocItemId::FunctionId(t) if item_name == name => Some(*t),
+            _ => None,
+        })
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]

crates/hir_expand/src/name.rs

@@ -285,6 +285,33 @@ pub mod known {
         wrapping_add,
         wrapping_mul,
         wrapping_sub,
+        // known methods of lang items
+        add,
+        mul,
+        sub,
+        div,
+        rem,
+        shl,
+        shr,
+        bitxor,
+        bitor,
+        bitand,
+        add_assign,
+        mul_assign,
+        sub_assign,
+        div_assign,
+        rem_assign,
+        shl_assign,
+        shr_assign,
+        bitxor_assign,
+        bitor_assign,
+        bitand_assign,
+        eq,
+        ne,
+        ge,
+        gt,
+        le,
+        lt,
     );
 
     // self/Self cannot be used as an identifier

crates/hir_ty/src/infer.rs

@@ -16,11 +16,11 @@
 use std::ops::Index;
 use std::sync::Arc;
 
-use chalk_ir::{cast::Cast, DebruijnIndex, Mutability, Safety};
+use chalk_ir::{cast::Cast, DebruijnIndex, Mutability, Safety, Scalar};
 use hir_def::{
     body::Body,
     data::{ConstData, FunctionData, StaticData},
-    expr::{ArithOp, BinaryOp, BindingAnnotation, ExprId, PatId},
+    expr::{BindingAnnotation, ExprId, PatId},
     lang_item::LangItemTarget,
     path::{path, Path},
     resolver::{HasResolver, ResolveValueResult, Resolver, TypeNs, ValueNs},
@@ -134,11 +134,17 @@ pub struct TypeMismatch {
 #[derive(Clone, PartialEq, Eq, Debug)]
 struct InternedStandardTypes {
     unknown: Ty,
+    bool_: Ty,
+    unit: Ty,
 }
 
 impl Default for InternedStandardTypes {
     fn default() -> Self {
-        InternedStandardTypes { unknown: TyKind::Error.intern(&Interner) }
+        InternedStandardTypes {
+            unknown: TyKind::Error.intern(&Interner),
+            bool_: TyKind::Scalar(Scalar::Bool).intern(&Interner),
+            unit: TyKind::Tuple(0, Substitution::empty(&Interner)).intern(&Interner),
+        }
     }
 }
 /// Represents coercing a value to a different type of value.
@@ -751,28 +757,6 @@ impl<'a> InferenceContext<'a> {
         self.db.trait_data(trait_).associated_type_by_name(&name![Output])
     }
 
-    fn resolve_binary_op_output(&self, bop: &BinaryOp) -> Option<TypeAliasId> {
-        let lang_item = match bop {
-            BinaryOp::ArithOp(aop) => match aop {
-                ArithOp::Add => "add",
-                ArithOp::Sub => "sub",
-                ArithOp::Mul => "mul",
-                ArithOp::Div => "div",
-                ArithOp::Shl => "shl",
-                ArithOp::Shr => "shr",
-                ArithOp::Rem => "rem",
-                ArithOp::BitXor => "bitxor",
-                ArithOp::BitOr => "bitor",
-                ArithOp::BitAnd => "bitand",
-            },
-            _ => return None,
-        };
-
-        let trait_ = self.resolve_lang_item(lang_item)?.as_trait();
-
-        self.db.trait_data(trait_?).associated_type_by_name(&name![Output])
-    }
-
     fn resolve_boxed_box(&self) -> Option<AdtId> {
         let struct_ = self.resolve_lang_item("owned_box")?.as_struct()?;
         Some(struct_.into())
@@ -846,6 +830,10 @@ impl Expectation {
         }
     }
 
+    fn from_option(ty: Option<Ty>) -> Self {
+        ty.map_or(Expectation::None, Expectation::HasType)
+    }
+
     /// The following explanation is copied straight from rustc:
     /// Provides an expectation for an rvalue expression given an *optional*
     /// hint, which is not required for type safety (the resulting type might

crates/hir_ty/src/infer/expr.rs

@@ -8,10 +8,13 @@ use std::{
 
 use chalk_ir::{cast::Cast, fold::Shift, Mutability, TyVariableKind};
 use hir_def::{
-    expr::{Array, BinaryOp, Expr, ExprId, Literal, MatchGuard, Statement, UnaryOp},
+    expr::{
+        ArithOp, Array, BinaryOp, CmpOp, Expr, ExprId, Literal, MatchGuard, Ordering, Statement,
+        UnaryOp,
+    },
     path::{GenericArg, GenericArgs},
     resolver::resolver_for_expr,
-    AssocContainerId, FieldId, Lookup,
+    AssocContainerId, FieldId, FunctionId, Lookup,
 };
 use hir_expand::name::{name, Name};
 use stdx::always;
@@ -23,7 +26,7 @@ use crate::{
     infer::coerce::CoerceMany,
     lower::lower_to_chalk_mutability,
     mapping::from_chalk,
-    method_resolution, op,
+    method_resolution,
     primitive::{self, UintTy},
     static_lifetime, to_chalk_trait_id,
     traits::FnTrait,
@@ -669,34 +672,21 @@ impl<'a> InferenceContext<'a> {
                 }
             }
             Expr::BinaryOp { lhs, rhs, op } => match op {
-                Some(op) => {
-                    let lhs_expectation = match op {
-                        BinaryOp::LogicOp(..) => {
-                            Expectation::has_type(TyKind::Scalar(Scalar::Bool).intern(&Interner))
-                        }
-                        _ => Expectation::none(),
-                    };
-                    let lhs_ty = self.infer_expr(*lhs, &lhs_expectation);
-                    let lhs_ty = self.resolve_ty_shallow(&lhs_ty);
-                    let rhs_expectation = op::binary_op_rhs_expectation(*op, lhs_ty.clone());
-                    let rhs_ty =
-                        self.infer_expr_coerce(*rhs, &Expectation::has_type(rhs_expectation));
-                    let rhs_ty = self.resolve_ty_shallow(&rhs_ty);
-
-                    let ret = op::binary_op_return_ty(*op, lhs_ty.clone(), rhs_ty.clone());
-
-                    if ret.is_unknown() {
-                        cov_mark::hit!(infer_expr_inner_binary_operator_overload);
-
-                        self.resolve_associated_type_with_params(
-                            lhs_ty,
-                            self.resolve_binary_op_output(op),
-                            &[rhs_ty],
-                        )
-                    } else {
-                        ret
-                    }
+                Some(BinaryOp::Assignment { op: None }) => {
+                    let lhs_ty = self.infer_expr(*lhs, &Expectation::none());
+                    self.infer_expr_coerce(*rhs, &Expectation::has_type(lhs_ty));
+                    self.result.standard_types.unit.clone()
+                }
+                Some(BinaryOp::LogicOp(_)) => {
+                    let bool_ty = self.result.standard_types.bool_.clone();
+                    self.infer_expr_coerce(*lhs, &Expectation::HasType(bool_ty.clone()));
+                    let lhs_diverges = self.diverges;
+                    self.infer_expr_coerce(*rhs, &Expectation::HasType(bool_ty.clone()));
+                    // Depending on the LHS' value, the RHS can never execute.
+                    self.diverges = lhs_diverges;
+                    bool_ty
                 }
+                Some(op) => self.infer_overloadable_binop(*lhs, *op, *rhs, tgt_expr),
                 _ => self.err_ty(),
             },
             Expr::Range { lhs, rhs, range_type } => {
@@ -862,6 +852,62 @@ impl<'a> InferenceContext<'a> {
         ty
     }
 
+    fn infer_overloadable_binop(
+        &mut self,
+        lhs: ExprId,
+        op: BinaryOp,
+        rhs: ExprId,
+        tgt_expr: ExprId,
+    ) -> Ty {
+        let lhs_expectation = Expectation::none();
+        let lhs_ty = self.infer_expr(lhs, &lhs_expectation);
+        let rhs_ty = self.table.new_type_var();
+
+        let func = self.resolve_binop_method(op);
+        let func = match func {
+            Some(func) => func,
+            None => {
+                let rhs_ty = self.builtin_binary_op_rhs_expectation(op, lhs_ty.clone());
+                let rhs_ty = self.infer_expr_coerce(rhs, &Expectation::from_option(rhs_ty));
+                return self
+                    .builtin_binary_op_return_ty(op, lhs_ty, rhs_ty)
+                    .unwrap_or_else(|| self.err_ty());
+            }
+        };
+
+        let subst = TyBuilder::subst_for_def(self.db, func)
+            .push(lhs_ty.clone())
+            .push(rhs_ty.clone())
+            .build();
+        self.write_method_resolution(tgt_expr, func, subst.clone());
+
+        let method_ty = self.db.value_ty(func.into()).substitute(&Interner, &subst);
+        self.register_obligations_for_call(&method_ty);
+
+        self.infer_expr_coerce(rhs, &Expectation::has_type(rhs_ty.clone()));
+
+        let ret_ty = match method_ty.callable_sig(self.db) {
+            Some(sig) => sig.ret().clone(),
+            None => self.err_ty(),
+        };
+
+        let ret_ty = self.normalize_associated_types_in(ret_ty);
+
+        // FIXME: record autoref adjustments
+
+        // use knowledge of built-in binary ops, which can sometimes help inference
+        if let Some(builtin_rhs) = self.builtin_binary_op_rhs_expectation(op, lhs_ty.clone()) {
+            self.unify(&builtin_rhs, &rhs_ty);
+        }
+        if let Some(builtin_ret) =
+            self.builtin_binary_op_return_ty(op, lhs_ty.clone(), rhs_ty.clone())
+        {
+            self.unify(&builtin_ret, &ret_ty);
+        }
+
+        ret_ty
+    }
+
     fn infer_block(
         &mut self,
         expr: ExprId,
@@ -1136,4 +1182,141 @@ impl<'a> InferenceContext<'a> {
             }
         }
     }
+
+    fn builtin_binary_op_return_ty(&mut self, op: BinaryOp, lhs_ty: Ty, rhs_ty: Ty) -> Option<Ty> {
+        let lhs_ty = self.resolve_ty_shallow(&lhs_ty);
+        let rhs_ty = self.resolve_ty_shallow(&rhs_ty);
+        match op {
+            BinaryOp::LogicOp(_) | BinaryOp::CmpOp(_) => {
+                Some(TyKind::Scalar(Scalar::Bool).intern(&Interner))
+            }
+            BinaryOp::Assignment { .. } => Some(TyBuilder::unit()),
+            BinaryOp::ArithOp(ArithOp::Shl | ArithOp::Shr) => {
+                // all integer combinations are valid here
+                if matches!(
+                    lhs_ty.kind(&Interner),
+                    TyKind::Scalar(Scalar::Int(_) | Scalar::Uint(_))
+                        | TyKind::InferenceVar(_, TyVariableKind::Integer)
+                ) && matches!(
+                    rhs_ty.kind(&Interner),
+                    TyKind::Scalar(Scalar::Int(_) | Scalar::Uint(_))
+                        | TyKind::InferenceVar(_, TyVariableKind::Integer)
+                ) {
+                    Some(lhs_ty)
+                } else {
+                    None
+                }
+            }
+            BinaryOp::ArithOp(_) => match (lhs_ty.kind(&Interner), rhs_ty.kind(&Interner)) {
+                // (int, int) | (uint, uint) | (float, float)
+                (TyKind::Scalar(Scalar::Int(_)), TyKind::Scalar(Scalar::Int(_)))
+                | (TyKind::Scalar(Scalar::Uint(_)), TyKind::Scalar(Scalar::Uint(_)))
+                | (TyKind::Scalar(Scalar::Float(_)), TyKind::Scalar(Scalar::Float(_))) => {
+                    Some(rhs_ty)
+                }
+                // ({int}, int) | ({int}, uint)
+                (
+                    TyKind::InferenceVar(_, TyVariableKind::Integer),
+                    TyKind::Scalar(Scalar::Int(_) | Scalar::Uint(_)),
+                ) => Some(rhs_ty),
+                // (int, {int}) | (uint, {int})
+                (
+                    TyKind::Scalar(Scalar::Int(_) | Scalar::Uint(_)),
+                    TyKind::InferenceVar(_, TyVariableKind::Integer),
+                ) => Some(lhs_ty),
+                // ({float} | float)
+                (
+                    TyKind::InferenceVar(_, TyVariableKind::Float),
+                    TyKind::Scalar(Scalar::Float(_)),
+                ) => Some(rhs_ty),
+                // (float, {float})
+                (
+                    TyKind::Scalar(Scalar::Float(_)),
+                    TyKind::InferenceVar(_, TyVariableKind::Float),
+                ) => Some(lhs_ty),
+                // ({int}, {int}) | ({float}, {float})
+                (
+                    TyKind::InferenceVar(_, TyVariableKind::Integer),
+                    TyKind::InferenceVar(_, TyVariableKind::Integer),
+                )
+                | (
+                    TyKind::InferenceVar(_, TyVariableKind::Float),
+                    TyKind::InferenceVar(_, TyVariableKind::Float),
+                ) => Some(rhs_ty),
+                _ => None,
+            },
+        }
+    }
+
+    fn builtin_binary_op_rhs_expectation(&mut self, op: BinaryOp, lhs_ty: Ty) -> Option<Ty> {
+        Some(match op {
+            BinaryOp::LogicOp(..) => TyKind::Scalar(Scalar::Bool).intern(&Interner),
+            BinaryOp::Assignment { op: None } => lhs_ty,
+            BinaryOp::CmpOp(CmpOp::Eq { .. }) => match self
+                .resolve_ty_shallow(&lhs_ty)
+                .kind(&Interner)
+            {
+                TyKind::Scalar(_) | TyKind::Str => lhs_ty,
+                TyKind::InferenceVar(_, TyVariableKind::Integer | TyVariableKind::Float) => lhs_ty,
+                _ => return None,
+            },
+            BinaryOp::ArithOp(ArithOp::Shl | ArithOp::Shr) => return None,
+            BinaryOp::CmpOp(CmpOp::Ord { .. })
+            | BinaryOp::Assignment { op: Some(_) }
+            | BinaryOp::ArithOp(_) => match self.resolve_ty_shallow(&lhs_ty).kind(&Interner) {
+                TyKind::Scalar(Scalar::Int(_) | Scalar::Uint(_) | Scalar::Float(_)) => lhs_ty,
+                TyKind::InferenceVar(_, TyVariableKind::Integer | TyVariableKind::Float) => lhs_ty,
+                _ => return None,
+            },
+        })
+    }
+
+    fn resolve_binop_method(&self, op: BinaryOp) -> Option<FunctionId> {
+        let (name, lang_item) = match op {
+            BinaryOp::LogicOp(_) => return None,
+            BinaryOp::ArithOp(aop) => match aop {
+                ArithOp::Add => (name!(add), "add"),
+                ArithOp::Mul => (name!(mul), "mul"),
+                ArithOp::Sub => (name!(sub), "sub"),
+                ArithOp::Div => (name!(div), "div"),
+                ArithOp::Rem => (name!(rem), "rem"),
+                ArithOp::Shl => (name!(shl), "shl"),
+                ArithOp::Shr => (name!(shr), "shr"),
+                ArithOp::BitXor => (name!(bitxor), "bitxor"),
+                ArithOp::BitOr => (name!(bitor), "bitor"),
+                ArithOp::BitAnd => (name!(bitand), "bitand"),
+            },
+            BinaryOp::Assignment { op: Some(aop) } => match aop {
+                ArithOp::Add => (name!(add_assign), "add_assign"),
+                ArithOp::Mul => (name!(mul_assign), "mul_assign"),
+                ArithOp::Sub => (name!(sub_assign), "sub_assign"),
+                ArithOp::Div => (name!(div_assign), "div_assign"),
+                ArithOp::Rem => (name!(rem_assign), "rem_assign"),
+                ArithOp::Shl => (name!(shl_assign), "shl_assign"),
+                ArithOp::Shr => (name!(shr_assign), "shr_assign"),
+                ArithOp::BitXor => (name!(bitxor_assign), "bitxor_assign"),
+                ArithOp::BitOr => (name!(bitor_assign), "bitor_assign"),
+                ArithOp::BitAnd => (name!(bitand_assign), "bitand_assign"),
+            },
+            BinaryOp::CmpOp(cop) => match cop {
+                CmpOp::Eq { negated: false } => (name!(eq), "eq"),
+                CmpOp::Eq { negated: true } => (name!(ne), "eq"),
+                CmpOp::Ord { ordering: Ordering::Less, strict: false } => {
+                    (name!(le), "partial_ord")
+                }
+                CmpOp::Ord { ordering: Ordering::Less, strict: true } => (name!(lt), "partial_ord"),
+                CmpOp::Ord { ordering: Ordering::Greater, strict: false } => {
+                    (name!(ge), "partial_ord")
+                }
+                CmpOp::Ord { ordering: Ordering::Greater, strict: true } => {
+                    (name!(gt), "partial_ord")
+                }
+            },
+            BinaryOp::Assignment { op: None } => return None,
+        };
+
+        let trait_ = self.resolve_lang_item(lang_item)?.as_trait()?;
+
+        self.db.trait_data(trait_).method_by_name(&name)
+    }
 }