[Macros] Don't visit macro-generated extensions in 'visitAuxiliaryDecls'.

Macro-generated extensions are hoisted to file scope, because extensions are
not valid in nested scopes. Callers of 'visitAuxiliaryDecls' assume that the
auxiliary decls are in the same decl context as the original, which is clearly
not the case for extensions, and it leads to issues like visiting extension at
the wrong time during SILGen. The extensions are already added to the top-level
decls, so we don't need to visit them as auxiliary decls, and we can type-check
macro-expanded decls at the end of visitation in TypeCheckDeclPrimary.

Author: hborla

include/swift/AST/Decl.h

@@ -924,6 +924,12 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
       bool visitFreestandingExpanded = true
   ) const;
 
+  using ExtensionCallback = llvm::function_ref<void(ExtensionDecl *)>;
+
+  /// Iterate over each macro-expanded extension for this declaration,
+  /// invoking the given callback with each extension.
+  void forEachExpandedExtension(ExtensionCallback callback) const;
+
   using MacroCallback = llvm::function_ref<void(CustomAttr *, MacroDecl *)>;
 
   /// Iterate over each attached macro with the given role, invoking the

lib/AST/Decl.cpp

@@ -407,20 +407,6 @@ void Decl::visitAuxiliaryDecls(
     }
   }
 
-  if (auto *nominal = dyn_cast<NominalTypeDecl>(mutableThis)) {
-    auto conformanceBuffers =
-        evaluateOrDefault(ctx.evaluator,
-                          ExpandConformanceMacros{nominal},
-                          {});
-    for (auto bufferID : conformanceBuffers) {
-      auto startLoc = sourceMgr.getLocForBufferStart(bufferID);
-      auto *sourceFile = moduleDecl->getSourceFileContainingLocation(startLoc);
-      for (auto *extension : sourceFile->getTopLevelDecls()) {
-        callback(extension);
-      }
-    }
-  }
-
   if (visitFreestandingExpanded) {
     if (auto *med = dyn_cast<MacroExpansionDecl>(mutableThis)) {
       if (auto bufferID = evaluateOrDefault(
@@ -436,6 +422,29 @@ void Decl::visitAuxiliaryDecls(
   // FIXME: fold VarDecl::visitAuxiliaryDecls into this.
 }
 
+void
+Decl::forEachExpandedExtension(ExtensionCallback callback) const {
+  auto &ctx = getASTContext();
+  auto *mutableThis = const_cast<Decl *>(this);
+  auto *nominal = dyn_cast<NominalTypeDecl>(mutableThis);
+  if (!nominal)
+    return;
+
+  SourceManager &sourceMgr = ctx.SourceMgr;
+  auto *moduleDecl = getModuleContext();
+  auto conformanceBuffers =
+      evaluateOrDefault(ctx.evaluator,
+                        ExpandConformanceMacros{nominal},
+                        {});
+  for (auto bufferID : conformanceBuffers) {
+    auto startLoc = sourceMgr.getLocForBufferStart(bufferID);
+    auto *sourceFile = moduleDecl->getSourceFileContainingLocation(startLoc);
+    for (auto *extension : sourceFile->getTopLevelDecls()) {
+      callback(cast<ExtensionDecl>(extension));
+    }
+  }
+}
+
 void Decl::forEachAttachedMacro(MacroRole role,
                                 MacroCallback macroCallback) const {
   for (auto customAttrConst : getSemanticAttrs().getAttributes<CustomAttr>()) {

lib/AST/Module.cpp

@@ -4104,8 +4104,7 @@ void FileUnit::getTopLevelDeclsWithAuxiliaryDecls(
   getTopLevelDecls(nonExpandedDecls);
   for (auto *decl : nonExpandedDecls) {
     decl->visitAuxiliaryDecls([&](Decl *auxDecl) {
-      if (!isa<ExtensionDecl>(auxDecl))
-        results.push_back(auxDecl);
+      results.push_back(auxDecl);
     });
     results.push_back(decl);
   }

lib/IRGen/GenDecl.cpp

@@ -5578,11 +5578,6 @@ void IRGenModule::emitNestedTypeDecls(DeclRange members) {
       continue;
 
     member->visitAuxiliaryDecls([&](Decl *decl) {
-      // FIXME: Conformance macros can generate extension decls. These
-      // are visited as top-level decls; skip them here.
-      if (isa<ExtensionDecl>(decl))
-        return;
-
       emitNestedTypeDecls({decl, nullptr});
     });
     switch (member->getKind()) {

lib/Index/Index.cpp

@@ -1148,6 +1148,9 @@ void IndexSwiftASTWalker::visitModule(ModuleDecl &Mod) {
     if (!handleSourceOrModuleFile(*SrcFile))
       return;
     walk(*SrcFile);
+    if (auto *synthesizedSF = SrcFile->getSynthesizedFile()) {
+      walk(synthesizedSF);
+    }
   } else {
     IsModuleFile = true;
     isSystemModule = Mod.isNonUserModule();

lib/SILGen/SILGen.cpp

@@ -2201,10 +2201,6 @@ class SILGenModuleRAII {
     for (auto *D : sf->getTopLevelDecls()) {
       // Emit auxiliary decls.
       D->visitAuxiliaryDecls([&](Decl *auxiliaryDecl) {
-        // Skip extensions decls; they are visited below.
-        if (isa<ExtensionDecl>(auxiliaryDecl))
-          return;
-
         FrontendStatsTracer StatsTracer(SGM.getASTContext().Stats,
                                         "SILgen-decl", auxiliaryDecl);
         SGM.visit(auxiliaryDecl);

lib/SILGen/SILGenType.cpp

@@ -1101,11 +1101,6 @@ class SILGenType : public TypeMemberVisitor<SILGenType> {
     SGM.emitLazyConformancesForType(theType);
 
     forEachMemberToLower(theType, [&](Decl *member) {
-      // FIXME: Conformance macros can generate extension decls. These
-      // are visited as top-level decls; skip them here.
-      if (isa<ExtensionDecl>(member))
-        return;
-
       visit(member);
     });
 

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -1936,6 +1936,11 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
                           "`" + VD->getBaseName().userFacingName().str() + "`");
       }
     }
+
+    // Visit macro-generated extensions.
+    decl->forEachExpandedExtension([&](ExtensionDecl *extension) {
+      DeclVisitor<DeclChecker>::visit(extension);
+    });
   }
 
 

lib/Sema/TypeCheckMacros.cpp

@@ -1442,7 +1442,13 @@ swift::expandConformances(CustomAttr *attr, MacroDecl *macro,
     extension->setExtendedNominal(nominal);
     nominal->addExtension(extension);
 
-    // Make it accessible to getTopLevelDecls()
+    // Most other macro-generated declarations are visited through calling
+    // 'visitAuxiliaryDecls' on the original declaration the macro is attached
+    // to. We don't do this for macro-generated extensions, because the
+    // extension is not a peer of the original declaration. Instead of
+    // requiring all callers of 'visitAuxiliaryDecls' to understand the
+    // hoisting behavior of macro-generated extensions, we make the
+    // extension accessible through 'getTopLevelDecls()'.
     if (auto file = dyn_cast<FileUnit>(
             decl->getDeclContext()->getModuleScopeContext()))
       file->getOrCreateSynthesizedFile().addTopLevelDecl(extension);

test/Index/index_macros.swift

@@ -81,6 +81,12 @@ struct AddOne {
 // CHECK: [[@LINE-5]]:1 | function/Swift | freeLog() | [[FREE_LOG_USR]] | Ref,Call,Impl,RelCall,RelCont
 // CHECK-NEXT: RelCall,RelCont | instance-method/Swift | freeFunc() | [[FREE_FUNC_USR]]
 
+func testExpr() {
+  #freestandingExpr
+  // CHECK: [[@LINE-1]]:3 | function/Swift | exprLog() | [[EXPR_LOG_USR]] | Ref,Call,Impl,RelCall,RelCont
+  // CHECK-NEXT: RelCall,RelCont | function/Swift | testExpr()
+}
+
 // CHECK: [[@LINE+4]]:40 | macro/Swift | Peer() | [[PEER_USR]] | Ref
 // CHECK: [[@LINE+3]]:23 | macro/Swift | MemberAttribute() | [[MEMBER_ATTRIBUTE_USR]] | Ref
 // CHECK: [[@LINE+2]]:15 | macro/Swift | Member() | [[MEMBER_USR]] | Ref
@@ -112,11 +118,6 @@ struct TestAttached {
 // CHECK: [[@LINE-24]]:39 | function/Swift | peerLog() | [[PEER_LOG_USR]] | Ref,Call,Impl,RelCall,RelCont
 // CHECK-NEXT: RelCall,RelCont | instance-method/Swift | peerFunc() | [[PEER_FUNC_USR]]
 
-// `Conformance` adds `TestProto` as a conformance on an extension of `TestAttached`
-// CHECK: [[@LINE-28]]:1 | extension/ext-struct/Swift | TestAttached | {{.*}} | Def,Impl
-// CHECK: [[@LINE-29]]:1 | protocol/Swift | TestProto | [[PROTO_USR]] | Ref,Impl,RelBase
-// CHECK-NEXT: RelBase | extension/ext-struct/Swift | TestAttached
-
 // CHECK: [[@LINE+1]]:8 | struct/Swift | Outer | [[OUTER_USR:.*]] | Def
 struct Outer {
   // CHECK: [[@LINE+1]]:4 | macro/Swift | PeerMember() | [[PEER_MEMBER_USR]] | Ref
@@ -137,16 +138,19 @@ struct Outer {
 // CHECK: [[@LINE-6]]:16 | function/Swift | memberLog() | [[MEMBER_LOG_USR]] | Ref,Call,Impl,RelCall,RelCont
 // CHECK-NEXT: RelCall,RelCont | instance-method/Swift | memberFunc() | [[INNER_FUNC_USR]]
 
+
+// Expanded extensions are visited last
+
+// `Conformance` adds `TestProto` as a conformance on an extension of `TestAttached`
+// CHECK: [[@LINE-51]]:1 | extension/ext-struct/Swift | TestAttached | {{.*}} | Def,Impl
+// CHECK: [[@LINE-52]]:1 | protocol/Swift | TestProto | [[PROTO_USR]] | Ref,Impl,RelBase
+// CHECK-NEXT: RelBase | extension/ext-struct/Swift | TestAttached
+
 // `Conformance` adds `TestProto` as a conformance on an extension of `TestInner`
-// CHECK: [[@LINE-10]]:3 | extension/ext-struct/Swift | TestInner | {{.*}} | Def,Impl
-// CHECK: [[@LINE-11]]:3 | protocol/Swift | TestProto | [[PROTO_USR]] | Ref,Impl,RelBase
+// CHECK: [[@LINE-18]]:3 | extension/ext-struct/Swift | TestInner | {{.*}} | Def,Impl
+// CHECK: [[@LINE-19]]:3 | protocol/Swift | TestProto | [[PROTO_USR]] | Ref,Impl,RelBase
 // CHECK-NEXT: RelBase | extension/ext-struct/Swift | TestInner
 
-func testExpr() {
-  #freestandingExpr
-  // CHECK: [[@LINE-1]]:3 | function/Swift | exprLog() | [[EXPR_LOG_USR]] | Ref,Call,Impl,RelCall,RelCont
-  // CHECK-NEXT: RelCall,RelCont | function/Swift | testExpr()
-}
 
 //--- IndexMacros.swift
 import SwiftSyntax

test/Macros/macro_expand_conformances.swift

@@ -63,6 +63,12 @@ requireHashable(S2())
 
 requireEquatable(E.Nested())
 
+extension E {
+  @Equatable struct NestedInExtension {}
+}
+
+requireEquatable(E.NestedInExtension())
+
 #if TEST_DIAGNOSTICS
 requireEquatable(PublicEquatable())