Merge pull request #64765 from rxwei/code-item-macro

rxwei · web-flow · commit 33b49d4885b3 · 2023-04-04T13:26:47.000-07:00
[Macros] Code item macros
diff --git a/include/swift/AST/DiagnosticsParse.def b/include/swift/AST/DiagnosticsParse.def
@@ -2049,8 +2049,8 @@ ERROR(macro_role_attr_expected_kind,PointsToFirstBadToken,
       "expected %select{a freestanding|an attached}0 macro role such as "
       "%select{'expression'|'accessor'}0", (bool))
 ERROR(macro_role_syntax_mismatch,PointsToFirstBadToken,
-      "expected %select{a freestanding|an attached}0 macro cannot have "
-      "the %1 role", (bool, Identifier))
+      "%select{a freestanding|an attached}0 macro cannot have the %1 role",
+      (bool, Identifier))
 ERROR(macro_attribute_unknown_label,PointsToFirstBadToken,
       "@%select{freestanding|attached}0 has no argument with label %1",
       (bool, Identifier))
diff --git a/include/swift/AST/DiagnosticsSema.def b/include/swift/AST/DiagnosticsSema.def
@@ -7032,6 +7032,11 @@ ERROR(invalid_macro_introduced_name,none,
 ERROR(global_freestanding_macro_script,none,
       "global freestanding macros not yet supported in script mode",
       ())
+ERROR(invalid_macro_role_for_macro_syntax,none,
+      "invalid macro role for %{a freestanding|an attached}0 macro",
+      (unsigned))
+ERROR(macro_cannot_introduce_names,none,
+      "'%0' macros are not allowed to introduce names", (StringRef))
 
 //------------------------------------------------------------------------------
 // MARK: Move Only Errors
diff --git a/include/swift/AST/MacroDeclaration.h b/include/swift/AST/MacroDeclaration.h
@@ -55,6 +55,9 @@ enum class MacroRole: uint32_t {
   /// An attached macro that adds conformances to the declaration the
   /// macro is attached to.
   Conformance = 0x40,
+  /// A freestanding macro that expands to expressions, statements and
+  /// declarations in a code block.
+  CodeItem = 0x80,
 };
 
 /// The contexts in which a particular macro declaration can be used.
diff --git a/include/swift/Basic/Features.def b/include/swift/Basic/Features.def
@@ -112,6 +112,7 @@ EXPERIMENTAL_FEATURE(VariadicGenerics, false)
 EXPERIMENTAL_FEATURE(NamedOpaqueTypes, false)
 EXPERIMENTAL_FEATURE(FlowSensitiveConcurrencyCaptures, false)
 EXPERIMENTAL_FEATURE(FreestandingMacros, true)
+EXPERIMENTAL_FEATURE(CodeItemMacros, true)
 
 // FIXME: MoveOnlyClasses is not intended to be in production,
 // but our tests currently rely on it, and we want to run those
diff --git a/lib/AST/ASTMangler.cpp b/lib/AST/ASTMangler.cpp
@@ -3965,6 +3965,7 @@ void ASTMangler::appendMacroExpansionOperator(
   switch (role) {
   case MacroRole::Expression:
   case MacroRole::Declaration:
+  case MacroRole::CodeItem:
     appendOperator("fMf", Index(discriminator));
     break;
 
diff --git a/lib/AST/ASTPrinter.cpp b/lib/AST/ASTPrinter.cpp
@@ -2838,6 +2838,14 @@ static bool usesFeatureFreestandingMacros(Decl *decl) {
   return macro->getMacroRoles().contains(MacroRole::Declaration);
 }
 
+static bool usesFeatureCodeItemMacros(Decl *decl) {
+  auto macro = dyn_cast<MacroDecl>(decl);
+  if (!macro)
+    return false;
+
+  return macro->getMacroRoles().contains(MacroRole::CodeItem);
+}
+
 static bool usesFeatureAttachedMacros(Decl *decl) {
   auto macro = dyn_cast<MacroDecl>(decl);
   if (!macro)
diff --git a/lib/AST/ASTScopeCreation.cpp b/lib/AST/ASTScopeCreation.cpp
@@ -261,6 +261,7 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
     switch (*macroRole) {
     case MacroRole::Expression:
     case MacroRole::Declaration:
+    case MacroRole::CodeItem:
     case MacroRole::Accessor:
     case MacroRole::MemberAttribute:
     case MacroRole::Conformance:
diff --git a/lib/AST/ASTWalker.cpp b/lib/AST/ASTWalker.cpp
@@ -442,6 +442,9 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   bool visitMacroExpansionDecl(MacroExpansionDecl *MED) {
+#ifndef NDEBUG
+    PrettyStackTraceDecl debugStack("walking into", MED);
+#endif
     bool shouldWalkArguments, shouldWalkExpansion;
     std::tie(shouldWalkArguments, shouldWalkExpansion) =
         Walker.shouldWalkMacroArgumentsAndExpansion();
@@ -1341,6 +1344,14 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     std::tie(shouldWalkArguments, shouldWalkExpansion) =
         Walker.shouldWalkMacroArgumentsAndExpansion();
 
+    if (auto *substituteDecl = E->getSubstituteDecl()) {
+      if (doIt(substituteDecl))
+        return nullptr;
+      // Visiting the substitute macro expansion decl will visit the same
+      // argument list. Skip visiting it again.
+      shouldWalkArguments = false;
+    }
+
     if (shouldWalkArguments && E->getArgs()) {
       ArgumentList *args = doIt(E->getArgs());
       if (!args) return nullptr;
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -10069,6 +10069,9 @@ StringRef swift::getMacroRoleString(MacroRole role) {
 
   case MacroRole::Conformance:
     return "conformance";
+
+  case MacroRole::CodeItem:
+    return "codeItem";
   }
 }
 
@@ -10110,7 +10113,8 @@ StringRef swift::getMacroIntroducedDeclNameString(
 static MacroRoles freestandingMacroRoles =
   (MacroRoles() |
    MacroRole::Expression |
-   MacroRole::Declaration);
+   MacroRole::Declaration |
+   MacroRole::CodeItem);
 static MacroRoles attachedMacroRoles = (MacroRoles() |
                                         MacroRole::Accessor |
                                         MacroRole::MemberAttribute |
@@ -10316,6 +10320,7 @@ void MacroDecl::getIntroducedNames(MacroRole role, ValueDecl *attachedTo,
   case MacroRole::Declaration:
   case MacroRole::Member:
   case MacroRole::Peer:
+  case MacroRole::CodeItem:
     names.push_back(MacroDecl::getUniqueNamePlaceholder(getASTContext()));
     break;
 
diff --git a/lib/ASTGen/Sources/ASTGen/Macros.swift b/lib/ASTGen/Sources/ASTGen/Macros.swift
@@ -606,6 +606,21 @@ func expandFreestandingMacroInProcess(
       evaluatedSyntax = Syntax(CodeBlockItemListSyntax(
         decls.map { CodeBlockItemSyntax(item: .decl($0)) }))
 
+    case let codeItemMacro as CodeItemMacro.Type:
+      func expandCodeItemMacro<Node: FreestandingMacroExpansionSyntax>(
+        _ node: Node
+      ) throws -> [CodeBlockItemSyntax] {
+        return try codeItemMacro.expansion(
+          of: sourceManager.detach(
+            node,
+            foldingWith: OperatorTable.standardOperators
+          ),
+          in: context
+        )
+      }
+      let items = try _openExistential(parentExpansion, do: expandCodeItemMacro)
+      evaluatedSyntax = Syntax(CodeBlockItemListSyntax(items))
+
     default:
       print("not an expression macro or a declaration macro")
       return nil
diff --git a/lib/Parse/ParseDecl.cpp b/lib/Parse/ParseDecl.cpp
@@ -2204,6 +2204,7 @@ static Optional<MacroRole> getMacroRole(
   auto role = llvm::StringSwitch<Optional<MacroRole>>(roleName->str())
       .Case("declaration", MacroRole::Declaration)
       .Case("expression", MacroRole::Expression)
+      .Case("codeItem", MacroRole::CodeItem)
       .Case("accessor", MacroRole::Accessor)
       .Case("memberAttribute", MacroRole::MemberAttribute)
       .Case("member", MacroRole::Member)
diff --git a/lib/SILGen/SILGenExpr.cpp b/lib/SILGen/SILGenExpr.cpp
@@ -6165,10 +6165,18 @@ RValue RValueEmitter::visitMacroExpansionExpr(MacroExpansionExpr *E,
     MacroScope scope(SGF, CleanupLocation(rewritten), E, name.str(),
                      E->getMacroRef().getDecl());
     return visit(rewritten, C);
-  } else {
-    assert(E->getSubstituteDecl());
+  }
+  else if (auto *MED = E->getSubstituteDecl()) {
+    Mangle::ASTMangler mangler;
+    MED->forEachExpandedExprOrStmt([&](ASTNode node) {
+      if (auto *expr = node.dyn_cast<Expr *>())
+        visit(expr, C);
+      else if (auto *stmt = node.dyn_cast<Stmt *>())
+        SGF.emitStmt(stmt);
+    });
     return RValue();
   }
+  return RValue();
 }
 
 RValue SILGenFunction::emitRValue(Expr *E, SGFContext C) {
diff --git a/lib/Sema/CSApply.cpp b/lib/Sema/CSApply.cpp
@@ -5394,8 +5394,8 @@ namespace {
           }
         }
         // For a non-expression macro, expand it as a declaration.
-        else if (macro->getMacroRoles().contains(MacroRole::Declaration)) {
-          auto &ctx = cs.getASTContext();
+        else if (macro->getMacroRoles().contains(MacroRole::Declaration) ||
+                 macro->getMacroRoles().contains(MacroRole::CodeItem)) {
           if (!E->getSubstituteDecl()) {
             auto *med = E->createSubstituteDecl();
             TypeChecker::typeCheckDecl(med);
diff --git a/lib/Sema/TypeCheckAttr.cpp b/lib/Sema/TypeCheckAttr.cpp
@@ -160,7 +160,6 @@ class AttributeChecker : public AttributeVisitor<AttributeChecker> {
   IGNORED_ATTR(Preconcurrency)
   IGNORED_ATTR(BackDeployed)
   IGNORED_ATTR(Documentation)
-  IGNORED_ATTR(MacroRole)
   IGNORED_ATTR(LexicalLifetimes)
 #undef IGNORED_ATTR
 
@@ -342,6 +341,8 @@ class AttributeChecker : public AttributeVisitor<AttributeChecker> {
   void visitSendableAttr(SendableAttr *attr);
 
   void visitRuntimeMetadataAttr(RuntimeMetadataAttr *attr);
+
+  void visitMacroRoleAttr(MacroRoleAttr *attr);
 };
 
 } // end anonymous namespace
@@ -6989,6 +6990,59 @@ void AttributeChecker::visitRuntimeMetadataAttr(RuntimeMetadataAttr *attr) {
   }
 }
 
+void AttributeChecker::visitMacroRoleAttr(MacroRoleAttr *attr) {
+  switch (attr->getMacroSyntax()) {
+  case MacroSyntax::Freestanding: {
+    switch (attr->getMacroRole()) {
+    case MacroRole::Expression:
+      if (!attr->getNames().empty())
+        diagnoseAndRemoveAttr(attr, diag::macro_cannot_introduce_names,
+                              getMacroRoleString(attr->getMacroRole()));
+      break;
+    case MacroRole::Declaration:
+      // TODO: Check names
+      break;
+    case MacroRole::CodeItem:
+      if (!attr->getNames().empty())
+        diagnoseAndRemoveAttr(attr, diag::macro_cannot_introduce_names,
+                              getMacroRoleString(attr->getMacroRole()));
+      break;
+    default:
+      diagnoseAndRemoveAttr(attr, diag::invalid_macro_role_for_macro_syntax,
+                            /*freestanding*/0);
+      break;
+    }
+    break;
+  }
+  case MacroSyntax::Attached: {
+    switch (attr->getMacroRole()) {
+    case MacroRole::Accessor:
+      // TODO: Check property observer names?
+      break;
+    case MacroRole::MemberAttribute:
+      if (!attr->getNames().empty())
+        diagnoseAndRemoveAttr(attr, diag::macro_cannot_introduce_names,
+                              getMacroRoleString(attr->getMacroRole()));
+      break;
+    case MacroRole::Member:
+      break;
+    case MacroRole::Peer:
+      break;
+    case MacroRole::Conformance:
+      if (!attr->getNames().empty())
+        diagnoseAndRemoveAttr(attr, diag::macro_cannot_introduce_names,
+                              getMacroRoleString(attr->getMacroRole()));
+      break;
+    default:
+      diagnoseAndRemoveAttr(attr, diag::invalid_macro_role_for_macro_syntax,
+                            /*attached*/1);
+      break;
+    }
+    break;
+  }
+  }
+}
+
 namespace {
 
 class ClosureAttributeChecker
diff --git a/lib/Sema/TypeCheckDeclPrimary.cpp b/lib/Sema/TypeCheckDeclPrimary.cpp
@@ -2037,7 +2037,10 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
   void visitMacroExpansionDecl(MacroExpansionDecl *MED) {
     // Assign a discriminator.
     (void)MED->getDiscriminator();
-    // Expansion already visited as auxiliary decls.
+    // Decls in expansion already visited as auxiliary decls.
+    MED->forEachExpandedExprOrStmt([&](ASTNode node) {
+      TypeChecker::typeCheckASTNode(node, MED->getDeclContext());
+    });
   }
 
   void visitBoundVariable(VarDecl *VD) {
@@ -3789,12 +3792,10 @@ ExpandMacroExpansionDeclRequest::evaluate(Evaluator &evaluator,
   // If it's not a declaration macro or a code item macro, it must have been
   // parsed as an expression macro, and this decl is just its substitute decl.
   // So there's no thing to be done here.
-  if (!roles.contains(MacroRole::Declaration))
+  if (!roles.contains(MacroRole::Declaration) &&
+      !roles.contains(MacroRole::CodeItem))
     return None;
 
-  // Otherwise, we treat it as a declaration macro.
-  assert(roles.contains(MacroRole::Declaration));
-
   // For now, restrict global freestanding macros in script mode.
   if (dc->isModuleScopeContext() &&
       dc->getParentSourceFile()->isScriptMode()) {
diff --git a/lib/Sema/TypeCheckEffects.cpp b/lib/Sema/TypeCheckEffects.cpp
@@ -413,6 +413,8 @@ class EffectsHandlingWalker : public ASTWalker {
     } else if (auto patternBinding = dyn_cast<PatternBindingDecl>(D)) {
       if (patternBinding->isAsyncLet())
         recurse = asImpl().checkAsyncLet(patternBinding);
+    } else if (auto macroExpansionDecl = dyn_cast<MacroExpansionDecl>(D)) {
+      recurse = ShouldRecurse;
     } else {
       recurse = ShouldNotRecurse;
     }
@@ -444,6 +446,8 @@ class EffectsHandlingWalker : public ASTWalker {
       recurse = asImpl().checkDeclRef(declRef);
     } else if (auto interpolated = dyn_cast<InterpolatedStringLiteralExpr>(E)) {
       recurse = asImpl().checkInterpolatedStringLiteral(interpolated);
+    } else if (auto macroExpansionExpr = dyn_cast<MacroExpansionExpr>(E)) {
+      recurse = ShouldRecurse;
     }
     // Error handling validation (via checkTopLevelEffects) happens after
     // type checking. If an unchecked expression is still around, the code was
diff --git a/lib/Sema/TypeCheckMacros.cpp b/lib/Sema/TypeCheckMacros.cpp
@@ -885,10 +885,13 @@ swift::expandFreestandingMacro(MacroExpansionDecl *med) {
   NullTerminatedStringRef evaluatedSource;
 
   MacroDecl *macro = cast<MacroDecl>(med->getMacroRef().getDecl());
-  assert(macro->getMacroRoles().contains(MacroRole::Declaration));
+  auto macroRoles = macro->getMacroRoles();
+  assert(macroRoles.contains(MacroRole::Declaration) ||
+         macroRoles.contains(MacroRole::CodeItem));
 
   if (isFromExpansionOfMacro(sourceFile, macro, MacroRole::Expression) ||
-      isFromExpansionOfMacro(sourceFile, macro, MacroRole::Declaration)) {
+      isFromExpansionOfMacro(sourceFile, macro, MacroRole::Declaration) ||
+      isFromExpansionOfMacro(sourceFile, macro, MacroRole::CodeItem)) {
     med->diagnose(diag::macro_recursive, macro->getName());
     return None;
   }
@@ -934,12 +937,20 @@ swift::expandFreestandingMacro(MacroExpansionDecl *med) {
       return None;
     }
 
-    // Make sure freestanding macros are enabled before we expand.
-    if (!ctx.LangOpts.hasFeature(Feature::FreestandingMacros) &&
-        !macro->getMacroRoles().contains(MacroRole::Expression)) {
-      med->diagnose(
-          diag::macro_experimental, "freestanding", "FreestandingMacros");
-      return None;
+    // Currently only expression macros are enabled by default. Declaration
+    // macros need the `FreestandingMacros` feature flag, and code item macros
+    // need both `FreestandingMacros` and `CodeItemMacros`.
+    if (!macroRoles.contains(MacroRole::Expression)) {
+      if (!ctx.LangOpts.hasFeature(Feature::FreestandingMacros)) {
+        med->diagnose(diag::macro_experimental, "freestanding",
+                      "FreestandingMacros");
+        return None;
+      }
+      if (!macroRoles.contains(MacroRole::Declaration) &&
+          !ctx.LangOpts.hasFeature(Feature::CodeItemMacros)) {
+        med->diagnose(diag::macro_experimental, "code item", "CodeItemMacros");
+        return None;
+      }
     }
 
 #if SWIFT_SWIFT_PARSER
@@ -1284,6 +1295,7 @@ evaluateAttachedMacro(MacroDecl *macro, Decl *attachedTo, CustomAttr *attr,
 
   case MacroRole::Expression:
   case MacroRole::Declaration:
+  case MacroRole::CodeItem:
     llvm_unreachable("freestanding macro in attached macro evaluation");
   }
 
diff --git a/lib/Serialization/Deserialization.cpp b/lib/Serialization/Deserialization.cpp
@@ -2671,6 +2671,7 @@ getActualMacroRole(uint8_t context) {
   CASE(Member)
   CASE(Peer)
   CASE(Conformance)
+  CASE(CodeItem)
 #undef CASE
   }
   return None;
diff --git a/lib/Serialization/ModuleFormat.h b/lib/Serialization/ModuleFormat.h
@@ -619,6 +619,7 @@ enum class MacroRole : uint8_t {
   Member,
   Peer,
   Conformance,
+  CodeItem,
 };
 using MacroRoleField = BCFixed<3>;
 
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
@@ -2239,6 +2239,7 @@ static uint8_t getRawStableMacroRole(swift::MacroRole context) {
   CASE(Member)
   CASE(Peer)
   CASE(Conformance)
+  CASE(CodeItem)
   }
 #undef CASE
   llvm_unreachable("bad result declaration macro kind");
diff --git a/test/Macros/Inputs/syntax_macro_definitions.swift b/test/Macros/Inputs/syntax_macro_definitions.swift
diff --git a/test/Macros/macro_expand.swift b/test/Macros/macro_expand.swift
diff --git a/test/Macros/macro_expand_codeitems.swift b/test/Macros/macro_expand_codeitems.swift
diff --git a/test/Macros/macros_diagnostics.swift b/test/Macros/macros_diagnostics.swift
