Fix inference for shift operators

Fixes #2602.

3 files changed, 35 insertions, 5 deletions

diff --git a/crates/ra_hir_ty/src/infer/expr.rs b/crates/ra_hir_ty/src/infer/expr.rs
index d6a17e469..31259a01d 100644
--- a/crates/ra_hir_ty/src/infer/expr.rs
+++ b/crates/ra_hir_ty/src/infer/expr.rs
@@ -386,11 +386,11 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                     let lhs_ty = self.infer_expr(*lhs, &lhs_expectation);
                     // FIXME: find implementation of trait corresponding to operation
                     // symbol and resolve associated `Output` type
-                    let rhs_expectation = op::binary_op_rhs_expectation(*op, lhs_ty);
+                    let rhs_expectation = op::binary_op_rhs_expectation(*op, lhs_ty.clone());
                     let rhs_ty = self.infer_expr(*rhs, &Expectation::has_type(rhs_expectation));
 
                     // FIXME: similar as above, return ty is often associated trait type
-                    op::binary_op_return_ty(*op, rhs_ty)
+                    op::binary_op_return_ty(*op, lhs_ty, rhs_ty)
                 }
                 _ => Ty::Unknown,
             },
diff --git a/crates/ra_hir_ty/src/op.rs b/crates/ra_hir_ty/src/op.rs
index 09c47a76d..ae253ca04 100644
--- a/crates/ra_hir_ty/src/op.rs
+++ b/crates/ra_hir_ty/src/op.rs
@@ -1,13 +1,21 @@
-//! FIXME: write short doc here
-use hir_def::expr::{BinaryOp, CmpOp};
+//! Helper functions for binary operator type inference.
+use hir_def::expr::{ArithOp, BinaryOp, CmpOp};
 
 use super::{InferTy, Ty, TypeCtor};
 use crate::ApplicationTy;
 
-pub(super) fn binary_op_return_ty(op: BinaryOp, rhs_ty: Ty) -> Ty {
+pub(super) fn binary_op_return_ty(op: BinaryOp, lhs_ty: Ty, rhs_ty: Ty) -> Ty {
     match op {
         BinaryOp::LogicOp(_) | BinaryOp::CmpOp(_) => Ty::simple(TypeCtor::Bool),
         BinaryOp::Assignment { .. } => Ty::unit(),
+        BinaryOp::ArithOp(ArithOp::Shl) | BinaryOp::ArithOp(ArithOp::Shr) => match lhs_ty {
+            Ty::Apply(ApplicationTy { ctor, .. }) => match ctor {
+                TypeCtor::Int(..) | TypeCtor::Float(..) => lhs_ty,
+                _ => Ty::Unknown,
+            },
+            Ty::Infer(InferTy::IntVar(..)) | Ty::Infer(InferTy::FloatVar(..)) => lhs_ty,
+            _ => Ty::Unknown,
+        },
         BinaryOp::ArithOp(_) => match rhs_ty {
             Ty::Apply(ApplicationTy { ctor, .. }) => match ctor {
                 TypeCtor::Int(..) | TypeCtor::Float(..) => rhs_ty,
@@ -36,6 +44,7 @@ pub(super) fn binary_op_rhs_expectation(op: BinaryOp, lhs_ty: Ty) -> Ty {
                 _ => Ty::Unknown,
             }
         }
+        BinaryOp::ArithOp(ArithOp::Shl) | BinaryOp::ArithOp(ArithOp::Shr) => Ty::Unknown,
         BinaryOp::CmpOp(CmpOp::Ord { .. })
         | BinaryOp::Assignment { op: Some(_) }
         | BinaryOp::ArithOp(_) => match lhs_ty {
diff --git a/crates/ra_hir_ty/src/tests/simple.rs b/crates/ra_hir_ty/src/tests/simple.rs
index f7e042c12..b7204ec00 100644
--- a/crates/ra_hir_ty/src/tests/simple.rs
+++ b/crates/ra_hir_ty/src/tests/simple.rs
@@ -614,6 +614,27 @@ fn test() -> bool {
 }
 
 #[test]
+fn infer_shift_op() {
+    assert_snapshot!(
+        infer(r#"
+fn test() {
+    1u32 << 5u8;
+    1u32 >> 5u8;
+}
+"#),
+        @r###"
+    [11; 48) '{     ...5u8; }': ()
+    [17; 21) '1u32': u32
+    [17; 28) '1u32 << 5u8': u32
+    [25; 28) '5u8': u8
+    [34; 38) '1u32': u32
+    [34; 45) '1u32 >> 5u8': u32
+    [42; 45) '5u8': u8
+    "###
+    );
+}
+
+#[test]
 fn infer_field_autoderef() {
     assert_snapshot!(
         infer(r#"