[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/AST/TypeRefinementContext.cpp

@@ -53,6 +53,13 @@ TypeRefinementContext::createRoot(SourceFile *SF,
   if (auto parentExpansion = SF->getMacroExpansion()) {
     if (auto parentTRC =
             SF->getEnclosingSourceFile()->getTypeRefinementContext()) {
+      // FIXME: Skip `MacroExpansionDecl` context?
+      // Expanded decl might _not_ have the attributes on 'MacroExpansionDecl'
+      // copied so it might want wider availability than the expansion itself.
+      // E.g. '@available(macOS 99, *) #foo' expands to 'func foo() {}' (instead
+      // of '@available(macOS 99, *) func foo() {}'). In such cases, the
+      // attribute should not affect the expanded 'foo'?
+
       auto charRange = Ctx.SourceMgr.getRangeForBuffer(*SF->getBufferID());
       range = SourceRange(charRange.getStart(), charRange.getEnd());
       parentContext = parentTRC->findMostRefinedSubContext(

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
@@ -3192,7 +3193,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

@@ -20,6 +20,7 @@
 // RUN:   -parse-as-library \
 // RUN:   -dump-macro-expansions \
 // RUN:   -plugin-path %t/plugins \
+// RUN:   -typo-correction-limit 0 \
 // RUN:   %t/src/test.swift
 
 
@@ -57,6 +58,36 @@ public struct LocalFuncAndVarMacro: DeclarationMacro {
   }
 }
 
+public struct FuncWithClosureMacro: DeclarationMacro {
+  public static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+
+    return ["func LocalFunc() {}", "var LocalVar: Int { 1 }"]
+  }
+}
+
+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 +128,20 @@ 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()
+  }
+}
+