AST: Fix availability checking for macros as default arguments.

Adoption of the `#isolation` macro as a default argument in the standard
library exposed that the availbility of contents of default argument macro
expansions were being mis-diagnosed. The nodes for auxilary files generated for
default argument expansions were not being added to the enclosing file's
`TypeRefinementContext` hierarchy, so the contents of the expansion would be
diagnosed as if it were located in a context with the minimum availability for
the compilation.

Resolves rdar://125812256

Author: tshortli

include/swift/AST/SourceFile.h

@@ -624,6 +624,11 @@ class SourceFile final : public FileUnit {
   /// enclosed in.
   SourceFile *getEnclosingSourceFile() const;
 
+  /// If this file has an enclosing source file (because it is the result of
+  /// expanding a macro or default argument), returns the node in the enclosing
+  /// file that this file's contents were expanded from.
+  ASTNode getNodeInEnclosingSourceFile() const;
+
   /// If this buffer corresponds to a file on disk, returns the path.
   /// Otherwise, return an empty string.
   StringRef getFilename() const;

include/swift/AST/TypeRefinementContext.h

@@ -193,8 +193,8 @@ class TypeRefinementContext : public ASTAllocated<TypeRefinementContext> {
 public:
 
   /// Create the root refinement context for the given SourceFile.
-  static TypeRefinementContext *createRoot(SourceFile *SF,
-                                           const AvailabilityContext &Info);
+  static TypeRefinementContext *
+  createForSourceFile(SourceFile *SF, const AvailabilityContext &Info);
 
   /// Create a refinement context for the given declaration.
   static TypeRefinementContext *createForDecl(ASTContext &Ctx, Decl *D,

lib/AST/Module.cpp

@@ -1179,9 +1179,7 @@ ASTNode SourceFile::getMacroExpansion() const {
   if (Kind != SourceFileKind::MacroExpansion)
     return nullptr;
 
-  auto genInfo =
-      *getASTContext().SourceMgr.getGeneratedSourceInfo(*getBufferID());
-  return ASTNode::getFromOpaqueValue(genInfo.astNode);
+  return getNodeInEnclosingSourceFile();
 }
 
 SourceRange SourceFile::getMacroInsertionRange() const {
@@ -1233,6 +1231,16 @@ SourceFile *SourceFile::getEnclosingSourceFile() const {
   return getParentModule()->getSourceFileContainingLocation(sourceLoc);
 }
 
+ASTNode SourceFile::getNodeInEnclosingSourceFile() const {
+  if (Kind != SourceFileKind::MacroExpansion &&
+      Kind != SourceFileKind::DefaultArgument)
+    return nullptr;
+
+  auto genInfo =
+      *getASTContext().SourceMgr.getGeneratedSourceInfo(*getBufferID());
+  return ASTNode::getFromOpaqueValue(genInfo.astNode);
+}
+
 void ModuleDecl::lookupClassMember(ImportPath::Access accessPath,
                                    DeclName name,
                                    SmallVectorImpl<ValueDecl*> &results) const {

lib/AST/TypeRefinementContext.cpp

@@ -42,24 +42,38 @@ TypeRefinementContext::TypeRefinementContext(ASTContext &Ctx, IntroNode Node,
 }
 
 TypeRefinementContext *
-TypeRefinementContext::createRoot(SourceFile *SF,
-                                  const AvailabilityContext &Info) {
+TypeRefinementContext::createForSourceFile(SourceFile *SF,
+                                           const AvailabilityContext &Info) {
   assert(SF);
 
   ASTContext &Ctx = SF->getASTContext();
 
   SourceRange range;
   TypeRefinementContext *parentContext = nullptr;
   AvailabilityContext availabilityContext = Info;
-  if (auto parentExpansion = SF->getMacroExpansion()) {
+  switch (SF->Kind) {
+  case SourceFileKind::MacroExpansion:
+  case SourceFileKind::DefaultArgument: {
+    // Look up the parent context in the enclosing file that this file's
+    // root context should be nested under.
     if (auto parentTRC =
             SF->getEnclosingSourceFile()->getTypeRefinementContext()) {
       auto charRange = Ctx.SourceMgr.getRangeForBuffer(*SF->getBufferID());
       range = SourceRange(charRange.getStart(), charRange.getEnd());
-      parentContext = parentTRC->findMostRefinedSubContext(
-          parentExpansion.getStartLoc(), Ctx);
-      availabilityContext = parentContext->getAvailabilityInfo();
+      auto originalNode = SF->getNodeInEnclosingSourceFile();
+      parentContext =
+          parentTRC->findMostRefinedSubContext(originalNode.getStartLoc(), Ctx);
+      if (parentContext)
+        availabilityContext = parentContext->getAvailabilityInfo();
     }
+    break;
+  }
+  case SourceFileKind::Library:
+  case SourceFileKind::Main:
+  case SourceFileKind::Interface:
+    break;
+  case SourceFileKind::SIL:
+    llvm_unreachable("unexpected SourceFileKind");
   }
 
   return new (Ctx)

lib/Sema/TypeCheckAvailability.cpp

@@ -1279,7 +1279,7 @@ void TypeChecker::buildTypeRefinementContextHierarchy(SourceFile &SF) {
     // the source file are guaranteed to be executing on at least the minimum
     // platform version for inlining.
     auto MinPlatformReq = AvailabilityContext::forInliningTarget(Context);
-    RootTRC = TypeRefinementContext::createRoot(&SF, MinPlatformReq);
+    RootTRC = TypeRefinementContext::createForSourceFile(&SF, MinPlatformReq);
     SF.setTypeRefinementContext(RootTRC);
   }
 

test/Concurrency/isolation_macro_availability.swift

@@ -0,0 +1,25 @@
+// RUN: %target-typecheck-verify-swift
+
+// REQUIRES: concurrency
+// REQUIRES: swift_swift_parser
+// REQUIRES: VENDOR=apple
+
+@available(SwiftStdlib 5.1, *)
+func isolatedFunc(isolation: isolated (any Actor)? = #isolation) {}
+
+func test() { // expected-note 3 {{add @available attribute to enclosing global function}}
+  _ = #isolation // expected-error {{'isolation()' is only available in}} expected-note {{add 'if #available' version check}}
+  isolatedFunc() // expected-error {{'isolatedFunc(isolation:)' is only available in}} expected-note {{add 'if #available' version check}}
+  // expected-error@-1 {{'isolation()' is only available in}}
+
+  if #available(SwiftStdlib 5.1, *) {
+    _ = #isolation
+    isolatedFunc()
+  }
+}
+
+@available(SwiftStdlib 5.1, *)
+func testAvailable5_1() {
+  _ = #isolation
+  isolatedFunc()
+}