Register Dependency Arcs for Extension Members

This was a hole in the existing dependency tracking infrastructure. David managed to discover a way to exploit this bug to get a miscompile in rdar://74583179. There, members of extensions were not counted towards the interface hash of a type and so mutating them could lead to e.g. the wrong declaration being selected in an overload set.

To start tracking extensions, we need to add two new kinds of arcs:
1) A nominal arc to the extension
2) A member arc to the extension

Unfortunately, extensions are also unique in Swift in that they do not have a name to allow us to unique them. Luckily, we do have a way of identifying extensions: their fingerprint. These arcs are therefore emitted with the extended nominal type and the fingerprint of the extension as their context. This effectively invents a new nominal type for every extension.

Author: CodaFi

include/swift/AST/FineGrainedDependencies.h

@@ -426,14 +426,14 @@ class DependencyKey {
   private:
     const NodeKind kind;
     const DeclAspect aspect;
-    const NominalTypeDecl *context;
+    const DeclContext *context;
     StringRef name;
 
   private:
     // A private copy constructor so our clients are forced to use the
     // move-only builder interface.
     explicit Builder(NodeKind kind, DeclAspect aspect,
-                     const NominalTypeDecl *context, StringRef name)
+                     const DeclContext *context, StringRef name)
         : kind(kind), aspect(aspect), context(context), name(name) {}
 
   public:

lib/AST/FrontendSourceFileDepGraphFactory.cpp

@@ -56,9 +56,18 @@ using namespace fine_grained_dependencies;
 // MARK: Helpers for key construction that must be in frontend
 //==============================================================================
 
-static std::string mangleTypeAsContext(const NominalTypeDecl *NTD) {
+static std::string identifierForContext(const DeclContext *DC) {
+  if (!DC) return "";
+
   Mangle::ASTMangler Mangler;
-  return !NTD ? "" : Mangler.mangleTypeAsContextUSR(NTD);
+  if (const auto *context = dyn_cast<NominalTypeDecl>(DC)) {
+    return Mangler.mangleTypeAsContextUSR(context);
+  }
+
+  const auto *ext = cast<ExtensionDecl>(DC);
+  auto fp = ext->getBodyFingerprint().getValueOr(Fingerprint::ZERO());
+  auto typeStr = Mangler.mangleTypeAsContextUSR(ext->getExtendedNominal());
+  return (typeStr + "@" + fp.getRawValue()).str();
 }
 
 //==============================================================================
@@ -69,7 +78,7 @@ DependencyKey DependencyKey::Builder::build() && {
   return DependencyKey{
     kind,
     aspect,
-    context ? mangleTypeAsContext(context) : "",
+    identifierForContext(context),
     name.str()
   };
 }
@@ -98,10 +107,15 @@ DependencyKey::Builder DependencyKey::Builder::withContext(const Decl *D) && {
   switch (kind) {
   case NodeKind::nominal:
   case NodeKind::potentialMember:
-  case NodeKind::member:
+  case NodeKind::member: {
     /// nominal and potential member dependencies are created from a Decl and
     /// use the context field.
-    return Builder{kind, aspect, cast<NominalTypeDecl>(D), name};
+    const DeclContext *context = dyn_cast<NominalTypeDecl>(D);
+    if (!context) {
+      context = cast<ExtensionDecl>(D);
+    }
+    return Builder{kind, aspect, context, name};
+  }
   case NodeKind::topLevel:
   case NodeKind::dynamicLookup:
   case NodeKind::externalDepend:
@@ -403,6 +417,7 @@ void FrontendSourceFileDepGraphFactory::addAllDefinedDecls() {
   addAllDefinedDeclsOfAGivenType<NodeKind::topLevel>(declFinder.topNominals);
   addAllDefinedDeclsOfAGivenType<NodeKind::topLevel>(declFinder.topValues);
   addAllDefinedDeclsOfAGivenType<NodeKind::nominal>(declFinder.allNominals);
+  addAllDefinedDeclsOfAGivenType<NodeKind::nominal>(declFinder.extensions);
   addAllDefinedDeclsOfAGivenType<NodeKind::potentialMember>(
       declFinder.potentialMemberHolders);
   addAllDefinedDeclsOfAGivenType<NodeKind::member>(
@@ -460,11 +475,23 @@ class UsedDeclEnumerator {
         return;
       }
 
-      auto key =
+      auto nominalKey =
+          DependencyKey::Builder(NodeKind::nominal, DeclAspect::interface)
+              .withContext(subject->getSelfNominalTypeDecl())
+              .build();
+      enumerateUse(nominalKey, enumerator);
+
+      auto declKey =
           DependencyKey::Builder(NodeKind::nominal, DeclAspect::interface)
               .withContext(subject->getAsDecl())
               .build();
-      enumerateUse(key, enumerator);
+
+      // If the subject of this dependency is not the nominal type itself,
+      // record another arc for the extension this member came from.
+      if (nominalKey != declKey) {
+        assert(isa<ExtensionDecl>(subject));
+        enumerateUse(declKey, enumerator);
+      }
     });
   }
 
@@ -571,6 +598,7 @@ void ModuleDepGraphFactory::addAllDefinedDecls() {
   addAllDefinedDeclsOfAGivenType<NodeKind::topLevel>(declFinder.topNominals);
   addAllDefinedDeclsOfAGivenType<NodeKind::topLevel>(declFinder.topValues);
   addAllDefinedDeclsOfAGivenType<NodeKind::nominal>(declFinder.allNominals);
+  addAllDefinedDeclsOfAGivenType<NodeKind::nominal>(declFinder.extensions);
   addAllDefinedDeclsOfAGivenType<NodeKind::potentialMember>(
       declFinder.potentialMemberHolders);
   addAllDefinedDeclsOfAGivenType<NodeKind::member>(

lib/AST/NameLookupRequests.cpp

@@ -305,6 +305,13 @@ void DirectLookupRequest::writeDependencySink(
     evaluator::DependencyCollector &tracker,
     const TinyPtrVector<ValueDecl *> &result) const {
   auto &desc = std::get<0>(getStorage());
+  // Add used members from the perspective of
+  // 1) The decl context they are found in
+  // 2) The decl context of the request
+  // This gets us a dependency not just on `Foo.bar` but on `extension Foo.bar`.
+  for (const auto *member : result) {
+    tracker.addUsedMember(member->getDeclContext(), desc.Name.getBaseName());
+  }
   tracker.addUsedMember(desc.DC, desc.Name.getBaseName());
 }