[Macros] Parse generic arguments in macro expansions.

DougGregor · DougGregor · commit 787af41765cd · 2022-11-17T16:46:43.000-08:00
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -8340,14 +8340,28 @@ class MacroExpansionDecl : public Decl {
   SourceLoc PoundLoc;
   DeclNameRef Macro;
   DeclNameLoc MacroLoc;
+  SourceLoc LeftAngleLoc, RightAngleLoc;
+  ArrayRef<TypeRepr *> GenericArgs;
   ArgumentList *ArgList;
   Decl *Rewritten;
 
 public:
   MacroExpansionDecl(DeclContext *dc, SourceLoc poundLoc, DeclNameRef macro,
-                     DeclNameLoc macroLoc, ArgumentList *args)
+                     DeclNameLoc macroLoc,
+                     SourceLoc leftAngleLoc,
+                     ArrayRef<TypeRepr *> genericArgs,
+                     SourceLoc rightAngleLoc,
+                     ArgumentList *args)
       : Decl(DeclKind::MacroExpansion, dc), PoundLoc(poundLoc),
-        Macro(macro), MacroLoc(macroLoc), ArgList(args), Rewritten(nullptr) {}
+        Macro(macro), MacroLoc(macroLoc),
+        LeftAngleLoc(leftAngleLoc), RightAngleLoc(rightAngleLoc),
+        GenericArgs(genericArgs), ArgList(args), Rewritten(nullptr) {}
+
+  ArrayRef<TypeRepr *> getGenericArgs() const { return GenericArgs; }
+
+  SourceRange getGenericArgsRange() const {
+    return SourceRange(LeftAngleLoc, RightAngleLoc);
+  }
 
   SourceRange getSourceRange() const;
   SourceLoc getLocFromSource() const { return PoundLoc; }
diff --git a/include/swift/AST/Expr.h b/include/swift/AST/Expr.h
@@ -6016,16 +6016,24 @@ class MacroExpansionExpr final : public Expr {
   SourceLoc PoundLoc;
   DeclNameRef MacroName;
   DeclNameLoc MacroNameLoc;
+  SourceLoc LeftAngleLoc, RightAngleLoc;
+  ArrayRef<TypeRepr *> GenericArgs;
   ArgumentList *ArgList;
   Expr *Rewritten;
 
 public:
   explicit MacroExpansionExpr(SourceLoc poundLoc, DeclNameRef macroName,
                               DeclNameLoc macroNameLoc,
+                              SourceLoc leftAngleLoc,
+                              ArrayRef<TypeRepr *> genericArgs,
+                              SourceLoc rightAngleLoc,
                               ArgumentList *argList, bool isImplicit = false,
                               Type ty = Type())
       : Expr(ExprKind::MacroExpansion, isImplicit, ty), PoundLoc(poundLoc),
-        MacroName(macroName), MacroNameLoc(macroNameLoc), ArgList(argList),
+        MacroName(macroName), MacroNameLoc(macroNameLoc),
+        LeftAngleLoc(leftAngleLoc), RightAngleLoc(rightAngleLoc),
+        GenericArgs(genericArgs),
+        ArgList(argList),
         Rewritten(nullptr) { }
 
   DeclNameRef getMacroName() const { return MacroName; }
@@ -6034,15 +6042,18 @@ class MacroExpansionExpr final : public Expr {
   Expr *getRewritten() const { return Rewritten; }
   void setRewritten(Expr *rewritten) { Rewritten = rewritten; }
 
+  ArrayRef<TypeRepr *> getGenericArgs() const { return GenericArgs; }
+
+  SourceRange getGenericArgsRange() const {
+    return SourceRange(LeftAngleLoc, RightAngleLoc);
+  }
+
   ArgumentList *getArgs() const { return ArgList; }
   void setArgs(ArgumentList *newArgs) { ArgList = newArgs; }
 
   SourceLoc getLoc() const { return PoundLoc; }
 
-  SourceRange getSourceRange() const {
-    return SourceRange(
-        PoundLoc, ArgList ? ArgList->getEndLoc() : MacroNameLoc.getEndLoc());
-  }
+  SourceRange getSourceRange() const;
 
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::MacroExpansion;
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -9644,6 +9644,13 @@ SourceRange MacroDecl::getSourceRange() const {
 }
 
 SourceRange MacroExpansionDecl::getSourceRange() const {
-  return SourceRange(PoundLoc,
-                     ArgList ? ArgList->getEndLoc() : MacroLoc.getEndLoc());
+  SourceLoc endLoc;
+  if (ArgList)
+    endLoc = ArgList->getEndLoc();
+  else if (RightAngleLoc.isValid())
+    endLoc = RightAngleLoc;
+  else
+    endLoc = MacroLoc.getEndLoc();
+
+  return SourceRange(PoundLoc, endLoc);
 }
diff --git a/lib/AST/Expr.cpp b/lib/AST/Expr.cpp
@@ -2470,6 +2470,18 @@ TypeJoinExpr *TypeJoinExpr::create(ASTContext &ctx, DeclRefExpr *var,
   return new (mem) TypeJoinExpr(var, elements);
 }
 
+SourceRange MacroExpansionExpr::getSourceRange() const {
+  SourceLoc endLoc;
+  if (ArgList)
+    endLoc = ArgList->getEndLoc();
+  else if (RightAngleLoc.isValid())
+    endLoc = RightAngleLoc;
+  else
+    endLoc = MacroNameLoc.getEndLoc();
+
+  return SourceRange(PoundLoc, endLoc);
+}
+
 void swift::simple_display(llvm::raw_ostream &out, const ClosureExpr *CE) {
   if (!CE) {
     out << "(null)";
diff --git a/lib/Parse/ParseDecl.cpp b/lib/Parse/ParseDecl.cpp
@@ -8987,28 +8987,44 @@ Parser::parseDeclMacroExpansion(ParseDeclOptions flags,
   if (!macroNameRef)
     return makeParserError();
 
+  ParserStatus status;
+  SourceLoc leftAngleLoc, rightAngleLoc;
+  SmallVector<TypeRepr *, 8> genericArgs;
+  if (canParseAsGenericArgumentList()) {
+    auto genericArgsStatus = parseGenericArguments(
+        genericArgs, leftAngleLoc, rightAngleLoc);
+    status |= genericArgsStatus;
+    if (genericArgsStatus.isErrorOrHasCompletion())
+      diagnose(leftAngleLoc, diag::while_parsing_as_left_angle_bracket);
+  }
+
   ArgumentList *argList = nullptr;
   if (Tok.isFollowingLParen()) {
     auto result = parseArgumentList(tok::l_paren, tok::r_paren,
                                     /*isExprBasic*/ false,
                                     /*allowTrailingClosure*/ true);
+    status |= result;
     if (result.hasCodeCompletion())
       return makeParserCodeCompletionResult<MacroExpansionDecl>();
-    if (result.isParseError())
-      return makeParserError();
-    argList = result.get();
+    argList = result.getPtrOrNull();
   } else if (Tok.is(tok::l_brace)) {
     SmallVector<Argument, 2> trailingClosures;
-    auto status = parseTrailingClosures(/*isExprBasic*/ false,
-                                        macroNameLoc.getSourceRange(),
-                                        trailingClosures);
-    if (status.isError() || trailingClosures.empty())
-      return makeParserError();
-    argList = ArgumentList::createParsed(Context, SourceLoc(),
-                                         trailingClosures, SourceLoc(),
-                                         /*trailingClosureIdx*/ 0);
+    auto closuresStatus = parseTrailingClosures(/*isExprBasic*/ false,
+                                                macroNameLoc.getSourceRange(),
+                                                trailingClosures);
+    status |= closuresStatus;
+
+    if (!trailingClosures.empty()) {
+      argList = ArgumentList::createParsed(Context, SourceLoc(),
+                                           trailingClosures, SourceLoc(),
+                                           /*trailingClosureIdx*/ 0);
+    }
   }
 
-  return makeParserResult(new (Context) MacroExpansionDecl(
-      CurDeclContext, poundLoc, macroNameRef, macroNameLoc, argList));
+  return makeParserResult(
+      status,
+      new (Context) MacroExpansionDecl(
+        CurDeclContext, poundLoc, macroNameRef, macroNameLoc,
+        leftAngleLoc, Context.AllocateCopy(genericArgs), rightAngleLoc,
+        argList));
 }
diff --git a/lib/Parse/ParseExpr.cpp b/lib/Parse/ParseExpr.cpp
@@ -3411,31 +3411,45 @@ ParserResult<Expr> Parser::parseExprMacroExpansion(bool isExprBasic) {
   if (!macroNameRef)
     return makeParserError();
 
+  ParserStatus status;
+  SourceLoc leftAngleLoc, rightAngleLoc;
+  SmallVector<TypeRepr *, 8> genericArgs;
+  if (canParseAsGenericArgumentList()) {
+    auto genericArgsStatus = parseGenericArguments(
+        genericArgs, leftAngleLoc, rightAngleLoc);
+    status |= genericArgsStatus;
+    if (genericArgsStatus.isErrorOrHasCompletion())
+      diagnose(leftAngleLoc, diag::while_parsing_as_left_angle_bracket);
+  }
+
   ArgumentList *argList = nullptr;
   if (Tok.isFollowingLParen()) {
     auto result = parseArgumentList(tok::l_paren, tok::r_paren, isExprBasic,
                                 /*allowTrailingClosure*/ true);
     if (result.hasCodeCompletion())
       return makeParserCodeCompletionResult<Expr>();
-    if (result.isParseError())
-      return makeParserError();
-    argList = result.get();
+    argList = result.getPtrOrNull();
+    status |= result;
   } else if (Tok.is(tok::l_brace) &&
              isValidTrailingClosure(isExprBasic, *this)) {
     SmallVector<Argument, 2> trailingClosures;
-    auto status = parseTrailingClosures(isExprBasic,
-                                        macroNameLoc.getSourceRange(),
-                                        trailingClosures);
-    if (status.isError() || trailingClosures.empty())
-      return makeParserError();
-    argList = ArgumentList::createParsed(Context, SourceLoc(),
-                                         trailingClosures, SourceLoc(),
-                                         /*trailingClosureIdx*/ 0);
+    auto closureStatus = parseTrailingClosures(isExprBasic,
+                                               macroNameLoc.getSourceRange(),
+                                               trailingClosures);
+    status |= closureStatus;
+
+    if (!trailingClosures.empty()) {
+      argList = ArgumentList::createParsed(Context, SourceLoc(),
+                                           trailingClosures, SourceLoc(),
+                                           /*trailingClosureIdx*/ 0);
+    }
   }
 
   return makeParserResult(
+      status,
       new (Context) MacroExpansionExpr(
-          poundLoc, macroNameRef, macroNameLoc, argList));
+          poundLoc, macroNameRef, macroNameLoc, leftAngleLoc,
+          Context.AllocateCopy(genericArgs), rightAngleLoc, argList));
 }
 
 /// parseExprCollection - Parse a collection literal expression.
diff --git a/lib/Sema/CSApply.cpp b/lib/Sema/CSApply.cpp
@@ -2986,8 +2986,8 @@ namespace {
         if (auto newExpr = expandMacroExpr(dc, expr, macroRef, expandedType)) {
           auto expansion = new (ctx) MacroExpansionExpr(
               expr->getStartLoc(), DeclNameRef(macro->getName()),
-              DeclNameLoc(expr->getLoc()), nullptr, /*isImplicit=*/true,
-              expandedType);
+              DeclNameLoc(expr->getLoc()), SourceLoc(), { }, SourceLoc(),
+              nullptr, /*isImplicit=*/true, expandedType);
           expansion->setRewritten(newExpr);
           cs.cacheExprTypes(expansion);
           return expansion;
diff --git a/test/Macros/macros.swift b/test/Macros/macros.swift
@@ -12,4 +12,9 @@ func test(a: Int, b: Int) {
   // CHECK: macro_expansion_expr type='(() -> Bool, String)'{{.*}}name=stringify
   // CHECK-NEXT: argument_list
   // CHECK: tuple_expr type='(() -> Bool, String)' location=Macro expansion of #stringify
+
+  let (b2, s3) = #stringify<Double>(1 + 2)
+  // CHECK: macro_expansion_expr type='(Int, String)'{{.*}}name=stringify
+  // CHECK-NEXT: argument_list
+  // CHECK: tuple_expr type='(Int, String)' location=Macro expansion of #stringify
 }
