[CSSyntacticElement] Fix solution applicator to identify context correctly

`TapExpr` do not form their own declaration context which means
that using `context.getAsDeclContext` is incorrect for them.

Author: xedin

lib/Sema/CSSyntacticElement.cpp

@@ -1645,7 +1645,6 @@ class SyntacticElementSolutionApplication
 protected:
   Solution &solution;
   SyntacticElementContext context;
-  Type resultType;
   RewriteTargetFn rewriteTarget;
 
   /// All `func`s declared in the body of the closure.
@@ -1658,24 +1657,39 @@ class SyntacticElementSolutionApplication
   SyntacticElementSolutionApplication(Solution &solution,
                                       SyntacticElementContext context,
                                       RewriteTargetFn rewriteTarget)
-      : solution(solution), context(context), rewriteTarget(rewriteTarget) {
-    if (auto fn = AnyFunctionRef::fromDeclContext(context.getAsDeclContext())) {
+      : solution(solution), context(context), rewriteTarget(rewriteTarget) {}
+
+  virtual ~SyntacticElementSolutionApplication() {}
+
+private:
+  Type getContextualResultType() const {
+    // Taps do not have a contextual result type.
+    if (context.is<TapExpr *>()) {
+      return Type();
+    }
+
+    auto fn = context.getAsAnyFunctionRef();
+
+    if (context.is<SingleValueStmtExpr *>()) {
+      // if/switch expressions can have `return` inside.
+      fn = AnyFunctionRef::fromDeclContext(context.getAsDeclContext());
+    }
+
+    if (fn) {
       if (auto transform = solution.getAppliedBuilderTransform(*fn)) {
-        resultType = solution.simplifyType(transform->bodyResultType);
+        return solution.simplifyType(transform->bodyResultType);
       } else if (auto *closure =
                      getAsExpr<ClosureExpr>(fn->getAbstractClosureExpr())) {
-        resultType = solution.getResolvedType(closure)
-                         ->castTo<FunctionType>()
-                         ->getResult();
+        return solution.getResolvedType(closure)
+            ->castTo<FunctionType>()
+            ->getResult();
       } else {
-        resultType = fn->getBodyResultType();
+        return fn->getBodyResultType();
       }
     }
-  }
-
-  virtual ~SyntacticElementSolutionApplication() {}
 
-private:
+    return Type();
+  }
 
   ASTNode visit(Stmt *S, bool performSyntacticDiagnostics = true) {
     auto rewritten = ASTVisitor::visit(S);
@@ -2031,17 +2045,18 @@ class SyntacticElementSolutionApplication
     auto closure = context.getAsAbstractClosureExpr();
     if (closure && !closure.get()->hasSingleExpressionBody() &&
         closure.get()->getBody() == braceStmt) {
+      auto resultType = getContextualResultType();
       if (resultType->getOptionalObjectType() &&
           resultType->lookThroughAllOptionalTypes()->isVoid() &&
           !braceStmt->getLastElement().isStmt(StmtKind::Return)) {
-        return addImplicitVoidReturn(braceStmt);
+        return addImplicitVoidReturn(braceStmt, resultType);
       }
     }
 
     return braceStmt;
   }
 
-  ASTNode addImplicitVoidReturn(BraceStmt *braceStmt) {
+  ASTNode addImplicitVoidReturn(BraceStmt *braceStmt, Type contextualResultTy) {
     auto &cs = solution.getConstraintSystem();
     auto &ctx = cs.getASTContext();
 
@@ -2056,7 +2071,7 @@ class SyntacticElementSolutionApplication
     // number of times.
     {
       SyntacticElementTarget target(resultExpr, context.getAsDeclContext(),
-                                    CTP_ReturnStmt, resultType,
+                                    CTP_ReturnStmt, contextualResultTy,
                                     /*isDiscarded=*/false);
       cs.setTargetFor(returnStmt, target);
 
@@ -2077,6 +2092,8 @@ class SyntacticElementSolutionApplication
   ASTNode visitReturnStmt(ReturnStmt *returnStmt) {
     auto &cs = solution.getConstraintSystem();
 
+    auto resultType = getContextualResultType();
+
     if (!returnStmt->hasResult()) {
       // If contextual is not optional, there is nothing to do here.
       if (resultType->isVoid())