[TBDGen] Add support for Objective-C Categories.

Emit Objective-C Categories for extensions that have the @objc attribute
directly (or indirectly via one of its methods, subscripts, etc) attached.

Also associate and emit all methods for that category into the API JSON file.

This fixes rdar://94734748.

lib/TBDGen/APIGen.cpp

@@ -42,12 +42,22 @@ ObjCInterfaceRecord *API::addObjCClass(StringRef name, APILinkage linkage,
   return interface;
 }
 
-void API::addObjCMethod(ObjCInterfaceRecord *cls, StringRef name, APILoc loc,
+ObjCCategoryRecord *API::addObjCCategory(StringRef name, APILinkage linkage,
+                                         APILoc loc, APIAccess access,
+                                         APIAvailability availability,
+                                         StringRef interface) {
+  auto *category = new (allocator)
+      ObjCCategoryRecord(name, linkage, loc, access, availability, interface);
+  categories.push_back(category);
+  return category;
+}
+
+void API::addObjCMethod(ObjCContainerRecord *record, StringRef name, APILoc loc,
                         APIAccess access, bool isInstanceMethod,
                         bool isOptional, APIAvailability availability) {
   auto method = new (allocator) ObjCMethodRecord(
       name, loc, access, isInstanceMethod, isOptional, availability);
-  cls->methods.push_back(method);
+  record->methods.push_back(method);
 }
 
 static void serialize(llvm::json::OStream &OS, APIAccess access) {
@@ -151,6 +161,30 @@ static void serialize(llvm::json::OStream &OS,
   });
 }
 
+static void serialize(llvm::json::OStream &OS,
+                      const ObjCCategoryRecord &record) {
+  OS.object([&]() {
+    OS.attribute("name", record.name);
+    serialize(OS, record.access);
+    serialize(OS, record.loc);
+    serialize(OS, record.linkage);
+    serialize(OS, record.availability);
+    OS.attribute("interface", record.interface);
+    OS.attributeArray("instanceMethods", [&]() {
+      for (auto &method : record.methods) {
+        if (method->isInstanceMethod)
+          serialize(OS, *method);
+      }
+    });
+    OS.attributeArray("classMethods", [&]() {
+      for (auto &method : record.methods) {
+        if (!method->isInstanceMethod)
+          serialize(OS, *method);
+      }
+    });
+  });
+}
+
 void API::writeAPIJSONFile(llvm::raw_ostream &os, bool PrettyPrint) {
   unsigned indentSize = PrettyPrint ? 2 : 0;
   llvm::json::OStream JSON(os, indentSize);
@@ -167,6 +201,11 @@ void API::writeAPIJSONFile(llvm::raw_ostream &os, bool PrettyPrint) {
       for (const auto *i : interfaces)
         serialize(JSON, *i);
     });
+    JSON.attributeArray("categories", [&]() {
+      llvm::sort(categories, sortAPIRecords);
+      for (const auto *c : categories)
+        serialize(JSON, *c);
+    });
     JSON.attribute("version", "1.0");
   });
 }

lib/TBDGen/APIGen.h

@@ -152,6 +152,16 @@ struct ObjCInterfaceRecord : ObjCContainerRecord {
         superClassName(superClassName.data(), superClassName.size()) {}
 };
 
+struct ObjCCategoryRecord : ObjCContainerRecord {
+  std::string interface;
+
+  ObjCCategoryRecord(StringRef name, APILinkage linkage, APILoc loc,
+                     APIAccess access, APIAvailability availability,
+                     StringRef interface)
+      : ObjCContainerRecord(name, linkage, loc, access, availability),
+        interface(interface.data(), interface.size()) {}
+};
+
 class API {
 public:
   API(const llvm::Triple &triple) : target(triple) {}
@@ -167,11 +177,16 @@ public:
                                     APIAvailability availability,
                                     StringRef superClassName);
 
-  void addObjCMethod(ObjCInterfaceRecord *cls, StringRef name, APILoc loc,
+  ObjCCategoryRecord *addObjCCategory(StringRef name, APILinkage linkage,
+                                      APILoc loc, APIAccess access,
+                                      APIAvailability availability,
+                                      StringRef interface);
+
+  void addObjCMethod(ObjCContainerRecord *record, StringRef name, APILoc loc,
                      APIAccess access, bool isInstanceMethod, bool isOptional,
                      APIAvailability availability);
 
-  void writeAPIJSONFile(llvm::raw_ostream &os, bool PrettyPrint = false);
+  void writeAPIJSONFile(raw_ostream &os, bool PrettyPrint = false);
 
 private:
   const llvm::Triple target;
@@ -179,6 +194,7 @@ private:
   llvm::BumpPtrAllocator allocator;
   std::vector<GlobalRecord*> globals;
   std::vector<ObjCInterfaceRecord*> interfaces;
+  std::vector<ObjCCategoryRecord *> categories;
 };
 
 } // end namespace apigen

lib/TBDGen/TBDGen.cpp

@@ -747,6 +747,15 @@ void TBDGenVisitor::visitAbstractFunctionDecl(AbstractFunctionDecl *AFD) {
   if (AFD->hasAsync()) {
     addAsyncFunctionPointerSymbol(SILDeclRef(AFD));
   }
+
+  // Skip non objc compatible methods or non-public methods.
+  if (isa<DestructorDecl>(AFD) || !AFD->isObjC() ||
+      AFD->getFormalAccess() != AccessLevel::Public)
+    return;
+  if (auto *CD = dyn_cast<ClassDecl>(AFD->getDeclContext()))
+    recorder.addObjCMethod(CD, SILDeclRef(AFD));
+  else if (auto *ED = dyn_cast<ExtensionDecl>(AFD->getDeclContext()))
+    recorder.addObjCMethod(ED, SILDeclRef(AFD));
 }
 
 void TBDGenVisitor::visitFuncDecl(FuncDecl *FD) {
@@ -956,30 +965,9 @@ void TBDGenVisitor::visitClassDecl(ClassDecl *CD) {
       }
 
       TBD.addMethodDescriptor(method);
-
-      if (auto methodOrCtorOrDtor = method.getDecl()) {
-        // Skip non objc compatible methods or non-public methods.
-        if (!methodOrCtorOrDtor->isObjC() ||
-            methodOrCtorOrDtor->getFormalAccess() != AccessLevel::Public)
-          return;
-
-        // only handle FuncDecl here. Initializers are handled in
-        // visitConstructorDecl.
-        if (isa<FuncDecl>(methodOrCtorOrDtor))
-          recorder.addObjCMethod(CD, method);
-      }
     }
 
-    void addMethodOverride(SILDeclRef baseRef, SILDeclRef derivedRef) {
-      if (auto methodOrCtorOrDtor = derivedRef.getDecl()) {
-        if (!methodOrCtorOrDtor->isObjC() ||
-            methodOrCtorOrDtor->getFormalAccess() != AccessLevel::Public)
-          return;
-
-        if (isa<FuncDecl>(methodOrCtorOrDtor))
-          recorder.addObjCMethod(CD, derivedRef);
-      }
-    }
+    void addMethodOverride(SILDeclRef baseRef, SILDeclRef derivedRef) {}
 
     void addPlaceholder(MissingMemberDecl *) {}
 
@@ -1001,10 +989,6 @@ void TBDGenVisitor::visitConstructorDecl(ConstructorDecl *CD) {
       addAsyncFunctionPointerSymbol(
           SILDeclRef(CD, SILDeclRef::Kind::Initializer));
     }
-    if (auto parentClass = CD->getParent()->getSelfClassDecl()) {
-      if (parentClass->isObjC() || CD->isObjC())
-        recorder.addObjCMethod(parentClass, SILDeclRef(CD));
-    }
   }
 
   visitAbstractFunctionDecl(CD);
@@ -1397,8 +1381,11 @@ public:
     addOrGetObjCInterface(decl);
   }
 
-  void addObjCMethod(const ClassDecl *cls,
-                     SILDeclRef method) override {
+  void addObjCCategory(const ExtensionDecl *decl) override {
+    addOrGetObjCCategory(decl);
+  }
+
+  void addObjCMethod(const GenericContext *ctx, SILDeclRef method) override {
     SmallString<128> buffer;
     StringRef name = getSelectorName(method, buffer);
     apigen::APIAvailability availability;
@@ -1413,12 +1400,23 @@ public:
         access = apigen::APIAccess::Private;
     }
 
-    auto *clsRecord = addOrGetObjCInterface(cls);
-    api.addObjCMethod(clsRecord, name, moduleLoc, access, isInstanceMethod,
+    apigen::ObjCContainerRecord *record = nullptr;
+    if (auto *cls = dyn_cast<ClassDecl>(ctx))
+      record = addOrGetObjCInterface(cls);
+    else if (auto *ext = dyn_cast<ExtensionDecl>(ctx))
+      record = addOrGetObjCCategory(ext);
+
+    if (record)
+      api.addObjCMethod(record, name, moduleLoc, access, isInstanceMethod,
                         false, availability);
   }
 
 private:
+  /// Follow the naming schema that IRGen uses for Categories (see
+  /// ClassDataBuilder).
+  using CategoryNameKey = std::pair<const ClassDecl *, const ModuleDecl *>;
+  llvm::DenseMap<CategoryNameKey, unsigned> CategoryCounts;
+
   apigen::APIAvailability getAvailability(const Decl *decl) {
     bool unavailable = false;
     std::string introduced, obsoleted;
@@ -1475,11 +1473,46 @@ private:
     return cls;
   }
 
+  void buildCategoryName(const ExtensionDecl *ext, const ClassDecl *cls,
+                         SmallVectorImpl<char> &s) {
+    llvm::raw_svector_ostream os(s);
+    ModuleDecl *module = ext->getParentModule();
+    os << module->getName();
+    unsigned categoryCount = CategoryCounts[{cls, module}]++;
+    if (categoryCount > 0)
+      os << categoryCount;
+  }
+
+  apigen::ObjCCategoryRecord *addOrGetObjCCategory(const ExtensionDecl *decl) {
+    auto entry = categoryMap.find(decl);
+    if (entry != categoryMap.end())
+      return entry->second;
+
+    SmallString<128> interfaceBuffer;
+    SmallString<128> nameBuffer;
+    ClassDecl *cls = decl->getSelfClassDecl();
+    auto interface = cls->getObjCRuntimeName(interfaceBuffer);
+    buildCategoryName(decl, cls, nameBuffer);
+    apigen::APIAvailability availability = getAvailability(decl);
+    apigen::APIAccess access =
+        decl->isSPI() ? apigen::APIAccess::Private : apigen::APIAccess::Public;
+    apigen::APILinkage linkage =
+        decl->getMaxAccessLevel() == AccessLevel::Public
+            ? apigen::APILinkage::Exported
+            : apigen::APILinkage::Internal;
+    auto category = api.addObjCCategory(nameBuffer, linkage, moduleLoc, access,
+                                        availability, interface);
+    categoryMap.try_emplace(decl, category);
+    return category;
+  }
+
   apigen::API &api;
   ModuleDecl *module;
   apigen::APILoc moduleLoc;
 
   llvm::DenseMap<const ClassDecl*, apigen::ObjCInterfaceRecord*> classMap;
+  llvm::DenseMap<const ExtensionDecl *, apigen::ObjCCategoryRecord *>
+      categoryMap;
 };
 
 apigen::API APIGenRequest::evaluate(Evaluator &evaluator,

lib/TBDGen/TBDGenVisitor.h

@@ -68,7 +68,8 @@ public:
   virtual void addSymbol(StringRef name, llvm::MachO::SymbolKind kind,
                          SymbolSource source) {}
   virtual void addObjCInterface(const ClassDecl *decl) {}
-  virtual void addObjCMethod(const ClassDecl *cls, SILDeclRef method) {}
+  virtual void addObjCCategory(const ExtensionDecl *decl) {}
+  virtual void addObjCMethod(const GenericContext *ctx, SILDeclRef method) {}
 };
 
 class SimpleAPIRecorder final : public APIRecorder {

test/APIJSON/apigen.swift

@@ -364,12 +364,12 @@ public func myFunction2() {}
 // CHECK-NEXT:     "super": "NSObject",
 // CHECK-NEXT:     "instanceMethods": [
 // CHECK-NEXT:       {
-// CHECK-NEXT:         "name": "init",
+// CHECK-NEXT:         "name": "method1",
 // CHECK-NEXT:         "access": "public",
 // CHECK-NEXT:         "file": "/@input/MyModule.swiftinterface"
 // CHECK-NEXT:       },
 // CHECK-NEXT:       {
-// CHECK-NEXT:         "name": "method1",
+// CHECK-NEXT:         "name": "init",
 // CHECK-NEXT:         "access": "public",
 // CHECK-NEXT:         "file": "/@input/MyModule.swiftinterface"
 // CHECK-NEXT:       }
@@ -416,12 +416,12 @@ public func myFunction2() {}
 // CHECK-NEXT:     "super": "_TtC8MyModule4Test",
 // CHECK-NEXT:     "instanceMethods": [
 // CHECK-NEXT:       {
-// CHECK-NEXT:         "name": "init",
+// CHECK-NEXT:         "name": "method1",
 // CHECK-NEXT:         "access": "public",
 // CHECK-NEXT:         "file": "/@input/MyModule.swiftinterface"
 // CHECK-NEXT:       },
 // CHECK-NEXT:       {
-// CHECK-NEXT:         "name": "method1",
+// CHECK-NEXT:         "name": "init",
 // CHECK-NEXT:         "access": "public",
 // CHECK-NEXT:         "file": "/@input/MyModule.swiftinterface"
 // CHECK-NEXT:       }
@@ -429,4 +429,5 @@ public func myFunction2() {}
 // CHECK-NEXT:     "classMethods": []
 // CHECK-NEXT:   }
 // CHECK-NEXT: ],
+// CHECK-NEXT: "categories": [],
 // CHECK-NEXT: "version": "1.0"

test/APIJSON/extension.swift

@@ -0,0 +1,108 @@
+// REQUIRES: objc_interop, OS=macosx
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk-nosource -I %t) %s -typecheck -emit-module-interface-path %t/MyModule.swiftinterface -enable-library-evolution -module-name MyModule -swift-version 5 
+// RUN: %target-swift-api-extract -o - -pretty-print %t/MyModule.swiftinterface -module-name MyModule -module-cache-path %t | %FileCheck %s
+
+import Foundation
+
+// This should create an ObjC Category and a method with custom name.
+extension NSDictionary {
+  @objc
+  public subscript(key: Any) -> Any? {
+    @objc(__custom_name:)
+    get { return nil }
+  }
+}
+
+// This shouldn't create an interface.
+public class A {}
+
+// This shouldn't create a category.
+extension A {
+  public func run() {}
+}
+
+// This creates an interface.
+public class B: NSObject {}
+
+// This creates a category.
+@objc
+extension B {
+  public func run() {}
+}
+
+// This shouldn't create a category.
+extension B {
+  public func noop() {}
+}
+
+// This creates a category with index 1.
+@objc
+extension B {
+  public func fun() {}
+}
+
+// CHECK:      "interfaces": [
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "name": "_TtC8MyModule1B",
+// CHECK-NEXT:      "access": "public",
+// CHECK-NEXT:      "file": "/@input/MyModule.swiftinterface",
+// CHECK-NEXT:      "linkage": "exported",
+// CHECK-NEXT:      "super": "NSObject",
+// CHECK-NEXT:      "instanceMethods": [
+// CHECK-NEXT:        {
+// CHECK-NEXT:          "name": "init",
+// CHECK-NEXT:          "access": "public",
+// CHECK-NEXT:          "file": "/@input/MyModule.swiftinterface"
+// CHECK-NEXT:        }
+// CHECK-NEXT:      ],
+// CHECK-NEXT:      "classMethods": []
+// CHECK-NEXT:    }
+// CHECK-NEXT:  ],
+// CHECK-NEXT:  "categories": [
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "name": "MyModule",
+// CHECK-NEXT:      "access": "public",
+// CHECK-NEXT:      "file": "/@input/MyModule.swiftinterface",
+// CHECK-NEXT:      "linkage": "exported",
+// CHECK-NEXT:      "interface": "NSDictionary",
+// CHECK-NEXT:      "instanceMethods": [
+// CHECK-NEXT:        {
+// CHECK-NEXT:          "name": "__custom_name:",
+// CHECK-NEXT:          "access": "public",
+// CHECK-NEXT:          "file": "/@input/MyModule.swiftinterface"
+// CHECK-NEXT:        }
+// CHECK-NEXT:      ],
+// CHECK-NEXT:      "classMethods": []
+// CHECK-NEXT:    },
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "name": "MyModule",
+// CHECK-NEXT:      "access": "public",
+// CHECK-NEXT:      "file": "/@input/MyModule.swiftinterface",
+// CHECK-NEXT:      "linkage": "exported",
+// CHECK-NEXT:      "interface": "_TtC8MyModule1B",
+// CHECK-NEXT:      "instanceMethods": [
+// CHECK-NEXT:        {
+// CHECK-NEXT:          "name": "run",
+// CHECK-NEXT:          "access": "public",
+// CHECK-NEXT:          "file": "/@input/MyModule.swiftinterface"
+// CHECK-NEXT:        }
+// CHECK-NEXT:      ],
+// CHECK-NEXT:      "classMethods": []
+// CHECK-NEXT:    },
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "name": "MyModule1",
+// CHECK-NEXT:      "access": "public",
+// CHECK-NEXT:      "file": "/@input/MyModule.swiftinterface",
+// CHECK-NEXT:      "linkage": "exported",
+// CHECK-NEXT:      "interface": "_TtC8MyModule1B",
+// CHECK-NEXT:      "instanceMethods": [
+// CHECK-NEXT:        {
+// CHECK-NEXT:          "name": "fun",
+// CHECK-NEXT:          "access": "public",
+// CHECK-NEXT:          "file": "/@input/MyModule.swiftinterface"
+// CHECK-NEXT:        }
+// CHECK-NEXT:      ],
+// CHECK-NEXT:      "classMethods": []
+// CHECK-NEXT:    }
+// CHECK-NEXT:  ],

test/APIJSON/spi.swift

@@ -92,6 +92,7 @@ public class MyClass2 : NSObject {
 // CHECK-NEXT:       "classMethods": []
 // CHECK-NEXT:     }
 // CHECK-NEXT:   ],
+// CHECK-NEXT:   "categories": [],
 // CHECK-NEXT:   "version": "1.0"
 // CHECK-NEXT: }
 
@@ -264,12 +265,12 @@ public class MyClass2 : NSObject {
 // CHECK-SPI-NEXT:       "super": "NSObject",
 // CHECK-SPI-NEXT:       "instanceMethods": [
 // CHECK-SPI-NEXT:         {
-// CHECK-SPI-NEXT:           "name": "init",
+// CHECK-SPI-NEXT:           "name": "spiMethod",
 // CHECK-SPI-NEXT:           "access": "private",
 // CHECK-SPI-NEXT:           "file": "/@input/MyModule.swiftmodule"
 // CHECK-SPI-NEXT:         },
 // CHECK-SPI-NEXT:         {
-// CHECK-SPI-NEXT:           "name": "spiMethod",
+// CHECK-SPI-NEXT:           "name": "init",
 // CHECK-SPI-NEXT:           "access": "private",
 // CHECK-SPI-NEXT:           "file": "/@input/MyModule.swiftmodule"
 // CHECK-SPI-NEXT:         }
@@ -284,18 +285,19 @@ public class MyClass2 : NSObject {
 // CHECK-SPI-NEXT:       "super": "NSObject",
 // CHECK-SPI-NEXT:       "instanceMethods": [
 // CHECK-SPI-NEXT:         {
-// CHECK-SPI-NEXT:           "name": "init",
-// CHECK-SPI-NEXT:           "access": "public",
+// CHECK-SPI-NEXT:           "name": "spiMethod",
+// CHECK-SPI-NEXT:           "access": "private",
 // CHECK-SPI-NEXT:           "file": "/@input/MyModule.swiftmodule"
 // CHECK-SPI-NEXT:         },
 // CHECK-SPI-NEXT:         {
-// CHECK-SPI-NEXT:           "name": "spiMethod",
-// CHECK-SPI-NEXT:           "access": "private",
+// CHECK-SPI-NEXT:           "name": "init",
+// CHECK-SPI-NEXT:           "access": "public",
 // CHECK-SPI-NEXT:           "file": "/@input/MyModule.swiftmodule"
 // CHECK-SPI-NEXT:         }
 // CHECK-SPI-NEXT:       ],
 // CHECK-SPI-NEXT:       "classMethods": []
 // CHECK-SPI-NEXT:     }
 // CHECK-SPI-NEXT:   ],
+ // CHECK-SPI-NEXT:  "categories": [],
 // CHECK-SPI-NEXT:   "version": "1.0"
 // CHECK-SPI-NEXT: }

test/APIJSON/struct.swift

@@ -48,5 +48,6 @@ public struct TestStruct {
 // CHECK-NEXT:     }
 // CHECK-NEXT:   ],
 // CHECK-NEXT:   "interfaces": [],
+// CHECK-NEXT:   "categories": [],
 // CHECK-NEXT:   "version": "1.0"
 // CHECK-NEXT: }