[Macros] Handle macro expansion for function-like macros.

DougGregor · DougGregor · commit d6d2318dc9a1 · 2022-10-27T17:15:30.000-07:00
diff --git a/include/swift/AST/Expr.h b/include/swift/AST/Expr.h
@@ -6013,20 +6013,23 @@ class TypeJoinExpr final : public Expr,
 
 class MacroExpansionExpr final : public Expr {
 private:
-  Expr *Macro;
-  Expr *Rewritten;
-  ArgumentList *ArgList;
   SourceLoc PoundLoc;
+  DeclNameRef MacroName;
+  DeclNameLoc MacroNameLoc;
+  ArgumentList *ArgList;
+  Expr *Rewritten;
 
 public:
-  explicit MacroExpansionExpr(SourceLoc poundLoc, Expr *macro,
+  explicit MacroExpansionExpr(SourceLoc poundLoc, DeclNameRef macroName,
+                              DeclNameLoc macroNameLoc,
                               ArgumentList *argList, bool isImplicit = false,
                               Type ty = Type())
-      : Expr(ExprKind::MacroExpansion, isImplicit, ty), Macro(macro),
-        Rewritten(nullptr), ArgList(argList), PoundLoc(poundLoc) {}
+      : Expr(ExprKind::MacroExpansion, isImplicit, ty), PoundLoc(poundLoc),
+        MacroName(macroName), MacroNameLoc(macroNameLoc), ArgList(argList),
+        Rewritten(nullptr) { }
 
-  Expr *getMacro() const { return Macro; }
-  void setMacro(Expr *macro) { Macro = macro; }
+  DeclNameRef getMacroName() const { return MacroName; }
+  DeclNameLoc getMacroNameLoc() const { return MacroNameLoc; }
 
   Expr *getRewritten() const { return Rewritten; }
   void setRewritten(Expr *rewritten) { Rewritten = rewritten; }
@@ -6038,7 +6041,7 @@ class MacroExpansionExpr final : public Expr {
 
   SourceRange getSourceRange() const {
     return SourceRange(
-        PoundLoc, ArgList ? ArgList->getEndLoc() : Macro->getEndLoc());
+        PoundLoc, ArgList ? ArgList->getEndLoc() : MacroNameLoc.getEndLoc());
   }
 
   static bool classof(const Expr *E) {
diff --git a/lib/AST/ASTDumper.cpp b/lib/AST/ASTDumper.cpp
@@ -2976,8 +2976,7 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
 
   void visitMacroExpansionExpr(MacroExpansionExpr *E) {
     printCommon(E, "macro_expansion_expr");
-    OS << '\n';
-    printRec(E->getMacro());
+    PrintWithColorRAII(OS, IdentifierColor) << " name=" << E->getMacroName();
     if (E->getArgs()) {
       OS << '\n';
       printArgumentList(E->getArgs());
diff --git a/lib/AST/ASTWalker.cpp b/lib/AST/ASTWalker.cpp
@@ -1247,8 +1247,6 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   Expr *visitMacroExpansionExpr(MacroExpansionExpr *E) {
-    auto *macro = doIt(E->getMacro());
-    if (!macro) return nullptr;
     Expr *rewritten = nullptr;
     if (E->getRewritten()) {
       rewritten = doIt(E->getRewritten());
@@ -1259,7 +1257,6 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       args = doIt(E->getArgs());
       if (!args) return nullptr;
     }
-    E->setMacro(macro);
     E->setRewritten(rewritten);
     E->setArgs(args);
     return E;
diff --git a/lib/Parse/ParseExpr.cpp b/lib/Parse/ParseExpr.cpp
@@ -3643,8 +3643,6 @@ ParserResult<Expr> Parser::parseExprMacroExpansion(bool isExprBasic) {
       DeclNameOptions());
   if (!macroNameRef)
     return makeParserError();
-  auto *macroExpr = new (Context) UnresolvedDeclRefExpr(
-      macroNameRef, DeclRefKind::Ordinary, macroNameLoc);
 
   ArgumentList *argList = nullptr;
   if (Tok.isFollowingLParen()) {
@@ -3669,7 +3667,8 @@ ParserResult<Expr> Parser::parseExprMacroExpansion(bool isExprBasic) {
   }
 
   return makeParserResult(
-      new (Context) MacroExpansionExpr(poundLoc, macroExpr, argList));
+      new (Context) MacroExpansionExpr(
+          poundLoc, macroNameRef, macroNameLoc, argList));
 }
 
 /// parseExprCollection - Parse a collection literal expression.
diff --git a/lib/Sema/CSApply.cpp b/lib/Sema/CSApply.cpp
@@ -2977,9 +2977,12 @@ namespace {
         auto kind = MagicIdentifierLiteralExpr::getKindString(expr->getKind())
             .drop_front();
         auto expandedType = solution.simplifyType(solution.getType(expr));
+        cs.setType(expr, expandedType);
+
         if (auto newExpr = expandMacroExpr(dc, expr, kind, expandedType)) {
           auto expansion = new (ctx) MacroExpansionExpr(
-              expr->getStartLoc(), expr, nullptr, /*isImplicit=*/true,
+              expr->getStartLoc(), DeclNameRef(ctx.getIdentifier(kind)),
+              DeclNameLoc(expr->getLoc()), nullptr, /*isImplicit=*/true,
               expandedType);
           expansion->setRewritten(newExpr);
           cs.cacheExprTypes(expansion);
@@ -5385,7 +5388,22 @@ namespace {
     }
 
     Expr *visitMacroExpansionExpr(MacroExpansionExpr *E) {
-      // TODO: Expand macro.
+#if SWIFT_SWIFT_PARSER
+      auto &ctx = cs.getASTContext();
+      if (ctx.LangOpts.hasFeature(Feature::Macros)) {
+        auto macroIdent = E->getMacroName().getBaseIdentifier();
+        auto expandedType = solution.simplifyType(solution.getType(E));
+        cs.setType(E, expandedType);
+        if (auto newExpr = expandMacroExpr(
+                dc, E, macroIdent.str(), expandedType)) {
+          E->setRewritten(newExpr);
+          cs.cacheExprTypes(E);
+          return E;
+        }
+
+        // Fall through to use old implementation.
+      }
+#endif
       return E;
     }
 
diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
@@ -3631,12 +3631,31 @@ namespace {
 #if SWIFT_SWIFT_PARSER
       auto &ctx = CS.getASTContext();
       if (ctx.LangOpts.hasFeature(Feature::Macros)) {
-        auto *UDRE = dyn_cast<UnresolvedDeclRefExpr>(expr->getMacro());
-        if (!UDRE)
-          return Type();
-        
-        auto macroIdent = UDRE->getName().getBaseIdentifier();
-        return CS.getTypeOfMacroReference(macroIdent.str(), expr);
+        auto macroIdent = expr->getMacroName().getBaseIdentifier();
+        auto refType = CS.getTypeOfMacroReference(macroIdent.str(), expr);
+        if (expr->getArgs()) {
+          CS.associateArgumentList(CS.getConstraintLocator(expr), expr->getArgs());
+          // FIXME: Do we have object-like vs. function-like macros?
+          if (auto fnType = dyn_cast<FunctionType>(refType.getPointer())) {
+            SmallVector<AnyFunctionType::Param, 8> params;
+            getMatchingParams(expr->getArgs(), params);
+
+            Type resultType = CS.createTypeVariable(
+                CS.getConstraintLocator(expr, ConstraintLocator::FunctionResult),
+                TVO_CanBindToNoEscape);
+
+            CS.addConstraint(
+                ConstraintKind::ApplicableFunction,
+                FunctionType::get(params, resultType),
+                fnType,
+                CS.getConstraintLocator(
+                  expr, ConstraintLocator::ApplyFunction));
+
+            return resultType;
+          }
+        }
+
+        return refType;
       }
 #endif
       return Type();
diff --git a/test/Macros/builtin_macros.swift b/test/Macros/builtin_macros.swift
@@ -2,22 +2,19 @@
 // REQUIRES: OS=macosx
 // REQUIRES: asserts
 
+
 // CHECK: macro_expansion_expr implicit type='String'
-// CHECK-NEXT: magic_identifier_literal_expr{{.*}}kind=#function
 // CHECK-NEXT: string_literal_expr{{.*}}Macro expansion of #function in{{.*}}value="MacrosTest"
 print(#function)
 
 func f(a: Int, b: Int) {
   print(#function, #line, #column)
   // CHECK: macro_expansion_expr implicit type='String'
-  // CHECK-NEXT: magic_identifier_literal_expr{{.*}}kind=#function
   // CHECK-NEXT: string_literal_expr{{.*}}Macro expansion of #function in{{.*}}value="f(a:b:)"
 
   // CHECK: macro_expansion_expr implicit type='Int'
-  // CHECK-NEXT: magic_identifier_literal_expr{{.*}}kind=#line
   // CHECK-NEXT: integer_literal_expr{{.*}}Macro expansion of #line in{{.*}}value=11
 
   // CHECK: macro_expansion_expr implicit type='Int'
-  // CHECK-NEXT: magic_identifier_literal_expr{{.*}}kind=#column
   // CHECK-NEXT: integer_literal_expr{{.*}}Macro expansion of #column in{{.*}}value=27
 }
diff --git a/test/Macros/macros.swift b/test/Macros/macros.swift
@@ -0,0 +1,10 @@
+// RUN: %target-swift-frontend -enable-experimental-feature Macros -dump-ast %s -module-name MacrosTest 2>&1 | %FileCheck %s
+// REQUIRES: OS=macosx
+// REQUIRES: asserts
+
+let a = 1, b = 1
+let s = #stringify(a + b)
+// CHECK: macro_expansion_expr type='(Int, String)'{{.*}}name=stringify
+// CHECK-NEXT: argument_list
+// CHECK: tuple_expr type='(Int, String)' location=Macro expansion of #stringify
+
