[Serialization] Serialize doc comments for extensions. Need this for rdar://25157796

We did not serialize them because getting USR for extensions is tricky (USRs are
usually for value decls). This commit starts to make up an USR for an extension by combining
the extended nominal's USR with the USR of the first value member of the extension. We use
this made-up USR to associate doc comments when (de)serializing them.

include/swift/AST/USRGeneration.h

@@ -18,6 +18,7 @@
 namespace swift {
 class AbstractStorageDecl;
 class ValueDecl;
+class ExtensionDecl;
 enum class AccessorKind;
 
 namespace ide {
@@ -31,6 +32,10 @@ bool printDeclUSR(const ValueDecl *D, raw_ostream &OS);
 bool printAccessorUSR(const AbstractStorageDecl *D, AccessorKind AccKind,
                       llvm::raw_ostream &OS);
 
+/// Prints out the extension USR for the given extension Decl.
+/// \returns true if it failed, false on success.
+bool printExtensionUSR(const ExtensionDecl *ED, raw_ostream &OS);
+
 } // namespace ide
 } // namespace swift
 

lib/AST/USRGeneration.cpp

@@ -114,3 +114,19 @@ bool ide::printAccessorUSR(const AbstractStorageDecl *D, AccessorKind AccKind,
   return false;
 }
 
+bool ide::printExtensionUSR(const ExtensionDecl *ED, raw_ostream &OS) {
+  if (ED->getExtendedType().isNull())
+    return true;
+
+  // We make up a unique usr for each extension by combining the usr of the
+  // extended type and the first value member of the extension.
+  if (printDeclUSR(ED->getExtendedType()->getAnyNominal(), OS))
+    return true;
+  for (auto D : ED->getMembers()) {
+    if (auto VD = dyn_cast<ValueDecl>(D)) {
+      return printDeclUSR(VD, OS);
+    }
+  }
+  return true;
+}
+

lib/Serialization/ModuleFile.cpp

@@ -1525,6 +1525,17 @@ Optional<CommentInfo> ModuleFile::getCommentForDecl(const Decl *D) const {
   if (D->isImplicit())
     return None;
 
+  if (auto *ED = dyn_cast<ExtensionDecl>(D)) {
+    // Compute the USR.
+    llvm::SmallString<128> USRBuffer;
+    {
+      llvm::raw_svector_ostream OS(USRBuffer);
+      if (ide::printExtensionUSR(ED, OS))
+        return None;
+    }
+     return getCommentForDeclByUSR(USRBuffer.str());
+  }
+
   auto *VD = dyn_cast<ValueDecl>(D);
   if (!VD)
     return None;

lib/Serialization/Serialization.cpp

@@ -3600,8 +3600,8 @@ class DeclGroupNameContext {
     const bool Enable;
     GroupNameCollector(bool Enable) : Enable(Enable) {}
     virtual ~GroupNameCollector() = default;
-    virtual StringRef getGroupNameInternal(const ValueDecl *VD) = 0;
-    StringRef getGroupName(const ValueDecl *VD) {
+    virtual StringRef getGroupNameInternal(const Decl *VD) = 0;
+    StringRef getGroupName(const Decl *VD) {
       return Enable ? getGroupNameInternal(VD) : StringRef(NullGroupName);
     };
   };
@@ -3609,7 +3609,7 @@ class DeclGroupNameContext {
   // FIXME: Implement better name collectors.
   struct GroupNameCollectorFromFileName : public GroupNameCollector {
     GroupNameCollectorFromFileName(bool Enable) : GroupNameCollector(Enable) {}
-    StringRef getGroupNameInternal(const ValueDecl *VD) override {
+    StringRef getGroupNameInternal(const Decl *VD) override {
       auto PathOp = VD->getDeclContext()->getParentSourceFile()->getBufferID();
       if (!PathOp.hasValue())
         return NullGroupName;
@@ -3627,7 +3627,7 @@ class DeclGroupNameContext {
     GroupNameCollectorFromJson(StringRef RecordPath, ASTContext &Ctx) :
       GroupNameCollector(!RecordPath.empty()), RecordPath(RecordPath),
       Ctx(Ctx) {}
-    StringRef getGroupNameInternal(const ValueDecl *VD) override {
+    StringRef getGroupNameInternal(const Decl *VD) override {
       // We need the file path, so there has to be a location.
       if (VD->getLoc().isInvalid())
         return NullGroupName;
@@ -3663,7 +3663,7 @@ class DeclGroupNameContext {
 public:
   DeclGroupNameContext(StringRef RecordPath, ASTContext &Ctx) :
     pNameCollector(new GroupNameCollectorFromJson(RecordPath, Ctx)) {}
-  uint32_t getGroupSequence(const ValueDecl *VD) {
+  uint32_t getGroupSequence(const Decl *VD) {
     return Map.insert(std::make_pair(pNameCollector->getGroupName(VD),
                                      Map.size())).first->second;
   }
@@ -3721,7 +3721,29 @@ static void writeDeclCommentTable(
       return StringRef(Mem, String.size());
     }
 
+    void writeDocForExtensionDecl(ExtensionDecl *ED) {
+      RawComment Raw = ED->getRawComment();
+      if (Raw.Comments.empty() && !GroupContext.isEnable())
+        return;
+      // Compute USR.
+      {
+        USRBuffer.clear();
+        llvm::raw_svector_ostream OS(USRBuffer);
+        if (ide::printExtensionUSR(ED, OS))
+          return;
+      }
+      generator.insert(copyString(USRBuffer.str()),
+                       { ED->getBriefComment(), Raw,
+                         GroupContext.getGroupSequence(ED),
+                         SourceOrder ++ });
+    }
+
     bool walkToDeclPre(Decl *D) override {
+      if (auto *ED = dyn_cast<ExtensionDecl>(D)) {
+        writeDocForExtensionDecl(ED);
+        return true;
+      }
+
       auto *VD = dyn_cast<ValueDecl>(D);
       if (!VD)
         return true;

test/IDE/comment_attach.swift

@@ -300,6 +300,7 @@ func unterminatedBlockDocComment() {}
 // CHECK-NEXT: comment_attach.swift:196:7: Var/decl_protocol_1.propertyWithGetSet RawComment=[/// propertyWithGetSet Aaa.\n]
 // CHECK-NEXT: comment_attach.swift:196:33: Func/decl_protocol_1.<getter for decl_protocol_1.propertyWithGetSet> RawComment=none
 // CHECK-NEXT: comment_attach.swift:196:37: Func/decl_protocol_1.<setter for decl_protocol_1.propertyWithGetSet> RawComment=none
+// CHECK-NEXT: comment_attach.swift:203:1: Extension/ RawComment=[/// decl_extension_1 Aaa.\n]
 // CHECK-NEXT: comment_attach.swift:207:6: Func/emptyBlockDocComment RawComment=[/***/]
 // CHECK-NEXT: comment_attach.swift:210:6: Func/weirdBlockDocComment RawComment=[/**/]
 // CHECK-NEXT: comment_attach.swift:217:6: Func/docCommentWithGybLineNumber RawComment=[/// docCommentWithGybLineNumber Aaa.\n/// Bbb.\n/// Ccc.\n]

test/Serialization/comments.swift

@@ -45,6 +45,16 @@ class first_decl_class_1 {
   func decl_func_3() {}
 }
 
+/// Comment for bar1
+extension first_decl_class_1 {
+  func bar1(){}
+}
+
+/// Comment for bar2
+extension first_decl_class_1 {
+  func bar2(){}
+}
+
 // FIRST: Class/first_decl_generic_class_1 RawComment=[/// first_decl_generic_class_1 Aaa.\n]
 // FIRST: Destructor/first_decl_generic_class_1.deinit RawComment=[/// deinit of first_decl_generic_class_1 Aaa.\n]
 // FIRST: Class/first_decl_class_1 RawComment=[/// first_decl_class_1 Aaa.\n]
@@ -53,4 +63,5 @@ class first_decl_class_1 {
 // FIRST: Func/first_decl_class_1.decl_func_3 RawComment=[/// decl_func_3 Aaa.\n/** Bbb. */]
 
 // SECOND: Class/second_decl_class_1 RawComment=[/// second_decl_class_1 Aaa.\n]
-
+// SECOND: Extension/ RawComment=[/// Comment for bar1\n] BriefComment=[Comment for bar1]
+// SECOND: Extension/ RawComment=[/// Comment for bar2\n] BriefComment=[Comment for bar2]

tools/swift-ide-test/swift-ide-test.cpp

@@ -2113,6 +2113,21 @@ public:
       OS << Decl::getKindName(VD->getKind()) << "/";
       printDeclName(VD);
 
+      OS << " ";
+      printRawComment(D->getRawComment());
+      OS << " ";
+      printBriefComment(D->getBriefComment());
+      OS << " ";
+      printDocComment(D);
+      OS << "\n";
+    } else if (D->getKind() == DeclKind::Extension) {
+      SourceLoc Loc = D->getLoc();
+      if (Loc.isValid()) {
+        auto LineAndColumn = SM.getLineAndColumn(Loc);
+        OS << getBufferIdentifier(D->getLoc())
+        << ":" << LineAndColumn.first << ":" << LineAndColumn.second << ": ";
+      }
+      OS << Decl::getKindName(D->getKind()) << "/";
       OS << " ";
       printRawComment(D->getRawComment());
       OS << " ";