From d68f0df018bacefd3ea649c58468a17e5b2c3660 Mon Sep 17 00:00:00 2001
From: yukang <moorekang@gmail.com>
Date: Thu, 2 Feb 2023 01:34:53 +0800
Subject: [PATCH] fix #106138, extend Infer ty for binary operators

---
 .../src/transform/validate.rs                 | 18 +-----
 compiler/rustc_hir_typeck/src/expr.rs         |  7 +++
 tests/ui/inference/issue-106138.rs            | 60 +++++++++++++++++++
 3 files changed, 69 insertions(+), 16 deletions(-)
 create mode 100644 tests/ui/inference/issue-106138.rs

diff --git a/compiler/rustc_const_eval/src/transform/validate.rs b/compiler/rustc_const_eval/src/transform/validate.rs
index dd168a9ac3cd3..80dd4c9ecb699 100644
--- a/compiler/rustc_const_eval/src/transform/validate.rs
+++ b/compiler/rustc_const_eval/src/transform/validate.rs
@@ -11,7 +11,7 @@ use rustc_middle::mir::{
     traversal, AggregateKind, BasicBlock, BinOp, Body, BorrowKind, CastKind, CopyNonOverlapping,
     Local, Location, MirPass, MirPhase, NonDivergingIntrinsic, Operand, Place, PlaceElem, PlaceRef,
     ProjectionElem, RetagKind, RuntimePhase, Rvalue, SourceScope, Statement, StatementKind,
-    Terminator, TerminatorKind, UnOp, START_BLOCK,
+    Terminator, TerminatorKind, START_BLOCK,
 };
 use rustc_middle::ty::{self, InstanceDef, ParamEnv, Ty, TyCtxt, TypeVisitable};
 use rustc_mir_dataflow::impls::MaybeStorageLive;
@@ -573,21 +573,7 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                     _ => self.fail(location, format!("There is no checked version of {:?}", op)),
                 }
             }
-            Rvalue::UnaryOp(op, operand) => {
-                let a = operand.ty(&self.body.local_decls, self.tcx);
-                match op {
-                    UnOp::Neg => {
-                        check_kinds!(a, "Cannot negate type {:?}", ty::Int(..) | ty::Float(..))
-                    }
-                    UnOp::Not => {
-                        check_kinds!(
-                            a,
-                            "Cannot binary not type {:?}",
-                            ty::Int(..) | ty::Uint(..) | ty::Bool
-                        );
-                    }
-                }
-            }
+            Rvalue::UnaryOp(..) => {}
             Rvalue::ShallowInitBox(operand, _) => {
                 let a = operand.ty(&self.body.local_decls, self.tcx);
                 check_kinds!(a, "Cannot shallow init type {:?}", ty::RawPtr(..));
diff --git a/compiler/rustc_hir_typeck/src/expr.rs b/compiler/rustc_hir_typeck/src/expr.rs
index bc7474cdfcf3d..837b659f8b870 100644
--- a/compiler/rustc_hir_typeck/src/expr.rs
+++ b/compiler/rustc_hir_typeck/src/expr.rs
@@ -383,6 +383,13 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         let mut oprnd_t = self.check_expr_with_expectation(&oprnd, expected_inner);
 
         if !oprnd_t.references_error() {
+            match (unop, oprnd_t.kind()) {
+                (hir::UnOp::Not | hir::UnOp::Neg, ty::Infer(ty::TyVar(_))) => {
+                    return oprnd_t;
+                }
+                _ => {}
+            }
+
             oprnd_t = self.structurally_resolved_type(expr.span, oprnd_t);
             match unop {
                 hir::UnOp::Deref => {
diff --git a/tests/ui/inference/issue-106138.rs b/tests/ui/inference/issue-106138.rs
new file mode 100644
index 0000000000000..0c75b10b3a859
--- /dev/null
+++ b/tests/ui/inference/issue-106138.rs
@@ -0,0 +1,60 @@
+// check-pass
+
+fn myfunction(x: &Vec<bool>, y: &Vec<i32> ) {
+    let one = |i, a: &Vec<bool>| {
+        a[i]  // ok
+    };
+
+    let two = |i, a: &Vec<bool>| {
+        !a[i] // cannot infer type
+    };
+
+    let three = |i, b: &Vec<i32>| {
+        -b[i] // ok
+    };
+
+    let r = one(0, x);
+    assert_eq!(r, x[0]);
+    let r = two(0, x);
+    assert_eq!(r, !x[0]);
+    let r = three(0, y);
+    assert_eq!(r, -y[0]);
+}
+
+fn bools(x: &Vec<bool>) {
+    let binary = |i, a: &Vec<bool>| {
+        a[i] && a[i+1] // ok
+    };
+
+    let unary = |i, a: &Vec<bool>| {
+        !a[i] // cannot infer type
+    };
+
+    let r = binary(0, x);
+    assert_eq!(r, x[0] && x[1]);
+
+    let r = unary(0, x);
+    assert_eq!(r, !x[0]);
+}
+
+fn ints(x: &Vec<i32>) {
+    let binary = |i, a: &Vec<i32>| {
+        a[i] + a[i+1] // ok
+    };
+    let unary = |i, a: &Vec<i32>| {
+        -a[i] // cannot infer type
+    };
+
+    let r = binary(0, x);
+    assert_eq!(r, x[0] + x[1]);
+    let r = unary(0, x);
+    assert_eq!(r, -x[0]);
+}
+
+fn main() {
+    let x = vec![true, false];
+    let y = vec![1, 2, 3];
+    myfunction(&x, &y);
+    bools(&x);
+    ints(&y);
+}