[SymbolGraphGen] Un-revert #78959 and clean up usage of DenseMap

lib/SymbolGraphGen/Symbol.cpp

@@ -37,19 +37,21 @@ using namespace swift;
 using namespace symbolgraphgen;
 
 Symbol::Symbol(SymbolGraph *Graph, const ExtensionDecl *ED,
-               const NominalTypeDecl *SynthesizedBaseTypeDecl, Type BaseType)
+               const ValueDecl *SynthesizedBaseTypeDecl, Type BaseType)
     : Symbol::Symbol(Graph, nullptr, ED, SynthesizedBaseTypeDecl, BaseType) {}
 
 Symbol::Symbol(SymbolGraph *Graph, const ValueDecl *VD,
-               const NominalTypeDecl *SynthesizedBaseTypeDecl, Type BaseType)
+               const ValueDecl *SynthesizedBaseTypeDecl, Type BaseType)
     : Symbol::Symbol(Graph, VD, nullptr, SynthesizedBaseTypeDecl, BaseType) {}
 
 Symbol::Symbol(SymbolGraph *Graph, const ValueDecl *VD, const ExtensionDecl *ED,
-               const NominalTypeDecl *SynthesizedBaseTypeDecl, Type BaseType)
+               const ValueDecl *SynthesizedBaseTypeDecl, Type BaseType)
     : Graph(Graph), D(VD), BaseType(BaseType),
       SynthesizedBaseTypeDecl(SynthesizedBaseTypeDecl) {
-  if (!BaseType && SynthesizedBaseTypeDecl)
-    BaseType = SynthesizedBaseTypeDecl->getDeclaredInterfaceType();
+  if (!BaseType && SynthesizedBaseTypeDecl) {
+    if (const auto *NTD = dyn_cast<NominalTypeDecl>(SynthesizedBaseTypeDecl))
+      BaseType = NTD->getDeclaredInterfaceType();
+  }
   if (D == nullptr) {
     D = ED;
   }

lib/SymbolGraphGen/Symbol.h

@@ -35,10 +35,10 @@ class Symbol {
   /// Either a ValueDecl* or ExtensionDecl*.
   const Decl *D;
   Type BaseType;
-  const NominalTypeDecl *SynthesizedBaseTypeDecl;
+  const ValueDecl *SynthesizedBaseTypeDecl;
 
   Symbol(SymbolGraph *Graph, const ValueDecl *VD, const ExtensionDecl *ED,
-         const NominalTypeDecl *SynthesizedBaseTypeDecl,
+         const ValueDecl *SynthesizedBaseTypeDecl,
          Type BaseTypeForSubstitution = Type());
 
   swift::DeclName getName(const Decl *D) const;
@@ -87,11 +87,11 @@ class Symbol {
 
 public:
   Symbol(SymbolGraph *Graph, const ExtensionDecl *ED,
-         const NominalTypeDecl *SynthesizedBaseTypeDecl,
+         const ValueDecl *SynthesizedBaseTypeDecl,
          Type BaseTypeForSubstitution = Type());
 
   Symbol(SymbolGraph *Graph, const ValueDecl *VD,
-         const NominalTypeDecl *SynthesizedBaseTypeDecl,
+         const ValueDecl *SynthesizedBaseTypeDecl,
          Type BaseTypeForSubstitution = Type());
 
   void serialize(llvm::json::OStream &OS) const;
@@ -108,7 +108,7 @@ class Symbol {
     return BaseType;
   }
 
-  const NominalTypeDecl *getSynthesizedBaseTypeDecl() const {
+  const ValueDecl *getSynthesizedBaseTypeDecl() const {
     return SynthesizedBaseTypeDecl;
   }
 
@@ -175,27 +175,28 @@ using SymbolGraph = swift::symbolgraphgen::SymbolGraph;
 
 template <> struct DenseMapInfo<Symbol> {
   static inline Symbol getEmptyKey() {
-    return Symbol {
-      DenseMapInfo<SymbolGraph *>::getEmptyKey(),
-      DenseMapInfo<const swift::ValueDecl *>::getEmptyKey(),
-      DenseMapInfo<const swift::NominalTypeDecl *>::getTombstoneKey(),
-      DenseMapInfo<swift::Type>::getEmptyKey(),
+    return Symbol{
+        DenseMapInfo<SymbolGraph *>::getEmptyKey(),
+        DenseMapInfo<const swift::ValueDecl *>::getEmptyKey(),
+        DenseMapInfo<const swift::ValueDecl *>::getTombstoneKey(),
+        DenseMapInfo<swift::Type>::getEmptyKey(),
     };
   }
   static inline Symbol getTombstoneKey() {
-    return Symbol {
-      DenseMapInfo<SymbolGraph *>::getTombstoneKey(),
-      DenseMapInfo<const swift::ValueDecl *>::getTombstoneKey(),
-      DenseMapInfo<const swift::NominalTypeDecl *>::getTombstoneKey(),
-      DenseMapInfo<swift::Type>::getTombstoneKey(),
+    return Symbol{
+        DenseMapInfo<SymbolGraph *>::getTombstoneKey(),
+        DenseMapInfo<const swift::ValueDecl *>::getTombstoneKey(),
+        DenseMapInfo<const swift::ValueDecl *>::getTombstoneKey(),
+        DenseMapInfo<swift::Type>::getTombstoneKey(),
     };
   }
   static unsigned getHashValue(const Symbol S) {
     unsigned H = 0;
     H ^= DenseMapInfo<SymbolGraph *>::getHashValue(S.getGraph());
     H ^=
         DenseMapInfo<const swift::Decl *>::getHashValue(S.getLocalSymbolDecl());
-    H ^= DenseMapInfo<const swift::NominalTypeDecl *>::getHashValue(S.getSynthesizedBaseTypeDecl());
+    H ^= DenseMapInfo<const swift::ValueDecl *>::getHashValue(
+        S.getSynthesizedBaseTypeDecl());
     H ^= DenseMapInfo<swift::Type>::getHashValue(S.getBaseType());
     return H;
   }

lib/SymbolGraphGen/SymbolGraph.cpp

@@ -233,6 +233,18 @@ void SymbolGraph::recordEdge(Symbol Source,
 
 void SymbolGraph::recordMemberRelationship(Symbol S) {
   const auto *DC = S.getLocalSymbolDecl()->getDeclContext();
+  const ValueDecl *ParentDecl = DC->getSelfNominalTypeDecl();
+  if (!ParentDecl) {
+    // If we couldn't look up the type the member is declared on (e.g.
+    // because the member is declared in an extension whose extended type
+    // doesn't exist), don't record a memberOf relationship.
+    return;
+  }
+  if (const auto *PublicDecl = Walker.PublicPrivateTypeAliases.lookup(ParentDecl)) {
+    // If our member target is a private type that has a public type alias,
+    // point the membership to that type alias instead.
+    ParentDecl = PublicDecl;
+  }
   switch (DC->getContextKind()) {
     case DeclContextKind::GenericTypeDecl:
     case DeclContextKind::ExtensionDecl:
@@ -251,13 +263,6 @@ void SymbolGraph::recordMemberRelationship(Symbol S) {
         return;
       }
 
-      if (DC->getSelfNominalTypeDecl() == nullptr) {
-        // If we couldn't look up the type the member is declared on (e.g.
-        // because the member is declared in an extension whose extended type
-        // doesn't exist), don't record a memberOf relationship.
-        return;
-      }
-
       // If this is an extension to an external type, we use the extension
       // symbol itself as the target.
       if (auto const *Extension =
@@ -269,8 +274,7 @@ void SymbolGraph::recordMemberRelationship(Symbol S) {
         }
       }
 
-      return recordEdge(S,
-                        Symbol(this, DC->getSelfNominalTypeDecl(), nullptr),
+      return recordEdge(S, Symbol(this, ParentDecl, nullptr),
                         RelationshipKind::MemberOf());
     case swift::DeclContextKind::AbstractClosureExpr:
     case swift::DeclContextKind::SerializedAbstractClosure:
@@ -323,7 +327,16 @@ void SymbolGraph::recordConformanceSynthesizedMemberRelationships(Symbol S) {
   bool dropSynthesizedMembers = !Walker.Options.EmitSynthesizedMembers ||
                                 Walker.Options.SkipProtocolImplementations;
 
-  const auto D = S.getLocalSymbolDecl();
+  const auto *D = S.getLocalSymbolDecl();
+
+  // If this symbol is a public type alias to a private symbol, collect
+  // synthesized members for the underlying type.
+  if (const auto *TD = dyn_cast<TypeAliasDecl>(D)) {
+    const auto *NTD = TD->getUnderlyingType()->getAnyNominal();
+    if (NTD && Walker.PublicPrivateTypeAliases.lookup(NTD) == D)
+        D = NTD;
+  }
+
   const NominalTypeDecl *OwningNominal = nullptr;
   if (const auto *ThisNominal = dyn_cast<NominalTypeDecl>(D)) {
     OwningNominal = ThisNominal;
@@ -376,9 +389,11 @@ void SymbolGraph::recordConformanceSynthesizedMemberRelationships(Symbol S) {
           // that protocol would otherwise be hidden.
           if (auto *Nominal = Info.Ext->getExtendedNominal()) {
             if (dropSynthesizedMembers &&
-                !isImplicitlyPrivate(
-                    Nominal, /*IgnoreContext =*/Nominal->getModuleContext() ==
-                                 StdlibModule))
+                !isImplicitlyPrivate(Nominal, /*IgnoreContext =*/
+                                     [&](const Decl *P) {
+                                       return Nominal->getModuleContext() ==
+                                              StdlibModule;
+                                     }))
               continue;
           } else if (dropSynthesizedMembers) {
             continue;
@@ -393,10 +408,12 @@ void SymbolGraph::recordConformanceSynthesizedMemberRelationships(Symbol S) {
               // There can be synthesized members on effectively private
               // protocols or things that conform to them. We don't want to
               // include those.
-              if (isImplicitlyPrivate(SynthMember,
-                                      /*IgnoreContext =*/
-                                      SynthMember->getModuleContext() ==
-                                          StdlibModule)) {
+              if (isImplicitlyPrivate(
+                      SynthMember,
+                      /*IgnoreContext =*/
+                      [&](const Decl *P) {
+                        return SynthMember->getModuleContext() == StdlibModule;
+                      })) {
                 continue;
               }
 
@@ -405,22 +422,29 @@ void SymbolGraph::recordConformanceSynthesizedMemberRelationships(Symbol S) {
                 continue;
               }
 
-              Symbol Source(this, SynthMember, OwningNominal);
+              const ValueDecl *BaseDecl = OwningNominal;
+              if (const auto *PublicDecl = Walker.PublicPrivateTypeAliases.lookup(BaseDecl))
+                BaseDecl = PublicDecl;
+
+              Symbol Source(this, SynthMember, BaseDecl);
 
               if (auto *InheritedDecl = Source.getInheritedDecl()) {
                 if (auto *ParentDecl =
                         InheritedDecl->getDeclContext()->getAsDecl()) {
                   if (dropSynthesizedMembers &&
                       !isImplicitlyPrivate(
                           ParentDecl,
-                          /*IgnoreContext =*/ParentDecl->getModuleContext() ==
-                              StdlibModule)) {
+                          /*IgnoreContext =*/
+                          [&](const Decl *P) {
+                            return ParentDecl->getModuleContext() ==
+                                   StdlibModule;
+                          })) {
                     continue;
                   }
                 }
               }
 
-              auto ExtendedSG = Walker.getModuleSymbolGraph(OwningNominal);
+              auto ExtendedSG = Walker.getModuleSymbolGraph(BaseDecl);
 
               ExtendedSG->Nodes.insert(Source);
 
@@ -433,7 +457,15 @@ void SymbolGraph::recordConformanceSynthesizedMemberRelationships(Symbol S) {
 
 void
 SymbolGraph::recordInheritanceRelationships(Symbol S) {
-  const auto D = S.getLocalSymbolDecl();
+  const auto *D = S.getLocalSymbolDecl();
+
+  // If this is a public type alias for a private symbol, gather inheritance
+  // for the underlying type instead.
+  if (const auto *TD = dyn_cast<TypeAliasDecl>(D)) {
+    const auto *NTD = TD->getUnderlyingType()->getAnyNominal();
+    if (NTD && Walker.PublicPrivateTypeAliases.lookup(NTD) == D)
+      D = NTD;
+  }
 
   ClassDecl *Super = nullptr;
   if (auto *CD = dyn_cast<ClassDecl>(D))
@@ -442,8 +474,7 @@ SymbolGraph::recordInheritanceRelationships(Symbol S) {
     Super = PD->getSuperclassDecl();
 
   if (Super) {
-    recordEdge(Symbol(this, cast<ValueDecl>(D), nullptr),
-               Symbol(this, Super, nullptr),
+    recordEdge(S, Symbol(this, Super, nullptr),
                RelationshipKind::InheritsFrom());
   }
 }
@@ -522,7 +553,16 @@ void SymbolGraph::recordOptionalRequirementRelationships(Symbol S) {
 }
 
 void SymbolGraph::recordConformanceRelationships(Symbol S) {
-  const auto D = S.getLocalSymbolDecl();
+  const auto *D = S.getLocalSymbolDecl();
+
+  // If this is a public type alias for a private symbol, gather conformances
+  // for the underlying type instead.
+  if (const auto *TD = dyn_cast<TypeAliasDecl>(D)) {
+    const auto *NTD = TD->getUnderlyingType()->getAnyNominal();
+    if (NTD && Walker.PublicPrivateTypeAliases.lookup(NTD) == D)
+      D = NTD;
+  }
+
   if (const auto *NTD = dyn_cast<NominalTypeDecl>(D)) {
     if (auto *PD = dyn_cast<ProtocolDecl>(NTD)) {
       for (auto *inherited : PD->getAllInheritedProtocols()) {
@@ -700,8 +740,8 @@ const ValueDecl *getProtocolRequirement(const ValueDecl *VD) {
 
 }
 
-bool SymbolGraph::isImplicitlyPrivate(const Decl *D,
-                                      bool IgnoreContext) const {
+bool SymbolGraph::isImplicitlyPrivate(
+    const Decl *D, llvm::function_ref<bool(const Decl *)> IgnoreContext) const {
   // Don't record unconditionally private declarations
   if (D->isPrivateSystemDecl(/*treatNonBuiltinProtocolsAsPublic=*/false)) {
     return true;
@@ -809,16 +849,15 @@ bool SymbolGraph::isImplicitlyPrivate(const Decl *D,
     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 = D->getDeclContext()) {
     if (const auto *Parent = DC->getAsDecl()) {
+      if (IgnoreContext && IgnoreContext(Parent))
+        return false;
+
       // Exception: Children of underscored protocols should be considered
       // public, even though the protocols themselves aren't. This way,
       // synthesized copies of those symbols are correctly added to the public
@@ -851,7 +890,11 @@ bool SymbolGraph::isUnconditionallyUnavailableOnAllPlatforms(const Decl *D) cons
 }
 
 /// Returns `true` if the symbol should be included as a node in the graph.
-bool SymbolGraph::canIncludeDeclAsNode(const Decl *D) const {
+bool SymbolGraph::canIncludeDeclAsNode(const Decl *D,
+                                       const Decl *PublicAncestorDecl) const {
+  if (PublicAncestorDecl && D == PublicAncestorDecl)
+    return true;
+
   // If this decl isn't in this module or module that this module imported with `@_exported`, don't record it,
   // as it will appear elsewhere in its module's symbol graph.
 
@@ -873,6 +916,8 @@ bool SymbolGraph::canIncludeDeclAsNode(const Decl *D) const {
   } else {
     return false;
   }
-  return !isImplicitlyPrivate(cast<ValueDecl>(D)) 
-    && !isUnconditionallyUnavailableOnAllPlatforms(cast<ValueDecl>(D));
+  return !isImplicitlyPrivate(
+             cast<ValueDecl>(D), /*IgnoreContext=*/
+             [&](const Decl *P) { return P == PublicAncestorDecl; }) &&
+         !isUnconditionallyUnavailableOnAllPlatforms(cast<ValueDecl>(D));
 }

lib/SymbolGraphGen/SymbolGraph.h

@@ -220,10 +220,12 @@ struct SymbolGraph {
   /// implicitly internal/private, such as underscore prefixes,
   /// and checking every named parent context as well.
   ///
-  /// \param IgnoreContext If `true`, don't consider
-  /// the context of the declaration to determine whether it is implicitly private.
-  bool isImplicitlyPrivate(const Decl *D,
-                           bool IgnoreContext = false) const;
+  /// \param IgnoreContext A function ref that receives the parent decl
+  /// and returns whether or not the context should be ignored when determining
+  /// privacy.
+  bool isImplicitlyPrivate(
+      const Decl *D,
+      llvm::function_ref<bool(const Decl *)> IgnoreContext = nullptr) const;
 
   /// Returns `true` if the declaration has an availability attribute
   /// that marks it as unconditionally unavailable on all platforms (i.e., where
@@ -232,7 +234,11 @@ struct SymbolGraph {
 
   /// Returns `true` if the declaration should be included as a node
   /// in the graph.
-  bool canIncludeDeclAsNode(const Decl *D) const;
+  ///
+  /// If `PublicAncestorDecl` is set and is an ancestor of `D`, that declaration
+  /// is considered to be public, regardless of its surrounding context.
+  bool canIncludeDeclAsNode(const Decl *D,
+                            const Decl *PublicAncestorDecl = nullptr) const;
 
   /// Returns `true` if the declaration is a requirement of a protocol
   /// or is a default implementation of a protocol