AST: Fix AvailabilityContext::forDeclSignature().

For declarations that aren't in source files there isn't an `AvailabilityScope`
tree to consult to find their `AvailabilityContext`. The fallback mechanism
used by `AvailabilityContext::forLocation()` walks up the `DeclContext`
hierarchy to build up the context for implicit code but this isn't entirely
correct since it skips over certain kinds of decls (like property
declarations). Instead, we need to walk up the parent decl hierarchy that is
used for availability inference to ensure that no decls with relevant
availability attributes are skipped.

Author: tshortli

lib/AST/AvailabilityContext.cpp

@@ -18,6 +18,7 @@
 #include "swift/AST/AvailabilityScope.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/Module.h"
+#include "swift/AST/SourceFile.h"
 #include "swift/Basic/Assertions.h"
 
 using namespace swift;
@@ -155,6 +156,30 @@ AvailabilityContext::forDeploymentTarget(const ASTContext &ctx) {
       AvailabilityRange::forDeploymentTarget(ctx), ctx);
 }
 
+static AvailabilityContext
+contextForLocationInSourceFile(SourceLoc loc, SourceFile *sf,
+                               AvailabilityContext baseAvailability,
+                               const AvailabilityScope **refinedScope) {
+  if (!sf || loc.isInvalid())
+    return baseAvailability;
+
+  auto &ctx = sf->getASTContext();
+  auto *rootScope = AvailabilityScope::getOrBuildForSourceFile(*sf);
+  if (!rootScope)
+    return baseAvailability;
+
+  auto *scope = rootScope->findMostRefinedSubContext(loc, ctx);
+  if (!scope)
+    return baseAvailability;
+
+  if (refinedScope)
+    *refinedScope = scope;
+
+  auto availability = scope->getAvailabilityContext();
+  availability.constrainWithContext(baseAvailability, ctx);
+  return availability;
+}
+
 AvailabilityContext
 AvailabilityContext::forLocation(SourceLoc loc, const DeclContext *declContext,
                                  const AvailabilityScope **refinedScope) {
@@ -191,37 +216,25 @@ AvailabilityContext::forLocation(SourceLoc loc, const DeclContext *declContext,
     declContext = decl->getDeclContext();
   }
 
-  if (!sf || loc.isInvalid())
-    return baseAvailability;
-
-  auto *rootScope = AvailabilityScope::getOrBuildForSourceFile(*sf);
-  if (!rootScope)
-    return baseAvailability;
-
-  auto *scope = rootScope->findMostRefinedSubContext(loc, ctx);
-  if (!scope)
-    return baseAvailability;
-
-  if (refinedScope)
-    *refinedScope = scope;
-
-  auto availability = scope->getAvailabilityContext();
-  availability.constrainWithContext(baseAvailability, ctx);
-  return availability;
+  return contextForLocationInSourceFile(loc, sf, baseAvailability,
+                                        refinedScope);
 }
 
 AvailabilityContext AvailabilityContext::forDeclSignature(const Decl *decl) {
-  // For decls with valid source locations, query the availability scope tree.
+  // For decls with valid source locations in source files, we can query the
+  // availability scope tree.
   auto loc = decl->getLoc();
-  if (loc.isValid())
-    return forLocation(loc, decl->getInnermostDeclContext());
-
-  // Otherwise, walk the decl hierarchy to compute availability. This can't be
-  // delegated to `AvailabilityContext::forLocation()` since it walks up the
-  // `DeclContext` hierachy for invalid source locations and that may skip
-  // some declarations with availability attributes.
   auto &ctx = decl->getASTContext();
   auto availability = forInliningTarget(ctx);
+
+  if (loc.isValid()) {
+    if (auto sf = decl->getDeclContext()
+                      ->getParentModule()
+                      ->getSourceFileContainingLocation(loc))
+      return contextForLocationInSourceFile(loc, sf, availability, nullptr);
+  }
+
+  // Otherwise, just walk the decl hierarchy to compute availability.
   while (decl) {
     availability.constrainWithDecl(decl);
     decl = decl->parentDeclForAvailability();