Serialize symbol graphs for extended modules separately

When a module extends a type from another module, serialize those symbols into
separated files dedicated to those extended modules. This makes it easier to
ingest and categorize those symbols under the extended module if desired.

rdar://58941718

Author: bitjammer

include/swift/Serialization/SerializedModuleLoader.h

@@ -403,6 +403,8 @@ class SerializedASTFile final : public LoadedFile {
                    SmallVectorImpl<GenericSignature> &genericSignatures)
                 override;
 
+  StringRef getTargetTriple() const;
+
   static bool classof(const FileUnit *file) {
     return file->getKind() == FileUnitKind::SerializedAST;
   }

include/swift/SymbolGraphGen/SymbolGraphGen.h

@@ -20,8 +20,8 @@ class ModuleDecl;
 namespace symbolgraphgen {
 
 struct SymbolGraphOptions {
-  /// The path to output the symbol graph JSON.
-  StringRef OutputPath;
+  /// The directory to output the symbol graph JSON files.
+  StringRef OutputDir;
 
   /// The target of the module.
   llvm::Triple Target;

lib/Serialization/ModuleFile.h

@@ -823,6 +823,10 @@ class ModuleFile
     return ModuleInputBuffer->getBufferIdentifier();
   }
 
+  StringRef getTargetTriple() const {
+    return TargetTriple;
+  }
+
   /// AST-verify imported decls.
   ///
   /// Has no effect in NDEBUG builds.

lib/Serialization/SerializedModuleLoader.cpp

@@ -1181,6 +1181,10 @@ StringRef SerializedASTFile::getFilename() const {
   return File.getModuleFilename();
 }
 
+StringRef SerializedASTFile::getTargetTriple() const {
+  return File.getTargetTriple();
+}
+
 const clang::Module *SerializedASTFile::getUnderlyingClangModule() const {
   if (auto *UnderlyingModule = File.getUnderlyingModule())
     return UnderlyingModule->findUnderlyingClangModule();

lib/SymbolGraphGen/FormatVersion.h

@@ -14,7 +14,7 @@
 #define SWIFT_SYMBOLGRAPHGEN_FORMATVERSION_H
 
 #define SWIFT_SYMBOLGRAPH_FORMAT_MAJOR 0
-#define SWIFT_SYMBOLGRAPH_FORMAT_MINOR 2
+#define SWIFT_SYMBOLGRAPH_FORMAT_MINOR 3
 #define SWIFT_SYMBOLGRAPH_FORMAT_PATCH 0
 
 #endif // SWIFT_SYMBOLGRAPHGEN_FORMATVERSION_H

lib/SymbolGraphGen/Symbol.cpp

@@ -256,7 +256,11 @@ void Symbol::serializeSwiftExtensionMixin(llvm::json::OStream &OS) const {
   if (const auto *Extension
           = dyn_cast_or_null<ExtensionDecl>(VD->getInnermostDeclContext())) {
     OS.attributeObject("swiftExtension", [&](){
-      OS.attribute("definedInModule", Graph.M.getNameStr());
+      if (const auto *ExtendedNominal = Extension->getExtendedNominal()) {
+        if (const auto *ExtendedModule = ExtendedNominal->getModuleContext()) {
+          OS.attribute("extendedModule", ExtendedModule->getNameStr());
+        }
+      }
       auto Generics = Extension->getGenericSignature();
       if (Generics && !Generics->getRequirements().empty()) {
         OS.attributeArray("constraints", [&](){

lib/SymbolGraphGen/SymbolGraph.cpp

@@ -16,6 +16,8 @@
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/USRGeneration.h"
 #include "swift/Basic/Version.h"
+#include "swift/ClangImporter/ClangModule.h"
+#include "swift/Serialization/SerializedModuleLoader.h"
 
 #include "DeclarationFragmentPrinter.h"
 #include "FormatVersion.h"
@@ -25,10 +27,16 @@
 using namespace swift;
 using namespace symbolgraphgen;
 
-SymbolGraph::SymbolGraph(ModuleDecl &M, llvm::Triple Target,
+SymbolGraph::SymbolGraph(ModuleDecl &M,
+                         Optional<ModuleDecl *> ExtendedModule,
+                         llvm::Triple Target,
                          markup::MarkupContext &Ctx,
                          Optional<llvm::VersionTuple> ModuleVersion)
-: M(M), Target(Target), Ctx(Ctx), ModuleVersion(ModuleVersion) {}
+: M(M),
+  ExtendedModule(ExtendedModule),
+  Target(Target),
+  Ctx(Ctx),
+  ModuleVersion(ModuleVersion) {}
 
 // MARK: - Utilities
 
@@ -98,11 +106,27 @@ PrintOptions SymbolGraph::getDeclarationFragmentsPrintOptions() const {
   return Opts;
 }
 
-// MARK: - Relationships
+// MARK: - Symbols (Nodes)
+
+void SymbolGraph::recordNode(const ValueDecl *VD) {
+  Nodes.insert(VD);
+
+  // Record all of the possible relationships (edges) originating
+  // with this declaration.
+  recordMemberRelationship(VD);
+  recordConformanceRelationships(VD);
+  recordInheritanceRelationships(VD);
+  recordDefaultImplementationRelationships(VD);
+  recordOverrideRelationship(VD);
+  recordRequirementRelationships(VD);
+  recordOptionalRequirementRelationships(VD);
+}
+
+// MARK: - Relationships (Edges)
 
 void SymbolGraph::recordEdge(const ValueDecl *Source,
-                                      const ValueDecl *Target,
-                                      RelationshipKind Kind) {
+                             const ValueDecl *Target,
+                             RelationshipKind Kind) {
   if (Target->isPrivateStdlibDecl(
       /*treatNonBuiltinProtocolsAsPublic = */false)) {
     return;
@@ -237,7 +261,36 @@ void SymbolGraph::serialize(llvm::json::OStream &OS) {
     OS.attributeObject("module", [&](){
       OS.attribute("name", M.getNameStr());
       AttributeRAII Platform("platform", OS);
-      symbolgraphgen::serialize(Target, OS);
+
+      auto *MainFile = M.getFiles().front();
+      switch (MainFile->getKind()) {
+          case FileUnitKind::Builtin:
+            llvm_unreachable("Unexpected module kind: Builtin");
+          case FileUnitKind::DWARFModule:
+            llvm_unreachable("Unexpected module kind: DWARFModule");
+          case FileUnitKind::Source:
+            llvm_unreachable("Unexpected module kind: Source");
+            break;
+          case FileUnitKind::SerializedAST: {
+            auto SerializedAST = cast<SerializedASTFile>(MainFile);
+            auto Target = llvm::Triple(SerializedAST->getTargetTriple());
+            symbolgraphgen::serialize(Target, OS);
+            break;
+          }
+          case FileUnitKind::ClangModule: {
+            auto ClangModule = cast<ClangModuleUnit>(MainFile);
+            if (const auto *Overlay = ClangModule->getOverlayModule()) {
+              auto &OverlayMainFile =
+                  Overlay->getMainFile(FileUnitKind::SerializedAST);
+              auto SerializedAST = cast<SerializedASTFile>(OverlayMainFile);
+              auto Target = llvm::Triple(SerializedAST.getTargetTriple());
+              symbolgraphgen::serialize(Target, OS);
+            } else {
+              symbolgraphgen::serialize(Target, OS);
+            }
+            break;
+        }
+      }
     });
 
     if (ModuleVersion) {

lib/SymbolGraphGen/SymbolGraph.h

@@ -31,6 +31,11 @@ struct SymbolGraph {
   */
   ModuleDecl &M;
 
+  /**
+   The module whose types were extended in `M`.
+   */
+  Optional<ModuleDecl *> ExtendedModule;
+  
   /**
    The module's target triple.
   */
@@ -60,7 +65,9 @@ struct SymbolGraph {
   /// A cache of USRs for declarations.
   llvm::DenseMap<const ValueDecl *, StringRef> USRCache;
 
-  SymbolGraph(ModuleDecl &M, llvm::Triple Target,
+  SymbolGraph(ModuleDecl &M,
+              Optional<ModuleDecl *> ExtendedModule,
+              llvm::Triple Target,
               markup::MarkupContext &Ctx,
               Optional<llvm::VersionTuple> ModuleVersion = None);
 
@@ -76,6 +83,13 @@ struct SymbolGraph {
   /// Get the base print options for declaration fragments.
   PrintOptions getDeclarationFragmentsPrintOptions() const;
 
+  // MARK: - Symbols (Nodes)
+
+  /**
+   Record a symbol as a node in the graph.
+   */
+  void recordNode(const ValueDecl *VD);
+
   // MARK: - Relationships (Edges)
 
   /**

lib/SymbolGraphGen/SymbolGraphASTWalker.cpp

@@ -13,6 +13,7 @@
 #include "llvm/ADT/StringSwitch.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/Module.h"
+#include "swift/Serialization/SerializedModuleLoader.h"
 #include "swift/SymbolGraphGen/SymbolGraphGen.h"
 
 #include "SymbolGraphASTWalker.h"
@@ -24,7 +25,7 @@ SymbolGraphASTWalker::SymbolGraphASTWalker(ModuleDecl &M,
                                            const SymbolGraphOptions &Options)
   : Options(Options),
     M(M),
-    Graph(M, Options.Target, Ctx) {}
+    Graph(M, None, Options.Target, Ctx) {}
 
 /// Returns `true` if the symbol should be included as a node in the graph.
 bool SymbolGraphASTWalker::shouldIncludeNode(const Decl *D) const {
@@ -86,6 +87,21 @@ bool SymbolGraphASTWalker::shouldIncludeNode(const Decl *D) const {
   return ShouldInclude;
 }
 
+/// Get a "sub" symbol graph for the parent module of a type that the main module `M` is extending.
+SymbolGraph &SymbolGraphASTWalker::getExtendedModuleSymbolGraph(ModuleDecl *M) {
+  auto Found = ExtendedModuleGraphs.find(M);
+  if (Found != ExtendedModuleGraphs.end()) {
+    return *Found->getSecond();
+  }
+  auto *Memory = Ctx.allocate(sizeof(SymbolGraph), alignof(SymbolGraph));  
+  auto *SG = new (Memory) SymbolGraph(Graph.M,
+                                      Optional<ModuleDecl *>(M),
+                                      Options.Target,
+                                      Ctx);
+  ExtendedModuleGraphs.insert({M, SG});
+  return *SG;
+}
+
 bool SymbolGraphASTWalker::walkToDeclPre(Decl *D, CharSourceRange Range) {
 
     switch (D->getKind()) {
@@ -111,17 +127,23 @@ bool SymbolGraphASTWalker::walkToDeclPre(Decl *D, CharSourceRange Range) {
   }
 
   auto *VD = cast<ValueDecl>(D);
-  Graph.Nodes.insert(VD);
-  
-  // Record all of the possible relationships (edges) originating
-  // with this declaration.
-  Graph.recordMemberRelationship(VD);
-  Graph.recordConformanceRelationships(VD);
-  Graph.recordInheritanceRelationships(VD);
-  Graph.recordDefaultImplementationRelationships(VD);
-  Graph.recordOverrideRelationship(VD);
-  Graph.recordRequirementRelationships(VD);
-  Graph.recordOptionalRequirementRelationships(VD);
-  
+
+  // If this symbol extends a type from another module, record it in that
+  // module's symbol graph, which will be emitted separately.
+  if (const auto *Extension
+          = dyn_cast_or_null<ExtensionDecl>(VD->getInnermostDeclContext())) {
+    if (const auto *ExtendedNominal = Extension->getExtendedNominal()) {
+      auto ExtendedModule = ExtendedNominal->getModuleContext();
+      if (ExtendedModule != &M) {
+        auto &SG = getExtendedModuleSymbolGraph(ExtendedModule);
+        SG.recordNode(VD);
+        return true;
+      }
+    }
+  }
+
+  // Otherwise, record this in the main module `M`'s symbol graph.
+  Graph.recordNode(VD);
+
   return true;
 }

lib/SymbolGraphGen/SymbolGraphASTWalker.h

@@ -13,6 +13,8 @@
 #ifndef SWIFT_SYMBOLGRAPHGEN_SYMBOLGRAPHASTWALKER_H
 #define SWIFT_SYMBOLGRAPHGEN_SYMBOLGRAPHASTWALKER_H
 
+#include "llvm/ADT/DenseMap.h"
+#include "swift/AST/Module.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/IDE/SourceEntityWalker.h"
 #include "swift/Markup/Markup.h"
@@ -48,16 +50,24 @@ struct SymbolGraphASTWalker : public SourceEntityWalker {
   /// The symbol graph for a module.
   SymbolGraph Graph;
 
-  // MARK: -
+  /// A map of modules whose types were extended by the main module of interest `M`.
+  llvm::DenseMap<ModuleDecl *, SymbolGraph *> ExtendedModuleGraphs;
+
+  // MARK: - Initialization
 
   SymbolGraphASTWalker(ModuleDecl &M, const SymbolGraphOptions &Options);
   virtual ~SymbolGraphASTWalker() {}
 
-  // MARK: - 
+  // MARK: - Utilities
 
   /// Returns `true` if the symbol should be included as a node in the graph.
   bool shouldIncludeNode(const Decl *D) const;
 
+  /// Get a "sub" symbol graph for the parent module of a type that the main module `M` is extending.
+  SymbolGraph &getExtendedModuleSymbolGraph(ModuleDecl *M);
+
+  // MARK: - SourceEntityWalker
+
   virtual bool walkToDeclPre(Decl *D, CharSourceRange Range);
 };