Add support to copy objc selector

lib/AST/Decl.cpp

@@ -7968,7 +7968,6 @@ AbstractStorageDecl::getObjCGetterSelector(Identifier preferredName) const {
     name = preferredName;
   return VarDecl::getDefaultObjCGetterSelector(ctx, name);
 }
-
 ObjCSelector
 AbstractStorageDecl::getObjCSetterSelector(Identifier preferredName) const {
    auto abiRole = ABIRoleInfo(this);
@@ -10551,6 +10550,7 @@ AbstractFunctionDecl::getBodyFingerprintIncludingLocalTypeMembers() const {
 }
 
 ObjCSelector
+//TODO: if it's a function, this is what we need
 AbstractFunctionDecl::getObjCSelector(DeclName preferredName,
                                       bool skipIsObjCResolution) const {
    auto abiRole = ABIRoleInfo(this);
@@ -10573,7 +10573,8 @@ AbstractFunctionDecl::getObjCSelector(DeclName preferredName,
   StringRef baseNameStr;
   if (auto destructor = dyn_cast<DestructorDecl>(this)) {
     return destructor->getObjCSelector();
-  } else if (auto func = dyn_cast<FuncDecl>(this)) {
+  }
+  if (auto func = dyn_cast<FuncDecl>(this)) {
     // Otherwise cast this to be able to access getName()
     baseNameStr = func->getBaseIdentifier().str();
   } else if (isa<ConstructorDecl>(this)) {

test/refactoring/CopyObjCSelector/basic.swift

@@ -0,0 +1,24 @@
+import Foundation
+
+class MyClass: NSObject {
+    @objc func simpleMethod() {
+        print("simple")
+    }
+
+    @objc func methodWithParameters(param1: Int, param2: String) {
+        print("params")
+    }
+
+    @objc(customSelector:with:)
+    func methodWithCustomSelector(foo: Int, bar: String) {
+        print("custom")
+    }
+}
+
+// RUN: %refactor -source-filename %s -pos=4:16 | %FileCheck %s -check-prefix=CHECK-SIMPLE
+// RUN: %refactor -source-filename %s -pos=8:16 | %FileCheck %s -check-prefix=CHECK-PARAMS
+// RUN: %refactor -source-filename %s -pos=13:10 | %FileCheck %s -check-prefix=CHECK-CUSTOM
+
+// CHECK-SIMPLE: Copy Objective-C Selector
+// CHECK-PARAMS: Copy Objective-C Selector
+// CHECK-CUSTOM: Copy Objective-C Selector

tools/SourceKit/include/SourceKit/Core/LangSupport.h

@@ -1249,6 +1249,13 @@ class LangSupport {
                                    SourceKitCancellationToken CancellationToken,
                                    CategorizedEditsReceiver Receiver) = 0;
 
+  virtual void getObjCSelector(StringRef PrimaryFilePath,
+                               StringRef InputBufferName,
+                               unsigned Offset,
+                               ArrayRef<const char *> Args,
+                               SourceKitCancellationToken CancellationToken,
+                               std::function<void(const RequestResult<std::string> &)> Receiver) = 0;
+
   virtual void collectExpressionTypes(
       StringRef PrimaryFilePath, StringRef InputBufferName,
       ArrayRef<const char *> Args, ArrayRef<const char *> ExpectedProtocols,

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -729,6 +729,13 @@ class SwiftLangSupport : public LangSupport {
                            SourceKitCancellationToken CancellationToken,
                            CategorizedEditsReceiver Receiver) override;
 
+  void getObjCSelector(StringRef PrimaryFilePath,
+                       StringRef InputBufferName,
+                       unsigned Offset,
+                       ArrayRef<const char *> Args,
+                       SourceKitCancellationToken CancellationToken,
+                       std::function<void(const RequestResult<std::string> &)> Receiver) override;
+
   void getDocInfo(llvm::MemoryBuffer *InputBuf,
                   StringRef ModuleName,
                   ArrayRef<const char *> Args,

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -2888,6 +2888,87 @@ void SwiftLangSupport::semanticRefactoring(
                                    llvm::vfs::getRealFileSystem());
 }
 
+void SwiftLangSupport::getObjCSelector(
+    StringRef PrimaryFilePath, StringRef InputBufferName,
+    unsigned Offset, ArrayRef<const char *> Args,
+    SourceKitCancellationToken CancellationToken,
+    std::function<void(const RequestResult<std::string> &)> Receiver) {
+  std::string Error;
+  SwiftInvocationRef Invok =
+      ASTMgr->getTypecheckInvocation(Args, PrimaryFilePath, Error);
+  if (!Invok) {
+    LOG_WARN_FUNC("failed to create an ASTInvocation: " << Error);
+    Receiver(RequestResult<std::string>::fromError(Error));
+    return;
+  }
+
+  class ObjCSelectorConsumer : public SwiftASTConsumer {
+    std::string InputBufferName;
+    unsigned Offset;
+    std::function<void(const RequestResult<std::string> &)> Receiver;
+
+  public:
+    ObjCSelectorConsumer(StringRef InputBufferName, unsigned Offset,
+                         std::function<void(const RequestResult<std::string> &)> Receiver)
+        : InputBufferName(InputBufferName.str()), Offset(Offset),
+          Receiver(std::move(Receiver)) {}
+
+    void handlePrimaryAST(ASTUnitRef AstUnit) override {
+      auto &CompIns = AstUnit->getCompilerInstance();
+      auto &SM = CompIns.getSourceMgr();
+
+      SourceFile *SF = retrieveInputFile(InputBufferName, CompIns);
+      if (!SF) {
+        Receiver(RequestResult<std::string>::fromError("Unable to find input file"));
+        return;
+      }
+
+      unsigned BufferID = SF->getBufferID();
+      SourceLoc Loc = SM.getLocForOffset(BufferID, Offset);
+
+      auto CursorInfo = evaluateOrDefault(
+          SF->getASTContext().evaluator,
+          CursorInfoRequest{CursorInfoOwner(SF, Loc)},
+          new ResolvedCursorInfo());
+
+      if (!CursorInfo || CursorInfo->isInvalid()) {
+        Receiver(RequestResult<std::string>::fromError("Invalid cursor position"));
+        return;
+      }
+
+      if (auto *ValueRefInfo = dyn_cast<ResolvedValueRefCursorInfo>(CursorInfo.get())) {
+        if (auto *VD = ValueRefInfo->getValueD()) {
+          if (auto *AFD = dyn_cast<AbstractFunctionDecl>(VD)) {
+            if (AFD->isObjC()) {
+              auto selector = AFD->getObjCSelector();
+              SmallString<64> scratch;
+              std::string selectorString = selector.getString(scratch).str();
+              Receiver(RequestResult<std::string>::fromResult(selectorString));
+              return;
+            }
+          }
+        }
+      }
+
+      Receiver(RequestResult<std::string>::fromError(
+          "Cursor is not on an Objective-C method"));
+    }
+
+    void cancelled() override {
+      Receiver(RequestResult<std::string>::cancelled());
+    }
+
+    void failed(StringRef Error) override {
+      Receiver(RequestResult<std::string>::fromError(Error));
+    }
+  };
+
+  auto Consumer = std::make_shared<ObjCSelectorConsumer>(InputBufferName, Offset, Receiver);
+  getASTManager()->processASTAsync(Invok, std::move(Consumer), nullptr,
+                                   CancellationToken,
+                                   llvm::vfs::getRealFileSystem());
+}
+
 void SwiftLangSupport::collectExpressionTypes(
     StringRef PrimaryFilePath, StringRef InputBufferName,
     ArrayRef<const char *> Args, ArrayRef<const char *> ExpectedProtocols,

tools/SourceKit/tools/sourcekitd/lib/Service/Requests.cpp

@@ -1767,6 +1767,51 @@ handleRequestSemanticRefactoring(const RequestDict &Req,
   });
 }
 
+static void
+handleRequestGetObjCSelector(const RequestDict &Req,
+                             SourceKitCancellationToken CancellationToken,
+                             ResponseReceiver Rec) {
+  if (checkVFSNotSupported(Req, Rec))
+    return;
+
+  handleSemanticRequest(Req, Rec, [Req, CancellationToken, Rec]() {
+    std::optional<StringRef> PrimaryFilePath =
+        getPrimaryFilePathForRequestOrEmitError(Req, Rec);
+    if (!PrimaryFilePath)
+      return;
+
+    StringRef InputBufferName = getInputBufferNameForRequest(Req, Rec);
+
+    SmallVector<const char *, 8> Args;
+    if (getCompilerArgumentsForRequestOrEmitError(Req, Args, Rec))
+      return;
+
+    int64_t Offset = 0;
+    if (!Req.getInt64(KeyOffset, Offset, /*isOptional=*/false)) {
+      return Rec(createErrorRequestInvalid("'key.offset' is required"));
+    }
+
+    LangSupport &Lang = getGlobalContext().getSwiftLangSupport();
+
+    return Lang.getObjCSelector(
+        *PrimaryFilePath, InputBufferName, Offset, Args,
+        CancellationToken,
+        [Rec](const RequestResult<std::string> &Result) {
+          if (Result.isCancelled()) {
+            return Rec(createErrorRequestCancelled());
+          }
+          if (Result.isError()) {
+            return Rec(createErrorRequestFailed(Result.getError()));
+          }
+
+          ResponseBuilder RespBuilder;
+          auto Dict = RespBuilder.getDictionary();
+          Dict.set(KeyText, Result.value());
+          Rec(RespBuilder.createResponse());
+        });
+  });
+}
+
 static void
 handleRequestCollectExpressionType(const RequestDict &Req,
                                    SourceKitCancellationToken CancellationToken,
@@ -2278,6 +2323,7 @@ void handleRequestImpl(sourcekitd_object_t ReqObj,
   HANDLE_REQUEST(RequestCursorInfo, handleRequestCursorInfo)
   HANDLE_REQUEST(RequestRangeInfo, handleRequestRangeInfo)
   HANDLE_REQUEST(RequestSemanticRefactoring, handleRequestSemanticRefactoring)
+  HANDLE_REQUEST(RequestGetObjCSelector, handleRequestGetObjCSelector)
 
   HANDLE_REQUEST(RequestCollectExpressionType,
                  handleRequestCollectExpressionType)

utils/gyb_sourcekit_support/UIDs.py

@@ -292,6 +292,7 @@ def __init__(self, internal_name, external_name):
     REQUEST('FindLocalRenameRanges',
             'source.request.find-local-rename-ranges'),
     REQUEST('SemanticRefactoring', 'source.request.semantic.refactoring'),
+    REQUEST('GetObjCSelector', 'source.request.objc.selector'),
     REQUEST('EnableCompileNotifications',
             'source.request.enable-compile-notifications'),
     REQUEST('TestNotification', 'source.request.test_notification'),