AST/Sema: Parse custom availability domains

include/swift/AST/AvailabilityDomain.h

@@ -26,6 +26,7 @@
 
 namespace swift {
 class ASTContext;
+class DeclContext;
 
 /// Represents a dimension of availability (e.g. macOS platform or Swift
 /// language mode).
@@ -142,7 +143,7 @@ class AvailabilityDomain final {
 
   /// Returns the built-in availability domain identified by the given string.
   static std::optional<AvailabilityDomain>
-  builtinDomainForString(StringRef string);
+  builtinDomainForString(StringRef string, const DeclContext *declContext);
 
   Kind getKind() const {
     if (auto inlineDomain = getInlineDomain())

include/swift/Basic/Features.def

@@ -447,6 +447,9 @@ SUPPRESSIBLE_EXPERIMENTAL_FEATURE(ABIAttribute, true)
 /// calling context.
 EXPERIMENTAL_FEATURE(NonIsolatedAsyncInheritsIsolationFromContext, false)
 
+/// Allow custom availability domains to be defined and referenced.
+SUPPRESSIBLE_EXPERIMENTAL_FEATURE(CustomAvailability, true)
+
 #undef EXPERIMENTAL_FEATURE_EXCLUDED_FROM_MODULE_INTERFACE
 #undef EXPERIMENTAL_FEATURE
 #undef UPCOMING_FEATURE

include/swift/Frontend/FrontendOptions.h

@@ -573,6 +573,18 @@ class FrontendOptions {
   /// All block list configuration files to be honored in this compilation.
   std::vector<std::string> BlocklistConfigFilePaths;
 
+  struct CustomAvailabilityDomains {
+    /// Domains defined with `-define-enabled-availability-domain=`.
+    llvm::SmallVector<std::string> EnabledDomains;
+    /// Domains defined with `-define-disabled-availability-domain=`.
+    llvm::SmallVector<std::string> DisabledDomains;
+    /// Domains defined with `-define-dynamic-availability-domain=`.
+    llvm::SmallVector<std::string> DynamicDomains;
+  };
+
+  /// The collection of AvailabilityDomain definitions specified as arguments.
+  CustomAvailabilityDomains AvailabilityDomains;
+
 private:
   static bool canActionEmitDependencies(ActionType);
   static bool canActionEmitReferenceDependencies(ActionType);

include/swift/Option/Options.td

@@ -561,6 +561,21 @@ def unavailable_decl_optimization_EQ : Joined<["-"], "unavailable-decl-optimizat
            "value may be 'none' (no optimization) or 'complete' (code is not "
            "generated at all unavailable declarations)">;
 
+def define_enabled_availability_domain : Separate<["-"], "define-enabled-availability-domain">,
+    Flags<[HelpHidden, FrontendOption, NoInteractiveOption, ModuleInterfaceOptionIgnorable]>,
+    HelpText<"Defines a custom availability domain that is available at compile time">,
+    MetaVarName<"<domain>">;
+
+def define_disabled_availability_domain : Separate<["-"], "define-disabled-availability-domain">,
+    Flags<[HelpHidden, FrontendOption, NoInteractiveOption, ModuleInterfaceOptionIgnorable]>,
+    HelpText<"Defines a custom availability domain that is unavailable at compile time">,
+    MetaVarName<"<domain>">;
+
+def define_dynamic_availability_domain : Separate<["-"], "define-dynamic-availability-domain">,
+    Flags<[HelpHidden, FrontendOption, NoInteractiveOption, ModuleInterfaceOptionIgnorable]>,
+    HelpText<"Defines a custom availability domain that can be enabled or disabled at runtime">,
+    MetaVarName<"<domain>">;
+
 def experimental_package_bypass_resilience : Flag<["-"], "experimental-package-bypass-resilience">,
   Flags<[FrontendOption]>,
   HelpText<"Deprecated; has no effect">;

lib/AST/ASTPrinter.cpp

@@ -3239,6 +3239,14 @@ suppressingFeatureAddressableTypes(PrintOptions &options,
   action();
 }
 
+static void
+suppressingFeatureCustomAvailability(PrintOptions &options,
+                                     llvm::function_ref<void()> action) {
+  // FIXME: [availability] Save and restore a bit controlling whether
+  // @available attributes for custom domains are printed.
+  action();
+}
+
 /// Suppress the printing of a particular feature.
 static void suppressingFeature(PrintOptions &options, Feature feature,
                                llvm::function_ref<void()> action) {

lib/AST/AvailabilityDomain.cpp

@@ -18,7 +18,11 @@
 using namespace swift;
 
 std::optional<AvailabilityDomain>
-AvailabilityDomain::builtinDomainForString(StringRef string) {
+AvailabilityDomain::builtinDomainForString(StringRef string,
+                                           const DeclContext *declContext) {
+  // This parameter is used in downstream forks, do not remove.
+  (void)declContext;
+
   auto domain = llvm::StringSwitch<std::optional<AvailabilityDomain>>(string)
                     .Case("*", AvailabilityDomain::forUniversal())
                     .Case("swift", AvailabilityDomain::forSwiftLanguage())

lib/AST/FeatureSet.cpp

@@ -407,6 +407,12 @@ static bool usesFeatureCoroutineAccessors(Decl *decl) {
   }
 }
 
+static bool usesFeatureCustomAvailability(Decl *decl) {
+  // FIXME: [availability] Check whether @available attributes for custom
+  // domains are attached to the decl.
+  return false;
+}
+
 // ----------------------------------------------------------------------------
 // MARK: - FeatureSet
 // ----------------------------------------------------------------------------

lib/Frontend/ArgsToFrontendOptionsConverter.cpp

@@ -419,6 +419,10 @@ bool ArgsToFrontendOptionsConverter::convert(
   }
 
   Opts.DisableSandbox = Args.hasArg(OPT_disable_sandbox);
+
+  if (computeAvailabilityDomains())
+    return true;
+
   return false;
 }
 
@@ -548,6 +552,39 @@ void ArgsToFrontendOptionsConverter::computeDumpScopeMapLocations() {
     Diags.diagnose(SourceLoc(), diag::error_no_source_location_scope_map);
 }
 
+bool ArgsToFrontendOptionsConverter::computeAvailabilityDomains() {
+  using namespace options;
+
+  bool hadError = false;
+  llvm::SmallSet<std::string, 4> seenDomains;
+
+  for (const Arg *A :
+       Args.filtered_reverse(OPT_define_enabled_availability_domain,
+                             OPT_define_disabled_availability_domain,
+                             OPT_define_dynamic_availability_domain)) {
+    std::string domain = A->getValue();
+    if (!seenDomains.insert(domain).second)
+      continue;
+
+    if (!Lexer::isIdentifier(domain)) {
+      Diags.diagnose(SourceLoc(), diag::error_invalid_arg_value,
+                     A->getAsString(Args), A->getValue());
+      hadError = true;
+      continue;
+    }
+
+    auto &option = A->getOption();
+    if (option.matches(OPT_define_enabled_availability_domain))
+      Opts.AvailabilityDomains.EnabledDomains.emplace_back(domain);
+    else if (option.matches(OPT_define_disabled_availability_domain))
+      Opts.AvailabilityDomains.DisabledDomains.emplace_back(domain);
+    else if (option.matches(OPT_define_dynamic_availability_domain))
+      Opts.AvailabilityDomains.DynamicDomains.emplace_back(domain);
+  }
+
+  return hadError;
+}
+
 FrontendOptions::ActionType
 ArgsToFrontendOptionsConverter::determineRequestedAction(const ArgList &args) {
   using namespace options;

lib/Frontend/ArgsToFrontendOptionsConverter.h

@@ -48,6 +48,7 @@ class ArgsToFrontendOptionsConverter {
   void computePlaygroundOptions();
   void computePrintStatsOptions();
   void computeTBDOptions();
+  bool computeAvailabilityDomains();
 
   bool setUpImmediateArgs();
 

lib/Parse/ParseDecl.cpp

@@ -537,7 +537,8 @@ ParserResult<AvailableAttr> Parser::parseExtendedAvailabilitySpecList(
     }
   }
 
-  if (!AnyAnnotations) {
+  if (!AnyAnnotations &&
+      !Context.LangOpts.hasFeature(Feature::CustomAvailability)) {
     diagnose(Tok.getLoc(), diag::attr_expected_comma, AttrName,
              /*isDeclModifier*/ false);
   }

lib/Sema/TypeCheckAttr.cpp

@@ -8271,7 +8271,6 @@ SemanticAvailableAttrRequest::evaluate(swift::Evaluator &evaluator,
   auto &diags = decl->getASTContext().Diags;
   auto attrLoc = attr->getLocation();
   auto attrName = attr->getAttrName();
-  auto attrKind = attr->getKind();
   auto domainLoc = attr->getDomainLoc();
   auto introducedVersion = attr->getRawIntroduced();
   auto deprecatedVersion = attr->getRawDeprecated();
@@ -8285,7 +8284,8 @@ SemanticAvailableAttrRequest::evaluate(swift::Evaluator &evaluator,
 
     // Attempt to resolve the domain specified for the attribute and diagnose
     // if no domain is found.
-    domain = AvailabilityDomain::builtinDomainForString(*string);
+    auto declContext = decl->getInnermostDeclContext();
+    domain = AvailabilityDomain::builtinDomainForString(*string, declContext);
     if (!domain) {
       if (auto suggestion = closestCorrectedPlatformString(*string)) {
         diags
@@ -8305,7 +8305,7 @@ SemanticAvailableAttrRequest::evaluate(swift::Evaluator &evaluator,
   auto domainName = domain->getNameForAttributePrinting();
 
   if (domain->isSwiftLanguage() || domain->isPackageDescription()) {
-    switch (attrKind) {
+    switch (attr->getKind()) {
     case AvailableAttr::Kind::Deprecated:
       diags.diagnose(attrLoc,
                      diag::attr_availability_expected_deprecated_version,
@@ -8320,7 +8320,8 @@ SemanticAvailableAttrRequest::evaluate(swift::Evaluator &evaluator,
     case AvailableAttr::Kind::NoAsync:
       diags.diagnose(attrLoc, diag::attr_availability_cannot_be_used_for_domain,
                      "noasync", attrName, domainName);
-      break;
+      return std::nullopt;
+
     case AvailableAttr::Kind::Default:
       break;
     }

test/attr/attr_availability_custom_domains.swift

@@ -0,0 +1,32 @@
+// RUN: %target-typecheck-verify-swift \
+// RUN:  -enable-experimental-feature CustomAvailability \
+// RUN:  -define-enabled-availability-domain EnabledDomain \
+// RUN:  -define-enabled-availability-domain RedefinedDomain \
+// RUN:  -define-disabled-availability-domain DisabledDomain \
+// RUN:  -define-dynamic-availability-domain DynamicDomain \
+// RUN:  -define-disabled-availability-domain RedefinedDomain
+
+// REQUIRES: swift_feature_CustomAvailability
+
+@available(EnabledDomain) // expected-warning {{unknown platform 'EnabledDomain' for attribute 'available'}}
+func availableInEnabledDomain() { }
+
+@available(DisabledDomain, unavailable) // expected-warning {{unknown platform 'DisabledDomain' for attribute 'available'}}
+func availableInDisabledDomain() { }
+
+@available(RedefinedDomain, deprecated, message: "Use something else") // expected-warning {{unknown platform 'RedefinedDomain' for attribute 'available'}}
+func availableInRedefinedDomain() { }
+
+@available(DynamicDomain) // expected-warning {{unknown platform 'DynamicDomain' for attribute 'available'}}
+func availableInDynamicDomain() { }
+
+@available(UnknownDomain) // expected-warning {{unknown platform 'UnknownDomain' for attribute 'available'}}
+func availableInUnknownDomain() { }
+
+func test() {
+  availableInEnabledDomain()
+  availableInDisabledDomain()
+  availableInRedefinedDomain()
+  availableInDynamicDomain()
+  availableInUnknownDomain()
+}

test/attr/attr_availability_noasync.swift

@@ -18,13 +18,11 @@ func asyncReplacement() async -> Int { }
 @available(*, noasync, renamed: "IOActor.readString()")
 func readStringFromIO() -> String {}
 
-// expected-warning@+2 {{'noasync' cannot be used in 'available' attribute for platform 'swift'}}
-// expected-warning@+1 {{expected 'introduced', 'deprecated', or 'obsoleted' in 'available' attribute for platform 'swift'}}
+// expected-warning@+1 {{'noasync' cannot be used in 'available' attribute for platform 'swift'}}
 @available(swift, noasync)
 func swiftNoAsync() { }
 
-// expected-warning@+2 {{'noasync' cannot be used in 'available' attribute for platform '_PackageDescription'}}
-// expected-warning@+1 {{expected 'introduced', 'deprecated', or 'obsoleted' in 'available' attribute for platform '_PackageDescription'}}
+// expected-warning@+1 {{'noasync' cannot be used in 'available' attribute for platform '_PackageDescription'}}
 @available(_PackageDescription, noasync)
 func packageDescriptionNoAsync() { }