Merge pull request #66387 from DougGregor/silgen-freestanding-local-vars

DougGregor · web-flow · commit 5ea559435d4c · 2023-06-07T06:49:23.000-07:00
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -911,7 +911,13 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
   ///
   /// Auxiliary declarations can be property wrapper backing variables,
   /// backing variables for 'lazy' vars, or peer macro expansions.
-  void visitAuxiliaryDecls(AuxiliaryDeclCallback callback) const;
+  ///
+  /// When \p visitFreestandingExpanded is true (the default), this will also
+  /// visit the declarations produced by a freestanding macro expansion.
+  void visitAuxiliaryDecls(
+      AuxiliaryDeclCallback callback,
+      bool visitFreestandingExpanded = true
+  ) const;
 
   using MacroCallback = llvm::function_ref<void(CustomAttr *, MacroDecl *)>;
 
@@ -8621,8 +8627,9 @@ class MacroExpansionDecl : public Decl, public FreestandingMacroExpansion {
     return getExpansionInfo()->getSourceRange();
   }
   SourceLoc getLocFromSource() const { return getExpansionInfo()->SigilLoc; }
-  using ExprOrStmtExpansionCallback = llvm::function_ref<void(ASTNode)>;
-  void forEachExpandedExprOrStmt(ExprOrStmtExpansionCallback) const;
+
+  /// Enumerate the nodes produced by expanding this macro expansion.
+  void forEachExpandedNode(llvm::function_ref<void(ASTNode)> callback) const;
 
   /// Returns a discriminator which determines this macro expansion's index
   /// in the sequence of macro expansions within the current function.
diff --git a/lib/AST/ASTWalker.cpp b/lib/AST/ASTWalker.cpp
@@ -455,22 +455,20 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       else
         return true;
     }
-    // Visit auxiliary decls, which may be decls from macro expansions.
     bool alreadyFailed = false;
     if (shouldWalkExpansion) {
-      MED->visitAuxiliaryDecls([&](Decl *decl) {
-        if (alreadyFailed) return;
-        if (!isa<VarDecl>(decl))
-          alreadyFailed = inherited::visit(decl);
-      });
-      MED->forEachExpandedExprOrStmt([&](ASTNode expandedNode) {
+      MED->forEachExpandedNode([&](ASTNode expandedNode) {
         if (alreadyFailed) return;
         if (auto *expr = expandedNode.dyn_cast<Expr *>()) {
           if (!doIt(expr))
             alreadyFailed = true;
         } else if (auto *stmt = expandedNode.dyn_cast<Stmt *>()) {
           if (!doIt(stmt))
             alreadyFailed = true;
+        } else {
+          auto decl = expandedNode.get<Decl *>();
+          if (!isa<VarDecl>(decl))
+            alreadyFailed = inherited::visit(decl);
         }
       });
     }
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -381,7 +381,10 @@ DeclAttributes Decl::getSemanticAttrs() const {
   return getAttrs();
 }
 
-void Decl::visitAuxiliaryDecls(AuxiliaryDeclCallback callback) const {
+void Decl::visitAuxiliaryDecls(
+    AuxiliaryDeclCallback callback,
+    bool visitFreestandingExpanded
+) const {
   auto &ctx = getASTContext();
   auto *mutableThis = const_cast<Decl *>(this);
   SourceManager &sourceMgr = ctx.SourceMgr;
@@ -414,13 +417,15 @@ void Decl::visitAuxiliaryDecls(AuxiliaryDeclCallback callback) const {
     }
   }
 
-  else if (auto *med = dyn_cast<MacroExpansionDecl>(mutableThis)) {
-    if (auto bufferID = evaluateOrDefault(
-            ctx.evaluator, ExpandMacroExpansionDeclRequest{med}, {})) {
-      auto startLoc = sourceMgr.getLocForBufferStart(*bufferID);
-      auto *sourceFile = moduleDecl->getSourceFileContainingLocation(startLoc);
-      for (auto *decl : sourceFile->getTopLevelDecls())
-        callback(decl);
+  if (visitFreestandingExpanded) {
+    if (auto *med = dyn_cast<MacroExpansionDecl>(mutableThis)) {
+      if (auto bufferID = evaluateOrDefault(
+              ctx.evaluator, ExpandMacroExpansionDeclRequest{med}, {})) {
+        auto startLoc = sourceMgr.getLocForBufferStart(*bufferID);
+        auto *sourceFile = moduleDecl->getSourceFileContainingLocation(startLoc);
+        for (auto *decl : sourceFile->getTopLevelDecls())
+          callback(decl);
+      }
     }
   }
 
@@ -10637,8 +10642,9 @@ unsigned MacroExpansionDecl::getDiscriminator() const {
   return getRawDiscriminator();
 }
 
-void MacroExpansionDecl::forEachExpandedExprOrStmt(
-    ExprOrStmtExpansionCallback callback) const {
+void MacroExpansionDecl::forEachExpandedNode(
+    llvm::function_ref<void(ASTNode)> callback
+) const {
   auto mutableThis = const_cast<MacroExpansionDecl *>(this);
   auto bufferID = evaluateOrDefault(
       getASTContext().evaluator,
@@ -10650,8 +10656,7 @@ void MacroExpansionDecl::forEachExpandedExprOrStmt(
   auto startLoc = sourceMgr.getLocForBufferStart(*bufferID);
   auto *sourceFile = moduleDecl->getSourceFileContainingLocation(startLoc);
   for (auto node : sourceFile->getTopLevelItems())
-    if (node.is<Expr *>() || node.is<Stmt *>())
-      callback(node);
+    callback(node);
 }
 
 NominalTypeDecl *
diff --git a/lib/SILGen/SILGenDecl.cpp b/lib/SILGen/SILGenDecl.cpp
@@ -1625,11 +1625,13 @@ void SILGenFunction::visitVarDecl(VarDecl *D) {
 }
 
 void SILGenFunction::visitMacroExpansionDecl(MacroExpansionDecl *D) {
-  D->forEachExpandedExprOrStmt([&](ASTNode node) {
+  D->forEachExpandedNode([&](ASTNode node) {
     if (auto *expr = node.dyn_cast<Expr *>())
       emitIgnoredExpr(expr);
     else if (auto *stmt = node.dyn_cast<Stmt *>())
       emitStmt(stmt);
+    else
+      visit(node.get<Decl *>());
   });
 }
 
diff --git a/lib/SILGen/SILGenExpr.cpp b/lib/SILGen/SILGenExpr.cpp
@@ -6294,11 +6294,13 @@ RValue RValueEmitter::visitMacroExpansionExpr(MacroExpansionExpr *E,
   }
   else if (auto *MED = E->getSubstituteDecl()) {
     Mangle::ASTMangler mangler;
-    MED->forEachExpandedExprOrStmt([&](ASTNode node) {
+    MED->forEachExpandedNode([&](ASTNode node) {
       if (auto *expr = node.dyn_cast<Expr *>())
         visit(expr, C);
       else if (auto *stmt = node.dyn_cast<Stmt *>())
         SGF.emitStmt(stmt);
+      else
+        SGF.visit(node.get<Decl *>());
     });
     return RValue();
   }
diff --git a/lib/Sema/TypeCheckDeclPrimary.cpp b/lib/Sema/TypeCheckDeclPrimary.cpp
@@ -1850,7 +1850,7 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
     if (!isa<ClassDecl>(decl->getDeclContext())) {
       decl->visitAuxiliaryDecls([&](Decl *auxiliaryDecl) {
         this->visit(auxiliaryDecl);
-      });
+      }, /*visitFreestandingExpanded=*/false);
     }
 
     if (auto *Stats = getASTContext().Stats)
@@ -2066,8 +2066,7 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
   void visitMacroExpansionDecl(MacroExpansionDecl *MED) {
     // Assign a discriminator.
     (void)MED->getDiscriminator();
-    // Decls in expansion already visited as auxiliary decls.
-    MED->forEachExpandedExprOrStmt([&](ASTNode node) {
+    MED->forEachExpandedNode([&](ASTNode node) {
       TypeChecker::typeCheckASTNode(node, MED->getDeclContext());
     });
   }
diff --git a/test/Macros/Inputs/syntax_macro_definitions.swift b/test/Macros/Inputs/syntax_macro_definitions.swift
@@ -299,6 +299,25 @@ public struct DefineDeclsWithKnownNamesMacro: DeclarationMacro {
   }
 }
 
+public struct VarDeclMacro: CodeItemMacro {
+  public static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [CodeBlockItemSyntax] {
+    let name = context.makeUniqueName("fromMacro")
+    return [
+      "let \(name) = 23",
+      "use(\(name))",
+      """
+      if true {
+        let \(name) = "string"
+        use(\(name))
+      }
+      """
+    ]
+  }
+}
+
 public struct WarningMacro: ExpressionMacro {
    public static func expansion(
      of macro: some FreestandingMacroExpansionSyntax,
diff --git a/test/Macros/macro_expand_codeitems.swift b/test/Macros/macro_expand_codeitems.swift
@@ -33,3 +33,10 @@ func testFreestandingMacroExpansion() {
   #codeItems
 }
 testFreestandingMacroExpansion()
+
+@freestanding(codeItem) macro varDecl() = #externalMacro(module: "MacroDefinition", type: "VarDeclMacro")
+
+func testVarDecl() {
+  func use<T>(_ t: T) {}
+  #varDecl()
+}
diff --git a/test/Macros/top_level_freestanding.swift b/test/Macros/top_level_freestanding.swift
@@ -50,7 +50,7 @@ func lookupGlobalFreestandingExpansion() {
 
 #anonymousTypes(public: true) { "hello" }
 
-// CHECK-SIL: sil @$s9MacroUser03$s9A71User33_082AE7CFEFA6960C804A9FE7366EB5A0Ll14anonymousTypesfMf0_4namefMu_C5helloSSyF
+// CHECK-SIL: sil @$s9MacroUser03$s9A70User33_082AE7CFEFA6960C804A9FE7366EB5A0Ll14anonymousTypesfMf_4namefMu_C5helloSSyF
 
 @main
 struct Main {
diff --git a/test/SourceKit/Macros/macro_basic.swift b/test/SourceKit/Macros/macro_basic.swift
@@ -131,7 +131,7 @@ macro anonymousTypes(_: () -> String) = #externalMacro(module: "MacroDefinition"
 // RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=57:1 %s -- ${COMPILER_ARGS[@]} -parse-as-library | %FileCheck -check-prefix=EXPAND_MACRO_DECL %s
 // RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=57:2 %s -- ${COMPILER_ARGS[@]} -parse-as-library | %FileCheck -check-prefix=EXPAND_MACRO_DECL %s
 // EXPAND_MACRO_DECL: source.edit.kind.active:
-// EXPAND_MACRO_DECL-NEXT: 57:1-57:28 (@__swiftmacro_9MacroUser33_70D4178875715FB9B8B50C58F66F8D53Ll14anonymousTypesfMf0_.swift) "class $s9MacroUser33_70D4178875715FB9B8B50C58F66F8D53Ll14anonymousTypesfMf0_4namefMu_ {
+// EXPAND_MACRO_DECL-NEXT: 57:1-57:28 (@__swiftmacro_9MacroUser33_70D4178875715FB9B8B50C58F66F8D53Ll14anonymousTypesfMf_.swift) "class $s9MacroUser33_70D4178875715FB9B8B50C58F66F8D53Ll14anonymousTypesfMf_4namefMu_ {
 // EXPAND_MACRO_DECL-NEXT:   func hello() -> String {
 // EXPAND_MACRO_DECL-NEXT:     "hello"
 // EXPAND_MACRO_DECL-NEXT:   }
@@ -140,15 +140,15 @@ macro anonymousTypes(_: () -> String) = #externalMacro(module: "MacroDefinition"
 // EXPAND_MACRO_DECL-NEXT:      return Self.self
 // EXPAND_MACRO_DECL-NEXT:   }
 // EXPAND_MACRO_DECL-NEXT: }
-// EXPAND_MACRO_DECL-NEXT: enum $s9MacroUser33_70D4178875715FB9B8B50C58F66F8D53Ll14anonymousTypesfMf0_4namefMu0_ {
+// EXPAND_MACRO_DECL-NEXT: enum $s9MacroUser33_70D4178875715FB9B8B50C58F66F8D53Ll14anonymousTypesfMf_4namefMu0_ {
 // EXPAND_MACRO_DECL-NEXT:   case apple
 // EXPAND_MACRO_DECL-NEXT:   case banana
 // EXPAND_MACRO_DECL-EMPTY:
 // EXPAND_MACRO_DECL-NEXT:   func hello() -> String {
 // EXPAND_MACRO_DECL-NEXT:     "hello"
 // EXPAND_MACRO_DECL-NEXT:   }
 // EXPAND_MACRO_DECL-NEXT: }
-// EXPAND_MACRO_DECL-NEXT: struct $s9MacroUser33_70D4178875715FB9B8B50C58F66F8D53Ll14anonymousTypesfMf0_4namefMu1_: Equatable {
+// EXPAND_MACRO_DECL-NEXT: struct $s9MacroUser33_70D4178875715FB9B8B50C58F66F8D53Ll14anonymousTypesfMf_4namefMu1_: Equatable {
 // EXPAND_MACRO_DECL-NEXT:   static func == (lhs: Self, rhs: Self) -> Bool {
 // EXPAND_MACRO_DECL-NEXT:     false
 // EXPAND_MACRO_DECL-NEXT:   }

Original file line number	Diff line number	Diff line change
`@@ -1625,11 +1625,13 @@ void SILGenFunction::visitVarDecl(VarDecl *D) {`
`1625`	`1625`	`}`
`1626`	`1626`
`1627`	`1627`	`void SILGenFunction::visitMacroExpansionDecl(MacroExpansionDecl *D) {`
`1628`		`- D->forEachExpandedExprOrStmt([&](ASTNode node) {`
	`1628`	`+ D->forEachExpandedNode([&](ASTNode node) {`
`1629`	`1629`	`if (auto expr = node.dyn_cast<Expr >())`
`1630`	`1630`	`emitIgnoredExpr(expr);`
`1631`	`1631`	`else if (auto stmt = node.dyn_cast<Stmt >())`
`1632`	`1632`	`emitStmt(stmt);`
	`1633`	`+ else`
	`1634`	`+ visit(node.get<Decl *>());`
`1633`	`1635`	`});`
`1634`	`1636`	`}`
`1635`	`1637`