1 files changed, 42 insertions, 22 deletions

diff --git a/lib/AST/RawCommentList.cpp b/lib/AST/RawCommentList.cpp
index df53b7fa10047..83e8a0b942a40 100644
--- a/lib/AST/RawCommentList.cpp
+++ b/lib/AST/RawCommentList.cpp
@@ -275,27 +275,25 @@ void RawCommentList::addComment(const RawComment &RC,
   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.
@@ -318,21 +316,43 @@ void RawCommentList::addComment(const RawComment &RC,
       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;
+}
 
-  std::merge(Comments.begin(), Comments.end(),
-             DeserializedComments.begin(), DeserializedComments.end(),
-             std::back_inserter(MergedComments),
-             BeforeThanCompare<RawComment>(SourceMgr));
-  std::swap(Comments, MergedComments);
+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::string RawComment::getFormattedText(const SourceManager &SourceMgr,