[ConstraintSystem] Shrink: If root expression is assignment always consider it as a candidate

If assignment expression is not considered as a top-level candidate
it would mean that other candidates would be allowed to produce
types inconsistent with destination type of the assignment.

Resolves: rdar://problem/51413254
(cherry picked from commit 93ab669)

Author: xedin

lib/Sema/CSSolver.cpp

@@ -829,8 +829,12 @@ void ConstraintSystem::shrink(Expr *expr) {
           return expr;
         }
 
-        // Or it's a function application with other candidates present.
-        if (isa<ApplyExpr>(expr)) {
+        // Or it's a function application or assignment with other candidates
+        // present. Assignment should be easy to solve because we'd get a
+        // contextual type from the destination expression, otherwise shrink
+        // might produce incorrect results without considering aforementioned
+        // destination type.
+        if (isa<ApplyExpr>(expr) || isa<AssignExpr>(expr)) {
           Candidates.push_back(Candidate(CS, PrimaryExpr));
           return expr;
         }

test/Constraints/generics.swift

@@ -750,3 +750,26 @@ func test_generic_subscript_with_missing_arg() {
     _ = s[0] // expected-error {{generic parameter 'U' could not be inferred}}
   }
 }
+
+// rdar://problem/51413254
+
+infix operator ==>
+
+struct Key {
+  init(_ key: String) {}
+}
+
+func ==> (lhs: Any, rhs: Key) throws -> Any {
+  return 0
+}
+
+func ==> <A>(lhs: Any, rhs: Key) throws -> A {
+  fatalError()
+}
+
+struct R_51413254 {
+  var str: String = ""
+  mutating func test(_ anyDict: Any) throws {
+    self.str = try anyDict ==> Key("a") // Ok
+  }
+}