From a5eb83d74b0a9aa5822e88c1b67010ac65368869 Mon Sep 17 00:00:00 2001
From: Eric Astor <epastor@google.com>
Date: Wed, 5 Mar 2025 17:34:30 +0000
Subject: [PATCH] [ms] [llvm-ml] Allow optional parenthesized arguments for
 macros

We match ML64.EXE, which allows optional parentheses around a macro's arguments.
---
 llvm/lib/MC/MCParser/AsmLexer.cpp   |  4 +++-
 llvm/lib/MC/MCParser/MasmParser.cpp | 11 +++++------
 llvm/test/tools/llvm-ml/macro.asm   | 19 +++++++++++++++++++
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/MC/MCParser/AsmLexer.cpp b/llvm/lib/MC/MCParser/AsmLexer.cpp
index 41ff94d125a16..23836438027c0 100644
--- a/llvm/lib/MC/MCParser/AsmLexer.cpp
+++ b/llvm/lib/MC/MCParser/AsmLexer.cpp
@@ -698,8 +698,10 @@ size_t AsmLexer::peekTokens(MutableArrayRef<AsmToken> Buf,
 
     Buf[ReadCount] = Token;
 
-    if (Token.is(AsmToken::Eof))
+    if (Token.is(AsmToken::Eof)) {
+      ReadCount++;
       break;
+    }
   }
 
   SetError(SavedErrLoc, SavedErr);
diff --git a/llvm/lib/MC/MCParser/MasmParser.cpp b/llvm/lib/MC/MCParser/MasmParser.cpp
index 4ef781c54f11d..15452c8c0822e 100644
--- a/llvm/lib/MC/MCParser/MasmParser.cpp
+++ b/llvm/lib/MC/MCParser/MasmParser.cpp
@@ -1986,7 +1986,10 @@ bool MasmParser::parseStatement(ParseStatementInfo &Info,
 
   // If macros are enabled, check to see if this is a macro instantiation.
   if (const MCAsmMacro *M = getContext().lookupMacro(IDVal.lower())) {
-    return handleMacroEntry(M, IDLoc);
+    AsmToken::TokenKind ArgumentEndTok = parseOptionalToken(AsmToken::LParen)
+                                             ? AsmToken::RParen
+                                             : AsmToken::EndOfStatement;
+    return handleMacroEntry(M, IDLoc, ArgumentEndTok);
   }
 
   // Otherwise, we have a normal instruction or directive.
@@ -2830,7 +2833,7 @@ bool MasmParser::handleMacroEntry(const MCAsmMacro *M, SMLoc NameLoc,
   }
 
   MCAsmMacroArguments A;
-  if (parseMacroArguments(M, A, ArgumentEndTok))
+  if (parseMacroArguments(M, A, ArgumentEndTok) || parseToken(ArgumentEndTok))
     return true;
 
   // Macro instantiation is lexical, unfortunately. We construct a new buffer
@@ -2914,10 +2917,6 @@ bool MasmParser::handleMacroInvocation(const MCAsmMacro *M, SMLoc NameLoc) {
       eatToEndOfStatement();
   }
 
-  // Consume the right-parenthesis on the other side of the arguments.
-  if (parseRParen())
-    return true;
-
   // Exit values may require lexing, unfortunately. We construct a new buffer to
   // hold the exit value.
   std::unique_ptr<MemoryBuffer> MacroValue =
diff --git a/llvm/test/tools/llvm-ml/macro.asm b/llvm/test/tools/llvm-ml/macro.asm
index 3c81b89414fc7..3373ff4837f7c 100644
--- a/llvm/test/tools/llvm-ml/macro.asm
+++ b/llvm/test/tools/llvm-ml/macro.asm
@@ -100,6 +100,25 @@ substitution_test_uppercase PROC
   ret
 substitution_test_uppercase ENDP
 
+substitution_test_with_parentheses PROC
+; CHECK-LABEL: substitution_test_with_parentheses:
+
+  SubstitutionMacro(2, 8)
+; CHECK: mov eax, 2
+; CHECK-NEXT: mov eax, 2
+; CHECK-NEXT: mov eax, 2
+; CHECK-NEXT: mov eax, 2
+; CHECK: mov eax, dword ptr [rip + xa1]
+; CHECK-NEXT: mov eax, dword ptr [rip + x2]
+; CHECK-NEXT: mov eax, dword ptr [rip + x2]
+; CHECK: mov eax, 8
+; CHECK-NEXT: mov eax, 8
+; CHECK-NEXT: mov eax, 8
+; CHECK-NEXT: mov eax, 8
+
+  ret
+substitution_test_with_parentheses ENDP
+
 AmbiguousSubstitutionMacro MACRO x, y
   x&y BYTE 0
 ENDM