[Macros] Ensure that the type refinement context accounts for macros

Make sure that we include all of the AST nodes from a macro expansion
buffer when building the type refinement context, not just the
declarations.

Author: DougGregor

lib/AST/TypeRefinementContext.cpp

@@ -46,10 +46,25 @@ TypeRefinementContext::createRoot(SourceFile *SF,
   assert(SF);
 
   ASTContext &Ctx = SF->getASTContext();
+
+  SourceRange range;
+  TypeRefinementContext *parentContext = nullptr;
+  AvailabilityContext availabilityContext = Info;
+  if (auto parentExpansion = SF->getMacroExpansion()) {
+    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.SourceMgr);
+      availabilityContext = parentContext->getAvailabilityInfo();
+    }
+  }
+
   return new (Ctx)
-      TypeRefinementContext(Ctx, SF,
-                            /*Parent=*/nullptr, SourceRange(),
-                            Info, AvailabilityContext::alwaysAvailable());
+      TypeRefinementContext(Ctx, SF, parentContext, range,
+                            availabilityContext,
+                            AvailabilityContext::alwaysAvailable());
 }
 
 TypeRefinementContext *
@@ -147,12 +162,34 @@ TypeRefinementContext::createForWhileStmtBody(ASTContext &Ctx, WhileStmt *S,
       Ctx, S, Parent, S->getBody()->getSourceRange(), Info, /* ExplicitInfo */Info);
 }
 
+/// Determine whether the child location is somewhere within the parent
+/// range.
+static bool rangeContainsTokenLocWithGeneratedSource(
+    SourceManager &sourceMgr, SourceRange parentRange, SourceLoc childLoc) {
+  auto parentBuffer = sourceMgr.findBufferContainingLoc(parentRange.Start);
+  auto childBuffer = sourceMgr.findBufferContainingLoc(childLoc);
+  while (parentBuffer != childBuffer) {
+    auto info = sourceMgr.getGeneratedSourceInfo(childBuffer);
+    if (!info)
+      return false;
+
+    childLoc = info->originalSourceRange.getStart();
+    if (childLoc.isInvalid())
+      return false;
+
+    childBuffer = sourceMgr.findBufferContainingLoc(childLoc);
+  }
+
+  return sourceMgr.rangeContainsTokenLoc(parentRange, childLoc);
+}
+
 TypeRefinementContext *
 TypeRefinementContext::findMostRefinedSubContext(SourceLoc Loc,
                                                  SourceManager &SM) {
   assert(Loc.isValid());
 
-  if (SrcRange.isValid() && !SM.rangeContainsTokenLoc(SrcRange, Loc))
+  if (SrcRange.isValid() &&
+      !rangeContainsTokenLocWithGeneratedSource(SM, SrcRange, Loc))
     return nullptr;
 
   // For the moment, we perform a linear search here, but we can and should

lib/Sema/TypeCheckAvailability.cpp

@@ -1232,8 +1232,13 @@ void TypeChecker::buildTypeRefinementContextHierarchy(SourceFile &SF) {
   // Build refinement contexts, if necessary, for all declarations starting
   // with StartElem.
   TypeRefinementContextBuilder Builder(RootTRC, Context);
-  for (auto D : SF.getTopLevelDecls()) {
-    Builder.build(D);
+  for (auto item : SF.getTopLevelItems()) {
+    if (auto decl = item.dyn_cast<Decl *>())
+      Builder.build(decl);
+    else if (auto expr = item.dyn_cast<Expr *>())
+      Builder.build(expr);
+    else if (auto stmt = item.dyn_cast<Stmt *>())
+      Builder.build(stmt);
   }
 }
 
@@ -1272,7 +1277,11 @@ AvailabilityContext
 TypeChecker::overApproximateAvailabilityAtLocation(SourceLoc loc,
                                                    const DeclContext *DC,
                                                    const TypeRefinementContext **MostRefined) {
-  SourceFile *SF = DC->getParentSourceFile();
+  SourceFile *SF;
+  if (loc.isValid())
+    SF = DC->getParentModule()->getSourceFileContainingLocation(loc);
+  else
+    SF = DC->getParentSourceFile();
   auto &Context = DC->getASTContext();
 
   // If our source location is invalid (this may be synthesized code), climb

test/Macros/macro_availability_macosx.swift

@@ -1,6 +1,9 @@
-// REQUIRES: swift_swift_parser, OS=macosx
+// REQUIRES: swift_swift_parser, executable_test, OS=macosx
 
-// RUN: %target-typecheck-verify-swift -swift-version 5 -module-name MacrosTest -target %target-cpu-apple-macosx11
+// RUN: %empty-directory(%t)
+// RUN: %host-build-swift -swift-version 5 -emit-library -o %t/%target-library-name(MacroDefinition) -module-name=MacroDefinition %S/Inputs/syntax_macro_definitions.swift -g -no-toolchain-stdlib-rpath -swift-version 5
+
+// RUN: %target-typecheck-verify-swift -swift-version 5 -module-name MacrosTest -target %target-cpu-apple-macosx11  -load-plugin-library %t/%target-library-name(MacroDefinition)
 
 @available(macOS 12.0, *)
 struct X { }
@@ -11,3 +14,16 @@ struct X { }
 @available(macOS 12.0, *)
 @freestanding(expression) macro m2() -> X = #externalMacro(module: "A", type: "B")
 // expected-warning@-1{{external macro implementation type 'A.B' could not be found for macro 'm2()'}}
+
+@freestanding(expression) macro stringify<T>(_ value: T) -> (T, String) = #externalMacro(module: "MacroDefinition", type: "StringifyMacro")
+
+@available(macOS 12.0, *)
+func onlyInMacOS12() { }
+
+func test() {
+  _ = #stringify({
+      if #available(macOS 12.0, *) {
+        onlyInMacOS12()
+      }
+    })
+}