1 files changed, 30 insertions, 14 deletions

diff --git a/clang-tools-extra/clangd/CodeCompletionStrings.cpp b/clang-tools-extra/clangd/CodeCompletionStrings.cpp
index f918a54c41c..8cdaa3380b8 100644
--- a/clang-tools-extra/clangd/CodeCompletionStrings.cpp
+++ b/clang-tools-extra/clangd/CodeCompletionStrings.cpp
@@ -9,6 +9,7 @@
 
 #include "CodeCompletionStrings.h"
 #include "clang/AST/ASTContext.h"
+#include "clang/AST/DeclObjC.h"
 #include "clang/AST/RawCommentList.h"
 #include "clang/Basic/SourceManager.h"
 #include <utility>
@@ -123,17 +124,32 @@ void processSnippetChunks(const CodeCompletionString &CCS,
   }
 }
 
-std::string getFormattedComment(const ASTContext &Ctx, const RawComment &RC,
-                                bool CommentsFromHeaders) {
+bool canRequestComment(const ASTContext &Ctx, const NamedDecl &D,
+                       bool CommentsFromHeaders) {
+  if (CommentsFromHeaders)
+    return true;
   auto &SourceMgr = Ctx.getSourceManager();
-  // Parsing comments from invalid preamble can lead to crashes. So we only
-  // return comments from the main file when doing code completion. For
-  // indexing, we still read all the comments.
+  // Accessing comments for decls from  invalid preamble can lead to crashes.
+  // So we only return comments from the main file when doing code completion.
+  // For indexing, we still read all the comments.
   // FIXME: find a better fix, e.g. store file contents in the preamble or get
   // doc comments from the index.
-  if (!CommentsFromHeaders && !SourceMgr.isWrittenInMainFile(RC.getLocStart()))
-    return "";
-  return RC.getFormattedText(Ctx.getSourceManager(), Ctx.getDiagnostics());
+  auto canRequestForDecl = [&](const NamedDecl &D) -> bool {
+    for (auto *Redecl : D.redecls()) {
+      auto Loc = SourceMgr.getSpellingLoc(Redecl->getLocation());
+      if (!SourceMgr.isWrittenInMainFile(Loc))
+        return false;
+    }
+    return true;
+  };
+  // First, check the decl itself.
+  if (!canRequestForDecl(D))
+    return false;
+  // Completion also returns comments for properties, corresponding to ObjC
+  // methods.
+  const ObjCMethodDecl *M = dyn_cast<ObjCMethodDecl>(&D);
+  const ObjCPropertyDecl *PDecl = M ? M->findPropertyDecl() : nullptr;
+  return !PDecl || canRequestForDecl(*PDecl);
 }
 } // namespace
 
@@ -145,26 +161,26 @@ std::string getDocComment(const ASTContext &Ctx,
   // get this declaration, so we don't show documentation in that case.
   if (Result.Kind != CodeCompletionResult::RK_Declaration)
     return "";
-  auto Decl = Result.getDeclaration();
-  if (!Decl)
+  auto *Decl = Result.getDeclaration();
+  if (!Decl || !canRequestComment(Ctx, *Decl, CommentsFromHeaders))
     return "";
   const RawComment *RC = getCompletionComment(Ctx, Decl);
   if (!RC)
     return "";
-  return getFormattedComment(Ctx, *RC, CommentsFromHeaders);
+  return RC->getFormattedText(Ctx.getSourceManager(), Ctx.getDiagnostics());
 }
 
 std::string
 getParameterDocComment(const ASTContext &Ctx,
                        const CodeCompleteConsumer::OverloadCandidate &Result,
                        unsigned ArgIndex, bool CommentsFromHeaders) {
-  auto Func = Result.getFunction();
-  if (!Func)
+  auto *Func = Result.getFunction();
+  if (!Func || !canRequestComment(Ctx, *Func, CommentsFromHeaders))
     return "";
   const RawComment *RC = getParameterComment(Ctx, Result, ArgIndex);
   if (!RC)
     return "";
-  return getFormattedComment(Ctx, *RC, CommentsFromHeaders);
+  return RC->getFormattedText(Ctx.getSourceManager(), Ctx.getDiagnostics());
 }
 
 void getLabelAndInsertText(const CodeCompletionString &CCS, std::string *Label,