[Macros] Use MacroExpansionExpr when resolving attached and freestanding macros.

Author: hborla

include/swift/AST/Expr.h

@@ -6051,16 +6051,19 @@ class TypeJoinExpr final : public Expr,
   }
 };
 
+/// An invocation of a macro expansion, spelled with `#` for freestanding
+/// macros or `@` for attached macros.
 class MacroExpansionExpr final : public Expr {
 private:
   DeclContext *DC;
-  SourceLoc PoundLoc;
+  SourceLoc SigilLoc;
   DeclNameRef MacroName;
   DeclNameLoc MacroNameLoc;
   SourceLoc LeftAngleLoc, RightAngleLoc;
   ArrayRef<TypeRepr *> GenericArgs;
   ArgumentList *ArgList;
   Expr *Rewritten;
+  MacroRoles Roles;
 
   /// The referenced macro.
   ConcreteDeclRef macroRef;
@@ -6069,21 +6072,22 @@ class MacroExpansionExpr final : public Expr {
   enum : unsigned { InvalidDiscriminator = 0xFFFF };
 
   explicit MacroExpansionExpr(DeclContext *dc,
-                              SourceLoc poundLoc, DeclNameRef macroName,
+                              SourceLoc sigilLoc, DeclNameRef macroName,
                               DeclNameLoc macroNameLoc,
                               SourceLoc leftAngleLoc,
                               ArrayRef<TypeRepr *> genericArgs,
                               SourceLoc rightAngleLoc,
                               ArgumentList *argList,
+                              MacroRoles roles,
                               bool isImplicit = false,
                               Type ty = Type())
       : Expr(ExprKind::MacroExpansion, isImplicit, ty),
-        DC(dc), PoundLoc(poundLoc),
+        DC(dc), SigilLoc(sigilLoc),
         MacroName(macroName), MacroNameLoc(macroNameLoc),
         LeftAngleLoc(leftAngleLoc), RightAngleLoc(rightAngleLoc),
         GenericArgs(genericArgs),
         ArgList(argList),
-        Rewritten(nullptr) {
+        Rewritten(nullptr), Roles(roles) {
     Bits.MacroExpansionExpr.Discriminator = InvalidDiscriminator;
   }
 
@@ -6102,7 +6106,9 @@ class MacroExpansionExpr final : public Expr {
   ArgumentList *getArgs() const { return ArgList; }
   void setArgs(ArgumentList *newArgs) { ArgList = newArgs; }
 
-  SourceLoc getLoc() const { return PoundLoc; }
+  MacroRoles getMacroRoles() const { return Roles; }
+
+  SourceLoc getLoc() const { return SigilLoc; }
 
   ConcreteDeclRef getMacroRef() const { return macroRef; }
   void setMacroRef(ConcreteDeclRef ref) { macroRef = ref; }

include/swift/AST/TypeCheckRequests.h

@@ -3151,6 +3151,7 @@ class UnresolvedMacroReference {
     return pointer.getOpaqueValue();
   }
 
+  SourceLoc getSigilLoc() const;
   DeclNameRef getMacroName() const;
   DeclNameLoc getMacroNameLoc() const;
   SourceRange getGenericArgsRange() const;

lib/AST/Expr.cpp

@@ -2538,14 +2538,14 @@ TypeJoinExpr *TypeJoinExpr::create(ASTContext &ctx, DeclRefExpr *var,
 
 SourceRange MacroExpansionExpr::getSourceRange() const {
   SourceLoc endLoc;
-  if (ArgList)
+  if (ArgList && !ArgList->isImplicit())
     endLoc = ArgList->getEndLoc();
   else if (RightAngleLoc.isValid())
     endLoc = RightAngleLoc;
   else
     endLoc = MacroNameLoc.getEndLoc();
 
-  return SourceRange(PoundLoc, endLoc);
+  return SourceRange(SigilLoc, endLoc);
 }
 
 unsigned MacroExpansionExpr::getDiscriminator() const {

lib/AST/TypeCheckRequests.cpp

@@ -1676,6 +1676,16 @@ DeclNameRef UnresolvedMacroReference::getMacroName() const {
   llvm_unreachable("Unhandled case");
 }
 
+SourceLoc UnresolvedMacroReference::getSigilLoc() const {
+  if (auto *med = pointer.dyn_cast<MacroExpansionDecl *>())
+    return med->getPoundLoc();
+  if (auto *mee = pointer.dyn_cast<MacroExpansionExpr *>())
+    return mee->getLoc();
+  if (auto *attr = pointer.dyn_cast<CustomAttr *>())
+    return attr->getRangeWithAt().Start;
+  llvm_unreachable("Unhandled case");
+}
+
 DeclNameLoc UnresolvedMacroReference::getMacroNameLoc() const {
   if (auto *med = pointer.dyn_cast<MacroExpansionDecl *>())
     return med->getMacroLoc();

lib/Parse/ParseExpr.cpp

@@ -3368,7 +3368,8 @@ ParserResult<Expr> Parser::parseExprMacroExpansion(bool isExprBasic) {
       status,
       new (Context) MacroExpansionExpr(
           CurDeclContext, poundLoc, macroNameRef, macroNameLoc, leftAngleLoc,
-          Context.AllocateCopy(genericArgs), rightAngleLoc, argList));
+          Context.AllocateCopy(genericArgs), rightAngleLoc, argList,
+          MacroRole::Expression));
 }
 
 /// parseExprCollection - Parse a collection literal expression.

lib/Sema/CSApply.cpp

@@ -2985,7 +2985,7 @@ namespace {
           auto expansion = new (ctx) MacroExpansionExpr(
               dc, expr->getStartLoc(), DeclNameRef(macro->getName()),
               DeclNameLoc(expr->getLoc()), SourceLoc(), { }, SourceLoc(),
-              nullptr, /*isImplicit=*/true, expandedType);
+              nullptr, MacroRole::Expression, /*isImplicit=*/true, expandedType);
           expansion->setMacroRef(macroRef);
           expansion->setRewritten(newExpr);
           cs.cacheExprTypes(expansion);
@@ -5406,7 +5406,8 @@ namespace {
       ConcreteDeclRef macroRef = resolveConcreteDeclRef(macro, locator);
       E->setMacroRef(macroRef);
 
-      if (!cs.Options.contains(ConstraintSystemFlags::DisableMacroExpansions)) {
+      if (E->getMacroRoles().contains(MacroRole::Expression) &&
+          !cs.Options.contains(ConstraintSystemFlags::DisableMacroExpansions)) {
         if (auto newExpr = expandMacroExpr(dc, E, macroRef, expandedType)) {
           E->setRewritten(newExpr);
         }

lib/Sema/CSGen.cpp

@@ -1211,7 +1211,8 @@ namespace {
 
         auto macroIdent = ctx.getIdentifier(kind);
         auto macros = lookupMacros(
-            macroIdent, expr->getLoc(), FunctionRefKind::Unapplied);
+            macroIdent, expr->getLoc(), FunctionRefKind::Unapplied,
+            MacroRole::Expression);
         if (!macros.empty()) {
           // Introduce an overload set for the macro reference.
           auto locator = CS.getConstraintLocator(expr);
@@ -3775,10 +3776,11 @@ namespace {
     /// Lookup all macros with the given macro name.
     SmallVector<OverloadChoice, 1>
     lookupMacros(Identifier macroName, SourceLoc loc,
-                 FunctionRefKind functionRefKind) {
+                 FunctionRefKind functionRefKind,
+                 MacroRoles roles) {
       SmallVector<OverloadChoice, 1> choices;
       auto results = TypeChecker::lookupMacros(
-          CurDC, DeclNameRef(macroName), loc, MacroRole::Expression);
+          CurDC, DeclNameRef(macroName), loc, roles);
       for (const auto &result : results) {
         OverloadChoice choice = OverloadChoice(Type(), result, functionRefKind);
         choices.push_back(choice);
@@ -3809,7 +3811,7 @@ namespace {
                                                : FunctionRefKind::Unapplied;
       auto macros = lookupMacros(
           macroIdent, expr->getMacroNameLoc().getBaseNameLoc(),
-          functionRefKind);
+          functionRefKind, expr->getMacroRoles());
       if (macros.empty()) {
         ctx.Diags.diagnose(expr->getMacroNameLoc(), diag::macro_undefined,
                            macroIdent)

lib/Sema/ConstraintSystem.cpp

@@ -3590,7 +3590,7 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
     if (auto macro = dyn_cast<MacroDecl>(decl)) {
       // Macro can only be used in an expansion. If we end up here, it's
       // because we found a macro but are missing the leading '#'.
-      if (!(locator->isForMacroExpansion() || locator->getAnchor().isImplicit())) {
+      if (!locator->isForMacroExpansion()) {
         // Record a fix here
         (void)recordFix(MacroMissingPound::create(*this, macro, locator));
       }

lib/Sema/TypeCheckMacros.cpp

@@ -1258,36 +1258,34 @@ ResolveMacroRequest::evaluate(Evaluator &evaluator,
     return nullptr;
 
   // Extract macro arguments, or create an empty list.
-  auto *args = macroRef.getArgs();
-  if (!args)
-    args = ArgumentList::createImplicit(ctx, {});
-
-  // Form an `OverloadedDeclRefExpr` with the filtered lookup result above
-  // to ensure @freestanding macros are not considered in overload resolution.
-  FunctionRefKind functionRefKind = FunctionRefKind::SingleApply;
-  SmallVector<ValueDecl *> valueDecls;
-  for (auto *macro : foundMacros)
-    valueDecls.push_back(macro);
-  Expr *callee = new (ctx) OverloadedDeclRefExpr(
-      valueDecls, macroRef.getMacroNameLoc(), functionRefKind,
-      /*implicit*/true);
-  auto genArgs = macroRef.getGenericArgs();
-  if (!genArgs.empty()) {
-    auto genArgsRange = macroRef.getGenericArgsRange();
-    callee = UnresolvedSpecializeExpr::create(
-        ctx, callee, genArgsRange.Start, genArgs, genArgsRange.End);
+  auto *argList = macroRef.getArgs();
+  if (!argList)
+    argList = ArgumentList::createImplicit(ctx, {});
+
+  // If we already have a MacroExpansionExpr, use that. Otherwise,
+  // create one.
+  MacroExpansionExpr *macroExpansion;
+  if (auto *expr = macroRef.getExpr()) {
+    macroExpansion = expr;
+  } else {
+    SourceRange genericArgsRange = macroRef.getGenericArgsRange();
+    macroExpansion = new (ctx) MacroExpansionExpr(
+      dc, macroRef.getSigilLoc(), macroRef.getMacroName(),
+      macroRef.getMacroNameLoc(), genericArgsRange.Start,
+      macroRef.getGenericArgs(), genericArgsRange.End,
+      argList, roles);
   }
-  auto *call = CallExpr::createImplicit(ctx, callee, args);
 
-  Expr *result = call;
+  Expr *result = macroExpansion;
   TypeChecker::typeCheckExpression(result, dc);
 
-  if (auto *fn = dyn_cast<DeclRefExpr>(call->getFn()))
-    if (auto *macro = dyn_cast<MacroDecl>(fn->getDecl()))
-      return macro;
+  auto macroDeclRef = macroExpansion->getMacroRef();
+  if (auto *macroDecl = dyn_cast_or_null<MacroDecl>(macroDeclRef.getDecl()))
+    return macroDecl;
 
   // If we couldn't resolve a macro decl, the attribute is invalid.
   if (auto *attr = macroRef.getAttr())
     attr->setInvalid();
+
   return nullptr;
 }

test/Macros/attached_macros_diags.swift

@@ -40,7 +40,7 @@ struct SkipNestedType {
 }
 
 struct TestMacroArgs {
-  @m1("extra arg") struct Args1 {} // expected-error{{argument passed to call that takes no arguments}}
+  @m1("extra arg") struct Args1 {} // expected-error{{argument passed to macro expansion that takes no arguments}}
 
   @m2(10) struct Args2 {}