[CodeComplete] Properly handle if/switch exprs

Run PreCheckFunctionBodyRequest to ensure we insert
an implicit return for an if/switch if needed, and
ensure we don't try and type-check an element in a
SingleValueStmtExpr separately, as it should be
type-checked as a whole by the constraint system.
This ensures we can propagate a contextual type from
outside an if/switch expression for code completion.

Author: hamishknight

lib/Sema/TypeCheckStmt.cpp

@@ -2396,6 +2396,13 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
         return false;
       }
     }
+    if (auto *AFD = dyn_cast<AbstractFunctionDecl>(DC)) {
+      if (AFD->hasBody() && !AFD->isBodyTypeChecked()) {
+        // Pre-check the function body if needed.
+        (void)evaluateOrDefault(evaluator, PreCheckFunctionBodyRequest{AFD},
+                                nullptr);
+      }
+    }
   }
 
   // Find innermost ASTNode at Loc from DC. Results the reference to the found
@@ -2519,6 +2526,11 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
       if (isa<TapExpr>(E))
         return Action::SkipChildren(E);
 
+      // Don't walk into SingleValueStmtExprs, they should be type-checked as
+      // a whole.
+      if (isa<SingleValueStmtExpr>(E))
+        return Action::SkipChildren(E);
+
       if (auto closure = dyn_cast<ClosureExpr>(E)) {
         // NOTE: When a client wants to type check a closure signature, it
         // requests with closure's 'getLoc()' location.
@@ -2599,11 +2611,6 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
         // thus the transform couldn't be applied. Perform code completion
         // pretending there was no result builder to recover.
       }
-    } else if (func->hasSingleExpressionBody() &&
-                func->getResultInterfaceType()->isVoid()) {
-       // The function returns void.  We don't need an explicit return, no matter
-       // what the type of the expression is.  Take the inserted return back out.
-      func->getBody()->setLastElement(func->getSingleExpressionBody());
     }
   }
 

test/IDE/complete_if_switch_expr.swift

@@ -0,0 +1,118 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-ide-test -batch-code-completion -source-filename %s -filecheck %raw-FileCheck -completion-output-dir %t
+
+enum E {
+  case e
+  case f(Int)
+}
+
+func ifExprDotReturn() -> E {
+  if .random() {
+    .#^DOT1?check=DOT^#
+  } else {
+    .e
+  }
+}
+
+func switchExprDotReturn() -> E {
+  switch Bool.random() {
+  case true:
+    .e
+  case false:
+    .#^DOT2?check=DOT^#
+  }
+}
+
+func ifExprDotClosureReturn() -> E {
+  let x: E = {
+    if .random() {
+      .e
+    } else {
+      .#^DOT3?check=DOT^#
+    }
+  }()
+  return x
+}
+
+func switchExprDotClosureReturn() -> E {
+  let x: E = {
+    switch Bool.random() {
+    case true:
+      .#^DOT4?check=DOT^#
+    case false:
+      .e
+    }
+  }()
+  return x
+}
+
+func ifExprBranchInferenceReturn1() -> E {
+  let fn = {
+    if .random() {
+      E.e
+    } else {
+      .#^DOT5?check=DOT^#
+    }
+  }
+  return fn()
+}
+
+func switchExprBranchInferenceReturn1() -> E {
+  let fn = {
+    switch Bool.random() {
+    case true:
+      E.e
+    case false:
+      .#^DOT6?check=DOT^#
+    }
+  }
+  return fn()
+}
+
+func ifExprBranchInferenceReturn2() -> E {
+  let fn = {
+    if .random() {
+      .#^DOT7?check=DOT^#
+    } else {
+      E.e
+    }
+  }
+  return fn()
+}
+
+func switchExprBranchInferenceReturn2() -> E {
+  let fn = {
+    switch Bool.random() {
+    case true:
+      .#^DOT8?check=DOT^#
+    case false:
+      E.e
+    }
+  }
+  return fn()
+}
+
+func ifExprBinding() -> E {
+  let x: E = 
+    if .random() {
+      .e
+    } else {
+      .#^DOT9?check=DOT^#
+    }
+  return x
+}
+
+func switchExprBinding() -> E {
+  let x: E =
+    switch Bool.random() {
+      case true:
+        .e
+      case false:
+        .#^DOT10?check=DOT^#
+    }
+  return x
+}
+
+// DOT:     Begin completions, 2 items
+// DOT-DAG: Decl[EnumElement]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: e[#E#]; name=e
+// DOT-DAG: Decl[EnumElement]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: f({#Int#})[#E#]; name=f()