[AST] Intro request to list declared and inherited SPI groups

Author: xymus

include/swift/AST/ASTTypeIDZone.def

@@ -54,3 +54,4 @@ SWIFT_TYPEID_NAMED(TypeAliasDecl *, TypeAliasDecl)
 SWIFT_TYPEID_NAMED(ValueDecl *, ValueDecl)
 SWIFT_TYPEID_NAMED(VarDecl *, VarDecl)
 SWIFT_TYPEID(FingerprintAndMembers)
+SWIFT_TYPEID(Identifier)

include/swift/AST/ASTTypeIDs.h

@@ -62,6 +62,7 @@ struct TypeWitnessAndDecl;
 enum class AncestryFlags : uint8_t;
 enum class ImplicitMemberAction : uint8_t;
 struct FingerprintAndMembers;
+class Identifier;
 
 // Define the AST type zone (zone 1)
 #define SWIFT_TYPEID_ZONE AST

include/swift/AST/TypeCheckRequests.h

@@ -2025,6 +2025,26 @@ class PatternTypeRequest
   }
 };
 
+/// List SPI group ids declared on a decl.
+class SPIGroupsRequest :
+    public SimpleRequest<SPIGroupsRequest,
+                         llvm::ArrayRef<Identifier>(const Decl *),
+                         CacheKind::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  llvm::Expected<llvm::ArrayRef<Identifier>>
+  evaluate(Evaluator &evaluator, const Decl *decl) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
+
 /// Type-checks a `@differentiable` attribute and returns the resolved parameter
 /// indices on success. On failure, emits diagnostics and returns `nullptr`.
 ///

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -105,6 +105,9 @@ SWIFT_REQUEST(TypeChecker, IsDynamicRequest, bool(ValueDecl *),
               SeparatelyCached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, IsFinalRequest, bool(ValueDecl *), SeparatelyCached,
               NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, SPIGroupsRequest,
+              llvm::ArrayRef<Identifier>(Decl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, IsGetterMutatingRequest, bool(AbstractStorageDecl *),
               SeparatelyCached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, IsImplicitlyUnwrappedOptionalRequest,

lib/AST/Module.cpp

@@ -1820,22 +1820,67 @@ void SourceFile::lookupImportedSPIGroups(const ModuleDecl *importedModule,
 }
 
 bool SourceFile::isImportedAsSPI(const ValueDecl *targetDecl) const {
-  if (!targetDecl->getAttrs().hasAttribute<SPIAccessControlAttr>())
-    return false;
-
   auto targetModule = targetDecl->getModuleContext();
-  SmallVector<Identifier, 4> importedSpis;
-  lookupImportedSPIGroups(targetModule, importedSpis);
-
-  for (auto attr : targetDecl->getAttrs().getAttributes<SPIAccessControlAttr>())
-    for (auto declSPI : attr->getSPIGroups())
-      for (auto importedSPI : importedSpis)
-        if (importedSPI == declSPI)
-          return true;
+  SmallVector<Identifier, 4> importedSPIGroups;
+  lookupImportedSPIGroups(targetModule, importedSPIGroups);
+  if (importedSPIGroups.empty()) return false;
+
+  auto declSPIGroups = evaluateOrDefault(
+                        targetDecl->getASTContext().evaluator,
+                        SPIGroupsRequest{ targetDecl },
+                        ArrayRef<Identifier>());
+
+  // Note: If we reach a point where there are many SPI imports or SPI groups
+  // on decls we could optimize this further by using a set.
+  for (auto importedSPI : importedSPIGroups)
+    for (auto declSPI : declSPIGroups)
+      if (importedSPI == declSPI)
+        return true;
 
   return false;
 }
 
+llvm::Expected<llvm::ArrayRef<Identifier>>
+SPIGroupsRequest::evaluate(Evaluator &evaluator, const Decl *decl) const {
+  // Applies only to public ValueDecls and ExtensionDecls.
+  if (auto vd = dyn_cast<ValueDecl>(decl)) {
+    if (vd->getFormalAccess() < AccessLevel::Public)
+      return ArrayRef<Identifier>();
+  } else if (!isa<ExtensionDecl>(decl))
+    return ArrayRef<Identifier>();
+
+  // First, look for local attributes.
+  llvm::SetVector<Identifier> spiGroups;
+  for (auto attr : decl->getAttrs().getAttributes<SPIAccessControlAttr>())
+    for (auto spi : attr->getSPIGroups())
+      spiGroups.insert(spi);
+
+  auto &ctx = decl->getASTContext();
+  if (spiGroups.empty()) {
+
+    // Then in the extended nominal type.
+    if (auto extension = dyn_cast<ExtensionDecl>(decl)) {
+      auto extended = extension->getExtendedNominal();
+      auto extSPIs = evaluateOrDefault(ctx.evaluator,
+                                       SPIGroupsRequest{ extended },
+                                       ArrayRef<Identifier>());
+      if (!extSPIs.empty()) return extSPIs;
+    }
+
+    // And finally in the parent context.
+    auto parent = decl->getDeclContext();
+    if (auto parentD = parent->getAsDecl()) {
+      if (!isa<ModuleDecl>(parentD)) {
+        return evaluateOrDefault(ctx.evaluator,
+                                 SPIGroupsRequest{ parentD },
+                                 ArrayRef<Identifier>());
+      }
+    }
+  }
+
+  return ctx.AllocateCopy(spiGroups.getArrayRef());
+}
+
 bool SourceFile::shouldCrossImport() const {
   return Kind != SourceFileKind::SIL && Kind != SourceFileKind::Interface &&
          getASTContext().LangOpts.EnableCrossImportOverlays;