MCAsmParser: Support \+

In .macro, \+ expands to the per-macro invocation count.
https://sourceware.org/pipermail/binutils/2024-May/134009.html

\+ counts from 0 for .irp/.irpc/.rept .

Note: We currently prints \q for `.print "\q"` while gas doesn't. This
patch does not change this behavior.

4 files changed, 61 insertions, 23 deletions

diff --git a/llvm/include/llvm/MC/MCAsmMacro.h b/llvm/include/llvm/MC/MCAsmMacro.h
index e2989c09017a..bdf55345adeb 100644
--- a/llvm/include/llvm/MC/MCAsmMacro.h
+++ b/llvm/include/llvm/MC/MCAsmMacro.h
@@ -145,6 +145,7 @@ struct MCAsmMacro {
   MCAsmMacroParameters Parameters;
   std::vector<std::string> Locals;
   bool IsFunction = false;
+  unsigned Count = 0;
 
 public:
   MCAsmMacro(StringRef N, StringRef B, MCAsmMacroParameters P)
diff --git a/llvm/include/llvm/MC/MCContext.h b/llvm/include/llvm/MC/MCContext.h
index 3de41b6a6ea7..b0ac432a065b 100644
--- a/llvm/include/llvm/MC/MCContext.h
+++ b/llvm/include/llvm/MC/MCContext.h
@@ -856,7 +856,7 @@ public:
   void reportError(SMLoc L, const Twine &Msg);
   void reportWarning(SMLoc L, const Twine &Msg);
 
-  const MCAsmMacro *lookupMacro(StringRef Name) {
+  MCAsmMacro *lookupMacro(StringRef Name) {
     StringMap<MCAsmMacro>::iterator I = MacroMap.find(Name);
     return (I == MacroMap.end()) ? nullptr : &I->getValue();
   }
diff --git a/llvm/lib/MC/MCParser/AsmParser.cpp b/llvm/lib/MC/MCParser/AsmParser.cpp
index 33287c6529ca..56da5e2f9a40 100644
--- a/llvm/lib/MC/MCParser/AsmParser.cpp
+++ b/llvm/lib/MC/MCParser/AsmParser.cpp
@@ -295,7 +295,7 @@ private:
 
   void checkForBadMacro(SMLoc DirectiveLoc, StringRef Name, StringRef Body,
                         ArrayRef<MCAsmMacroParameter> Parameters);
-  bool expandMacro(raw_svector_ostream &OS, const MCAsmMacro &Macro,
+  bool expandMacro(raw_svector_ostream &OS, MCAsmMacro &Macro,
                    ArrayRef<MCAsmMacroParameter> Parameters,
                    ArrayRef<MCAsmMacroArgument> A, bool EnableAtPseudoVariable);
 
@@ -312,7 +312,7 @@ private:
   ///
   /// \param M The macro.
   /// \param NameLoc Instantiation location.
-  bool handleMacroEntry(const MCAsmMacro *M, SMLoc NameLoc);
+  bool handleMacroEntry(MCAsmMacro *M, SMLoc NameLoc);
 
   /// Handle exit from macro instantiation.
   void handleMacroExit();
@@ -1980,9 +1980,8 @@ bool AsmParser::parseStatement(ParseStatementInfo &Info,
 
   // If macros are enabled, check to see if this is a macro instantiation.
   if (areMacrosEnabled())
-    if (const MCAsmMacro *M = getContext().lookupMacro(IDVal)) {
+    if (MCAsmMacro *M = getContext().lookupMacro(IDVal))
       return handleMacroEntry(M, IDLoc);
-    }
 
   // Otherwise, we have a normal instruction or directive.
 
@@ -2495,7 +2494,7 @@ static bool isIdentifierChar(char c) {
          c == '.';
 }
 
-bool AsmParser::expandMacro(raw_svector_ostream &OS, const MCAsmMacro &Macro,
+bool AsmParser::expandMacro(raw_svector_ostream &OS, MCAsmMacro &Macro,
                             ArrayRef<MCAsmMacroParameter> Parameters,
                             ArrayRef<MCAsmMacroArgument> A,
                             bool EnableAtPseudoVariable) {
@@ -2560,14 +2559,18 @@ bool AsmParser::expandMacro(raw_svector_ostream &OS, const MCAsmMacro &Macro,
       }
       Pos += 2;
     } else {
+      // Check for \@ and \+ pseudo variables.
       unsigned I = Pos + 1;
-
-      // Check for the \@ pseudo-variable.
-      if (EnableAtPseudoVariable && Body[I] == '@' && I + 1 != End)
-        ++I;
-      else
-        while (isIdentifierChar(Body[I]) && I + 1 != End)
+      if (I + 1 != End) {
+        if (EnableAtPseudoVariable && Body[I] == '@') {
+          ++I;
+        } else if (Body[I] == '+') {
           ++I;
+        } else {
+          while (isIdentifierChar(Body[I]) && I + 1 != End)
+            ++I;
+        }
+      }
 
       const char *Begin = Body.data() + Pos + 1;
       StringRef Argument(Begin, I - (Pos + 1));
@@ -2576,6 +2579,9 @@ bool AsmParser::expandMacro(raw_svector_ostream &OS, const MCAsmMacro &Macro,
       if (Argument == "@") {
         OS << NumOfMacroInstantiations;
         Pos += 2;
+      } else if (Argument == "+") {
+        OS << Macro.Count++;
+        Pos += 2;
       } else {
         for (; Index < NParameters; ++Index)
           if (Parameters[Index].Name == Argument)
@@ -2860,7 +2866,7 @@ bool AsmParser::parseMacroArguments(const MCAsmMacro *M,
   return TokError("too many positional arguments");
 }
 
-bool AsmParser::handleMacroEntry(const MCAsmMacro *M, SMLoc NameLoc) {
+bool AsmParser::handleMacroEntry(MCAsmMacro *M, SMLoc NameLoc) {
   // Arbitrarily limit macro nesting depth (default matches 'as'). We can
   // eliminate this, although we should protect against infinite loops.
   unsigned MaxNestingDepth = AsmMacroMaxNestingDepth;
diff --git a/llvm/test/MC/AsmParser/macro-at-pseudo-variable.s b/llvm/test/MC/AsmParser/macro-at-pseudo-variable.s
index b0966e137170..a083b17aa54f 100644
--- a/llvm/test/MC/AsmParser/macro-at-pseudo-variable.s
+++ b/llvm/test/MC/AsmParser/macro-at-pseudo-variable.s
@@ -1,35 +1,44 @@
-# RUN: llvm-mc -triple i386-unknown-unknown %s | FileCheck %s
+# RUN: rm -rf %t && split-file %s %t && cd %t
+# RUN: llvm-mc -triple i386 a.s | FileCheck %s
+# RUN: llvm-mc -triple i386 b.s | FileCheck %s --check-prefix=CHECK2
 
+#--- a.s
 .macro A
   add  $1\@, %eax
+  add  $2\+, %eax
 .endm
 
 .macro B
   sub  $1\@, %eax
+  sub  $2\+, %eax
 .endm
 
   A
-# CHECK: addl  $10, %eax
+# CHECK:      addl  $10, %eax
+# CHECK-NEXT: addl  $20, %eax
   A
-# CHECK: addl  $11, %eax
+# CHECK:      addl  $11, %eax
+# CHECK-NEXT: addl  $21, %eax
   B
-# CHECK: subl  $12, %eax
+# CHECK:      subl  $12, %eax
+# CHECK-NEXT: subl  $20, %eax
   B
-# CHECK: subl  $13, %eax
+# CHECK:      subl  $13, %eax
+# CHECK-NEXT: subl  $21, %eax
 
 # The following uses of \@ are undocumented, but valid:
 .irpc foo,234
   add  $\foo\@, %eax
 .endr
-# CHECK: addl  $24, %eax
-# CHECK: addl  $34, %eax
-# CHECK: addl  $44, %eax
+# CHECK:      addl  $24, %eax
+# CHECK-NEXT: addl  $34, %eax
+# CHECK-NEXT: addl  $44, %eax
 
 .irp reg,%eax,%ebx
   sub  $2\@, \reg
 .endr
-# CHECK: subl  $24, %eax
-# CHECK: subl  $24, %ebx
+# CHECK:      subl  $24, %eax
+# CHECK-NEXT: subl  $24, %ebx
 
 # Test that .irp(c) and .rep(t) do not increase \@.
 # Only the use of A should increase \@, so we can test that it increases by 1
@@ -62,3 +71,25 @@
 
   A
 # CHECK: addl  $17, %eax
+
+#--- b.s
+.rept 2
+  .print "r\+"
+.endr
+.irpc foo,12
+  .print "\+i"
+.endr
+# CHECK2:      r0
+# CHECK2-NEXT: r1
+# CHECK2-NEXT: 0i
+# CHECK2-NEXT: 1i
+
+.rept 2
+  .rept 2
+    .print "n\+"
+  .endr
+.endr
+# CHECK2:      n0
+# CHECK2-NEXT: n0
+# CHECK2-NEXT: n1
+# CHECK2-NEXT: n1