[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

Author: xedin

lib/Sema/CSSolver.cpp

@@ -836,8 +836,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

@@ -770,3 +770,26 @@ func rdar_50007727() {
   // expected-error@-1 {{generic parameter 'T' could not be inferred in cast to 'A.B'}}
   // expected-note@-2 {{explicitly specify the generic arguments to fix this issue}} {{12-12=<Any>}}
 }
+
+// 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
+  }
+}