Merge pull request #77983 from hamishknight/ref-wrap

[CS] Correctly handle compound-applied functions with property wrappers

Author: hamishknight

include/swift/AST/Expr.h

@@ -552,6 +552,10 @@ class alignas(8) Expr : public ASTAllocated<Expr> {
   /// \c nullptr.
   ArgumentList *getArgs() const;
 
+  /// If the expression has a DeclNameLoc, returns it. Otherwise, returns
+  /// an nullp DeclNameLoc.
+  DeclNameLoc getNameLoc() const;
+
   /// Produce a mapping from each subexpression to its parent
   /// expression, with the provided expression serving as the root of
   /// the parent map.

lib/AST/Expr.cpp

@@ -859,6 +859,25 @@ ArgumentList *Expr::getArgs() const {
   return nullptr;
 }
 
+DeclNameLoc Expr::getNameLoc() const {
+  if (auto *DRE = dyn_cast<DeclRefExpr>(this))
+    return DRE->getNameLoc();
+  if (auto *UDRE = dyn_cast<UnresolvedDeclRefExpr>(this))
+    return UDRE->getNameLoc();
+  if (auto *ODRE = dyn_cast<OverloadedDeclRefExpr>(this))
+    return ODRE->getNameLoc();
+  if (auto *UDE = dyn_cast<UnresolvedDotExpr>(this))
+    return UDE->getNameLoc();
+  if (auto *UME = dyn_cast<UnresolvedMemberExpr>(this))
+    return UME->getNameLoc();
+  if (auto *MRE = dyn_cast<MemberRefExpr>(this))
+    return MRE->getNameLoc();
+  if (auto *DRME = dyn_cast<DynamicMemberRefExpr>(this))
+    return DRME->getNameLoc();
+
+  return DeclNameLoc();
+}
+
 llvm::DenseMap<Expr *, Expr *> Expr::getParentMap() {
   class RecordingTraversal : public ASTWalker {
   public:

lib/IDE/CursorInfo.cpp

@@ -214,23 +214,6 @@ class NodeFinder : ASTWalker {
     return Action::Continue();
   }
 
-  /// Retrieve the name location for an expression that supports cursor info.
-  DeclNameLoc getExprNameLoc(Expr *E) {
-    if (auto *DRE = dyn_cast<DeclRefExpr>(E))
-      return DRE->getNameLoc();
-    
-    if (auto *UDRE = dyn_cast<UnresolvedDeclRefExpr>(E))
-      return UDRE->getNameLoc();
-
-    if (auto *ODRE = dyn_cast<OverloadedDeclRefExpr>(E))
-      return ODRE->getNameLoc();
-
-    if (auto *UDE = dyn_cast<UnresolvedDotExpr>(E))
-      return UDE->getNameLoc();
-
-    return DeclNameLoc();
-  }
-
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
     if (auto closure = dyn_cast<ClosureExpr>(E)) {
       DeclContextStack.push_back(closure);
@@ -247,7 +230,7 @@ class NodeFinder : ASTWalker {
       }
     }
 
-    if (getExprNameLoc(E).getBaseNameLoc() != LocToResolve)
+    if (E->getNameLoc().getBaseNameLoc() != LocToResolve)
       return Action::Continue(E);
 
     assert(Result == nullptr);

lib/Sema/CSApply.cpp

@@ -764,10 +764,7 @@ namespace {
 
       if (auto *fnDecl = dyn_cast<AbstractFunctionDecl>(decl)) {
         if (AnyFunctionRef(fnDecl).hasExternalPropertyWrapperParameters() &&
-            // FIXME(FunctionRefInfo): This should just be `isUnapplied()`, see
-            // the FIXME in `unwrapPropertyWrapperParameterTypes`.
-            (declRefExpr->getFunctionRefInfo().isCompoundName() ||
-             declRefExpr->getFunctionRefInfo().isUnappliedBaseName())) {
+            declRefExpr->getFunctionRefInfo().isUnapplied()) {
           // We don't need to do any further adjustment once we've built the
           // curry thunk.
           return buildSingleCurryThunk(result, fnDecl,

lib/Sema/CSSimplify.cpp

@@ -1725,6 +1725,17 @@ static ConstraintSystem::TypeMatchResult matchCallArguments(
       continue;
     }
 
+    // See if we have a parameter label specified in the function's DeclNameLoc.
+    Identifier compoundParamLabel;
+    if (auto *E = getAsExpr(calleeLocator->getAnchor())) {
+      auto nameLoc = E->getNameLoc();
+      if (auto labelLoc = nameLoc.getArgumentLabelLoc(paramIdx)) {
+        auto &ctx = cs.getASTContext();
+        auto labelTok = Lexer::getTokenAtLocation(ctx.SourceMgr, labelLoc);
+        compoundParamLabel = ctx.getIdentifier(labelTok.getText());
+      }
+    }
+
     // Compare each of the bound arguments for this parameter.
     for (auto argIdx : parameterBindings[paramIdx]) {
       auto loc = locator.withPathElement(LocatorPathElt::ApplyArgToParam(
@@ -1813,14 +1824,17 @@ static ConstraintSystem::TypeMatchResult matchCallArguments(
         openedExistentials.push_back({openedTypeVar, opened});
       }
 
-      auto argLabel = argument.getLabel();
+      // If we have a compound function reference (e.g `fn($x:)`), respect
+      // the parameter label given. Otherwise look at the argument label.
+      auto wrapperArgLabel = compoundParamLabel.empty() ? argument.getLabel()
+                                                        : compoundParamLabel;
       if (paramInfo.hasExternalPropertyWrapper(paramIdx) ||
-          argLabel.hasDollarPrefix()) {
+          wrapperArgLabel.hasDollarPrefix()) {
         auto *param = getParameterAt(callee, paramIdx);
         assert(param);
         if (cs.applyPropertyWrapperToParameter(paramTy, argTy,
                                                const_cast<ParamDecl *>(param),
-                                               argLabel, subKind, loc)
+                                               wrapperArgLabel, subKind, loc)
                 .isFailure()) {
           return cs.getTypeMatchFailure(loc);
         }

lib/Sema/TypeOfReference.cpp

@@ -694,13 +694,8 @@ unwrapPropertyWrapperParameterTypes(ConstraintSystem &cs, AbstractFunctionDecl *
                                     FunctionRefInfo functionRefInfo, FunctionType *functionType,
                                     ConstraintLocatorBuilder locator) {
   // Only apply property wrappers to unapplied references to functions.
-  // FIXME(FunctionRefInfo): This should just be `isUnapplied()`, which would
-  // fix https://github.com/swiftlang/swift/issues/77823, but we also need to
-  // correctly handle the wrapping in matchCallArguments.
-  if (!(functionRefInfo.isCompoundName() ||
-        functionRefInfo.isUnappliedBaseName())) {
+  if (!functionRefInfo.isUnapplied())
     return functionType;
-  }
 
   // This transform is not applicable to pattern matching context.
   //
@@ -714,16 +709,8 @@ unwrapPropertyWrapperParameterTypes(ConstraintSystem &cs, AbstractFunctionDecl *
   SmallVector<AnyFunctionType::Param, 4> adjustedParamTypes;
 
   DeclNameLoc nameLoc;
-  auto *ref = getAsExpr(locator.getAnchor());
-  if (auto *declRef = dyn_cast<DeclRefExpr>(ref)) {
-    nameLoc = declRef->getNameLoc();
-  } else if (auto *dotExpr = dyn_cast<UnresolvedDotExpr>(ref)) {
-    nameLoc = dotExpr->getNameLoc();
-  } else if (auto *overloadedRef = dyn_cast<OverloadedDeclRefExpr>(ref)) {
-    nameLoc = overloadedRef->getNameLoc();
-  } else if (auto *memberExpr = dyn_cast<UnresolvedMemberExpr>(ref)) {
-    nameLoc = memberExpr->getNameLoc();
-  }
+  if (auto *ref = getAsExpr(locator.getAnchor()))
+    nameLoc = ref->getNameLoc();
 
   for (unsigned i : indices(*paramList)) {
     Identifier argLabel;

test/Sema/property_wrapper_parameter.swift

@@ -159,6 +159,40 @@ struct ProjectionWrapper<Value> {
   }
 }
 
+// https://github.com/swiftlang/swift/issues/77823
+// Make sure we correctly handle compound applied functions.
+func testCompoundApplication() {
+  func foo(@ProjectionWrapper x: Int) {}
+  struct HasProjectionWrapperMember {
+    static func foo(@ProjectionWrapper x: Int) {}
+  }
+
+  foo(x:)(0)
+  foo($x:)(ProjectionWrapper(wrappedValue: 0))
+
+  (foo($x:).self)(ProjectionWrapper(wrappedValue: 0))
+  HasProjectionWrapperMember.foo($x:)(ProjectionWrapper(wrappedValue: 0))
+
+  foo(x:)(ProjectionWrapper(wrappedValue: 0)) // expected-error {{cannot convert value of type 'ProjectionWrapper<Int>' to expected argument type 'Int'}}
+  foo(x:)(ProjectionWrapper(wrappedValue: "")) // expected-error {{cannot convert value of type 'ProjectionWrapper<String>' to expected argument type 'Int'}}
+  foo(x:)("") // expected-error {{cannot convert value of type 'String' to expected argument type 'Int'}}
+
+  foo($x:)(ProjectionWrapper(wrappedValue: "")) // expected-error {{cannot convert value of type 'String' to expected argument type 'Int'}}
+  foo($x:)(0) // expected-error {{cannot convert value of type 'Int' to expected argument type 'ProjectionWrapper<Int>'}}
+  foo($x:)("") // expected-error {{cannot convert value of type 'String' to expected argument type 'ProjectionWrapper<Int>'}}
+
+  func bar(x: Int) {} // expected-note 2{{parameter 'x' does not have an attached property wrapper}}
+  bar($x:)(0) // expected-error {{cannot use property wrapper projection argument}}
+  _ = bar($x:) // expected-error {{cannot use property wrapper projection argument}}
+
+  func baz(@ProjectionWrapper x: Int, @ProjectionWrapper y: Int) {}
+  baz($x:y:)(ProjectionWrapper(wrappedValue: 0), 0)
+  baz(x:$y:)(0, ProjectionWrapper(wrappedValue: 0))
+
+  let _: (ProjectionWrapper<Int>, Int) -> Void = baz($x:y:)
+  let _: (Int, ProjectionWrapper<Int>) -> Void = baz(x:$y:)
+}
+
 func testImplicitPropertyWrapper() {
   typealias PropertyWrapperTuple = (ProjectionWrapper<Int>, Int, ProjectionWrapper<Int>)
 

test/SourceKit/CursorInfo/issue-77981.swift

@@ -0,0 +1,14 @@
+struct S {
+  static func foo(_ x: Int) -> S {}
+  static func foo(_ x: String) -> S {}
+}
+
+// https://github.com/swiftlang/swift/issues/77981 - Make sure we can resolve
+// solver-based cursor info for UnresolvedMemberExprs.
+func bar() {
+  // RUN: %sourcekitd-test -req=cursor -pos=%(line + 1):15 %s -- %s | %FileCheck %s
+  let _: S = .foo()
+}
+
+// CHECK-DAG: source.lang.swift.ref.function.method.static (2:15-2:28)
+// CHECK-DAG: source.lang.swift.ref.function.method.static (3:15-3:31)