[CodeComplete] Handle assign expression in sequence completion

rdar://problem/41232519
rdar://problem/41159258

Author: rintaro

lib/IDE/CodeCompletion.cpp

@@ -3358,11 +3358,17 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
       SE->setElement(SE->getNumElements() - 2, nullptr);
       prepareForRetypechecking(SE);
 
-      // Reset any references to operators in types, so they are properly
-      // handled as operators by sequence folding.
-      //
-      // FIXME: Would be better to have some kind of 'OperatorRefExpr'?
       for (auto &element : sequence.drop_back(2)) {
+        // Unfold AssignExpr for re-typechecking sequence.
+        if (auto *AE = dyn_cast_or_null<AssignExpr>(element)) {
+          AE->setSrc(nullptr);
+          AE->setDest(nullptr);
+        }
+
+        // Reset any references to operators in types, so they are properly
+        // handled as operators by sequence folding.
+        //
+        // FIXME: Would be better to have some kind of 'OperatorRefExpr'?
         if (auto operatorRef = element->getMemberOperatorRef()) {
           operatorRef->setType(nullptr);
           element = operatorRef;
@@ -3396,20 +3402,20 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     }
   }
 
-  void flattenBinaryExpr(BinaryExpr *expr, SmallVectorImpl<Expr *> &sequence) {
-    auto LHS = expr->getArg()->getElement(0);
-    if (auto binexpr = dyn_cast<BinaryExpr>(LHS))
-      flattenBinaryExpr(binexpr, sequence);
-    else
-      sequence.push_back(LHS);
-
-    sequence.push_back(expr->getFn());
-
-    auto RHS = expr->getArg()->getElement(1);
-    if (auto binexpr = dyn_cast<BinaryExpr>(RHS))
-      flattenBinaryExpr(binexpr, sequence);
-    else
-      sequence.push_back(RHS);
+  void flattenBinaryExpr(Expr *expr, SmallVectorImpl<Expr *> &sequence) {
+    if (auto binExpr = dyn_cast<BinaryExpr>(expr)) {
+      flattenBinaryExpr(binExpr->getArg()->getElement(0), sequence);
+      sequence.push_back(binExpr->getFn());
+      flattenBinaryExpr(binExpr->getArg()->getElement(1), sequence);
+    } else if (auto assignExpr = dyn_cast<AssignExpr>(expr)) {
+      flattenBinaryExpr(assignExpr->getDest(), sequence);
+      sequence.push_back(assignExpr);
+      flattenBinaryExpr(assignExpr->getSrc(), sequence);
+      assignExpr->setDest(nullptr);
+      assignExpr->setSrc(nullptr);
+    } else {
+      sequence.push_back(expr);
+    }
   }
 
   void typeCheckLeadingSequence(SmallVectorImpl<Expr *> &sequence) {
@@ -3419,10 +3425,10 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     // Take advantage of the fact the type-checker leaves the types on the AST.
     if (!typeCheckExpression(const_cast<DeclContext *>(CurrDeclContext),
                              expr)) {
-      if (auto binexpr = dyn_cast<BinaryExpr>(expr)) {
+      if (isa<BinaryExpr>(expr) || isa<AssignExpr>(expr)) {
         // Rebuild the sequence from the type-checked version.
         sequence.clear();
-        flattenBinaryExpr(binexpr, sequence);
+        flattenBinaryExpr(expr, sequence);
         return;
       }
     }
@@ -3448,6 +3454,9 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     if (sequence.size() > 1)
       typeCheckLeadingSequence(sequence);
 
+    // Retrieve typechecked LHS.
+    LHS = sequence.back();
+
     // Create a single sequence expression, which we will modify for each
     // operator, filling in the operator and dummy right-hand side.
     sequence.push_back(nullptr); // operator

test/IDE/complete_crashes.swift

@@ -220,10 +220,8 @@ protocol Bar_38149042 {
 func foo_38149042(bar: Bar_38149042) {
   _ = bar.foo? #^RDAR_38149042^# .x
 }
-// RDAR_38149042: Begin completions, 4 items
+// RDAR_38149042: Begin completions, 2 items
 // RDAR_38149042-DAG: Decl[InstanceVar]/CurrNominal:                  .x[#Int#]; name=x
-// RDAR_38149042-DAG: Decl[InfixOperatorFunction]/OtherModule[Swift]: [' ']=== {#AnyObject?#}[#Bool#]; name==== AnyObject?
-// RDAR_38149042-DAG: Decl[InfixOperatorFunction]/OtherModule[Swift]: [' ']!== {#AnyObject?#}[#Bool#]; name=!== AnyObject?
 // RDAR_38149042-DAG: Keyword[self]/CurrNominal: .self[#Baz_38149042#]; name=self
 // RDAR_38149042: End completions
 
@@ -247,3 +245,45 @@ func foo_38272904(a: A_38272904) {
   bar_38272904(a: .foo() #^RDAR_38272904^#)
 }
 // RDAR_38272904: Begin completions
+
+
+// rdar://problem/41159258
+// RUN: %target-swift-ide-test -code-completion -source-filename=%s -code-completion-token=RDAR41159258_1 | %FileCheck %s -check-prefix=RDAR_41159258
+// RUN: %target-swift-ide-test -code-completion -source-filename=%s -code-completion-token=RDAR41159258_2 | %FileCheck %s -check-prefix=RDAR_41159258
+public func ==(lhs: RDAR41159258_MyResult1, rhs: RDAR41159258_MyResult1) -> Bool {
+  fatalError()
+}
+public func ==(lhs: RDAR41159258_MyResult2, rhs: RDAR41159258_MyResult2) -> Bool {
+  fatalError()
+}
+public enum RDAR41159258_MyResult1 {
+  case failure(Error)
+}
+public enum RDAR41159258_MyResult2 {
+  case failure(Error)
+}
+
+public struct RDAR41159258_MyError: Error {}
+
+func foo(x: RDAR41159258_MyResult1) {
+  let x: RDAR41159258_MyResult1
+  x = .failure(RDAR41159258_MyError()) #^RDAR41159258_1^#
+  let y: Bool
+  y = .failure(RDAR41159258_MyError()) #^RDAR41159258_2^#
+}
+// RDAR_41159258: Begin completions
+
+
+// rdar://problem/41232519
+// RUN: %target-swift-ide-test -code-completion -source-filename=%s -code-completion-token=RDAR41232519 | %FileCheck %s -check-prefix=RDAR_41232519
+public protocol IntProvider {
+  func nextInt() -> Int
+}
+
+public final class IntStore {
+  public var myInt: Int = 0
+  func readNextInt(from provider: IntProvider) {
+      myInt = provider.nextInt() #^RDAR41232519^#
+  }
+}
+// RDAR_41232519: Begin completions