From f8b1ca4992a22b4b65282c09dd6f07a1a2839070 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Fri, 5 Jul 2024 14:08:13 -0700
Subject: [PATCH] [MCParser] .altmacro: Support argument expansion not preceded
 by \

In the .altmacro mode, an argument can be expanded even if not preceded
by \
---
 llvm/lib/MC/MCParser/AsmParser.cpp    | 79 +++++++++++++++++----------
 llvm/test/MC/AsmParser/altmacro-arg.s | 22 ++++++++
 2 files changed, 71 insertions(+), 30 deletions(-)
 create mode 100644 llvm/test/MC/AsmParser/altmacro-arg.s

diff --git a/llvm/lib/MC/MCParser/AsmParser.cpp b/llvm/lib/MC/MCParser/AsmParser.cpp
index bf3061ad541e34..707edb0481a619 100644
--- a/llvm/lib/MC/MCParser/AsmParser.cpp
+++ b/llvm/lib/MC/MCParser/AsmParser.cpp
@@ -2500,7 +2500,34 @@ bool AsmParser::expandMacro(raw_svector_ostream &OS, MCAsmMacro &Macro,
                             ArrayRef<MCAsmMacroArgument> A,
                             bool EnableAtPseudoVariable) {
   unsigned NParameters = Parameters.size();
-  bool HasVararg = NParameters ? Parameters.back().Vararg : false;
+  auto expandArg = [&](unsigned Index) {
+    bool HasVararg = NParameters ? Parameters.back().Vararg : false;
+    bool VarargParameter = HasVararg && Index == (NParameters - 1);
+    for (const AsmToken &Token : A[Index])
+      // For altmacro mode, you can write '%expr'.
+      // The prefix '%' evaluates the expression 'expr'
+      // and uses the result as a string (e.g. replace %(1+2) with the
+      // string "3").
+      // Here, we identify the integer token which is the result of the
+      // absolute expression evaluation and replace it with its string
+      // representation.
+      if (AltMacroMode && Token.getString().front() == '%' &&
+          Token.is(AsmToken::Integer))
+        // Emit an integer value to the buffer.
+        OS << Token.getIntVal();
+      // Only Token that was validated as a string and begins with '<'
+      // is considered altMacroString!!!
+      else if (AltMacroMode && Token.getString().front() == '<' &&
+               Token.is(AsmToken::String)) {
+        OS << angleBracketString(Token.getStringContents());
+      }
+      // We expect no quotes around the string's contents when
+      // parsing for varargs.
+      else if (Token.isNot(AsmToken::String) || VarargParameter)
+        OS << Token.getString();
+      else
+        OS << Token.getStringContents();
+  };
 
   // A macro without parameters is handled differently on Darwin:
   // gas accepts no arguments and does no substitutions
@@ -2534,36 +2561,10 @@ bool AsmParser::expandMacro(raw_svector_ostream &OS, MCAsmMacro &Macro,
       for (; Index < NParameters; ++Index)
         if (Parameters[Index].Name == Argument)
           break;
-      if (Index == NParameters) {
+      if (Index == NParameters)
         OS << '\\' << Argument;
-      } else {
-        bool VarargParameter = HasVararg && Index == (NParameters - 1);
-        for (const AsmToken &Token : A[Index]) {
-          // For altmacro mode, you can write '%expr'.
-          // The prefix '%' evaluates the expression 'expr'
-          // and uses the result as a string (e.g. replace %(1+2) with the
-          // string "3").
-          // Here, we identify the integer token which is the result of the
-          // absolute expression evaluation and replace it with its string
-          // representation.
-          if (AltMacroMode && Token.getString().front() == '%' &&
-              Token.is(AsmToken::Integer))
-            // Emit an integer value to the buffer.
-            OS << Token.getIntVal();
-          // Only Token that was validated as a string and begins with '<'
-          // is considered altMacroString!!!
-          else if (AltMacroMode && Token.getString().front() == '<' &&
-                   Token.is(AsmToken::String)) {
-            OS << angleBracketString(Token.getStringContents());
-          }
-          // We expect no quotes around the string's contents when
-          // parsing for varargs.
-          else if (Token.isNot(AsmToken::String) || VarargParameter)
-            OS << Token.getString();
-          else
-            OS << Token.getStringContents();
-        }
-      }
+      else
+        expandArg(Index);
       continue;
     }
 
@@ -2595,6 +2596,24 @@ bool AsmParser::expandMacro(raw_svector_ostream &OS, MCAsmMacro &Macro,
       }
     }
 
+    if (AltMacroMode && isIdentifierChar(Body[I])) {
+      size_t Len = 1;
+      while (I + Len != End && isIdentifierChar(Body[I + Len]))
+        ++Len;
+      StringRef Argument(Body.data() + I, Len);
+      unsigned Index = 0;
+      for (; Index != NParameters; ++Index)
+        if (Parameters[Index].Name == Argument)
+          break;
+      if (Index != NParameters) {
+        expandArg(Index);
+        I += Len;
+        if (I != End && Body[I] == '&')
+          ++I;
+        continue;
+      }
+    }
+
     OS << Body[I];
     ++I;
   }
diff --git a/llvm/test/MC/AsmParser/altmacro-arg.s b/llvm/test/MC/AsmParser/altmacro-arg.s
new file mode 100644
index 00000000000000..262c5eac832e0c
--- /dev/null
+++ b/llvm/test/MC/AsmParser/altmacro-arg.s
@@ -0,0 +1,22 @@
+## Arguments can be expanded even if they are not preceded by \
+# RUN: llvm-mc -triple=x86_64 %s | FileCheck %s
+
+# CHECK:      1 1 1a
+# CHECK-NEXT: 1 2 1a 2b
+# CHECK-NEXT: \$b \$b
+.altmacro
+.irp ._a,1
+  .print "\._a \._a& ._a&a"
+  .irp $b,2
+    .print "\._a \$b ._a&a $b&b"
+  .endr
+  .print "\$b \$b&"
+.endr
+
+# CHECK:      1 1& ._a&a
+# CHECK-NEXT: \$b \$b&
+.noaltmacro
+.irp ._a,1
+  .print "\._a \._a& ._a&a"
+  .print "\$b \$b&"
+.endr