[Macros] Create type refinement context for MacroExpansionDecl

Author: rintaro

lib/AST/ASTWalker.cpp

@@ -1482,7 +1482,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
 
   bool shouldSkip(Decl *D) {
     if (!Walker.shouldWalkMacroArgumentsAndExpansion().second &&
-        D->isInMacroExpansionInContext())
+        D->isInMacroExpansionInContext() && !Walker.Parent.isNull())
       return true;
 
     if (auto *VD = dyn_cast<VarDecl>(D)) {

lib/AST/Decl.cpp

@@ -17,6 +17,7 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTMangler.h"
+#include "swift/AST/ASTPrinter.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/AccessRequests.h"
 #include "swift/AST/AccessScope.h"
@@ -10078,6 +10079,13 @@ void swift::simple_display(llvm::raw_ostream &out, const Decl *decl) {
       typeRepr->print(out);
     else
       ext->getSelfNominalTypeDecl()->dumpRef(out);
+  } else if (auto med = dyn_cast<MacroExpansionDecl>(decl)) {
+    out << '#' << med->getMacroName() << " in ";
+    printContext(out, med->getDeclContext());
+    if (med->getLoc().isValid()) {
+      out << '@';
+      med->getLoc().print(out, med->getASTContext().SourceMgr);
+    }
   } else {
     out << "(unknown decl)";
   }

lib/Sema/TypeCheckAvailability.cpp

@@ -473,6 +473,7 @@ class TypeRefinementContextBuilder : private ASTWalker {
 
 private:
   MacroWalking getMacroWalkingBehavior() const override {
+    // Expansion buffers will have their type refinement contexts built lazily.
     return MacroWalking::Arguments;
   }
 
@@ -558,7 +559,7 @@ class TypeRefinementContextBuilder : private ASTWalker {
   /// Returns a new context to be introduced for the declaration, or nullptr
   /// if no new context should be introduced.
   TypeRefinementContext *getNewContextForSignatureOfDecl(Decl *D) {
-    if (!isa<ValueDecl>(D) && !isa<ExtensionDecl>(D))
+    if (!isa<ValueDecl>(D) && !isa<ExtensionDecl>(D) && !isa<MacroExpansionDecl>(D))
       return nullptr;
 
     // Only introduce for an AbstractStorageDecl if it is not local. We
@@ -1430,7 +1431,9 @@ class InnermostAncestorFinder : private ASTWalker {
   Optional<ASTNode> getInnermostMatchingNode() { return InnermostMatchingNode; }
 
   MacroWalking getMacroWalkingBehavior() const override {
-    return MacroWalking::ArgumentsAndExpansion;
+    // This is SourceRange based finder. 'SM.rangeContains()' fails anyway when
+    // crossing source buffers.
+    return MacroWalking::Arguments;
   }
 
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
@@ -3192,7 +3195,8 @@ class ExprAvailabilityWalker : public ASTWalker {
   bool shouldWalkIntoTapExpression() override { return false; }
 
   MacroWalking getMacroWalkingBehavior() const override {
-    return MacroWalking::ArgumentsAndExpansion;
+    // Expanded source should be type checked and diagnosed separately.
+    return MacroWalking::Arguments;
   }
 
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override {

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -2060,10 +2060,6 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
   }
 
   void visitMacroExpansionDecl(MacroExpansionDecl *MED) {
-    // TODO: Type check attributes.
-    // Type checking arguments should reflect the attributes.
-    // e.g. '@available(macOS 999) #Future { newAPIFrom999() }'.
-
     // Assign a discriminator.
     (void)MED->getDiscriminator();
     // Decls in expansion already visited as auxiliary decls.

test/Macros/Inputs/top_level_freestanding_other.swift

@@ -11,9 +11,9 @@ var globalVar2 = { #stringify(1 + 1) }()
 func deprecated() -> Int { 0 }
 
 var globalVar3 = #stringify({ deprecated() })
-// expected-note@-1 2{{in expansion of macro 'stringify' here}}
+// expected-note@-1 {{in expansion of macro 'stringify' here}}
 // expected-warning@-2{{'deprecated()' is deprecated}}
 
 var globalVar4 = #stringify({ deprecated() })
-// expected-note@-1 2{{in expansion of macro 'stringify' here}}
+// expected-note@-1 {{in expansion of macro 'stringify' here}}
 // expected-warning@-2{{'deprecated()' is deprecated}}

test/Macros/macro_attribute_expansiondecl.swift

@@ -57,6 +57,26 @@ public struct LocalFuncAndVarMacro: DeclarationMacro {
   }
 }
 
+public struct FuncFromClosureMacro: DeclarationMacro {
+  public static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    guard
+      let closure = node.trailingClosure,
+      let arg1 = node.argumentList.first?.expression else {
+      return []
+    }
+
+    return ["""
+      func fromClosure() {
+        print(\(arg1))
+        \(closure.statements)
+      }
+      """]
+  }
+}
+
 //--- test.swift
 
 @freestanding(declaration, names: named(globalFunc), named(globalVar)) macro globalDecls() = #externalMacro(module: "MacroDefinition", type: "GlobalFuncAndVarMacro")
@@ -97,3 +117,44 @@ func testLocal() {
   }
 #endif
 }
+
+@freestanding(declaration, names: named(fromClosure)) macro funcFromClosureMacro(_: String, _: () -> Void) = #externalMacro(module: "MacroDefinition", type: "FuncFromClosureMacro")
+
+@available(macOS 99, *)
+func APIFrom99() -> String { "" }
+@available(macOS 999, *)
+func APIFrom999() -> String { "" }
+
+@available(macOS 99, *)
+#funcFromClosureMacro(APIFrom99()) {
+  _ = APIFrom99()
+  if #available(macOS 999, *) {
+    _ = APIFrom99()
+    _ = APIFrom999()
+  }
+}
+
+struct S1 {
+  @available(macOS 99, *)
+  #funcFromClosureMacro(APIFrom99()) {
+    _ = APIFrom99()
+    if #available(macOS 999, *) {
+      _ = APIFrom99()
+      _ = APIFrom999()
+    }
+  }
+}
+
+// FIXME: Diagnostics could be better.
+struct S2 { // expected-note 4 {{add @available attribute to enclosing struct}}
+  // expected-note@+3 6 {{in expansion of macro 'funcFromClosureMacro' here}}
+  // expected-error@+2 {{'APIFrom99()' is only available in macOS 99 or newer}}
+  // expected-error@+2 {{'APIFrom99()' is only available in macOS 99 or newer}} expected-note@+2 {{add 'if #available' version check}}
+  #funcFromClosureMacro(APIFrom99()) {
+    _ = APIFrom99()
+    if #available(macOS 999, *) {
+      _ = APIFrom99()
+      _ = APIFrom999()
+    }
+  }
+}