Merge pull request swiftlang#77517 from DougGregor/rtc-macro-expansion

Teach the TypeRefinementContext not to skip declarations within macro expansions

Author: DougGregor

include/swift/AST/ASTWalker.h

@@ -616,6 +616,11 @@ class ASTWalker {
     return MacroWalking::ArgumentsAndExpansion;
   }
 
+  /// This method determines whether the given declaration should be
+  /// considered to be in a macro expansion context. It can be configured
+  /// by subclasses.
+  virtual bool isDeclInMacroExpansion(Decl *decl) const;
+
   /// Determine whether we should walk macro arguments (as they appear in
   /// source) and the expansion (which is semantically part of the program).
   std::pair<bool, bool> shouldWalkMacroArgumentsAndExpansion() const {

lib/AST/ASTWalker.cpp

@@ -63,6 +63,10 @@ using namespace swift;
 
 void ASTWalker::anchor() {}
 
+bool ASTWalker::isDeclInMacroExpansion(Decl *decl) const {
+  return decl->isInMacroExpansionInContext();
+}
+
 namespace {
 
 /// Traversal - This class implements a simple expression/statement
@@ -1540,7 +1544,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
 
   bool shouldSkip(Decl *D) {
     if (!Walker.shouldWalkMacroArgumentsAndExpansion().second &&
-        D->isInMacroExpansionInContext() && !Walker.Parent.isNull())
+        Walker.isDeclInMacroExpansion(D) && !Walker.Parent.isNull())
       return true;
 
     if (auto *VD = dyn_cast<VarDecl>(D)) {

lib/AST/SwiftNameTranslation.cpp

@@ -63,9 +63,6 @@ getErrorDomainStringForObjC(const EnumDecl *ED) {
   // Should have already been diagnosed as diag::objc_enum_generic.
   assert(!ED->isGenericContext() && "Trying to bridge generic enum error to Obj-C");
 
-  // Clang decls have custom domains, but we shouldn't see them here anyway.
-  assert(!ED->getClangDecl() && "clang decls shouldn't be re-exported");
-
   SmallVector<const NominalTypeDecl *, 4> outerTypes;
   for (const NominalTypeDecl * D = ED;
        D != nullptr;

lib/Sema/TypeCheckAvailability.cpp

@@ -535,6 +535,69 @@ class TypeRefinementContextBuilder : private ASTWalker {
     return MacroWalking::Arguments;
   }
 
+  /// Check whether this declaration is in a source file buried within
+  /// a macro expansion of the
+  bool isDeclInMacroExpansion(Decl *decl) const override {
+    // If it's not in a macro expansion relative to its context, it's not
+    // considered to be in a macro expansion.
+    if (!decl->isInMacroExpansionInContext())
+      return false;
+
+    auto module = decl->getDeclContext()->getParentModule();
+    auto *declFile = module->getSourceFileContainingLocation(decl->getLoc());
+    if (!declFile)
+      return false;
+
+    // Look for a parent context that implies that we are producing a
+    // type refinement context for this expansion.
+    for (auto iter = ContextStack.rbegin(), endIter = ContextStack.rend();
+         iter != endIter; ++iter) {
+      const auto &context = *iter;
+      if (auto contextRTC = context.TRC) {
+        // If the context is the same source file, don't treat it as an
+        // expansion.
+        auto introNode = contextRTC->getIntroductionNode();
+        switch (auto reason = contextRTC->getReason()) {
+        case TypeRefinementContext::Reason::Root:
+          if (auto contextFile = introNode.getAsSourceFile())
+            if (declFile == contextFile)
+              return false;
+
+          break;
+
+        case TypeRefinementContext::Reason::Decl:
+        case TypeRefinementContext::Reason::DeclImplicit:
+          // If the context is a declaration, check whether the declaration
+          // is in the same source file as this declaration.
+          if (auto contextDecl = introNode.getAsDecl()) {
+            if (decl == contextDecl)
+              return false;
+
+            auto contextModule =
+                contextDecl->getDeclContext()->getParentModule();
+            SourceLoc contextDeclLoc = contextDecl->getLoc();
+            auto contextDeclFile =
+                contextModule->getSourceFileContainingLocation(contextDeclLoc);
+            if (declFile == contextDeclFile)
+              return false;
+          }
+          break;
+
+          case TypeRefinementContext::Reason::IfStmtThenBranch:
+          case TypeRefinementContext::Reason::IfStmtElseBranch:
+          case TypeRefinementContext::Reason::ConditionFollowingAvailabilityQuery:
+          case TypeRefinementContext::Reason::GuardStmtFallthrough:
+          case TypeRefinementContext::Reason::GuardStmtElseBranch:
+          case TypeRefinementContext::Reason::WhileStmtBody:
+            // Nothing to check here.
+            break;
+        }
+      }
+    }
+
+    return true;
+  }
+
   bool shouldSkipDecl(Decl *D) const {
     // Only visit a node that has a corresponding concrete syntax node if we are
     // already walking that concrete syntax node.

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -758,6 +758,39 @@ public struct AddArbitraryMembers: MemberMacro {
   }
 }
 
+public struct MemberThatCallsCodeMacro: MemberMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingMembersOf decl: some DeclGroupSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    guard case let .argumentList(arguments) = node.arguments,
+          let firstElement = arguments.first,
+          let stringLiteral = firstElement.expression.as(StringLiteralExprSyntax.self) else {
+      throw CustomError.message("Macro requires a string literal")
+    }
+
+    let codeString = try stringLiteral.segments.map { segment in
+      switch segment {
+      case .stringSegment(let string):
+        return string.content.text
+
+      case .expressionSegment(_):
+        throw CustomError.message("Macro cannot handle string interpolation")
+      }
+    }.joined(separator: "")
+    return [
+      """
+      static var synthesizedMember: String {
+        \(raw: codeString)
+
+        return "hello"
+      }
+      """,
+    ]
+  }
+}
+
 /// Implementation of the `wrapStoredProperties` macro, which can be
 /// used to apply an attribute to all of the stored properties of a type.
 ///

test/Macros/macro_availability_macosx.swift

@@ -17,6 +17,9 @@ struct X { }
 
 @freestanding(expression) macro stringify<T>(_ value: T) -> (T, String) = #externalMacro(module: "MacroDefinition", type: "StringifyMacro")
 
+@attached(member, names: named(synthesizedMember))
+macro MemberThatCallsCode(_ codeString: String) = #externalMacro(module: "MacroDefinition", type: "MemberThatCallsCodeMacro")
+
 @available(macOS 12.0, *)
 func onlyInMacOS12() { }
 
@@ -27,3 +30,11 @@ func test() {
       }
     })
 }
+
+@MemberThatCallsCode("""
+  if #available(macOS 12.0, *) {
+    onlyInMacOS12()
+  }
+  """)
+struct HasMembers {
+}