AST: Generalize availability fix-its to support custom availability domains.

Resolves rdar://156118254.

Author: tshortli

include/swift/AST/DiagnosticsSema.def

@@ -7070,9 +7070,6 @@ NOTE(availability_guard_with_version_check, none,
 
 NOTE(availability_add_attribute, none,
      "add '@available' attribute to enclosing %kindonly0", (const Decl *))
-FIXIT(insert_available_attr,
-      "@available(%0 %1, *)\n%2",
-      (StringRef, StringRef, StringRef))
 
 ERROR(availability_inout_accessor_only_in, none,
       "cannot pass as inout because %0 is only available in %1"

lib/Sema/TypeCheckAvailability.cpp

@@ -659,10 +659,18 @@ static void fixAvailabilityForDecl(
   StringRef OriginalIndent =
       Lexer::getIndentationForLine(Context.SourceMgr, InsertLoc);
 
+  llvm::SmallString<64> FixItBuffer;
+  llvm::raw_svector_ostream FixIt(FixItBuffer);
+
+  FixIt << "@available(" << Domain.getNameForAttributePrinting();
+  if (Domain.isVersioned())
+    FixIt << " " <<  RequiredAvailability.getVersionString();
+  if (Domain.isPlatform())
+    FixIt << ", *";
+  FixIt << ")\n" << OriginalIndent;
+
   D->diagnose(diag::availability_add_attribute, DeclForDiagnostic)
-      .fixItInsert(InsertLoc, diag::insert_available_attr,
-                   Domain.getNameForAttributePrinting(),
-                   RequiredAvailability.getVersionString(), OriginalIndent);
+      .fixItInsert(InsertLoc, FixIt.str());
 }
 
 /// In the special case of being in an existing, nontrivial availability scope
@@ -684,7 +692,6 @@ static bool fixAvailabilityByNarrowingNearbyVersionCheck(
   if (!scope)
     return false;
 
-  // FIXME: [availability] Support fixing availability for versionless domains.
   auto ExplicitAvailability =
       scope->getExplicitAvailabilityRange(Domain, Context);
   if (ExplicitAvailability && !RequiredAvailability.isAlwaysAvailable() &&
@@ -722,8 +729,9 @@ static bool fixAvailabilityByNarrowingNearbyVersionCheck(
 /// Emit a diagnostic note and Fix-It to add an if #available(...) { } guard
 /// that checks for the given version range around the given node.
 static void fixAvailabilityByAddingVersionCheck(
-    ASTNode NodeToWrap, const AvailabilityRange &RequiredAvailability,
-    SourceRange ReferenceRange, ASTContext &Context) {
+    ASTNode NodeToWrap, AvailabilityDomain Domain,
+    const AvailabilityRange &RequiredAvailability, SourceRange ReferenceRange,
+    ASTContext &Context) {
   // If this is an implicit variable that wraps an expression,
   // let's point to it's initializer. For example, result builder
   // transform captures expressions into implicit variables.
@@ -768,24 +776,32 @@ static void fixAvailabilityByAddingVersionCheck(
       StartAt += NewLine.length();
     }
 
-    PlatformKind Target = targetPlatform(Context.LangOpts);
+    AvailabilityDomain QueryDomain = Domain;
 
     // Runtime availability checks that specify app extension platforms don't
     // work, so only suggest checks against the base platform.
-    if (auto TargetRemovingAppExtension =
-            basePlatformForExtensionPlatform(Target))
-      Target = *TargetRemovingAppExtension;
+    if (auto CanonicalPlatform =
+            basePlatformForExtensionPlatform(QueryDomain.getPlatformKind())) {
+      QueryDomain = AvailabilityDomain::forPlatform(*CanonicalPlatform);
+    }
+
+    Out << "if #available(" << QueryDomain.getNameForAttributePrinting();
+    if (QueryDomain.isVersioned())
+      Out << " " << RequiredAvailability.getVersionString();
+    if (QueryDomain.isPlatform())
+      Out << ", *";
 
-    Out << "if #available(" << platformString(Target) << " "
-        << RequiredAvailability.getVersionString() << ", *) {\n";
+    Out << ") {\n";
 
     Out << OriginalIndent << ExtraIndent << GuardedText << "\n";
 
     // We emit an empty fallback case with a comment to encourage the developer
     // to think explicitly about whether fallback on earlier versions is needed.
     Out << OriginalIndent << "} else {\n";
-    Out << OriginalIndent << ExtraIndent << "// Fallback on earlier versions\n";
-    Out << OriginalIndent << "}";
+    Out << OriginalIndent << ExtraIndent << "// Fallback";
+    if (QueryDomain.isVersioned())
+      Out << " on earlier versions";
+    Out << "\n" << OriginalIndent << "}";
   }
 
   Context.Diags.diagnose(
@@ -803,10 +819,6 @@ static void fixAvailability(SourceRange ReferenceRange,
   if (ReferenceRange.isInvalid())
     return;
 
-  // FIXME: [availability] Support non-platform domains.
-  if (!Domain.isPlatform())
-    return;
-
   std::optional<ASTNode> NodeToWrapInVersionCheck;
   const Decl *FoundMemberDecl = nullptr;
   const Decl *FoundTypeLevelDecl = nullptr;
@@ -818,8 +830,8 @@ static void fixAvailability(SourceRange ReferenceRange,
   // Suggest wrapping in if #available(...) { ... } if possible.
   if (NodeToWrapInVersionCheck.has_value()) {
     fixAvailabilityByAddingVersionCheck(NodeToWrapInVersionCheck.value(),
-                                        RequiredAvailability, ReferenceRange,
-                                        Context);
+                                        Domain, RequiredAvailability,
+                                        ReferenceRange, Context);
   }
 
   // Suggest adding availability attributes.

test/Availability/availability_custom_domains.swift

@@ -29,22 +29,25 @@ func unavailableInDynamicDomain() { } // expected-note * {{'unavailableInDynamic
 @available(UnknownDomain) // expected-error {{unrecognized platform name 'UnknownDomain'}}
 func availableInUnknownDomain() { }
 
-func testDeployment() {
+func testDeployment() { // expected-note 2 {{add '@available' attribute to enclosing global function}}
   alwaysAvailable()
   availableInEnabledDomain() // expected-error {{'availableInEnabledDomain()' is only available in EnabledDomain}}
+  // expected-note@-1 {{add 'if #available' version check}}
   unavailableInEnabledDomain() // expected-error {{'unavailableInEnabledDomain()' is unavailable}}
   unavailableInDisabledDomain() // expected-error {{'unavailableInDisabledDomain()' is unavailable}}
   deprecatedInDynamicDomain() // expected-warning {{'deprecatedInDynamicDomain()' is deprecated: Use something else}}
   unavailableInDynamicDomain() // expected-error {{'unavailableInDynamicDomain()' is unavailable}}
   availableInDynamicDomain() // expected-error {{'availableInDynamicDomain()' is only available in DynamicDomain}}
+  // expected-note@-1 {{add 'if #available' version check}}
   availableInUnknownDomain()
 }
 
-func testIfAvailable(_ truthy: Bool) {
+func testIfAvailable(_ truthy: Bool) { // expected-note 9 {{add '@available' attribute to enclosing global function}}
   if #available(EnabledDomain) { // expected-note {{enclosing scope here}}
     availableInEnabledDomain()
     unavailableInEnabledDomain() // expected-error {{'unavailableInEnabledDomain()' is unavailable}}
     availableInDynamicDomain() // expected-error {{'availableInDynamicDomain()' is only available in DynamicDomain}}
+    // expected-note@-1 {{add 'if #available' version check}}
     unavailableInDynamicDomain() // expected-error {{'unavailableInDynamicDomain()' is unavailable}}
 
     if #available(DynamicDomain) {
@@ -56,6 +59,7 @@ func testIfAvailable(_ truthy: Bool) {
       availableInEnabledDomain()
       unavailableInEnabledDomain() // expected-error {{'unavailableInEnabledDomain()' is unavailable}}
       availableInDynamicDomain() // expected-error {{'availableInDynamicDomain()' is only available in DynamicDomain}}
+      // expected-note@-1 {{add 'if #available' version check}}
       unavailableInDynamicDomain()
     }
 
@@ -66,12 +70,15 @@ func testIfAvailable(_ truthy: Bool) {
       availableInEnabledDomain()
       unavailableInEnabledDomain() // expected-error {{'unavailableInEnabledDomain()' is unavailable}}
       availableInDynamicDomain() // expected-error {{'availableInDynamicDomain()' is only available in DynamicDomain}}
+      // expected-note@-1 {{add 'if #available' version check}}
       unavailableInDynamicDomain() // expected-error {{'unavailableInDynamicDomain()' is unavailable}}
     }
   } else {
     availableInEnabledDomain() // expected-error {{'availableInEnabledDomain()' is only available in EnabledDomain}}
+    // expected-note@-1 {{add 'if #available' version check}}
     unavailableInEnabledDomain()
     availableInDynamicDomain() // expected-error {{'availableInDynamicDomain()' is only available in DynamicDomain}}
+    // expected-note@-1 {{add 'if #available' version check}}
     unavailableInDynamicDomain() // expected-error {{'unavailableInDynamicDomain()' is unavailable}}
   }
 
@@ -84,8 +91,10 @@ func testIfAvailable(_ truthy: Bool) {
     // In this branch, we only know that one of the domains is unavailable,
     // but we don't know which.
     availableInEnabledDomain() // expected-error {{'availableInEnabledDomain()' is only available in EnabledDomain}}
+    // expected-note@-1 {{add 'if #available' version check}}
     unavailableInEnabledDomain() // expected-error {{'unavailableInEnabledDomain()' is unavailable}}
     availableInDynamicDomain() // expected-error {{'availableInDynamicDomain()' is only available in DynamicDomain}}
+    // expected-note@-1 {{add 'if #available' version check}}
     unavailableInDynamicDomain() // expected-error {{'unavailableInDynamicDomain()' is unavailable}}
   }
 
@@ -96,11 +105,13 @@ func testIfAvailable(_ truthy: Bool) {
     // In this branch, the state of EnabledDomain remains unknown since
     // execution will reach here if "truthy" is false.
     availableInEnabledDomain() // expected-error {{'availableInEnabledDomain()' is only available in EnabledDomain}}
+    // expected-note@-1 {{add 'if #available' version check}}
     unavailableInEnabledDomain() // expected-error {{'unavailableInEnabledDomain()' is unavailable}}
   }
 
   if #unavailable(EnabledDomain) {
     availableInEnabledDomain() // expected-error {{'availableInEnabledDomain()' is only available in EnabledDomain}}
+    // expected-note@-1 {{add 'if #available' version check}}
     unavailableInEnabledDomain()
   } else {
     availableInEnabledDomain()
@@ -113,7 +124,7 @@ func testIfAvailable(_ truthy: Bool) {
   }
 }
 
-func testWhileAvailable() {
+func testWhileAvailable() { // expected-note {{add '@available' attribute to enclosing global function}}
   while #available(EnabledDomain) { // expected-note {{enclosing scope here}}
     availableInEnabledDomain()
     unavailableInEnabledDomain() // expected-error {{'unavailableInEnabledDomain()' is unavailable}}
@@ -124,32 +135,36 @@ func testWhileAvailable() {
 
   while #unavailable(EnabledDomain) {
     availableInEnabledDomain() // expected-error {{'availableInEnabledDomain()' is only available in EnabledDomain}}
+    // expected-note@-1 {{add 'if #available' version check}}
     unavailableInEnabledDomain()
 
     if #available(EnabledDomain) {} // FIXME: [availability] Diagnose as unreachable
     if #unavailable(EnabledDomain) {} // FIXME: [availability] Diagnose as redundant
   }
 }
 
-func testGuardAvailable() {
+func testGuardAvailable() { // expected-note 3 {{add '@available' attribute to enclosing global function}}
   guard #available(EnabledDomain) else { // expected-note {{enclosing scope here}}
     availableInEnabledDomain() // expected-error {{'availableInEnabledDomain()' is only available in EnabledDomain}}
+    // expected-note@-1 {{add 'if #available' version check}}
     unavailableInEnabledDomain()
     availableInDynamicDomain() // expected-error {{'availableInDynamicDomain()' is only available in DynamicDomain}}
+    // expected-note@-1 {{add 'if #available' version check}}
 
     return
   }
 
   availableInEnabledDomain()
   unavailableInEnabledDomain() // expected-error {{'unavailableInEnabledDomain()' is unavailable}}
   availableInDynamicDomain() // expected-error {{'availableInDynamicDomain()' is only available in DynamicDomain}}
+  // expected-note@-1 {{add 'if #available' version check}}
 
   if #available(EnabledDomain) {} // expected-warning {{unnecessary check for 'EnabledDomain'; enclosing scope ensures guard will always be true}}
   if #unavailable(EnabledDomain) {} // FIXME: [availability] Diagnose as unreachable
 }
 
 @available(EnabledDomain)
-func testEnabledDomainAvailable() { // expected-note {{enclosing scope here}}
+func testEnabledDomainAvailable() { // expected-note {{add '@available' attribute to enclosing global function}} expected-note {{enclosing scope here}}
   availableInEnabledDomain()
   unavailableInEnabledDomain() // expected-error {{'unavailableInEnabledDomain()' is unavailable}}
 
@@ -160,12 +175,14 @@ func testEnabledDomainAvailable() { // expected-note {{enclosing scope here}}
   unavailableInDisabledDomain() // expected-error {{'unavailableInDisabledDomain()' is unavailable}}
   deprecatedInDynamicDomain() // expected-warning {{'deprecatedInDynamicDomain()' is deprecated: Use something else}}
   availableInDynamicDomain() // expected-error {{'availableInDynamicDomain()' is only available in DynamicDomain}}
+  // expected-note@-1 {{add 'if #available' version check}}
   availableInUnknownDomain()
 }
 
 @available(EnabledDomain, unavailable)
-func testEnabledDomainUnavailable() {
+func testEnabledDomainUnavailable() { // expected-note {{add '@available' attribute to enclosing global function}}
   availableInEnabledDomain() // expected-error {{'availableInEnabledDomain()' is only available in EnabledDomain}}
+  // expected-note@-1 {{add 'if #available' version check}}
   unavailableInEnabledDomain()
 
   if #available(EnabledDomain) {} // FIXME: [availability] Diagnose as unreachable
@@ -175,6 +192,7 @@ func testEnabledDomainUnavailable() {
   unavailableInDisabledDomain() // expected-error {{'unavailableInDisabledDomain()' is unavailable}}
   deprecatedInDynamicDomain() // expected-warning {{'deprecatedInDynamicDomain()' is deprecated: Use something else}}
   availableInDynamicDomain() // expected-error {{'availableInDynamicDomain()' is only available in DynamicDomain}}
+  // expected-note@-1 {{add 'if #available' version check}}
   availableInUnknownDomain()
 }
 
@@ -184,9 +202,11 @@ func testUniversallyUnavailable() {
   // FIXME: [availability] Diagnostic consistency: potentially unavailable declaration shouldn't be diagnosed
   // in contexts that are unavailable to broader domains
   availableInEnabledDomain() // expected-error {{'availableInEnabledDomain()' is only available in EnabledDomain}}
+  // expected-note@-1 {{add 'if #available' version check}}
   unavailableInDisabledDomain()
   deprecatedInDynamicDomain() // expected-warning {{'deprecatedInDynamicDomain()' is deprecated: Use something else}}
   availableInDynamicDomain() // expected-error {{'availableInDynamicDomain()' is only available in DynamicDomain}}
+  // expected-note@-1 {{add 'if #available' version check}}
   availableInUnknownDomain()
 
   if #available(EnabledDomain) {} // FIXME: [availability] Diagnose?
@@ -205,6 +225,12 @@ struct EnabledDomainAvailable {
   }
 }
 
+func testFixIts() {
+  // expected-note@-1 {{add '@available' attribute to enclosing global function}}{{1-1=@available(EnabledDomain)\n}}
+  availableInEnabledDomain() // expected-error {{'availableInEnabledDomain()' is only available in EnabledDomain}}
+  // expected-note@-1 {{add 'if #available' version check}}{{3-29=if #available(EnabledDomain) {\n      availableInEnabledDomain()\n  \} else {\n      // Fallback\n  \}}}
+}
+
 protocol P { }
 
 @available(EnabledDomain)

test/ClangImporter/Inputs/availability_custom_domains_other.swift

@@ -6,6 +6,7 @@ func availableInArctic() { }
 @available(Mediterranean)
 func availableInMediterranean() { }
 
-func testOtherClangDecls() {
+func testOtherClangDecls() { // expected-note {{add '@available' attribute to enclosing global function}}
   available_in_baltic() // expected-error {{'available_in_baltic()' is only available in Baltic}}
+  // expected-note@-1 {{add 'if #available' version check}}
 }

test/ClangImporter/availability_custom_domains.swift

@@ -23,11 +23,14 @@
 
 import Oceans // re-exports Rivers
 
-func testClangDecls() {
+func testClangDecls() { // expected-note 3 {{add '@available' attribute to enclosing global function}}
   available_in_arctic() // expected-error {{'available_in_arctic()' is only available in Arctic}}
+  // expected-note@-1 {{add 'if #available' version check}}
   unavailable_in_pacific() // expected-error {{'unavailable_in_pacific()' is unavailable}}
   available_in_colorado_river_delta() // expected-error {{'available_in_colorado_river_delta()' is only available in Pacific}}
+  // expected-note@-1 {{add 'if #available' version check}}
   available_in_colorado() // expected-error {{'available_in_colorado()' is only available in Colorado}}
+  // expected-note@-1 {{add 'if #available' version check}}
   available_in_baltic() // expected-error {{cannot find 'available_in_baltic' in scope}}
 }
 
@@ -47,12 +50,15 @@ func unavailableInColorado() { } // expected-note {{'unavailableInColorado()' ha
 @available(Baltic) // expected-error {{unrecognized platform name 'Baltic'}}
 func availableInBaltic() { } // expected-note {{did you mean 'availableInBaltic'}}
 
-func testSwiftDecls() {
+func testSwiftDecls() { // expected-note 3 {{add '@available' attribute to enclosing global function}}
   availableInBayBridge() // expected-error {{'availableInBayBridge()' is only available in BayBridge}}
+  // expected-note@-1 {{add 'if #available' version check}}
   unavailableInBayBridge() // expected-error {{'unavailableInBayBridge()' is unavailable}}
   availableInArctic()
   availableInPacific() // expected-error {{'availableInPacific()' is only available in Pacific}}
+  // expected-note@-1 {{add 'if #available' version check}}
   unavailableInColorado() // expected-error {{'unavailableInColorado()' is unavailable}}
   availableInBaltic()
   availableInMediterranean() // expected-error {{'availableInMediterranean()' is only available in Mediterranean}}
+  // expected-note@-1 {{add 'if #available' version check}}
 }

test/Serialization/availability_custom_domains.swift

@@ -31,7 +31,8 @@ public func unavailableInColorado() { }
 
 import lib
 
-func test() {
+func test() { // expected-note {{add '@available' attribute to enclosing global function}}
   availableInPacific() // expected-error {{'availableInPacific()' is only available in Pacific}}
+  // expected-note@-1 {{add 'if #available' version check}}
   unavailableInColorado() // expected-error {{'unavailableInColorado()' is unavailable}}
 }