[clang] Refactor how are doc comments attached to Decls

[jkorous]

Summary

• Add ASTContext::searchForDocComments so we don't have to load external comments in Sema::ActOnDocumentableDecls when trying to attach existing doc comments to newly parsed Decls.
• Keep comments ordered and cache their decomposed locations.

This noticeably improves performance of code completion when using large modules with doc comments.

[jkorous]

Ping.

[gribozavr]

Not a complete review, but sending out the comments I have so far.

clang/include/clang/AST/ASTContext.h

806

"attachCommentsToJustParsedDecls"?

clang/include/clang/AST/RawCommentList.h

202–209

"... in case there are no comments in ..."

219

s/FileId hash/FileId/

clang/lib/AST/ASTContext.cpp

268

Please add a blank line between functions.

270

Please un-abbrevieate "search".

Also maybe call it "DeclLoc"?

456

We don't care about VS2015 anymore. https://reviews.llvm.org/D64326

clang/lib/AST/RawCommentList.cpp

289–290

Please move this comment down, right before the if statement that implement this logic.

319–320

No need to parenthesize.

329

I don't understand what you mean by "slow path".

334

No need to parenthesize.

clang/lib/Sema/SemaDecl.cpp

12729–12730

s/D/d/

12731

Also macros. (However I do agree with you that we can ignore this case.)

12737

Would it be better to move this logic into searchForDocComments? It seems like it belongs there. That function already computes CommentsInThisFile etc.

clang/lib/Serialization/ASTReader.cpp

9733

Wouldn't it be faster to directly modify Context.Comments.OrderedComments?

[jkorous]

clang/include/clang/AST/ASTContext.h

806

Sounds good - let's just finish the discussion about moving the logic from Sema::ActOnDocumentableDecls to this method first.

clang/lib/AST/RawCommentList.cpp

329

The comment is no longer relevant. Thanks!

clang/lib/Sema/SemaDecl.cpp

12731

You're absolutely right! I added a FIXME.

12737

I was asking myself the same thing.

We can move the logic but then we'd have to change semantics of ASTContext::searchForDocComments - essentially documenting that it makes sense to be called only from Sema::ActOnDocumentableDecls context.

I don't have a strong preference but I'm worried that it would be kind of weird interface.

WDYT?

[jkorous]

Addressed feedback (comments, naming).

[gribozavr]

Finished the review. I like this approach a lot.

clang/lib/AST/ASTContext.cpp

448

cacheRawCommentForDecl

We don't know that it is a doc comment (it could be an ordinary comment when ParseAllComments is true), and "data" in the name does not explain anything.

463

Should it be setOriginalDecl(OriginalDeclForComment), like in ASTContext::getRawCommentForAnyRedecl?

Could you also call cacheDocCommentDataForDecl from ASTContext::getRawCommentForAnyRedecl?

468

cacheRawCommentForRedecls

486

Could you call cacheDocCommentDataForRedecls from ASTContext::getRawCommentForAnyRedecl?

clang/lib/Sema/SemaDecl.cpp

12737

I think changing the name and semantics like that would be an improvement in clarity. I don't think it makes sense to call ASTContext::searchForDocComments in any other case anyway -- users who want to retrieve the comment should be calling ASTContext::getRawCommentForAnyRedecl.

[jkorous]

Rewrite caching of Decl -> RawComment mapping.

[jkorous]

I still have to do the searchForDocComments -> attachCommentsToJustParsedDecls change.

clang/lib/AST/ASTContext.cpp

463

Well, turned out this was a spot on comment!

I started by explaining the use NoCommentInDecl and why this change is fine and realized we can save some memory and make the code (arguably) simpler. I ended up rewriting the caching completely.

[jkorous]

searchForDocComments() -> attachCommentsToJustParsedDecls()

[gribozavr]

clang/lib/AST/ASTContext.cpp

370

Is it possible for D to be a null pointer? adjustDeclToTemplate does a dyn_cast on D, which would crash if D was null.

396

return DeclRawComments.find(RedeclComment->second)->second; might be slightly better, because it will detect coding errors -- it will catch the case when DeclRawComments does not have an entry for that decl. The square bracket syntax will insert a null entry silently.

430

I like the new approach a lot. The last loop is tricky, but I do think it is an improvement.

439

Why not retrieve the canonical decl inside the cacheRawCommentForDecl from OriginalD?

466

Indentation is off.

505

Are there cases when canonical decl is null?..

clang/lib/Sema/SemaDecl.cpp

12747

I think this FIXME should be sunk down into attachCommentsToJustParsedDecls, because it relates to the implementation of that function.

clang/lib/Serialization/ASTReader.cpp

9733

So what about this point?

[jkorous]

clang/lib/AST/ASTContext.cpp

370

You're right. I moved the check above the adjustment and changed interface of adjustDeclToTemplate to make it obvious that != NULL is expected of input and guaranteed for output.

396

Good point!

430

I agree. I don't think we're doing any rocket science here so the code should be easy to understand.

439

You are right, this is premature optimization.

//  - The first element is also often called the canonical element. Every
//    element has a pointer to it so that "getCanonical" can be fast.
...
class Redeclarable

505

IIUC correctly there are no such cases at the moment just that the interface doesn't guarantee that. But it seems like we don't expect that possibility anywhere in clang.

clang/lib/Sema/SemaDecl.cpp

12747

I've put it as a requirement to attachCommentsToJustParsedDecls doc comment. Since we're violating it here I'd leave the FIXME also here.

Or did you mean that I shouldn't be expecting that in the first place?

clang/lib/Serialization/ASTReader.cpp

9733

Oh, absolutely. Thanks for reminding me!

[jkorous]

Addressed comments & clang-formatted the patch.

[gribozavr]

clang/lib/AST/ASTContext.cpp

345

Why add a null check?

352

Why null check?

370

But why even add a null check at all? The contract of this function is that D is never null.

396

No need to. Dereferencing the end iterator is undefined behavior, which is sufficient to find the problem in sanitizer builds.

398

Please don't add code after llvm_unreachable. We don't do defensive programming in LLVM/Clang.

clang/lib/Sema/SemaDecl.cpp

12747

What I meant is that the relevant logic (that is violating the expectation, and should be eventually fixed) is in attachCommentsToJustParsedDecls, therefore, the FIXME belongs there, not in the caller. There isn't much a caller can do about the problem. Also, when the problem will get fixed, I doubt someone will look at the callsite for FIXMEs.

clang/test/Index/comment-redeclarations.cpp

15 ↗

(On Diff #214226)

Please add a newline.

[jkorous]

clang/lib/AST/ASTContext.cpp

345

You're right, nullptr can't be returned here.

352

You're right, nullptr can't be returned here.

370

I am not sure what you mean. "D is never null" isn't explicitly stated in the ASTContext.h header - do you mean that it's implied by the interface & doc comment of ASTContext::getRawCommentForAnyRedecl?

It seems to me that we should either

• Explicitly mention this in the header file and add an assert here.

or

• Handle nullptr value here.

396

Arguably debug builds are more common than sanitized builds so I'd rather not rely on sanitizer being used.

clang/lib/Sema/SemaDecl.cpp

12747

Sorry, I don't understand what you mean by:

relevant logic (that is violating the expectation, and should be eventually fixed)

The way I understand it is that attachCommentsToJustParsedDecls states the expectation in its doc comment

/// Requirement: All \p Decls are in the same file.

and it's up to a caller to satisfy it, which isn't something Sema::ActOnDocumentableDecls does.

Are you saying that attachCommentsToJustParsedDecls should pretend it's more generic than that and admit it's cheating by FIXME in its implementation?

[jkorous]

Addressed comments

[gribozavr]

Very nice improvement! Thank you!

[ivanmurashko]

cfe/trunk/lib/Serialization/ASTReader.cpp

9735 ↗

(On Diff #214886)

The var has never been used and can be removed

clang/include/clang/AST/ASTContext.h

public:
   }
 
   /// All comments in this translation unit.
   RawCommentList Comments;
 
   /// True if comments are already loaded from ExternalASTSource.
   mutable bool CommentsLoaded = false;
 
-  class RawCommentAndCacheFlags {
-  public:
-    enum Kind {
-      /// We searched for a comment attached to the particular declaration, but
-      /// didn't find any.
-      ///
-      /// getRaw() == 0.
-      NoCommentInDecl = 0,
-
-      /// We have found a comment attached to this particular declaration.
-      ///
-      /// getRaw() != 0.
-      FromDecl,
-
-      /// This declaration does not have an attached comment, and we have
-      /// searched the redeclaration chain.
-      ///
-      /// If getRaw() == 0, the whole redeclaration chain does not have any
-      /// comments.
-      ///
-      /// If getRaw() != 0, it is a comment propagated from other
-      /// redeclaration.
-      FromRedecl
-    };
-
-    Kind getKind() const LLVM_READONLY {
-      return Data.getInt();
-    }
-
-    void setKind(Kind K) {
-      Data.setInt(K);
-    }
-
-    const RawComment *getRaw() const LLVM_READONLY {
-      return Data.getPointer();
-    }
-
-    void setRaw(const RawComment *RC) {
-      Data.setPointer(RC);
-    }
-
-    const Decl *getOriginalDecl() const LLVM_READONLY {
-      return OriginalDecl;
-    }
-
-    void setOriginalDecl(const Decl *Orig) {
-      OriginalDecl = Orig;
-    }
-
-  private:
-    llvm::PointerIntPair<const RawComment *, 2, Kind> Data;
-    const Decl *OriginalDecl;
-  };
+  /// Mapping from declaration to directly attached comment.
+  /// Raw comments are owned by Comments list.  This mapping is populated
+  /// lazily.
+  mutable llvm::DenseMap<const Decl *, const RawComment *> DeclRawComments;
 
-  /// Mapping from declarations to comments attached to any
-  /// redeclaration.
+  /// Mapping from canonical declaration to the first redeclaration in chain
+  /// that has a comment attached.
   ///
   /// Raw comments are owned by Comments list.  This mapping is populated
   /// lazily.
-  mutable llvm::DenseMap<const Decl *, RawCommentAndCacheFlags> RedeclComments;
+  mutable llvm::DenseMap<const Decl *, const Decl *> RedeclChainComments;
+
+  /// Keeps track of redeclaration chains that don't have any comment attached.
+  /// Mapping from canonical declaration to redeclaration chain that has no
+  /// comments attached to any redeclaration. Specifically it's mapping to
+  /// the last redeclaration we've checked.
+  ///
+  /// Shall not contain declarations that have comments attached to any
+  /// redeclaration in their chain.
+  mutable llvm::DenseMap<const Decl *, const Decl *> CommentlessRedeclChains;
 
   /// Mapping from declarations to parsed comments attached to any
   /// redeclaration.
   mutable llvm::DenseMap<const Decl *, comments::FullComment *> ParsedComments;
 
+  /// Attaches \p Comment to \p OriginalD and to its redeclaration chain
+  /// represented by \p CanonicalD and removes the redeclaration chain
+  /// from the set of commentless chains.
+  ///
+  /// \p CanonicalD is canonical declaration of \p OriginalD.
+  ///
+  /// Don't do anything if a comment has already been attached to \p OriginalD
+  /// or its redeclaration chain.
+  void cacheRawCommentForDecl(
+    const Decl &OriginalD,
+    const Decl &CanonicalD,
+    const RawComment &Comment) const;
+
+  /// \returns searches \p CommentsInFile for doc comment for \p D.
+  ///
+  /// \p RepresentativeLocForDecl is used as a location for searching doc
+  /// comments. \p CommentsInFile is a mapping offset -> comment of files in the
+  /// same file where \p RepresentativeLocForDecl is.
+  RawComment *getRawCommentForDeclNoCacheImpl(
+      const Decl *D, const SourceLocation RepresentativeLocForDecl,
+      const std::map<unsigned, RawComment *> &CommentsInFile) const;
+
   /// Return the documentation comment attached to a given declaration,
   /// without looking into cache.
   RawComment *getRawCommentForDeclNoCache(const Decl *D) const;
 
 public:
   RawCommentList &getRawCommentList() {
     return Comments;
   }
 
   void addComment(const RawComment &RC) {
     assert(LangOpts.RetainCommentsFromSystemHeaders ||
            !SourceMgr.isInSystemHeader(RC.getSourceRange().getBegin()));
     Comments.addComment(RC, LangOpts.CommentOpts, BumpAlloc);
   }
 
   /// Return the documentation comment attached to a given declaration.
   /// Returns nullptr if no comment is attached.
   ///
   /// \param OriginalDecl if not nullptr, is set to declaration AST node that
   /// had the comment, if the comment we found comes from a redeclaration.
   const RawComment *
   getRawCommentForAnyRedecl(const Decl *D,
                             const Decl **OriginalDecl = nullptr) const;
 
+  /// Searches existing comments for doc comments that should be attached to \p
+  /// Decls. If any doc comment is found, it is parsed.
+  void searchForDocComments(ArrayRef<Decl *> Decls, const Preprocessor *PP);
+
   /// Return parsed documentation comment attached to a given declaration.
   /// Returns nullptr if no comment is attached.
   ///
   /// \param PP the Preprocessor used with this TU.  Could be nullptr if
   /// preprocessor is not available.
   comments::FullComment *getCommentForDecl(const Decl *D,
                                            const Preprocessor *PP) const;
 

clang/include/clang/AST/RawCommentList.h

 //===--- RawCommentList.h - Classes for processing raw comments -*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_AST_RAWCOMMENTLIST_H
 #define LLVM_CLANG_AST_RAWCOMMENTLIST_H
 
 #include "clang/Basic/CommentOptions.h"
+#include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include <map>
 
 namespace clang {
 
 class ASTContext;
 class ASTReader;
 class Decl;
 class Preprocessor;
 

 /// sorted in order of appearance in the translation unit.
 class RawCommentList {
 public:
   RawCommentList(SourceManager &SourceMgr) : SourceMgr(SourceMgr) {}
 
   void addComment(const RawComment &RC, const CommentOptions &CommentOpts,
                   llvm::BumpPtrAllocator &Allocator);
 
-  ArrayRef<RawComment *> getComments() const {
-    return Comments;
-  }
+  /// \returns nullptr in case there are no comments in in \p File.
+  const std::map<unsigned, RawComment *> *getCommentsInFile(FileID File) const;
+
+  bool empty() const;
+
+  unsigned getCommentBeginLine(RawComment *C, FileID File,
+                               unsigned Offset) const;
+  unsigned getCommentEndOffset(RawComment *C) const;
 
 private:
   SourceManager &SourceMgr;
-  std::vector<RawComment *> Comments;
-
-  void addDeserializedComments(ArrayRef<RawComment *> DeserializedComments);
+  // mapping: FileId -> comment begin offset -> comment
+  llvm::DenseMap<FileID, std::map<unsigned, RawComment *>> OrderedComments;
+  mutable llvm::DenseMap<RawComment *, unsigned> CommentBeginLine;
+  mutable llvm::DenseMap<RawComment *, unsigned> CommentEndOffset;
+
+  /// Adds comments that were deserialized from an AST file.
+  /// \p DeserializedComments is a mapping: FileId -> comment begin offset
+  /// -> comment
+  void addDeserializedComments(
+      const llvm::DenseMap<FileID, std::map<unsigned, RawComment *>>
+          &DeserializedComments);
 
   friend class ASTReader;
+  friend class ASTWriter;
 };
 
 } // end namespace clang
 
 #endif

clang/lib/AST/ASTContext.cpp

 #include <utility>
 
 using namespace clang;
 
 enum FloatingRank {
   Float16Rank, HalfRank, FloatRank, DoubleRank, LongDoubleRank, Float128Rank
 };
 
-RawComment *ASTContext::getRawCommentForDeclNoCache(const Decl *D) const {
+/// \returns location that is relevant when searching for Doc comments related
+/// to \p D.
+static SourceLocation getDeclLocForCommentSearch(const Decl *D,
+                                                 SourceManager &SourceMgr) {
   assert(D);
 
-  // If we already tried to load comments but there are none,
-  // we won't find anything.
-  if (CommentsLoaded && Comments.getComments().empty())
-    return nullptr;
-
   // User can not attach documentation to implicit declarations.
   if (D->isImplicit())
-    return nullptr;
+    return {};
 
   // User can not attach documentation to implicit instantiations.
   if (const auto *FD = dyn_cast<FunctionDecl>(D)) {
     if (FD->getTemplateSpecializationKind() == TSK_ImplicitInstantiation)
-      return nullptr;
+      return {};
   }
 
   if (const auto *VD = dyn_cast<VarDecl>(D)) {
     if (VD->isStaticDataMember() &&
         VD->getTemplateSpecializationKind() == TSK_ImplicitInstantiation)
-      return nullptr;
+      return {};
   }
 
   if (const auto *CRD = dyn_cast<CXXRecordDecl>(D)) {
     if (CRD->getTemplateSpecializationKind() == TSK_ImplicitInstantiation)
-      return nullptr;
+      return {};
   }
 
   if (const auto *CTSD = dyn_cast<ClassTemplateSpecializationDecl>(D)) {
     TemplateSpecializationKind TSK = CTSD->getSpecializationKind();
     if (TSK == TSK_ImplicitInstantiation ||
         TSK == TSK_Undeclared)
-      return nullptr;
+      return {};
   }
 
   if (const auto *ED = dyn_cast<EnumDecl>(D)) {
     if (ED->getTemplateSpecializationKind() == TSK_ImplicitInstantiation)
-      return nullptr;
+      return {};
   }
   if (const auto *TD = dyn_cast<TagDecl>(D)) {
     // When tag declaration (but not definition!) is part of the
     // decl-specifier-seq of some other declaration, it doesn't get comment
     if (TD->isEmbeddedInDeclarator() && !TD->isCompleteDefinition())
-      return nullptr;
+      return {};
   }
   // TODO: handle comments for function parameters properly.
   if (isa<ParmVarDecl>(D))
-    return nullptr;
+    return {};
 
   // TODO: we could look up template parameter documentation in the template
   // documentation.
   if (isa<TemplateTypeParmDecl>(D) ||
       isa<NonTypeTemplateParmDecl>(D) ||
       isa<TemplateTemplateParmDecl>(D))
-    return nullptr;
+    return {};
 
   // Find declaration location.
   // For Objective-C declarations we generally don't expect to have multiple
   // declarators, thus use declaration starting location as the "declaration
   // location".
   // For all other declarations multiple declarators are used quite frequently,
   // so we use the location of the identifier as the "declaration location".
-  SourceLocation DeclLoc;
   if (isa<ObjCMethodDecl>(D) || isa<ObjCContainerDecl>(D) ||
       isa<ObjCPropertyDecl>(D) ||
       isa<RedeclarableTemplateDecl>(D) ||
       isa<ClassTemplateSpecializationDecl>(D))
-    DeclLoc = D->getBeginLoc();
+    return D->getBeginLoc();
   else {
-    DeclLoc = D->getLocation();
+    const SourceLocation DeclLoc = D->getLocation();
     if (DeclLoc.isMacroID()) {
       if (isa<TypedefDecl>(D)) {
         // If location of the typedef name is in a macro, it is because being
         // declared via a macro. Try using declaration's starting location as
         // the "declaration location".
-        DeclLoc = D->getBeginLoc();
+        return D->getBeginLoc();
       } else if (const auto *TD = dyn_cast<TagDecl>(D)) {
         // If location of the tag decl is inside a macro, but the spelling of
         // the tag name comes from a macro argument, it looks like a special
         // macro like NS_ENUM is being used to define the tag decl.  In that
         // case, adjust the source location to the expansion loc so that we can
         // attach the comment to the tag decl.
         if (SourceMgr.isMacroArgExpansion(DeclLoc) &&
             TD->isCompleteDefinition())
-          DeclLoc = SourceMgr.getExpansionLoc(DeclLoc);
+          return SourceMgr.getExpansionLoc(DeclLoc);
       }
     }
+    return DeclLoc;
   }
 
+  return {};
+}
+
+RawComment *ASTContext::getRawCommentForDeclNoCacheImpl(
+    const Decl *D, const SourceLocation RepresentativeLocForDecl,
+    const std::map<unsigned, RawComment *> &CommentsInTheFile) const {
   // If the declaration doesn't map directly to a location in a file, we
   // can't find the comment.
-  if (DeclLoc.isInvalid() || !DeclLoc.isFileID())
+  if (RepresentativeLocForDecl.isInvalid() ||
+      !RepresentativeLocForDecl.isFileID())
     return nullptr;
 
-  if (!CommentsLoaded && ExternalSource) {
-    ExternalSource->ReadComments();
-
-#ifndef NDEBUG
-    ArrayRef<RawComment *> RawComments = Comments.getComments();
-    assert(std::is_sorted(RawComments.begin(), RawComments.end(),
-                          BeforeThanCompare<RawComment>(SourceMgr)));
-#endif
-
-    CommentsLoaded = true;
-  }
-
-  ArrayRef<RawComment *> RawComments = Comments.getComments();
   // If there are no comments anywhere, we won't find anything.
-  if (RawComments.empty())
+  if (CommentsInTheFile.empty())
     return nullptr;
 
-  // Find the comment that occurs just after this declaration.
-  ArrayRef<RawComment *>::iterator Comment;
-  {
-    // When searching for comments during parsing, the comment we are looking
-    // for is usually among the last two comments we parsed -- check them
-    // first.
-    RawComment CommentAtDeclLoc(
-        SourceMgr, SourceRange(DeclLoc), LangOpts.CommentOpts, false);
-    BeforeThanCompare<RawComment> Compare(SourceMgr);
-    ArrayRef<RawComment *>::iterator MaybeBeforeDecl = RawComments.end() - 1;
-    bool Found = Compare(*MaybeBeforeDecl, &CommentAtDeclLoc);
-    if (!Found && RawComments.size() >= 2) {
-      MaybeBeforeDecl--;
-      Found = Compare(*MaybeBeforeDecl, &CommentAtDeclLoc);
-    }
-
-    if (Found) {
-      Comment = MaybeBeforeDecl + 1;
-      assert(Comment ==
-             llvm::lower_bound(RawComments, &CommentAtDeclLoc, Compare));
-    } else {
-      // Slow path.
-      Comment = llvm::lower_bound(RawComments, &CommentAtDeclLoc, Compare);
-    }
-  }
-
   // Decompose the location for the declaration and find the beginning of the
   // file buffer.
-  std::pair<FileID, unsigned> DeclLocDecomp = SourceMgr.getDecomposedLoc(DeclLoc);
+  const std::pair<FileID, unsigned> DeclLocDecomp =
+      SourceMgr.getDecomposedLoc(RepresentativeLocForDecl);
+
+  // Slow path.
+  auto OffsetCommentBehindDecl =
+      CommentsInTheFile.lower_bound(DeclLocDecomp.second);
 
   // First check whether we have a trailing comment.
-  if (Comment != RawComments.end() &&
-      ((*Comment)->isDocumentation() || LangOpts.CommentOpts.ParseAllComments)
-      && (*Comment)->isTrailingComment() &&
+  if (OffsetCommentBehindDecl != CommentsInTheFile.end()) {
+    RawComment *CommentBehindDecl = OffsetCommentBehindDecl->second;
+    if ((CommentBehindDecl->isDocumentation() ||
+         LangOpts.CommentOpts.ParseAllComments) &&
+        CommentBehindDecl->isTrailingComment() &&
         (isa<FieldDecl>(D) || isa<EnumConstantDecl>(D) || isa<VarDecl>(D) ||
          isa<ObjCMethodDecl>(D) || isa<ObjCPropertyDecl>(D))) {
-    std::pair<FileID, unsigned> CommentBeginDecomp
-      = SourceMgr.getDecomposedLoc((*Comment)->getSourceRange().getBegin());
+
       // Check that Doxygen trailing comment comes after the declaration, starts
       // on the same line and in the same file as the declaration.
-    if (DeclLocDecomp.first == CommentBeginDecomp.first &&
-        SourceMgr.getLineNumber(DeclLocDecomp.first, DeclLocDecomp.second)
-          == SourceMgr.getLineNumber(CommentBeginDecomp.first,
-                                     CommentBeginDecomp.second)) {
-      (**Comment).setAttached();
-      return *Comment;
+      if (SourceMgr.getLineNumber(DeclLocDecomp.first, DeclLocDecomp.second) ==
+          Comments.getCommentBeginLine(CommentBehindDecl, DeclLocDecomp.first,
+                                       OffsetCommentBehindDecl->first)) {
+        return CommentBehindDecl;
+      }
     }
   }
 
   // The comment just after the declaration was not a trailing comment.
   // Let's look at the previous comment.
-  if (Comment == RawComments.begin())
+  if (OffsetCommentBehindDecl == CommentsInTheFile.begin())
     return nullptr;
-  --Comment;
+
+  auto OffsetCommentBeforeDecl = --OffsetCommentBehindDecl;
+  RawComment *CommentBeforeDecl = OffsetCommentBeforeDecl->second;
 
   // Check that we actually have a non-member Doxygen comment.
-  if (!((*Comment)->isDocumentation() ||
+  if (!(CommentBeforeDecl->isDocumentation() ||
         LangOpts.CommentOpts.ParseAllComments) ||
-      (*Comment)->isTrailingComment())
+      CommentBeforeDecl->isTrailingComment())
     return nullptr;
 
   // Decompose the end of the comment.
-  std::pair<FileID, unsigned> CommentEndDecomp
-    = SourceMgr.getDecomposedLoc((*Comment)->getSourceRange().getEnd());
-
-  // If the comment and the declaration aren't in the same file, then they
-  // aren't related.
-  if (DeclLocDecomp.first != CommentEndDecomp.first)
-    return nullptr;
+  const unsigned CommentEndOffset =
+      Comments.getCommentEndOffset(CommentBeforeDecl);
 
   // Get the corresponding buffer.
   bool Invalid = false;
   const char *Buffer = SourceMgr.getBufferData(DeclLocDecomp.first,
                                                &Invalid).data();
   if (Invalid)
     return nullptr;
 
   // Extract text between the comment and declaration.
-  StringRef Text(Buffer + CommentEndDecomp.second,
-                 DeclLocDecomp.second - CommentEndDecomp.second);
+  StringRef Text(Buffer + CommentEndOffset,
+                 DeclLocDecomp.second - CommentEndOffset);
 
   // There should be no other declarations or preprocessor directives between
   // comment and declaration.
   if (Text.find_first_of(";{}#@") != StringRef::npos)
     return nullptr;
 
-  (**Comment).setAttached();
-  return *Comment;
+  return CommentBeforeDecl;
+}
+
+RawComment *ASTContext::getRawCommentForDeclNoCache(const Decl *D) const {
+  const SourceLocation DeclLoc =
+      getDeclLocForCommentSearch(D, SourceMgr);
+
+  // If the declaration doesn't map directly to a location in a file, we
+  // can't find the comment.
+  if (DeclLoc.isInvalid() || !DeclLoc.isFileID())
+    return nullptr;
+
+  if (ExternalSource && !CommentsLoaded) {
+    ExternalSource->ReadComments();
+    CommentsLoaded = true;
+  }
+
+  if (Comments.empty())
+    return nullptr;
+
+  const FileID File = SourceMgr.getDecomposedLoc(DeclLoc).first;
+  const auto CommentsInThisFile = Comments.getCommentsInFile(File);
+  if (!CommentsInThisFile || CommentsInThisFile->empty())
+    return nullptr;
+
+  return getRawCommentForDeclNoCacheImpl(D, DeclLoc,
+                                         *CommentsInThisFile);
 }
 
 /// If we have a 'templated' declaration for a template, adjust 'D' to
 /// refer to the actual template.
 /// If we have an implicit instantiation, adjust 'D' to refer to template.
 static const Decl *adjustDeclToTemplate(const Decl *D) {
   if (const auto *FD = dyn_cast<FunctionDecl>(D)) {
     // Is this function declaration part of a function template?
if (const auto *CTSD = dyn_cast<ClassTemplateSpecializationDecl>(CRD)) {
       if (CTSD->getSpecializationKind() != TSK_ImplicitInstantiation)
         return D;
       llvm::PointerUnion<ClassTemplateDecl *,
                          ClassTemplatePartialSpecializationDecl *>
           PU = CTSD->getSpecializedTemplateOrPartial();
       return PU.is<ClassTemplateDecl*>() ?
           static_cast<const Decl*>(PU.get<ClassTemplateDecl *>()) :
           static_cast<const Decl*>(
               PU.get<ClassTemplatePartialSpecializationDecl *>());
     }
 
     // Class is instantiated from a member definition of a class template?
     if (const MemberSpecializationInfo *Info =
                    CRD->getMemberSpecializationInfo())
       return Info->getInstantiatedFrom();
 
     return D;
   }
   if (const auto *ED = dyn_cast<EnumDecl>(D)) {
     // Enum is instantiated from a member definition of a class template?
     if (const EnumDecl *MemberDecl = ED->getInstantiatedFromMemberEnum())
       return MemberDecl;
 
     return D;
   }
   // FIXME: Adjust alias templates?
   return D;
 }
 
 const RawComment *ASTContext::getRawCommentForAnyRedecl(
                                                 const Decl *D,
                                                 const Decl **OriginalDecl) const {
   D = adjustDeclToTemplate(D);
+  if (!D) {
+    if (OriginalDecl)
+      OriginalDecl = nullptr;
+    return nullptr;
+  }
 
-  // Check whether we have cached a comment for this declaration already.
+  // Any comment directly attached to D?
   {
-    llvm::DenseMap<const Decl *, RawCommentAndCacheFlags>::iterator Pos =
-        RedeclComments.find(D);
-    if (Pos != RedeclComments.end()) {
-      const RawCommentAndCacheFlags &Raw = Pos->second;
-      if (Raw.getKind() != RawCommentAndCacheFlags::NoCommentInDecl) {
-        if (OriginalDecl)
-          *OriginalDecl = Raw.getOriginalDecl();
-        return Raw.getRaw();
-      }
+    auto DeclComment = DeclRawComments.find(D);
+    if (DeclComment != DeclRawComments.end()) {
+        *OriginalDecl = D;
+      return DeclComment->second;
     }
   }
 
-  // Search for comments attached to declarations in the redeclaration chain.
-  const RawComment *RC = nullptr;
-  const Decl *OriginalDeclForRC = nullptr;
-  for (auto I : D->redecls()) {
-    llvm::DenseMap<const Decl *, RawCommentAndCacheFlags>::iterator Pos =
-        RedeclComments.find(I);
-    if (Pos != RedeclComments.end()) {
-      const RawCommentAndCacheFlags &Raw = Pos->second;
-      if (Raw.getKind() != RawCommentAndCacheFlags::NoCommentInDecl) {
-        RC = Raw.getRaw();
-        OriginalDeclForRC = Raw.getOriginalDecl();
-        break;
-      }
-    } else {
-      RC = getRawCommentForDeclNoCache(I);
-      OriginalDeclForRC = I;
-      RawCommentAndCacheFlags Raw;
-      if (RC) {
-        // Call order swapped to work around ICE in VS2015 RTM (Release Win32)
-        // https://connect.microsoft.com/VisualStudio/feedback/details/1741530
-        Raw.setKind(RawCommentAndCacheFlags::FromDecl);
-        Raw.setRaw(RC);
-      } else
-        Raw.setKind(RawCommentAndCacheFlags::NoCommentInDecl);
-      Raw.setOriginalDecl(I);
-      RedeclComments[I] = Raw;
-      if (RC)
-        break;
+  // Any comment attached to any redeclaration of D?
+  const Decl* CanonicalD = D->getCanonicalDecl();
+  if (!CanonicalD)
+    return nullptr;
+
+  {
+    auto RedeclComment = RedeclChainComments.find(CanonicalD);
+    if (RedeclComment != RedeclChainComments.end()) {
+      if (OriginalDecl)
+        *OriginalDecl = RedeclComment->second;
+      return DeclRawComments[RedeclComment->second];
     }
   }
 
-  // If we found a comment, it should be a documentation comment.
-  assert(!RC || RC->isDocumentation() || LangOpts.CommentOpts.ParseAllComments);
+  // Any redeclarations of D that we haven't checked for comments yet?
+  // We can't use DenseMap::iterator directly since it'd get invalid.
+  auto LastCheckedRedecl = [this, CanonicalD]() -> const Decl * {
+    auto LookupRes = CommentlessRedeclChains.find(CanonicalD);
+    if (LookupRes != CommentlessRedeclChains.end())
+      return LookupRes->second;
+    return nullptr;
+  }();
 
+  for (const auto Redecl : D->redecls()) {
+    assert(Redecl);
+    // Skip all redeclarations that have've been checked previously.
+    if (LastCheckedRedecl) {
+      if (LastCheckedRedecl == Redecl) {
+        LastCheckedRedecl = nullptr;
+      }
+      continue;
+    }
+    const RawComment* RedeclComment = getRawCommentForDeclNoCache(Redecl);
+    if (RedeclComment) {
+      cacheRawCommentForDecl(
+        *Redecl,
+        *CanonicalD,
+        *RedeclComment
+      );
       if (OriginalDecl)
-    *OriginalDecl = OriginalDeclForRC;
-
-  // Update cache for every declaration in the redeclaration chain.
-  RawCommentAndCacheFlags Raw;
-  Raw.setRaw(RC);
-  Raw.setKind(RawCommentAndCacheFlags::FromRedecl);
-  Raw.setOriginalDecl(OriginalDeclForRC);
+        *OriginalDecl = Redecl;
+      return RedeclComment;
+    }
+    CommentlessRedeclChains[CanonicalD] = Redecl;
+  }
 
-  for (auto I : D->redecls()) {
-    RawCommentAndCacheFlags &R = RedeclComments[I];
-    if (R.getKind() == RawCommentAndCacheFlags::NoCommentInDecl)
-      R = Raw;
-  }
+  if (OriginalDecl)
+    *OriginalDecl = nullptr;
+  return nullptr;
 
-  return RC;
+void ASTContext::cacheRawCommentForDecl(
+    const Decl &OriginalD,
+    const Decl &CanonicalD,
+    const RawComment &Comment) const {
+  assert(Comment.isDocumentation() ||
+         LangOpts.CommentOpts.ParseAllComments);
+  DeclRawComments.try_emplace(&OriginalD, &Comment);
+  RedeclChainComments.try_emplace(&CanonicalD, &OriginalD);
+  CommentlessRedeclChains.erase(&CanonicalD);
 }
 
 static void addRedeclaredMethods(const ObjCMethodDecl *ObjCMethod,
                    SmallVectorImpl<const NamedDecl *> &Redeclared) {
   const DeclContext *DC = ObjCMethod->getDeclContext();
   if (const auto *IMD = dyn_cast<ObjCImplDecl>(DC)) {
     const ObjCInterfaceDecl *ID = IMD->getClassInterface();
     if (!ID)
       return;
     // Add redeclared method here.
     for (const auto *Ext : ID->known_extensions()) {
       if (ObjCMethodDecl *RedeclaredMethod =
             Ext->getMethod(ObjCMethod->getSelector(),
                                   ObjCMethod->isInstanceMethod()))
         Redeclared.push_back(RedeclaredMethod);
     }
   }
 }
 
+void ASTContext::searchForDocComments(ArrayRef<Decl *> Decls,
+                                      const Preprocessor *PP) {
+  if (Comments.empty() || Decls.empty())
+    return;
+
+  const FileID File =
+      SourceMgr.getDecomposedLoc((*Decls.begin())->getLocation()).first;
+  const auto CommentsInThisFile = getRawCommentList().getCommentsInFile(File);
+  if (!CommentsInThisFile || CommentsInThisFile->empty())
+    return;
+
+  for (const Decl *D : Decls) {
+    assert(D);
+    if (D->isInvalidDecl())
+      continue;
+
+    D = adjustDeclToTemplate(D);
+
+    const SourceLocation DeclLoc = getDeclLocForCommentSearch(D, SourceMgr);
+
+    if (DeclLoc.isInvalid() || !DeclLoc.isFileID())
+      continue;
+
+    if (DeclRawComments.count(D) > 0)
+      continue;
+
+    RawComment *const DocComment =
+        getRawCommentForDeclNoCacheImpl(D, DeclLoc, *CommentsInThisFile);
+    if (DocComment) {
+      if (const Decl *CanonicalD = D->getCanonicalDecl()) {
+        cacheRawCommentForDecl(*D, *CanonicalD, *DocComment);
+        comments::FullComment *FC = DocComment->parse(*this, PP, D);
+        ParsedComments[CanonicalD] = FC;
+      }
+    }
+  }
+}
+
 comments::FullComment *ASTContext::cloneFullComment(comments::FullComment *FC,
                                                     const Decl *D) const {
   auto *ThisDeclInfo = new (*this) comments::DeclInfo;
   ThisDeclInfo->CommentDecl = D;
   ThisDeclInfo->IsFilled = false;
   ThisDeclInfo->fill();
   ThisDeclInfo->CommentDecl = FC->getDecl();
   if (!ThisDeclInfo->TemplateParameters)
     ThisDeclInfo->TemplateParameters = FC->getDeclInfo()->TemplateParameters;
   comments::FullComment *CFC =
     new (*this) comments::FullComment(FC->getBlocks(),
if (Pos != ParsedComments.end()) {
     if (Canonical != D) {
       comments::FullComment *FC = Pos->second;
       comments::FullComment *CFC = cloneFullComment(FC, D);
       return CFC;
     }
     return Pos->second;
   }
 
-  const Decl *OriginalDecl;
+  const Decl *OriginalDecl = nullptr;
 
   const RawComment *RC = getRawCommentForAnyRedecl(D, &OriginalDecl);
   if (!RC) {
     if (isa<ObjCMethodDecl>(D) || isa<FunctionDecl>(D)) {
       SmallVector<const NamedDecl*, 8> Overridden;
       const auto *OMD = dyn_cast<ObjCMethodDecl>(D);
       if (OMD && OMD->isPropertyAccessor())
         if (const ObjCPropertyDecl *PDecl = OMD->findPropertyDecl())
if (!RC) {
     }
     return nullptr;
   }
 
   // If the RawComment was attached to other redeclaration of this Decl, we
   // should parse the comment in context of that other Decl.  This is important
   // because comments can contain references to parameter names which can be
   // different across redeclarations.
-  if (D != OriginalDecl)
+  if (D != OriginalDecl && OriginalDecl)
     return getCommentForDecl(OriginalDecl, PP);
 
   comments::FullComment *FC = RC->parse(*this, PP, D);
   ParsedComments[Canonical] = FC;
   return FC;
 }
 
 void

clang/lib/AST/RawCommentList.cpp

 }
 
 void RawCommentList::addComment(const RawComment &RC,
                                 const CommentOptions &CommentOpts,
                                 llvm::BumpPtrAllocator &Allocator) {
   if (RC.isInvalid())
     return;
 
-  // Check if the comments are not in source order.
-  while (!Comments.empty() &&
-         !SourceMgr.isBeforeInTranslationUnit(Comments.back()->getBeginLoc(),
-                                              RC.getBeginLoc())) {
-    // If they are, just pop a few last comments that don't fit.
-    // This happens if an \#include directive contains comments.
-    Comments.pop_back();
-  }
-
   // Ordinary comments are not interesting for us.
   if (RC.isOrdinary() && !CommentOpts.ParseAllComments)
     return;
 
+  std::pair<FileID, unsigned> Loc =
+      SourceMgr.getDecomposedLoc(RC.getBeginLoc());
+
+  const FileID CommentFile = Loc.first;
+  const unsigned CommentOffset = Loc.second;
+
   // If this is the first Doxygen comment, save it (because there isn't
   // anything to merge it with).
-  if (Comments.empty()) {
-    Comments.push_back(new (Allocator) RawComment(RC));
+  if (OrderedComments[CommentFile].empty()) {
+    OrderedComments[CommentFile][CommentOffset] =
+        new (Allocator) RawComment(RC);
     return;
   }
 
-  const RawComment &C1 = *Comments.back();
+  const RawComment &C1 = *OrderedComments[CommentFile].rbegin()->second;
   const RawComment &C2 = RC;
 
   // Merge comments only if there is only whitespace between them.
   // Can't merge trailing and non-trailing comments unless the second is
   // non-trailing ordinary in the same column, as in the case:
   //   int x; // documents x
   //          // more text
   // versus:
   //   int x; // documents x
   //   int y; // documents y
   // or:
   //   int x; // documents x
   //   // documents y
   //   int y;
   // Merge comments if they are on same or consecutive lines.
   if ((C1.isTrailingComment() == C2.isTrailingComment() ||
        (C1.isTrailingComment() && !C2.isTrailingComment() &&
         isOrdinaryKind(C2.getKind()) &&
         commentsStartOnSameColumn(SourceMgr, C1, C2))) &&
       onlyWhitespaceBetween(SourceMgr, C1.getEndLoc(), C2.getBeginLoc(),
                             /*MaxNewlinesAllowed=*/1)) {
     SourceRange MergedRange(C1.getBeginLoc(), C2.getEndLoc());
-    *Comments.back() = RawComment(SourceMgr, MergedRange, CommentOpts, true);
+    *OrderedComments[CommentFile].rbegin()->second =
+        RawComment(SourceMgr, MergedRange, CommentOpts, true);
   } else {
-    Comments.push_back(new (Allocator) RawComment(RC));
+    OrderedComments[CommentFile][CommentOffset] =
+        new (Allocator) RawComment(RC);
   }
 }
 
-void RawCommentList::addDeserializedComments(ArrayRef<RawComment *> DeserializedComments) {
-  std::vector<RawComment *> MergedComments;
-  MergedComments.reserve(Comments.size() + DeserializedComments.size());
+const std::map<unsigned, RawComment *> *
+RawCommentList::getCommentsInFile(FileID File) const {
+  auto CommentsInFile = OrderedComments.find(File);
+  if (CommentsInFile == OrderedComments.end())
+    return nullptr;
+
+  return &CommentsInFile->second;
+}
+
+bool RawCommentList::empty() const { return OrderedComments.empty(); }
+
+unsigned RawCommentList::getCommentBeginLine(RawComment *C, FileID File,
+                                             unsigned Offset) const {
+  auto Cached = CommentBeginLine.find(C);
+  if (Cached != CommentBeginLine.end())
+    return Cached->second;
+  const unsigned Line = SourceMgr.getLineNumber(File, Offset);
+  CommentBeginLine[C] = Line;
+  return Line;
+}
+
+unsigned RawCommentList::getCommentEndOffset(RawComment *C) const {
+  auto Cached = CommentEndOffset.find(C);
+  if (Cached != CommentEndOffset.end())
+    return Cached->second;
+  const unsigned Offset =
+      SourceMgr.getDecomposedLoc(C->getSourceRange().getEnd()).second;
+  CommentEndOffset[C] = Offset;
+  return Offset;
+}
 
-  std::merge(Comments.begin(), Comments.end(),
-             DeserializedComments.begin(), DeserializedComments.end(),
-             std::back_inserter(MergedComments),
-             BeforeThanCompare<RawComment>(SourceMgr));
-  std::swap(Comments, MergedComments);
+void RawCommentList::addDeserializedComments(
+    const llvm::DenseMap<FileID, std::map<unsigned, RawComment *>>
+        &DeserializedComments) {
+  for (const auto &F : DeserializedComments)
+    OrderedComments[F.first].insert(F.second.begin(), F.second.end());
 }
 
 std::string RawComment::getFormattedText(const SourceManager &SourceMgr,
                                          DiagnosticsEngine &Diags) const {
   llvm::StringRef CommentText = getRawText(SourceMgr);
   if (CommentText.empty())
     return "";
 

clang/lib/Sema/SemaDecl.cpp

if (Group.size() >= 2) {
     //   'struct S *pS;'
     // FinalizeDeclaratorGroup adds these as separate declarations.
     Decl *MaybeTagDecl = Group[0];
     if (MaybeTagDecl && isa<TagDecl>(MaybeTagDecl)) {
       Group = Group.slice(1);
     }
   }
 
   // See if there are any new comments that are not attached to a decl.
-  ArrayRef<RawComment *> Comments = Context.getRawCommentList().getComments();
-  if (!Comments.empty() &&
-      !Comments.back()->isAttached()) {
-    // There is at least one comment that not attached to a decl.
+  // The location doesn't have to be precise - we care only about the file.
+  // FIMXE: We assume every Decl in the group is in the same file.
+  // This is false when preprocessor constructs the group from decls in
+  // different files (e. g. macros or #include).
+  const FileID File =
+      SourceMgr.getDecomposedLoc((*Group.begin())->getLocation()).first;
+  auto CommentsInThisFile = Context.getRawCommentList().getCommentsInFile(File);
+  if (CommentsInThisFile && !CommentsInThisFile->empty() &&
+      !CommentsInThisFile->rbegin()->second->isAttached()) {
+    // There is at least one comment not attached to a decl.
     // Maybe it should be attached to one of these decls?
     //
     // Note that this way we pick up not only comments that precede the
     // declaration, but also comments that *follow* the declaration -- thanks to
     // the lookahead in the lexer: we've consumed the semicolon and looked
     // ahead through comments.
-    for (unsigned i = 0, e = Group.size(); i != e; ++i)
-      Context.getCommentForDecl(Group[i], &PP);
+    Context.searchForDocComments(Group, &getPreprocessor());
   }
 }
 
 /// Common checks for a parameter-declaration that should apply to both function
 /// parameters and non-type template parameters.
 void Sema::CheckFunctionOrTemplateParamDeclarator(Scope *S, Declarator &D) {
   // Check that there are no default arguments inside the type of this
   // parameter.
   if (getLangOpts().CPlusPlus)

clang/lib/Serialization/ASTReader.cpp

while (true) {
         bool IsAlmostTrailingComment = Record[Idx++];
         Comments.push_back(new (Context) RawComment(
             SR, Kind, IsTrailingComment, IsAlmostTrailingComment));
         break;
       }
       }
     }
   NextCursor:
-    // De-serialized SourceLocations get negative FileIDs for other modules,
-    // potentially invalidating the original order. Sort it again.
-    llvm::sort(Comments, BeforeThanCompare<RawComment>(SourceMgr));
-    Context.Comments.addDeserializedComments(Comments);
+    llvm::DenseMap<FileID, std::map<unsigned, RawComment *>>
+        FileToOffsetToComment;
+    for (RawComment *C : Comments) {
+      SourceLocation CommentLoc = C->getBeginLoc();
+      if (CommentLoc.isValid()) {
+        std::pair<FileID, unsigned> Loc =
+            SourceMgr.getDecomposedLoc(CommentLoc);
+        if (Loc.first.isValid())
+          FileToOffsetToComment[Loc.first][Loc.second] = C;
+      }
+    }
+
+    Context.Comments.addDeserializedComments(FileToOffsetToComment);
   }
 }
 
 void ASTReader::visitInputFiles(serialization::ModuleFile &MF,
                                 bool IncludeSystem, bool Complain,
                     llvm::function_ref<void(const serialization::InputFile &IF,
                                             bool isSystem)> Visitor) {
   unsigned NumUserInputs = MF.NumUserInputFiles;

clang/lib/Serialization/ASTWriter.cpp

void ASTWriter::WriteFileDeclIDsMap() {
   Stream.EmitRecordWithBlob(AbbrevCode, Record, bytes(FileGroupedDeclIDs));
 }
 
 void ASTWriter::WriteComments() {
   Stream.EnterSubblock(COMMENTS_BLOCK_ID, 3);
   auto _ = llvm::make_scope_exit([this] { Stream.ExitBlock(); });
   if (!PP->getPreprocessorOpts().WriteCommentListToPCH)
     return;
-  ArrayRef<RawComment *> RawComments = Context->Comments.getComments();
   RecordData Record;
-  for (const auto *I : RawComments) {
+  for (const auto &FO : Context->Comments.OrderedComments) {
+    for (const auto &OC : FO.second) {
+      const RawComment *I = OC.second;
       Record.clear();
       AddSourceRange(I->getSourceRange(), Record);
       Record.push_back(I->getKind());
       Record.push_back(I->isTrailingComment());
       Record.push_back(I->isAlmostTrailingComment());
       Stream.EmitRecord(COMMENTS_RAW_COMMENT, Record);
     }
   }
+}
 
 //===----------------------------------------------------------------------===//
 // Global Method Pool and Selector Serialization
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 // Trait used for the on-disk hash table used in the method pool.