[SymbolGraph] Look for @_spi on extensions

Consider declarations inside `@_spi` extensions to be internal.

Clean up the "implicitly private" check to work for `Decl` and not just
`ValueDecl`, allowing it to be used directly on extensions instead of having to
look for extensions everywhere.

rdar://63361634

Author: bitjammer

Diff for: lib/SymbolGraphGen/SymbolGraph.cpp

@@ -566,71 +566,74 @@ SymbolGraph::serializeDeclarationFragments(StringRef Key, Type T,
   T->print(Printer, getDeclarationFragmentsPrintOptions());
 }
 
-bool SymbolGraph::isImplicitlyPrivate(const ValueDecl *VD,
+bool SymbolGraph::isImplicitlyPrivate(const Decl *D,
                                       bool IgnoreContext) const {
   // Don't record unconditionally private declarations
-  if (VD->isPrivateStdlibDecl(/*treatNonBuiltinProtocolsAsPublic=*/false)) {
+  if (D->isPrivateStdlibDecl(/*treatNonBuiltinProtocolsAsPublic=*/false)) {
     return true;
   }
 
   // Don't record effectively internal declarations if specified
   if (Walker.Options.MinimumAccessLevel > AccessLevel::Internal &&
-      VD->hasUnderscoredNaming()) {
+      D->hasUnderscoredNaming()) {
     return true;
   }
 
-  // Don't include declarations with the @_spi attribute for now.
-  if (VD->getAttrs().getAttribute(DeclAttrKind::DAK_SPIAccessControl)) {
-    return true;
+  // Don't include declarations with the @_spi attribute unless the
+  // access control filter is internal or below.
+  if (D->isSPI()) {
+    return Walker.Options.MinimumAccessLevel > AccessLevel::Internal;
   }
 
-  // Symbols must meet the minimum access level to be included in the graph.
-  if (VD->getFormalAccess() < Walker.Options.MinimumAccessLevel) {
-    return true;
+  if (const auto *Extension = dyn_cast<ExtensionDecl>(D)) {
+    if (const auto *Nominal = Extension->getExtendedNominal()) {
+      return isImplicitlyPrivate(Nominal, IgnoreContext);
+    }
   }
 
-  // Special cases below.
+  if (const auto *VD = dyn_cast<ValueDecl>(D)) {
+    // Symbols must meet the minimum access level to be included in the graph.
+    if (VD->getFormalAccess() < Walker.Options.MinimumAccessLevel) {
+      return true;
+    }
 
-  auto BaseName = VD->getBaseName().userFacingName();
+    // Special cases below.
 
-  // ${MODULE}Version{Number,String} in ${Module}.h
-  SmallString<32> VersionNameIdentPrefix { M.getName().str() };
-  VersionNameIdentPrefix.append("Version");
-  if (BaseName.startswith(VersionNameIdentPrefix.str())) {
-    return true;
-  }
+    auto BaseName = VD->getBaseName().userFacingName();
 
-  // Automatically mapped SIMD types
-  auto IsGlobalSIMDType = llvm::StringSwitch<bool>(BaseName)
-#define MAP_SIMD_TYPE(C_TYPE, _, __) \
-.Case("swift_" #C_TYPE "2", true) \
-.Case("swift_" #C_TYPE "3", true) \
-.Case("swift_" #C_TYPE "4", true)
-#include "swift/ClangImporter/SIMDMappedTypes.def"
-  .Case("SWIFT_TYPEDEFS", true)
-  .Case("char16_t", true)
-  .Case("char32_t", true)
-  .Default(false);
-
-  if (IsGlobalSIMDType) {
-    return true;
-  }
+    // ${MODULE}Version{Number,String} in ${Module}.h
+    SmallString<32> VersionNameIdentPrefix { M.getName().str() };
+    VersionNameIdentPrefix.append("Version");
+    if (BaseName.startswith(VersionNameIdentPrefix.str())) {
+      return true;
+    }
 
-  if (IgnoreContext) {
-    return false;
+    // Automatically mapped SIMD types
+    auto IsGlobalSIMDType = llvm::StringSwitch<bool>(BaseName)
+  #define MAP_SIMD_TYPE(C_TYPE, _, __) \
+  .Case("swift_" #C_TYPE "2", true) \
+  .Case("swift_" #C_TYPE "3", true) \
+  .Case("swift_" #C_TYPE "4", true)
+  #include "swift/ClangImporter/SIMDMappedTypes.def"
+    .Case("SWIFT_TYPEDEFS", true)
+    .Case("char16_t", true)
+    .Case("char32_t", true)
+    .Default(false);
+
+    if (IsGlobalSIMDType) {
+      return true;
+    }
+
+    if (IgnoreContext) {
+      return false;
+    }
   }
 
   // Check up the parent chain. Anything inside a privately named
   // thing is also private. We could be looking at the `B` of `_A.B`.
-  if (const auto *DC = VD->getDeclContext()) {
+  if (const auto *DC = D->getDeclContext()) {
     if (const auto *Parent = DC->getAsDecl()) {
-      if (const auto *ParentVD = dyn_cast<ValueDecl>(Parent)) {
-        return isImplicitlyPrivate(ParentVD, IgnoreContext);
-      } else if (const auto *Extension = dyn_cast<ExtensionDecl>(Parent)) {
-        if (const auto *Nominal = Extension->getExtendedNominal()) {
-          return isImplicitlyPrivate(Nominal, IgnoreContext);
-        }
-      }
+      return isImplicitlyPrivate(Parent, IgnoreContext);
     }
   }
   return false;

Diff for: lib/SymbolGraphGen/SymbolGraph.h

@@ -215,8 +215,8 @@ struct SymbolGraph {
   /// and checking every named parent context as well.
   ///
   /// \param IgnoreContext If `true`, don't consider
-  /// the context of the symbol to determine whether it is implicitly private.
-  bool isImplicitlyPrivate(const ValueDecl *VD,
+  /// the context of the declaration to determine whether it is implicitly private.
+  bool isImplicitlyPrivate(const Decl *D,
                            bool IgnoreContext = false) const;
 
   /// Returns `true` if the declaration should be included as a node

Diff for: lib/SymbolGraphGen/SymbolGraphASTWalker.cpp

@@ -46,16 +46,17 @@ SymbolGraph *SymbolGraphASTWalker::getModuleSymbolGraph(const Decl *D) {
   if (this->M.getNameStr().equals(M->getNameStr())) {
     return &MainGraph;
   }
-  auto Found = ExtendedModuleGraphs.find(M);
+  auto Found = ExtendedModuleGraphs.find(M->getNameStr());
   if (Found != ExtendedModuleGraphs.end()) {
-    return Found->getSecond();
+    return Found->getValue();
   }
   auto *Memory = Ctx.allocate(sizeof(SymbolGraph), alignof(SymbolGraph));  
   auto *SG = new (Memory) SymbolGraph(*this,
                                       MainGraph.M,
                                       Optional<ModuleDecl *>(M),
                                       Ctx);
-  ExtendedModuleGraphs.insert({M, SG});
+
+  ExtendedModuleGraphs.insert({M->getNameStr(), SG});
   return SG;
 }
 
@@ -108,8 +109,8 @@ bool SymbolGraphASTWalker::walkToDeclPre(Decl *D, CharSourceRange Range) {
   if (const auto *Extension = dyn_cast<ExtensionDecl>(D)) {
     const auto *ExtendedNominal = Extension->getExtendedNominal();
     auto ExtendedSG = getModuleSymbolGraph(ExtendedNominal);
-    // Ignore effecively private decls.
-    if (ExtendedSG->isImplicitlyPrivate(ExtendedNominal)) {
+    // Ignore effectively private decls.
+    if (ExtendedSG->isImplicitlyPrivate(Extension)) {
       return false;
     }
 

Diff for: lib/SymbolGraphGen/SymbolGraphASTWalker.h

@@ -51,7 +51,7 @@ struct SymbolGraphASTWalker : public SourceEntityWalker {
   SymbolGraph MainGraph;
 
   /// A map of modules whose types were extended by the main module of interest `M`.
-  llvm::DenseMap<ModuleDecl *, SymbolGraph *> ExtendedModuleGraphs;
+  llvm::StringMap<SymbolGraph *> ExtendedModuleGraphs;
 
   // MARK: - Initialization
 

Diff for: lib/SymbolGraphGen/SymbolGraphGen.cpp

@@ -73,8 +73,8 @@ symbolgraphgen::emitSymbolGraphForModule(ModuleDecl *M,
 
   Success |= serializeSymbolGraph(Walker.MainGraph, Options);
 
-  for (auto Pair : Walker.ExtendedModuleGraphs) {
-    Success |= serializeSymbolGraph(*Pair.getSecond(), Options);
+  for (const auto &Entry : Walker.ExtendedModuleGraphs) {
+    Success |= serializeSymbolGraph(*Entry.getValue(), Options);
   }
 
   return Success;

Diff for: test/SymbolGraph/Symbols/SkipsSPI.swift

@@ -57,3 +57,10 @@ extension StructShouldntAppear.InnerStructShouldntAppear {
   @_spi(OtherModule)
   public func extendedFunctionShouldntAppear() {}
 }
+
+@_spi(OtherModule)
+extension StructShouldAppear {
+  // Although StructShouldAppear is fair to include, this extension is
+  // tagged as SPI, so everything inside it should be considered SPI as well.
+  public func functionInSPIExtensionShouldntAppear() {}
+}