AST/Sema: Look up custom availability domains by name.

Store the domain definitions defined on the command line on the `ModuleDecl` for
the main module.

Author: tshortli

include/swift/AST/Module.h

@@ -19,6 +19,7 @@
 
 #include "swift/AST/AccessNotes.h"
 #include "swift/AST/AttrKind.h"
+#include "swift/AST/AvailabilityDomain.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/DeclContext.h"
 #include "swift/AST/Identifier.h"
@@ -41,8 +42,8 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MD5.h"
 #include <optional>
-#include <unordered_map>
 #include <set>
+#include <unordered_map>
 
 namespace clang {
   class Module;
@@ -51,20 +52,16 @@ namespace clang {
 namespace swift {
   enum class ArtificialMainKind : uint8_t;
   class ASTContext;
-  class ASTScope;
   class ASTWalker;
-  class AvailabilityScope;
-  class BraceStmt;
+  class CustomAvailabilityDomain;
   class Decl;
   class DeclAttribute;
   class TypeDecl;
   enum class DeclKind : uint8_t;
-  class ExtensionDecl;
   class DebuggerClient;
   class DeclName;
   class FileUnit;
   class FuncDecl;
-  class InfixOperatorDecl;
   enum class LibraryLevel : uint8_t;
   class LinkLibrary;
   class ModuleLoader;
@@ -74,20 +71,12 @@ namespace swift {
   class PostfixOperatorDecl;
   class PrefixOperatorDecl;
   class ProtocolConformance;
-  class ProtocolDecl;
   struct PrintOptions;
   class SourceLookupCache;
-  class Token;
-  class TupleType;
   class Type;
   class ValueDecl;
-  class VarDecl;
   class VisibleDeclConsumer;
 
-namespace ast_scope {
-class ASTSourceFileScope;
-}
-
 /// Discriminator for file-units.
 enum class FileUnitKind {
   /// For a .swift source file.
@@ -361,6 +350,10 @@ class ModuleDecl
   /// Used by the debugger to bypass resilient access to fields.
   bool BypassResilience = false;
 
+  using AvailabilityDomainMap =
+      llvm::SmallDenseMap<Identifier, CustomAvailabilityDomain *>;
+  AvailabilityDomainMap AvailabilityDomains;
+
 public:
   using PopulateFilesFn = llvm::function_ref<void(
       ModuleDecl *, llvm::function_ref<void(FileUnit *)>)>;
@@ -1233,6 +1226,15 @@ class ModuleDecl
   /// An empty `Version` is returned if the information is not available.
   version::Version getLanguageVersionBuiltWith() const;
 
+  /// Returns the custom availability domain defined by this module with the
+  /// given identifier, if one exists.
+  std::optional<AvailabilityDomain>
+  getAvailabilityDomainForIdentifier(Identifier identifier) const;
+
+  void setAvailabilityDomains(const AvailabilityDomainMap &&map) {
+    AvailabilityDomains = std::move(map);
+  }
+
   static bool classof(const DeclContext *DC) {
     if (auto D = DC->getAsDecl())
       return classof(D);

include/swift/AST/SourceFile.h

@@ -25,9 +25,11 @@
 #include "llvm/ADT/STLExtras.h"
 
 namespace swift {
-
+class ASTScope;
+class AvailabilityScope;
 class PersistentParserState;
 struct SourceFileExtras;
+class Token;
 
 /// Kind of import affecting how a decl can be reexported.
 ///

lib/AST/Module.cpp

@@ -4135,3 +4135,12 @@ version::Version ModuleDecl::getLanguageVersionBuiltWith() const {
 
   return version::Version();
 }
+
+std::optional<AvailabilityDomain>
+ModuleDecl::getAvailabilityDomainForIdentifier(Identifier identifier) const {
+  auto iter = AvailabilityDomains.find(identifier);
+  if (iter == AvailabilityDomains.end())
+    return std::nullopt;
+
+  return AvailabilityDomain::forCustom(iter->getSecond());
+}

lib/Frontend/Frontend.cpp

@@ -17,6 +17,7 @@
 
 #include "swift/Frontend/Frontend.h"
 #include "swift/AST/ASTContext.h"
+#include "swift/AST/AvailabilityDomain.h"
 #include "swift/AST/DiagnosticsFrontend.h"
 #include "swift/AST/DiagnosticsSema.h"
 #include "swift/AST/FileSystem.h"
@@ -29,24 +30,23 @@
 #include "swift/Basic/Platform.h"
 #include "swift/Basic/SourceManager.h"
 #include "swift/Basic/Statistic.h"
+#include "swift/DependencyScan/ModuleDependencyScanner.h"
 #include "swift/Frontend/CachingUtils.h"
 #include "swift/Frontend/CompileJobCacheKey.h"
 #include "swift/Frontend/ModuleInterfaceLoader.h"
 #include "swift/Parse/Lexer.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/Utils/Generics.h"
+#include "swift/Serialization/ScanningLoaders.h"
 #include "swift/Serialization/SerializationOptions.h"
 #include "swift/Serialization/SerializedModuleLoader.h"
-#include "swift/Serialization/ScanningLoaders.h"
-#include "swift/DependencyScan/ModuleDependencyScanner.h"
 #include "swift/Strings.h"
 #include "swift/Subsystems.h"
 #include "clang/AST/ASTContext.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/TargetParser/Triple.h"
 #include "llvm/CAS/ActionCache.h"
 #include "llvm/CAS/BuiltinUnifiedCASDatabases.h"
 #include "llvm/CAS/CASFileSystem.h"
@@ -55,8 +55,9 @@
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Process.h"
-#include "llvm/Support/VirtualOutputBackends.h"
 #include "llvm/Support/ThreadPool.h"
+#include "llvm/Support/VirtualOutputBackends.h"
+#include "llvm/TargetParser/Triple.h"
 #include <llvm/ADT/StringExtras.h>
 
 using namespace swift;
@@ -1392,6 +1393,28 @@ bool CompilerInstance::createFilesForMainModule(
   return false;
 }
 
+static void configureAvailabilityDomains(const ASTContext &ctx,
+                                         const FrontendOptions &opts,
+                                         ModuleDecl *mainModule) {
+  llvm::SmallDenseMap<Identifier, CustomAvailabilityDomain *> domainMap;
+  auto createAndInsertDomain = [&](const std::string &name,
+                                   CustomAvailabilityDomain::Kind kind) {
+    auto *domain = CustomAvailabilityDomain::create(
+        ctx, name, mainModule, CustomAvailabilityDomain::Kind::Enabled);
+    bool inserted = domainMap.insert({domain->getName(), domain}).second;
+    ASSERT(inserted); // Domains must be unique.
+  };
+
+  for (auto enabled : opts.AvailabilityDomains.EnabledDomains)
+    createAndInsertDomain(enabled, CustomAvailabilityDomain::Kind::Enabled);
+  for (auto disabled : opts.AvailabilityDomains.DisabledDomains)
+    createAndInsertDomain(disabled, CustomAvailabilityDomain::Kind::Disabled);
+  for (auto dynamic : opts.AvailabilityDomains.DynamicDomains)
+    createAndInsertDomain(dynamic, CustomAvailabilityDomain::Kind::Dynamic);
+
+  mainModule->setAvailabilityDomains(std::move(domainMap));
+}
+
 ModuleDecl *CompilerInstance::getMainModule() const {
   if (MainModule)
     return MainModule;
@@ -1441,6 +1464,9 @@ ModuleDecl *CompilerInstance::getMainModule() const {
     if (Invocation.getLangOptions().hasFeature(Feature::WarnUnsafe))
       MainModule->setStrictMemorySafety(true);
 
+    configureAvailabilityDomains(getASTContext(),
+                                 Invocation.getFrontendOptions(), MainModule);
+
     if (!Invocation.getFrontendOptions()
              .SwiftInterfaceCompilerVersion.empty()) {
       auto compilerVersion =

lib/Sema/TypeCheckAttr.cpp

@@ -8332,14 +8332,30 @@ ValueDecl *RenamedDeclRequest::evaluate(Evaluator &evaluator,
   return renamedDecl;
 }
 
+static std::optional<AvailabilityDomain>
+getAvailabilityDomainForName(StringRef name, const DeclContext *declContext) {
+  if (auto builtinDomain =
+          AvailabilityDomain::builtinDomainForString(name, declContext))
+    return builtinDomain;
+
+  auto &ctx = declContext->getASTContext();
+  auto identifier = ctx.getIdentifier(name);
+  if (auto customDomain =
+          ctx.MainModule->getAvailabilityDomainForIdentifier(identifier))
+    return customDomain;
+
+  return std::nullopt;
+}
+
 std::optional<SemanticAvailableAttr>
 SemanticAvailableAttrRequest::evaluate(swift::Evaluator &evaluator,
                                        const AvailableAttr *attr,
                                        const Decl *decl) const {
   if (attr->hasCachedDomain())
     return SemanticAvailableAttr(attr);
 
-  auto &diags = decl->getASTContext().Diags;
+  auto &ctx = decl->getASTContext();
+  auto &diags = ctx.Diags;
   auto attrLoc = attr->getLocation();
   auto attrName = attr->getAttrName();
   auto domainLoc = attr->getDomainLoc();
@@ -8350,22 +8366,22 @@ SemanticAvailableAttrRequest::evaluate(swift::Evaluator &evaluator,
   auto domain = attr->getCachedDomain();
 
   if (!domain) {
-    auto string = attr->getDomainString();
-    ASSERT(string);
+    auto domainName = attr->getDomainString();
+    ASSERT(domainName);
 
     // Attempt to resolve the domain specified for the attribute and diagnose
     // if no domain is found.
     auto declContext = decl->getInnermostDeclContext();
-    domain = AvailabilityDomain::builtinDomainForString(*string, declContext);
+    domain = getAvailabilityDomainForName(*domainName, declContext);
     if (!domain) {
-      if (auto suggestion = closestCorrectedPlatformString(*string)) {
+      if (auto suggestion = closestCorrectedPlatformString(*domainName)) {
         diags
             .diagnose(domainLoc, diag::attr_availability_suggest_platform,
-                      *string, attrName, *suggestion)
+                      *domainName, attrName, *suggestion)
             .fixItReplace(SourceRange(domainLoc), *suggestion);
       } else {
         diags.diagnose(attrLoc, diag::attr_availability_unknown_platform,
-                       *string, attrName);
+                       *domainName, attrName);
       }
       return std::nullopt;
     }

test/attr/attr_availability_custom_domains.swift

@@ -8,25 +8,25 @@
 
 // REQUIRES: swift_feature_CustomAvailability
 
-@available(EnabledDomain) // expected-warning {{unknown platform 'EnabledDomain' for attribute 'available'}}
+@available(EnabledDomain)
 func availableInEnabledDomain() { }
 
-@available(DisabledDomain, unavailable) // expected-warning {{unknown platform 'DisabledDomain' for attribute 'available'}}
-func availableInDisabledDomain() { }
+@available(DisabledDomain, unavailable)
+func unavailableInDisabledDomain() { } // expected-note {{'unavailableInDisabledDomain()' has been explicitly marked unavailable here}}
 
-@available(RedefinedDomain, deprecated, message: "Use something else") // expected-warning {{unknown platform 'RedefinedDomain' for attribute 'available'}}
-func availableInRedefinedDomain() { }
+@available(RedefinedDomain, deprecated, message: "Use something else")
+func deprecatedInRedefinedDomain() { }
 
-@available(DynamicDomain) // expected-warning {{unknown platform 'DynamicDomain' for attribute 'available'}}
+@available(DynamicDomain)
 func availableInDynamicDomain() { }
 
 @available(UnknownDomain) // expected-warning {{unknown platform 'UnknownDomain' for attribute 'available'}}
 func availableInUnknownDomain() { }
 
 func test() {
-  availableInEnabledDomain()
-  availableInDisabledDomain()
-  availableInRedefinedDomain()
-  availableInDynamicDomain()
-  availableInUnknownDomain()
+  availableInEnabledDomain() // FIXME: [availability] should be diagnosed
+  unavailableInDisabledDomain() // expected-error {{'unavailableInDisabledDomain()' is unavailable}}
+  deprecatedInRedefinedDomain() // expected-warning {{'deprecatedInRedefinedDomain()' is deprecated: Use something else}}
+  availableInDynamicDomain() // FIXME: [availability] should be diagnosed
+  availableInUnknownDomain() // Ok
 }