From 93ab669c35344789a8b69ebae3cc71396f4087bb Mon Sep 17 00:00:00 2001
From: Pavel Yaskevich <pyaskevich@apple.com>
Date: Mon, 24 Jun 2019 14:03:54 -0700
Subject: [PATCH] [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
---
 lib/Sema/CSSolver.cpp           |  8 ++++++--
 test/Constraints/generics.swift | 23 +++++++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/lib/Sema/CSSolver.cpp b/lib/Sema/CSSolver.cpp
index 6c2a81fa64660..df0450eb0dd08 100644
--- a/lib/Sema/CSSolver.cpp
+++ b/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;
         }
diff --git a/test/Constraints/generics.swift b/test/Constraints/generics.swift
index 5c266fb9fbe9f..62e619aae5bc4 100644
--- a/test/Constraints/generics.swift
+++ b/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
+  }
+}