[CSSimplify] Look through optionals in contextual type while resolving a closure

lib/Sema/CSSimplify.cpp

@@ -8681,6 +8681,22 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
   auto *closure = castToExpr<ClosureExpr>(closureLocator->getAnchor());
   auto *inferredClosureType = getClosureType(closure);
 
+  // Let's look through all optionals associated with contextual
+  // type to make it possible to infer parameter/result type of
+  // the closure faster e.g.:
+  //
+  // func test(_: ((Int) -> Void)?) {
+  //   ...
+  // }
+  //
+  // test { $0 + ... }
+  //
+  // In this case dropping optionality from contextual type
+  // `((Int) -> Void)?` allows `resolveClosure` to infer type
+  // of `$0` directly (via `getContextualParamAt`) instead of
+  // having to use type variable inference mechanism.
+  contextualType = contextualType->lookThroughAllOptionalTypes();
+
   auto getContextualParamAt =
       [&contextualType, &inferredClosureType](
           unsigned index) -> Optional<AnyFunctionType::Param> {

unittests/Sema/ConstraintSimplificationTests.cpp

@@ -11,6 +11,10 @@
 //===----------------------------------------------------------------------===//
 
 #include "SemaFixture.h"
+#include "swift/AST/Decl.h"
+#include "swift/AST/Expr.h"
+#include "swift/AST/ParameterList.h"
+#include "swift/AST/Types.h"
 #include "swift/Sema/ConstraintSystem.h"
 
 using namespace swift;
@@ -44,3 +48,81 @@ TEST_F(SemaTest, TestTrailingClosureMatchRecordingForIdenticalFunctions) {
       TrailingClosureMatching::Forward, {{0}, {1}}, None};
   ASSERT_EQ(choice->second, expected);
 }
+
+/// Emulates code like this:
+///
+/// func test(_: ((Int) -> Void)?) {}
+///
+/// test { $0 }
+///
+/// To make sure that closure resolution propagates contextual
+/// information into the body of the closure even when contextual
+/// type is wrapped in an optional.
+TEST_F(SemaTest, TestClosureInferenceFromOptionalContext) {
+  ConstraintSystem cs(DC, ConstraintSystemOptions());
+
+  DeclAttributes closureAttrs;
+
+  // Anonymous closure parameter
+  auto paramName = Context.getIdentifier("0");
+
+  auto *paramDecl =
+      new (Context) ParamDecl(/*specifierLoc=*/SourceLoc(),
+                              /*argumentNameLoc=*/SourceLoc(), paramName,
+                              /*parameterNameLoc=*/SourceLoc(), paramName, DC);
+
+  paramDecl->setSpecifier(ParamSpecifier::Default);
+
+  auto *closure = new (Context) ClosureExpr(
+      closureAttrs,
+      /*bracketRange=*/SourceRange(),
+      /*capturedSelfDecl=*/nullptr, ParameterList::create(Context, {paramDecl}),
+      /*asyncLoc=*/SourceLoc(),
+      /*throwsLoc=*/SourceLoc(),
+      /*arrowLoc=*/SourceLoc(),
+      /*inLoc=*/SourceLoc(),
+      /*explicitResultType=*/nullptr,
+      /*discriminator=*/0,
+      /*parent=*/DC);
+
+  closure->setImplicit();
+
+  closure->setBody(BraceStmt::create(Context, /*startLoc=*/SourceLoc(), {},
+                                     /*endLoc=*/SourceLoc()),
+                   /*isSingleExpression=*/false);
+
+  auto *closureLoc = cs.getConstraintLocator(closure);
+
+  auto *paramTy = cs.createTypeVariable(
+      cs.getConstraintLocator(closure, LocatorPathElt::TupleElement(0)),
+      /*options=*/TVO_CanBindToInOut);
+
+  auto *resultTy = cs.createTypeVariable(
+      cs.getConstraintLocator(closure, ConstraintLocator::ClosureResult),
+      /*options=*/0);
+
+  auto extInfo = FunctionType::ExtInfo();
+
+  auto defaultTy = FunctionType::get({FunctionType::Param(paramTy, paramName)},
+                                     resultTy, extInfo);
+
+  cs.setClosureType(closure, defaultTy);
+
+  auto *closureTy = cs.createTypeVariable(closureLoc, /*options=*/0);
+
+  cs.addUnsolvedConstraint(Constraint::create(
+      cs, ConstraintKind::DefaultClosureType, closureTy, defaultTy,
+      cs.getConstraintLocator(closure), /*referencedVars=*/{}));
+
+  auto contextualTy =
+      FunctionType::get({FunctionType::Param(getStdlibType("Int"))},
+                        Context.TheEmptyTupleType, extInfo);
+
+  // Try to resolve closure:
+  // - external type `paramTy` should get bound to `Int`.
+  // - result type should be bound to `Void`.
+  cs.resolveClosure(closureTy, OptionalType::get(contextualTy), closureLoc);
+
+  ASSERT_TRUE(cs.simplifyType(paramTy)->isEqual(getStdlibType("Int")));
+  ASSERT_TRUE(cs.simplifyType(resultTy)->isEqual(Context.TheEmptyTupleType));
+}