[ConstraintSystem] Store declaration context in which application occurs in ApplicableFunction constraint

This is required because `ApplicableFunction` constraint can
inject member reference constraints that require a declaration
context.

For example, `_ = { Double(...) }` would now produce a disjunction
for `Double.init` where overload choice declaration contexts point
to the closure instead of the enclosing context.

This addresses a long-standing FIXME in `simplifyApplicableFnConstraint`
and helps with disjunction optimizer because its correctness depends on
correct identification of declaration contexts where applications happen.

Author: xedin

include/swift/Sema/Constraint.h

@@ -380,10 +380,6 @@ class Constraint final : public llvm::ilist_node<Constraint>,
   /// The kind of function reference, for member references.
   unsigned TheFunctionRefInfo : 3;
 
-  /// The trailing closure matching for an applicable function constraint,
-  /// if any. 0 = None, 1 = Forward, 2 = Backward.
-  unsigned trailingClosureMatching : 2;
-
   union {
     struct {
       /// The first type.
@@ -439,6 +435,21 @@ class Constraint final : public llvm::ilist_node<Constraint>,
       /// Identifies whether result of this node is unused.
       bool IsDiscarded;
     } SyntacticElement;
+
+    struct {
+      /// The function type that is being applied where parameters
+      /// represent argument types passed to callee and result type
+      /// represents result type of the application.
+      FunctionType *AppliedFn;
+      /// The type being called, primarily a function type, but could
+      /// be a metatype, a tuple or a nominal type.
+      Type Callee;
+      /// The trailing closure matching for an applicable function constraint,
+      /// if any. 0 = None, 1 = Forward, 2 = Backward.
+      unsigned TrailingClosureMatching : 2;
+      /// The declaration context in which the application appears.
+      DeclContext *UseDC;
+    } Apply;
   };
 
   /// The locator that describes where in the expression this
@@ -495,6 +506,11 @@ class Constraint final : public llvm::ilist_node<Constraint>,
              ConstraintLocator *locator,
              SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
+  Constraint(FunctionType *appliedFn, Type calleeType,
+             unsigned trailingClosureMatching, DeclContext *useDC,
+             ConstraintLocator *locator,
+             SmallPtrSetImpl<TypeVariableType *> &typeVars);
+
   /// Retrieve the type variables buffer, for internal mutation.
   MutableArrayRef<TypeVariableType *> getTypeVariablesBuffer() {
     return { getTrailingObjects<TypeVariableType *>(), NumTypeVariables };
@@ -580,7 +596,7 @@ class Constraint final : public llvm::ilist_node<Constraint>,
 
   /// Create a new Applicable Function constraint.
   static Constraint *createApplicableFunction(
-      ConstraintSystem &cs, Type argumentFnType, Type calleeType,
+      ConstraintSystem &cs, FunctionType *argumentFnType, Type calleeType,
       std::optional<TrailingClosureMatching> trailingClosureMatching,
       DeclContext *useDC, ConstraintLocator *locator);
 
@@ -739,6 +755,9 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     case ConstraintKind::SyntacticElement:
       llvm_unreachable("closure body element constraint has no type operands");
 
+    case ConstraintKind::ApplicableFunction:
+      return Apply.AppliedFn;
+
     default:
       return Types.First;
     }
@@ -758,6 +777,9 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     case ConstraintKind::ValueWitness:
       return Member.Second;
 
+    case ConstraintKind::ApplicableFunction:
+      return Apply.Callee;
+
     default:
       return Types.Second;
     }
@@ -851,6 +873,21 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     return Member.UseDC;
   }
 
+  FunctionType *getAppliedFunctionType() const {
+    assert(Kind == ConstraintKind::ApplicableFunction);
+    return Apply.AppliedFn;
+  }
+
+  Type getCalleeType() const {
+    assert(Kind == ConstraintKind::ApplicableFunction);
+    return Apply.Callee;
+  }
+
+  DeclContext *getApplicationDC() const {
+    assert(Kind == ConstraintKind::ApplicableFunction);
+    return Apply.UseDC;
+  }
+
   ASTNode getSyntacticElement() const {
     assert(Kind == ConstraintKind::SyntacticElement);
     return SyntacticElement.Element;

include/swift/Sema/ConstraintSystem.h

@@ -4965,7 +4965,7 @@ class ConstraintSystem {
 
   /// Attempt to simplify the ApplicableFunction constraint.
   SolutionKind simplifyApplicableFnConstraint(
-      Type type1, Type type2,
+      FunctionType *appliedFn, Type calleeTy,
       std::optional<TrailingClosureMatching> trailingClosureMatching,
       DeclContext *useDC,
       TypeMatchOptions flags, ConstraintLocatorBuilder locator);

lib/Sema/CSSimplify.cpp

@@ -13107,16 +13107,12 @@ createImplicitRootForCallAsFunction(ConstraintSystem &cs, Type refType,
 }
 
 ConstraintSystem::SolutionKind ConstraintSystem::simplifyApplicableFnConstraint(
-    Type type1, Type type2,
+    FunctionType *func1, Type type2,
     std::optional<TrailingClosureMatching> trailingClosureMatching,
     DeclContext *useDC,
     TypeMatchOptions flags, ConstraintLocatorBuilder locator) {
   auto &ctx = getASTContext();
 
-  // By construction, the left hand side is a type that looks like the
-  // following: $T1 -> $T2.
-  auto func1 = type1->castTo<FunctionType>();
-
   // Before stripping lvalue-ness and optional types, save the original second
   // type for handling `func callAsFunction` and `@dynamicCallable`
   // applications. This supports the following cases:
@@ -13173,7 +13169,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyApplicableFnConstraint(
   auto formUnsolved = [&](bool activate = false) {
     if (flags.contains(TMF_GenerateConstraints)) {
       auto *application = Constraint::createApplicableFunction(
-          *this, type1, type2, trailingClosureMatching, useDC,
+          *this, func1, type2, trailingClosureMatching, useDC,
           getConstraintLocator(locator));
 
       addUnsolvedConstraint(application);
@@ -13203,7 +13199,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyApplicableFnConstraint(
   // If the types are obviously equivalent, we're done. This optimization
   // is not valid for operators though, where an inout parameter does not
   // have an explicit inout argument.
-  if (type1.getPointer() == desugar2) {
+  if (func1 == desugar2) {
     // Note that this could throw.
     recordPotentialThrowSite(
         PotentialThrowSite::Application, Type(desugar2), outerLocator);
@@ -13370,10 +13366,10 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyApplicableFnConstraint(
 
       auto applyLocator = getConstraintLocator(locator);
       auto forwardConstraint = Constraint::createApplicableFunction(
-          *this, type1, type2, TrailingClosureMatching::Forward, useDC,
+          *this, func1, type2, TrailingClosureMatching::Forward, useDC,
           applyLocator);
       auto backwardConstraint = Constraint::createApplicableFunction(
-          *this, type1, type2, TrailingClosureMatching::Backward, useDC,
+          *this, func1, type2, TrailingClosureMatching::Backward, useDC,
           applyLocator);
       addDisjunctionConstraint({forwardConstraint, backwardConstraint},
                                applyLocator);
@@ -13472,7 +13468,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyApplicableFnConstraint(
 
   // Handle applications of @dynamicCallable types.
   auto result = simplifyDynamicCallableApplicableFnConstraint(
-      type1, origType2, subflags, locator);
+      func1, origType2, subflags, locator);
 
   if (shouldAttemptFixes() && result == SolutionKind::Error) {
     // Skip this fix if the type is not yet resolved or
@@ -16204,9 +16200,9 @@ ConstraintSystem::simplifyConstraint(const Constraint &constraint) {
 
   case ConstraintKind::ApplicableFunction:
     return simplifyApplicableFnConstraint(
-        constraint.getFirstType(), constraint.getSecondType(),
+        constraint.getAppliedFunctionType(), constraint.getCalleeType(),
         constraint.getTrailingClosureMatching(),
-        /*FIXME*/DC, /*flags=*/std::nullopt,
+        constraint.getApplicationDC(), /*flags=*/std::nullopt,
         constraint.getLocator());
 
   case ConstraintKind::DynamicCallableApplicableFunction:

lib/Sema/Constraint.cpp

@@ -87,11 +87,9 @@ Constraint::Constraint(ConstraintKind Kind, Type First, Type Second,
     assert(!First.isNull());
     assert(!Second.isNull());
     break;
-  case ConstraintKind::ApplicableFunction:
   case ConstraintKind::DynamicCallableApplicableFunction:
     assert(First->is<FunctionType>()
            && "The left-hand side type should be a function type");
-    trailingClosureMatching = 0;
     break;
 
   case ConstraintKind::ValueMember:
@@ -120,6 +118,10 @@ Constraint::Constraint(ConstraintKind Kind, Type First, Type Second,
 
   case ConstraintKind::SyntacticElement:
     llvm_unreachable("Syntactic element constraint should use create()");
+
+  case ConstraintKind::ApplicableFunction:
+    llvm_unreachable(
+        "Application constraint should use create()");
   }
 
   std::uninitialized_copy(typeVars.begin(), typeVars.end(),
@@ -287,6 +289,27 @@ Constraint::Constraint(ASTNode node, ContextualTypeInfo context,
   std::copy(typeVars.begin(), typeVars.end(), getTypeVariablesBuffer().begin());
 }
 
+Constraint::Constraint(FunctionType *appliedFn, Type calleeType,
+                       unsigned trailingClosureMatching, DeclContext *useDC,
+                       ConstraintLocator *locator,
+                       SmallPtrSetImpl<TypeVariableType *> &typeVars)
+    : Kind(ConstraintKind::ApplicableFunction), HasFix(false),
+      HasRestriction(false), IsActive(false), IsDisabled(false),
+      IsDisabledForPerformance(false), RememberChoice(false), IsFavored(false),
+      IsIsolated(false), NumTypeVariables(typeVars.size()), Locator(locator) {
+  assert(appliedFn);
+  assert(calleeType);
+  assert(trailingClosureMatching >= 0 && trailingClosureMatching <= 2);
+  assert(useDC);
+
+  Apply.AppliedFn = appliedFn;
+  Apply.Callee = calleeType;
+  Apply.TrailingClosureMatching = trailingClosureMatching;
+  Apply.UseDC = useDC;
+
+  std::copy(typeVars.begin(), typeVars.end(), getTypeVariablesBuffer().begin());
+}
+
 ProtocolDecl *Constraint::getProtocol() const {
   assert((Kind == ConstraintKind::ConformsTo ||
           Kind == ConstraintKind::LiteralConformsTo ||
@@ -986,7 +1009,7 @@ Constraint *Constraint::createConjunction(
 }
 
 Constraint *Constraint::createApplicableFunction(
-    ConstraintSystem &cs, Type argumentFnType, Type calleeType,
+    ConstraintSystem &cs, FunctionType *argumentFnType, Type calleeType,
     std::optional<TrailingClosureMatching> trailingClosureMatching,
     DeclContext *useDC, ConstraintLocator *locator) {
   // Collect type variables.
@@ -996,30 +1019,29 @@ Constraint *Constraint::createApplicableFunction(
   if (calleeType->hasTypeVariable())
     calleeType->getTypeVariables(typeVars);
 
-  // Create the constraint.
-  auto size =
-    totalSizeToAlloc<TypeVariableType *, ConstraintFix *, OverloadChoice>(
-      typeVars.size(), /*hasFix=*/0, /*hasOverloadChoice=*/0);
-  void *mem = cs.getAllocator().Allocate(size, alignof(Constraint));
-  auto constraint = new (mem) Constraint(
-      ConstraintKind::ApplicableFunction, argumentFnType, calleeType, locator,
-      typeVars);
-
   // Encode the trailing closure matching.
+  unsigned rawTrailingClosureMatching = 0;
   if (trailingClosureMatching) {
     switch (*trailingClosureMatching) {
       case TrailingClosureMatching::Forward:
-        constraint->trailingClosureMatching = 1;
+        rawTrailingClosureMatching = 1;
         break;
 
       case TrailingClosureMatching::Backward:
-        constraint->trailingClosureMatching = 2;
+        rawTrailingClosureMatching = 2;
         break;
     }
-  } else {
-    constraint->trailingClosureMatching = 0;
   }
 
+    // Create the constraint.
+  auto size =
+    totalSizeToAlloc<TypeVariableType *, ConstraintFix *, OverloadChoice>(
+      typeVars.size(), /*hasFix=*/0, /*hasOverloadChoice=*/0);
+  void *mem = cs.getAllocator().Allocate(size, alignof(Constraint));
+  auto constraint = new (mem)
+      Constraint(argumentFnType, calleeType, rawTrailingClosureMatching, useDC,
+                 locator, typeVars);
+
   return constraint;
 }
 
@@ -1051,7 +1073,7 @@ Constraint *Constraint::createSyntacticElement(ConstraintSystem &cs,
 std::optional<TrailingClosureMatching>
 Constraint::getTrailingClosureMatching() const {
   assert(Kind == ConstraintKind::ApplicableFunction);
-  switch (trailingClosureMatching) {
+  switch (Apply.TrailingClosureMatching) {
   case 0:
     return std::nullopt;
   case 1: return TrailingClosureMatching::Forward;

unittests/Sema/ConstraintSimplificationTests.cpp

@@ -127,3 +127,98 @@ TEST_F(SemaTest, TestClosureInferenceFromOptionalContext) {
   ASSERT_TRUE(cs.simplifyType(paramTy)->isEqual(getStdlibType("Int")));
   ASSERT_TRUE(cs.simplifyType(resultTy)->isEqual(Context.TheEmptyTupleType));
 }
+
+/// Emulates code like this:
+///
+/// func test(_: (Int) -> Void) {}
+///
+/// test { Double($0) }
+///
+/// To make sure that constructor application sets correct
+/// declaration context for implicit `.init` member.
+TEST_F(SemaTest, TestInitializerUseDCIsSetCorrectlyInClosure) {
+  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(),
+      /*thrownType=*/nullptr,
+      /*arrowLoc=*/SourceLoc(),
+      /*inLoc=*/SourceLoc(),
+      /*explicitResultType=*/nullptr,
+      /*parent=*/DC);
+  closure->setDiscriminator(0);
+
+  closure->setImplicit();
+
+  // Double($0)
+  auto initCall = CallExpr::createImplicit(
+      Context, TypeExpr::createImplicit(getStdlibType("Double"), Context),
+      ArgumentList::forImplicitUnlabeled(
+          Context, {new (Context) DeclRefExpr(ConcreteDeclRef(paramDecl),
+                                              /*Loc*/ DeclNameLoc(),
+                                              /*Implicit=*/true)}));
+
+  closure->setBody(BraceStmt::createImplicit(Context, {initCall}));
+
+  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.setType(closure, closureTy);
+
+  cs.addUnsolvedConstraint(Constraint::create(
+      cs, ConstraintKind::FallbackType, closureTy, defaultTy,
+      cs.getConstraintLocator(closure), /*referencedVars=*/{}));
+
+  auto contextualTy =
+      FunctionType::get({FunctionType::Param(getStdlibType("Int"))},
+                        Context.TheEmptyTupleType, extInfo);
+
+  cs.resolveClosure(closureTy, contextualTy, closureLoc);
+
+  auto &graph = cs.getConstraintGraph();
+
+  for (const auto &component :
+       graph.computeConnectedComponents(cs.getTypeVariables())) {
+    for (auto *constraint : component.getConstraints()) {
+      if (constraint->getKind() != ConstraintKind::Disjunction)
+        continue;
+
+      ASSERT_TRUE(constraint->getLocator()
+                      ->isLastElement<LocatorPathElt::ConstructorMember>());
+
+      for (auto *choice : constraint->getNestedConstraints())
+        ASSERT_EQ(choice->getOverloadUseDC(), closure);
+    }
+  }
+}