Merge pull request #69340 from rintaro/macros-visitauxiliary-toplevel

[Macros] Improve visitation of auxiliary decls

Author: rintaro

include/swift/AST/Module.h

@@ -967,6 +967,10 @@ class ModuleDecl
   /// The order of the results is not guaranteed to be meaningful.
   void getTopLevelDecls(SmallVectorImpl<Decl*> &Results) const;
 
+  /// Finds all top-level decls of this module including auxiliary decls.
+  void
+  getTopLevelDeclsWithAuxiliaryDecls(SmallVectorImpl<Decl *> &Results) const;
+
   void getExportedPrespecializations(SmallVectorImpl<Decl *> &results) const;
 
   /// Finds top-level decls of this module filtered by their attributes.

lib/AST/ASTPrinter.cpp

@@ -1035,12 +1035,6 @@ class PrintAST : public ASTVisitor<PrintAST> {
         Options.TransformContext->isPrintingSynthesizedExtension() &&
         isa<ExtensionDecl>(D);
 
-    SWIFT_DEFER {
-      D->visitAuxiliaryDecls([&](Decl *auxDecl) {
-        visit(auxDecl);
-      });
-    };
-
     if (!shouldPrint(D, true) && !Synthesize)
       return false;
 
@@ -4510,13 +4504,20 @@ bool PrintAST::printASTNodes(const ArrayRef<ASTNode> &Elements,
                              bool NeedIndent) {
   IndentRAII IndentMore(*this, NeedIndent);
   bool PrintedSomething = false;
+
+  std::function<void(Decl *)> printDecl;
+  printDecl = [&](Decl *d) {
+    if (d->shouldPrintInContext(Options))
+      visit(d);
+    d->visitAuxiliaryDecls(printDecl);
+  };
+
   for (auto element : Elements) {
     PrintedSomething = true;
     Printer.printNewline();
     indent();
     if (auto decl = element.dyn_cast<Decl*>()) {
-      if (decl->shouldPrintInContext(Options))
-        visit(decl);
+      printDecl(decl);
     } else if (auto stmt = element.dyn_cast<Stmt*>()) {
       visit(stmt);
     } else {

lib/AST/Module.cpp

@@ -1331,6 +1331,11 @@ void ModuleDecl::getTopLevelDecls(SmallVectorImpl<Decl*> &Results) const {
   FORWARD(getTopLevelDecls, (Results));
 }
 
+void ModuleDecl::getTopLevelDeclsWithAuxiliaryDecls(
+    SmallVectorImpl<Decl *> &Results) const {
+  FORWARD(getTopLevelDeclsWithAuxiliaryDecls, (Results));
+}
+
 void ModuleDecl::dumpDisplayDecls() const {
   SmallVector<Decl *, 32> Decls;
   getDisplayDecls(Decls);
@@ -3095,7 +3100,9 @@ void SourceFile::print(raw_ostream &OS, const PrintOptions &PO) {
 void SourceFile::print(ASTPrinter &Printer, const PrintOptions &PO) {
   std::set<DeclKind> MajorDeclKinds = {DeclKind::Class, DeclKind::Enum,
     DeclKind::Extension, DeclKind::Protocol, DeclKind::Struct};
-  for (auto decl : getTopLevelDecls()) {
+  SmallVector<Decl *> topLevelDecls;
+  getTopLevelDeclsWithAuxiliaryDecls(topLevelDecls);
+  for (auto decl : topLevelDecls) {
     if (!decl->shouldPrintInContext(PO))
       continue;
     // For a major decl, we print an empty line before it.
@@ -4190,14 +4197,18 @@ void FileUnit::getTopLevelDeclsWhereAttributesMatch(
 
 void FileUnit::getTopLevelDeclsWithAuxiliaryDecls(
     SmallVectorImpl<Decl*> &results) const {
+
+  std::function<void(Decl *)> addResult;
+  addResult = [&](Decl *decl) {
+    results.push_back(decl);
+    decl->visitAuxiliaryDecls(addResult);
+  };
+
   SmallVector<Decl *, 32> nonExpandedDecls;
   nonExpandedDecls.reserve(results.capacity());
   getTopLevelDecls(nonExpandedDecls);
   for (auto *decl : nonExpandedDecls) {
-    decl->visitAuxiliaryDecls([&](Decl *auxDecl) {
-      results.push_back(auxDecl);
-    });
-    results.push_back(decl);
+    addResult(decl);
   }
 }
 

lib/AST/NameLookup.cpp

@@ -4005,16 +4005,12 @@ void FindLocalVal::visitBraceStmt(BraceStmt *S, bool isTopLevelCode) {
     }
   }
 
-  auto visitDecl = [&](Decl *D) {
+  std::function<void(Decl *)> visitDecl;
+  visitDecl = [&](Decl *D) {
     if (auto *VD = dyn_cast<ValueDecl>(D))
       checkValueDecl(VD, DeclVisibilityKind::LocalVariable);
-    D->visitAuxiliaryDecls([&](Decl *D) {
-      if (auto *VD = dyn_cast<ValueDecl>(D))
-        checkValueDecl(VD, DeclVisibilityKind::LocalVariable);
-      // FIXME: Recursively call `visitDecl` to handle nested macros.
-    });
+    D->visitAuxiliaryDecls(visitDecl);
   };
-
   for (auto elem : S->getElements()) {
     if (auto *E = elem.dyn_cast<Expr *>()) {
       // 'MacroExpansionExpr' at code-item position may introduce value decls.

lib/Frontend/ModuleInterfaceSupport.cpp

@@ -833,7 +833,7 @@ bool swift::emitSwiftInterface(raw_ostream &out,
   InheritedProtocolCollector::PerTypeMap inheritedProtocolMap;
 
   SmallVector<Decl *, 16> topLevelDecls;
-  M->getTopLevelDecls(topLevelDecls);
+  M->getTopLevelDeclsWithAuxiliaryDecls(topLevelDecls);
   for (const Decl *D : topLevelDecls) {
     InheritedProtocolCollector::collectProtocols(inheritedProtocolMap, D);
 

lib/Sema/LookupVisibleDecls.cpp

@@ -217,24 +217,17 @@ static void collectVisibleMemberDecls(const DeclContext *CurrDC, LookupState LS,
                                       Type BaseType,
                                       IterableDeclContext *Parent,
                                       SmallVectorImpl<ValueDecl *> &FoundDecls) {
-  auto check = [&](Decl *decl) {
-    auto *VD = dyn_cast<ValueDecl>(decl);
+  for (auto Member : Parent->getAllMembers()) {
+    auto *VD = dyn_cast<ValueDecl>(Member);
     if (!VD)
-      return;
+      continue;
     if (!isDeclVisibleInLookupMode(VD, LS, CurrDC))
-      return;
+      continue;
     if (!evaluateOrDefault(CurrDC->getASTContext().evaluator,
         IsDeclApplicableRequest(DeclApplicabilityOwner(CurrDC, BaseType, VD)),
                            false))
-      return;
+      continue;
     FoundDecls.push_back(VD);
-  };
-
-  for (auto Member : Parent->getAllMembers()) {
-    check(Member);
-    Member->visitAuxiliaryDecls([&](Decl *d) {
-      check(d);
-    });
   }
 }
 

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -1900,7 +1900,7 @@ public struct AddPeerStoredPropertyMacro: PeerMacro, Sendable {
     return [
       """
 
-      private var _foo: Int = 100
+      public var _foo: Int = 100
       """
     ]
   }

test/ModuleInterface/macros.swift

@@ -64,8 +64,35 @@ macro structWithUnqualifiedLookup() = #externalMacro(module: "MacroDefinition",
 
 let world = 17
 
-// CHECK-NOT: structWithUnqualifiedLookup
 public
 #structWithUnqualifiedLookup
-
+// CHECK-NOT: structWithUnqualifiedLookup
+// CHECK-NOT: struct StructWithUnqualifiedLookup
 // CHECK: struct StructWithUnqualifiedLookup
+// CHECK-NOT: struct StructWithUnqualifiedLookup
+
+@attached(peer, names: named(_foo))
+macro AddPeerStoredProperty() = #externalMacro(module: "MacroDefinition", type: "AddPeerStoredPropertyMacro")
+
+@AddPeerStoredProperty
+public var test: Int = 10
+// CHECK: var test
+// CHECK-NOT: var _foo
+// CHECK: var _foo
+// CHECK-NOT: var _foo
+
+// CHECK: struct TestStruct {
+public struct TestStruct {
+  public #structWithUnqualifiedLookup
+  // CHECK-NOT: structWithUnqualifiedLookup
+  // CHECK-NOT: struct StructWithUnqualifiedLookup
+  // CHECK: struct StructWithUnqualifiedLookup
+  // CHECK-NOT: struct StructWithUnqualifiedLookup
+
+  @AddPeerStoredProperty
+  public var test: Int = 10
+  // CHECK: var test
+  // CHECK-NOT: var _foo
+  // CHECK: var _foo
+  // CHECK-NOT: var _foo
+}

test/SourceKit/Macros/macro_basic.swift

@@ -60,6 +60,13 @@ macro anonymousTypes(_: () -> String) = #externalMacro(module: "MacroDefinition"
 @freestanding(expression) macro assert(_: String) = #externalMacro(module: "MacroDefinition", type: "AssertMacro")
 #assert("foobar")
 
+@attached(peer, names: named(_foo))
+macro AddPeerStoredProperty() = #externalMacro(module: "MacroDefinition", type: "AddPeerStoredPropertyMacro")
+struct S5 {
+  @AddPeerStoredProperty
+  var test: Int = 10
+}
+
 // REQUIRES: swift_swift_parser, executable_test, shell
 
 // RUN: %empty-directory(%t)
@@ -281,3 +288,8 @@ macro anonymousTypes(_: () -> String) = #externalMacro(module: "MacroDefinition"
 //##-- Expansion on "fails to typecheck" macro expression
 // RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=61:2 %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=ERRONEOUS_EXPAND %s
 // ERRONEOUS_EXPAND: 61:1-61:18 (@__swiftmacro_{{.+}}.swift) "assert("foobar")"
+
+//##-- Cursor-info on a decl where a peer macro attached.
+// RUN: %sourcekitd-test -req=cursor -pos=67:7 %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=CURSOR_ON_DECL_WITH_PEER %s
+// CURSOR_ON_DECL_WITH_PEER: <decl.var.instance><syntaxtype.keyword>var</syntaxtype.keyword> <decl.name>test</decl.name>: <decl.var.type><ref.struct usr="s:Si">Int</ref.struct></decl.var.type></decl.var.instance>
+// CURSOR_ON_DECL_WITH_PEER-NOT: _foo