Module selector fix-its in @_dynamicReplacement

Author: beccadax

lib/Sema/TypeCheckAttr.cpp

@@ -3684,6 +3684,28 @@ static void lookupReplacedDecl(DeclNameRef replacedDeclName,
                                      results);
 }
 
+static void
+diagnoseCandidatesEliminatedByModuleSelector(DeclNameRefWithLoc replacedDeclName,
+                                             DeclAttribute *attr,
+                                             const ValueDecl *replacement,
+                                             DiagnosticEngine &Diags) {
+  if (replacedDeclName.Name.getModuleSelector().empty())
+    return;
+
+  // Look up without the module selector
+  SmallVector<ValueDecl *, 4> results;
+  lookupReplacedDecl(DeclNameRef(replacedDeclName.Name.getFullName()),
+                     attr, replacement, results);
+
+  auto selectorLoc = replacedDeclName.Loc.getModuleSelectorLoc();
+
+  for (auto candidate : results)
+    Diags.diagnose(selectorLoc, diag::note_change_module_selector,
+                   candidate->getModuleContext()->getBaseIdentifier())
+      .fixItReplace(selectorLoc,
+                    candidate->getModuleContext()->getBaseIdentifier().str());
+}
+
 /// Remove any argument labels from the interface type of the given value that
 /// are extraneous from the type system's point of view, producing the
 /// type to compare against for the purposes of dynamic replacement.
@@ -3703,14 +3725,14 @@ static Type getDynamicComparisonType(ValueDecl *value) {
   return interfaceType->removeArgumentLabels(numArgumentLabels);
 }
 
-static FuncDecl *findSimilarAccessor(DeclNameRef replacedVarName,
+static FuncDecl *findSimilarAccessor(DeclNameRefWithLoc replacedVarName,
                                      const AccessorDecl *replacement,
                                      DeclAttribute *attr, ASTContext &ctx,
                                      bool forDynamicReplacement) {
 
   // Retrieve the replaced abstract storage decl.
   SmallVector<ValueDecl *, 4> results;
-  lookupReplacedDecl(replacedVarName, attr, replacement, results);
+  lookupReplacedDecl(replacedVarName.Name, attr, replacement, results);
 
   // Filter out any accessors that won't work.
   if (!results.empty()) {
@@ -3743,15 +3765,19 @@ static FuncDecl *findSimilarAccessor(DeclNameRef replacedVarName,
   if (results.empty()) {
     Diags.diagnose(attr->getLocation(),
                    diag::dynamic_replacement_accessor_not_found,
-                   replacedVarName);
+                   replacedVarName.Name);
     attr->setInvalid();
+
+    diagnoseCandidatesEliminatedByModuleSelector(replacedVarName, attr,
+                                                 replacement, Diags);
+
     return nullptr;
   }
 
   if (results.size() > 1) {
     Diags.diagnose(attr->getLocation(),
                    diag::dynamic_replacement_accessor_ambiguous,
-                   replacedVarName);
+                   replacedVarName.Name);
     for (auto result : results) {
       Diags.diagnose(result,
                      diag::dynamic_replacement_accessor_ambiguous_candidate,
@@ -3793,15 +3819,15 @@ static FuncDecl *findSimilarAccessor(DeclNameRef replacedVarName,
   return origAccessor;
 }
 
-static FuncDecl *findReplacedAccessor(DeclNameRef replacedVarName,
+static FuncDecl *findReplacedAccessor(DeclNameRefWithLoc replacedVarName,
                                       const AccessorDecl *replacement,
                                       DeclAttribute *attr,
                                       ASTContext &ctx) {
   return findSimilarAccessor(replacedVarName, replacement, attr, ctx,
                              /*forDynamicReplacement*/ true);
 }
 
-static FuncDecl *findTargetAccessor(DeclNameRef replacedVarName,
+static FuncDecl *findTargetAccessor(DeclNameRefWithLoc replacedVarName,
                                       const AccessorDecl *replacement,
                                       DeclAttribute *attr,
                                       ASTContext &ctx) {
@@ -3810,15 +3836,15 @@ static FuncDecl *findTargetAccessor(DeclNameRef replacedVarName,
 }
 
 static AbstractFunctionDecl *
-findSimilarFunction(DeclNameRef replacedFunctionName,
+findSimilarFunction(DeclNameRefWithLoc replacedFunctionName,
                     const AbstractFunctionDecl *base, DeclAttribute *attr,
                     DiagnosticEngine *Diags, bool forDynamicReplacement) {
 
   // Note: we might pass a constant attribute when typechecker is nullptr.
   // Any modification to attr must be guarded by a null check on TC.
   //
   SmallVector<ValueDecl *, 4> lookupResults;
-  lookupReplacedDecl(replacedFunctionName, attr, base, lookupResults);
+  lookupReplacedDecl(replacedFunctionName.Name, attr, base, lookupResults);
 
   SmallVector<AbstractFunctionDecl *, 4> candidates;
   for (auto *result : lookupResults) {
@@ -3839,7 +3865,9 @@ findSimilarFunction(DeclNameRef replacedFunctionName,
                       forDynamicReplacement
                           ? diag::dynamic_replacement_function_not_found
                           : diag::specialize_target_function_not_found,
-                      replacedFunctionName);
+                      replacedFunctionName.Name);
+      diagnoseCandidatesEliminatedByModuleSelector(replacedFunctionName, attr,
+                                                   base, *Diags);
     }
 
     attr->setInvalid();
@@ -3905,7 +3933,7 @@ findSimilarFunction(DeclNameRef replacedFunctionName,
                   forDynamicReplacement
                       ? diag::dynamic_replacement_function_of_type_not_found
                       : diag::specialize_target_function_of_type_not_found,
-                  replacedFunctionName,
+                  replacedFunctionName.Name,
                   base->getInterfaceType()->getCanonicalType());
 
   for (auto *result : matches) {
@@ -3921,15 +3949,15 @@ findSimilarFunction(DeclNameRef replacedFunctionName,
 }
 
 static AbstractFunctionDecl *
-findReplacedFunction(DeclNameRef replacedFunctionName,
+findReplacedFunction(DeclNameRefWithLoc replacedFunctionName,
                      const AbstractFunctionDecl *replacement,
                      DynamicReplacementAttr *attr, DiagnosticEngine *Diags) {
   return findSimilarFunction(replacedFunctionName, replacement, attr, Diags,
                              true /*forDynamicReplacement*/);
 }
 
 static AbstractFunctionDecl *
-findTargetFunction(DeclNameRef targetFunctionName,
+findTargetFunction(DeclNameRefWithLoc targetFunctionName,
                    const AbstractFunctionDecl *base,
                    SpecializeAttr * attr, DiagnosticEngine *diags) {
   return findSimilarFunction(targetFunctionName, base, attr, diags,
@@ -5587,13 +5615,15 @@ DynamicallyReplacedDeclRequest::evaluate(Evaluator &evaluator,
   }
 
   auto &Ctx = VD->getASTContext();
+  DeclNameRefWithLoc nameWithLoc{attr->getReplacedFunctionName(),
+                                 attr->getReplacedFunctionNameLoc(),
+                                 std::nullopt};
   if (auto *AD = dyn_cast<AccessorDecl>(VD)) {
-    return findReplacedAccessor(attr->getReplacedFunctionName(), AD, attr, Ctx);
+    return findReplacedAccessor(nameWithLoc, AD, attr, Ctx);
   }
 
   if (auto *AFD = dyn_cast<AbstractFunctionDecl>(VD)) {
-    return findReplacedFunction(attr->getReplacedFunctionName(), AFD,
-                                attr, &Ctx.Diags);
+    return findReplacedFunction(nameWithLoc, AFD, attr, &Ctx.Diags);
   }
 
   if (auto *SD = dyn_cast<AbstractStorageDecl>(VD)) {
@@ -5616,8 +5646,10 @@ SpecializeAttrTargetDeclRequest::evaluate(Evaluator &evaluator,
 
   auto &ctx = vd->getASTContext();
 
-  auto targetFunctionName = attr->getTargetFunctionName();
-  if (!targetFunctionName)
+  DeclNameRefWithLoc targetFunctionName{attr->getTargetFunctionName(),
+                                        attr->getTargetFunctionNameLoc(),
+                                        std::nullopt};
+  if (!targetFunctionName.Name)
     return nullptr;
 
   if (auto *ad = dyn_cast<AccessorDecl>(vd)) {

test/NameLookup/module_selector.swift

@@ -83,7 +83,8 @@ extension B: main::Equatable {
   // @_derivative(of:)
 
   @_dynamicReplacement(for: main::negate())
-  // FIXME improve: expected-error@-1 {{replaced function 'main::negate()' could not be found}}
+  // FIXME shouldn't this succeed?: expected-error@-1 {{replaced function 'main::negate()' could not be found}}
+  // FIXME: expected-note@-2 {{did you mean module 'ModuleSelectorTestingKit'?}} {{29-33=ModuleSelectorTestingKit}}
   mutating func myNegate() {
     let fn: (main::Int, main::Int) -> main::Int =
     // expected-error@-1 3{{type 'Int' is not imported through module 'main'}}
@@ -206,7 +207,8 @@ extension D: @retroactive Swift::Equatable {
   // @_derivative(of:)
 
   @_dynamicReplacement(for: Swift::negate())
-  // FIXME improve: expected-error@-1 {{replaced function 'Swift::negate()' could not be found}}
+  // expected-error@-1 {{replaced function 'Swift::negate()' could not be found}}
+  // expected-note@-2 {{did you mean module 'ModuleSelectorTestingKit'?}} {{29-34=ModuleSelectorTestingKit}}
   mutating func myNegate() {
   // FIXME improve: expected-note@-1 {{did you mean 'myNegate'?}}