[CodeComplete] Handle ternary expression in sequence completion

Follow-up to f577578. The same
treatment for ternary expression (IfExpr). Plus, fix a regression
introduced in f577578 where infix operators were disappeard from
results.

Author: rintaro

lib/IDE/CodeCompletion.cpp

@@ -3468,10 +3468,13 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
       prepareForRetypechecking(SE);
 
       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);
+        // Unfold expressions for re-typechecking sequence.
+        if (auto *assignExpr = dyn_cast_or_null<AssignExpr>(element)) {
+          assignExpr->setSrc(nullptr);
+          assignExpr->setDest(nullptr);
+        } else if (auto *ifExpr = dyn_cast_or_null<IfExpr>(element)) {
+          ifExpr->setCondExpr(nullptr);
+          ifExpr->setElseExpr(nullptr);
         }
 
         // Reset any references to operators in types, so they are properly
@@ -3522,6 +3525,12 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
       flattenBinaryExpr(assignExpr->getSrc(), sequence);
       assignExpr->setDest(nullptr);
       assignExpr->setSrc(nullptr);
+    } else if (auto ifExpr = dyn_cast<IfExpr>(expr)) {
+      flattenBinaryExpr(ifExpr->getCondExpr(), sequence);
+      sequence.push_back(ifExpr);
+      flattenBinaryExpr(ifExpr->getElseExpr(), sequence);
+      ifExpr->setCondExpr(nullptr);
+      ifExpr->setElseExpr(nullptr);
     } else if (auto tryExpr = dyn_cast<AnyTryExpr>(expr)) {
       // Strip out try expression. It doesn't affect completion.
       flattenBinaryExpr(tryExpr->getSubExpr(), sequence);

lib/Sema/TypeCheckConstraints.cpp

@@ -2106,11 +2106,20 @@ bool TypeChecker::typeCheckCompletionSequence(Expr *&expr, DeclContext *DC) {
   expr = foldSequence(SE, DC);
 
   // Find the code-completion expression and operator again.
-  BinaryExpr *exprAsBinOp = nullptr;
-  while (auto *binExpr = dyn_cast<BinaryExpr>(expr)) {
-    auto *RHS = binExpr->getArg()->getElement(1);
+  Expr *exprAsBinOp = nullptr;
+  while (true) {
+    Expr *RHS;
+    if (auto *binExpr = dyn_cast<BinaryExpr>(expr))
+      RHS = binExpr->getArg()->getElement(1);
+    else if (auto *assignExpr = dyn_cast<AssignExpr>(expr))
+      RHS = assignExpr->getSrc();
+    else if (auto *ifExpr = dyn_cast<IfExpr>(expr))
+      RHS = ifExpr->getElseExpr();
+    else
+      break;
+
     if (RHS == CCE) {
-      exprAsBinOp = binExpr;
+      exprAsBinOp = expr;
       break;
     }
     expr = RHS;
@@ -2119,7 +2128,7 @@ bool TypeChecker::typeCheckCompletionSequence(Expr *&expr, DeclContext *DC) {
     return true;
 
   // Ensure the output expression is up to date.
-  assert(exprAsBinOp == expr && "found wrong expr?");
+  assert(exprAsBinOp == expr && isa<BinaryExpr>(expr) && "found wrong expr?");
 
   // Add type variable for the code-completion expression.
   auto tvRHS =

test/IDE/complete_crashes.swift

@@ -220,9 +220,10 @@ protocol Bar_38149042 {
 func foo_38149042(bar: Bar_38149042) {
   _ = bar.foo? #^RDAR_38149042^# .x
 }
-
 // RDAR_38149042: Begin completions
 // 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: End completions
 
 // rdar://problem/38272904
@@ -274,18 +275,22 @@ func foo(x: RDAR41159258_MyResult1) {
 
 
 // rdar://problem/41232519
-// RUN: %target-swift-ide-test -code-completion -source-filename=%s -code-completion-token=RDAR41232519 | %FileCheck %s -check-prefix=RDAR_41232519
+// RUN: %target-swift-ide-test -code-completion -source-filename=%s -code-completion-token=RDAR41232519_1 | %FileCheck %s -check-prefix=RDAR_41232519
+// RUN: %target-swift-ide-test -code-completion -source-filename=%s -code-completion-token=RDAR41232519_2 | %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^#
+      myInt = provider.nextInt() #^RDAR41232519_1^#
+      _ = true ? 1 : provider.nextInt() #^RDAR41232519_2^#
   }
 }
 // RDAR_41232519: Begin completions
+// RDAR_41232519: Decl[InfixOperatorFunction]/OtherModule[Swift]: [' ']+ {#Int#}[#Int#]; name=+ Int
+// RDAR_41232519: End completions
 
 // rdar://problem/28188259
 // RUN: %target-swift-ide-test -code-completion -code-completion-token=RDAR_28188259 -source-filename=%s | %FileCheck %s -check-prefix=RDAR_28188259