[clang][ExtractAPI] Add support for Objective-C categories

Differential Revision: https://reviews.llvm.org/D152770

clang/include/clang/ExtractAPI/API.h

@@ -80,6 +80,7 @@ struct APIRecord {
     RK_ObjCInstanceMethod,
     RK_ObjCInterface,
     RK_ObjCCategory,
+    RK_ObjCCategoryModule,
     RK_ObjCProtocol,
     RK_MacroDefinition,
     RK_Typedef,
@@ -153,6 +154,9 @@ struct APIRecord {
         Comment(Comment), Declaration(Declaration), SubHeading(SubHeading),
         IsFromSystemHeader(IsFromSystemHeader), Kind(Kind) {}
 
+  APIRecord(RecordKind Kind, StringRef USR, StringRef Name)
+      : USR(USR), Name(Name), Kind(Kind) {}
+
   // Pure virtual destructor to make APIRecord abstract
   virtual ~APIRecord() = 0;
 };
@@ -643,6 +647,8 @@ struct CXXClassRecord : APIRecord {
 /// This holds information associated with Objective-C categories.
 struct ObjCCategoryRecord : ObjCContainerRecord {
   SymbolReference Interface;
+  /// Determine whether the Category is derived from external class interface.
+  bool IsFromExternalModule = false;
 
   ObjCCategoryRecord(StringRef USR, StringRef Name, PresumedLoc Loc,
                      AvailabilitySet Availabilities, const DocComment &Comment,
@@ -895,7 +901,7 @@ class APISet {
                   AvailabilitySet Availability, const DocComment &Comment,
                   DeclarationFragments Declaration,
                   DeclarationFragments SubHeading, SymbolReference Interface,
-                  bool IsFromSystemHeader);
+                  bool IsFromSystemHeader, bool IsFromExternalModule);
 
   /// Create and add an Objective-C interface record into the API set.
   ///

clang/include/clang/ExtractAPI/ExtractAPIVisitor.h

@@ -578,9 +578,17 @@ bool ExtractAPIVisitorBase<Derived>::VisitObjCCategoryDecl(
   SymbolReference Interface(InterfaceDecl->getName(),
                             API.recordUSR(InterfaceDecl));
 
+  bool IsFromExternalModule = true;
+  for (const auto &Interface : API.getObjCInterfaces()) {
+    if (InterfaceDecl->getName() == Interface.second.get()->Name) {
+      IsFromExternalModule = false;
+      break;
+    }
+  }
+
   ObjCCategoryRecord *ObjCCategoryRecord = API.addObjCCategory(
       Name, USR, Loc, AvailabilitySet(Decl), Comment, Declaration, SubHeading,
-      Interface, isInSystemHeader(Decl));
+      Interface, isInSystemHeader(Decl), IsFromExternalModule);
 
   getDerivedExtractAPIVisitor().recordObjCMethods(ObjCCategoryRecord,
                                                   Decl->methods());

clang/include/clang/ExtractAPI/Serialization/SerializerBase.h

@@ -39,6 +39,8 @@ template <typename Derived> class APISetVisitor {
 
     getDerived()->traverseObjCProtocols();
 
+    getDerived()->traverseObjCCategories();
+
     getDerived()->traverseMacroDefinitionRecords();
 
     getDerived()->traverseTypedefRecords();
@@ -84,6 +86,11 @@ template <typename Derived> class APISetVisitor {
       getDerived()->visitObjCContainerRecord(*Protocol.second);
   }
 
+  void traverseObjCCategories() {
+    for (const auto &Category : API.getObjCCategories())
+      getDerived()->visitObjCCategoryRecord(*Category.second);
+  }
+
   void traverseMacroDefinitionRecords() {
     for (const auto &Macro : API.getMacros())
       getDerived()->visitMacroDefinitionRecord(*Macro.second);
@@ -113,6 +120,9 @@ template <typename Derived> class APISetVisitor {
   /// Visit an Objective-C container record.
   void visitObjCContainerRecord(const ObjCContainerRecord &Record){};
 
+  /// Visit an Objective-C category record.
+  void visitObjCCategoryRecord(const ObjCCategoryRecord &Record){};
+
   /// Visit a macro definition record.
   void visitMacroDefinitionRecord(const MacroDefinitionRecord &Record){};
 

clang/include/clang/ExtractAPI/Serialization/SymbolGraphSerializer.h

@@ -21,6 +21,7 @@
 #include "clang/ExtractAPI/APIIgnoresList.h"
 #include "clang/ExtractAPI/Serialization/SerializerBase.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/VersionTuple.h"
 #include "llvm/Support/raw_ostream.h"
@@ -87,6 +88,10 @@ class SymbolGraphSerializer : public APISetVisitor<SymbolGraphSerializer> {
     /// The source symbol conforms to the target symbol.
     /// For example Objective-C protocol conformances.
     ConformsTo,
+
+    /// The source symbol is an extension to the target symbol.
+    /// For example Objective-C categories extending an external type.
+    ExtensionTo,
   };
 
   /// Get the string representation of the relationship kind.
@@ -147,6 +152,8 @@ class SymbolGraphSerializer : public APISetVisitor<SymbolGraphSerializer> {
 
   SymbolGraphSerializerOption Options;
 
+  llvm::StringSet<> visitedCategories;
+
 public:
   /// Visit a global function record.
   void visitGlobalFunctionRecord(const GlobalFunctionRecord &Record);
@@ -167,6 +174,9 @@ class SymbolGraphSerializer : public APISetVisitor<SymbolGraphSerializer> {
   /// Visit an Objective-C container record.
   void visitObjCContainerRecord(const ObjCContainerRecord &Record);
 
+  /// Visit an Objective-C category record.
+  void visitObjCCategoryRecord(const ObjCCategoryRecord &Record);
+
   /// Visit a macro definition record.
   void visitMacroDefinitionRecord(const MacroDefinitionRecord &Record);
 

clang/lib/ExtractAPI/API.cpp

@@ -209,15 +209,16 @@ ObjCCategoryRecord *APISet::addObjCCategory(
     StringRef Name, StringRef USR, PresumedLoc Loc,
     AvailabilitySet Availabilities, const DocComment &Comment,
     DeclarationFragments Declaration, DeclarationFragments SubHeading,
-    SymbolReference Interface, bool IsFromSystemHeader) {
+    SymbolReference Interface, bool IsFromSystemHeader,
+    bool IsFromExternalModule) {
   // Create the category record.
   auto *Record =
       addTopLevelRecord(USRBasedLookupTable, ObjCCategories, USR, Name, Loc,
                         std::move(Availabilities), Comment, Declaration,
                         SubHeading, Interface, IsFromSystemHeader);
 
-  // If this category is extending a known interface, associate it with the
-  // ObjCInterfaceRecord.
+  Record->IsFromExternalModule = IsFromExternalModule;
+
   auto It = ObjCInterfaces.find(Interface.USR);
   if (It != ObjCInterfaces.end())
     It->second->Categories.push_back(Record);

clang/lib/ExtractAPI/Serialization/SymbolGraphSerializer.cpp

@@ -328,7 +328,13 @@ serializeDeclarationFragments(const DeclarationFragments &DF) {
 ///     Objective-C methods). Can be used as sub-headings for documentation.
 Object serializeNames(const APIRecord &Record) {
   Object Names;
-  Names["title"] = Record.Name;
+  if (auto *CategoryRecord =
+          dyn_cast_or_null<const ObjCCategoryRecord>(&Record))
+    Names["title"] =
+        (CategoryRecord->Interface.Name + " (" + Record.Name + ")").str();
+  else
+    Names["title"] = Record.Name;
+
   serializeArray(Names, "subHeading",
                  serializeDeclarationFragments(Record.SubHeading));
   DeclarationFragments NavigatorFragments;
@@ -432,9 +438,12 @@ Object serializeSymbolKind(APIRecord::RecordKind RK, Language Lang) {
     Kind["displayName"] = "Class";
     break;
   case APIRecord::RK_ObjCCategory:
-    // We don't serialize out standalone Objective-C category symbols yet.
-    llvm_unreachable("Serializing standalone Objective-C category symbols is "
-                     "not supported.");
+    Kind["identifier"] = AddLangPrefix("class.extension");
+    Kind["displayName"] = "Class Extension";
+    break;
+  case APIRecord::RK_ObjCCategoryModule:
+    Kind["identifier"] = AddLangPrefix("module.extension");
+    Kind["displayName"] = "Module Extension";
     break;
   case APIRecord::RK_ObjCProtocol:
     Kind["identifier"] = AddLangPrefix("protocol");
@@ -563,14 +572,16 @@ bool generatePathComponents(
     if (!ParentRecord)
       ParentRecord = API.findRecordForUSR(CurrentParent->ParentUSR);
 
-    // If the parent is a category then we need to pretend this belongs to the
-    // associated interface.
+    // If the parent is a category extended from internal module then we need to
+    // pretend this belongs to the associated interface.
     if (auto *CategoryRecord =
             dyn_cast_or_null<ObjCCategoryRecord>(ParentRecord)) {
-      ParentRecord = API.findRecordForUSR(CategoryRecord->Interface.USR);
-      CurrentParentComponent = PathComponent(CategoryRecord->Interface.USR,
-                                             CategoryRecord->Interface.Name,
-                                             APIRecord::RK_ObjCInterface);
+      if (!CategoryRecord->IsFromExternalModule) {
+        ParentRecord = API.findRecordForUSR(CategoryRecord->Interface.USR);
+        CurrentParentComponent = PathComponent(CategoryRecord->Interface.USR,
+                                               CategoryRecord->Interface.Name,
+                                               APIRecord::RK_ObjCInterface);
+      }
     }
 
     // The parent record doesn't exist which means the symbol shouldn't be
@@ -709,6 +720,8 @@ StringRef SymbolGraphSerializer::getRelationshipString(RelationshipKind Kind) {
     return "inheritsFrom";
   case RelationshipKind::ConformsTo:
     return "conformsTo";
+  case RelationshipKind::ExtensionTo:
+    return "extensionTo";
   }
   llvm_unreachable("Unhandled relationship kind");
 }
@@ -820,6 +833,45 @@ void SymbolGraphSerializer::visitObjCContainerRecord(
   }
 }
 
+void SymbolGraphSerializer::visitObjCCategoryRecord(
+    const ObjCCategoryRecord &Record) {
+  if (!Record.IsFromExternalModule)
+    return;
+
+  // Check if the current Category' parent has been visited before, if so skip.
+  if (!(visitedCategories.contains(Record.Interface.Name) > 0)) {
+    visitedCategories.insert(Record.Interface.Name);
+    Object Obj;
+    serializeObject(Obj, "identifier",
+                    serializeIdentifier(Record, API.getLanguage()));
+    serializeObject(Obj, "kind",
+                    serializeSymbolKind(APIRecord::RK_ObjCCategoryModule,
+                                        API.getLanguage()));
+    Obj["accessLevel"] = "public";
+    Symbols.emplace_back(std::move(Obj));
+  }
+
+  Object Relationship;
+  Relationship["source"] = Record.USR;
+  Relationship["target"] = Record.Interface.USR;
+  Relationship["targetFallback"] = Record.Interface.Name;
+  Relationship["kind"] = getRelationshipString(RelationshipKind::ExtensionTo);
+  Relationships.emplace_back(std::move(Relationship));
+
+  auto ObjCCategory = serializeAPIRecord(Record);
+
+  if (!ObjCCategory)
+    return;
+
+  Symbols.emplace_back(std::move(*ObjCCategory));
+  serializeMembers(Record, Record.Methods);
+  serializeMembers(Record, Record.Properties);
+
+  // Surface the protocols of the category to the interface.
+  for (const auto &Protocol : Record.Protocols)
+    serializeRelationship(RelationshipKind::ConformsTo, Record, Protocol);
+}
+
 void SymbolGraphSerializer::visitMacroDefinitionRecord(
     const MacroDefinitionRecord &Record) {
   auto Macro = serializeAPIRecord(Record);
@@ -858,6 +910,9 @@ void SymbolGraphSerializer::serializeSingleRecord(const APIRecord *Record) {
   case APIRecord::RK_ObjCProtocol:
     visitObjCContainerRecord(*cast<ObjCProtocolRecord>(Record));
     break;
+  case APIRecord::RK_ObjCCategory:
+    visitObjCCategoryRecord(*cast<ObjCCategoryRecord>(Record));
+    break;
   case APIRecord::RK_MacroDefinition:
     visitMacroDefinitionRecord(*cast<MacroDefinitionRecord>(Record));
     break;
@@ -926,9 +981,6 @@ SymbolGraphSerializer::serializeSingleSymbolSGF(StringRef USR,
   if (!Record)
     return {};
 
-  if (isa<ObjCCategoryRecord>(Record))
-    return {};
-
   Object Root;
   APIIgnoresList EmptyIgnores;
   SymbolGraphSerializer Serializer(API, EmptyIgnores,