[MC] AsmLexer: add extensible identifier's character set support.

[vpykhtin]

Summary

Working on AMDGPU project I need to support assembler identifiers started with '&'. Looking into AsmLexer.cpp I found similar need to optionally support '@' inside identifiers and decided this is time to add generic support for identifier charset. I added configurable bitvector set for prefix and body identifier's characters.

[vpykhtin]

Kind reminder if someone can take a look at this.

[• rafael]

include/llvm/MC/MCParser/MCAsmLexer.h

238

Why do you need to make these virtual?

243

is..Contains is a strange name since it has two verbs.

lib/MC/MCParser/AsmLexer.cpp

33

Why?

[vpykhtin]

include/llvm/MC/MCParser/MCAsmLexer.h

238

Well not making it virtual would require bitvector sets to be part of this class. I'm not objecting though as it already done with SkipSpace and AllowAtInIdentifier.

243

What would be a better name here?

lib/MC/MCParser/AsmLexer.cpp

33

Well it based on my previuos experience on Windows where we had lexer using these routines eating up to 10% of scan time. Probably not so "generally" as I stated though. I'm not insisting on this particular change and can remove it.

[vpykhtin]

After a loooong time I would like to reanimate this review requiest.

Previously I incorrectly measured performance impact for this patch and obtained 30% performance gain - this result was incorrect. Current measurement on a large .s file shows no affect on parsing performance.

[vpykhtin]

ping

[vpykhtin]

ping.

[vpykhtin]

Last ping?

[grosbach]

include/llvm/MC/MCParser/MCAsmLexer.h

238

With the generalization, these can go away entirely, yes? Replace the callsites w/ the new API.

246

This should start with "is" not "Is" according the the coding guidelines.

251

Ditto.

258

This feels really weird. Wouldn't any callsites want to be using one of the other two? They'll know their context. I don't see any invocations of this method in the patch. Why is it needed at all?

lib/MC/MCParser/AsmLexer.cpp

576

Can you elaborate on this bit? Not sure I follow why this is so much more logic than previously.

lib/MC/MCParser/MCAsmLexer.cpp

37

Given the bimodal behaviour based on Value, this should probably just be two functions.

include/llvm/MC/MCParser/MCAsmLexer.h

 //===-- llvm/MC/MCAsmLexer.h - Abstract Asm Lexer Interface -----*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_MC_MCPARSER_MCASMLEXER_H
 #define LLVM_MC_MCPARSER_MCASMLEXER_H
 
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/DataTypes.h"
 #include "llvm/Support/SMLoc.h"
 #include <utility>
 
 namespace llvm {
class MCAsmLexer {
   SMLoc ErrLoc;
   std::string Err;
 
   MCAsmLexer(const MCAsmLexer &) = delete;
   void operator=(const MCAsmLexer &) = delete;
 protected: // Can only create subclasses.
   const char *TokStart;
   bool SkipSpace;
-  bool AllowAtInIdentifier;
   bool IsAtStartOfStatement;
+  BitVector IdPrefixCharSet;
+  BitVector IdBodyCharSet;
 
   MCAsmLexer();
 
   virtual AsmToken LexToken() = 0;
 
   void SetError(SMLoc errLoc, const std::string &err) {
     ErrLoc = errLoc;
     Err = err;
public:
   bool is(AsmToken::TokenKind K) const { return getTok().is(K); }
 
   /// Check if the current token has kind \p K.
   bool isNot(AsmToken::TokenKind K) const { return getTok().isNot(K); }
 
   /// Set whether spaces should be ignored by the lexer
   void setSkipSpace(bool val) { SkipSpace = val; }
 
-  bool getAllowAtInIdentifier() { return AllowAtInIdentifier; }
-  void setAllowAtInIdentifier(bool v) { AllowAtInIdentifier = v; }
+  bool getAllowAtInIdentifier() { return IsAllowedIDBodyChar('@'); }
+  void setAllowAtInIdentifier(bool v) { setIdentifierCharSet(v, "", "@"); }
+
+  /// allow/disallow an identifier to contain specified characters
+  void setIdentifierCharSet(bool Value,
+                            StringRef PfxCharSet,
+                            StringRef BodyCharSet);
+
+  /// test whether the specified character can start an identifier
+  bool IsAllowedIDPrefixChar(char C) const {
+    return IdPrefixCharSet.test((unsigned char)C);
+  }
+
+  /// test whether the specified character can follow identifier start char
+  bool IsAllowedIDBodyChar(char C) const {
+    return IdBodyCharSet.test((unsigned char)C);
+  }
+
+  /// test whether the specified character can be found in an identifier
+  bool isAllowedIDChar(char C) const {
+    return IsAllowedIDBodyChar(C) || IsAllowedIDPrefixChar(C);
+  }
 };
 
 } // End llvm namespace
 
 #endif

lib/MC/MCParser/AsmLexer.cpp

 #include <cstdio>
 #include <cstring>
 #include <string>
 #include <tuple>
 #include <utility>
 
 using namespace llvm;
 
 AsmLexer::AsmLexer(const MCAsmInfo &MAI)
     : MAI(MAI), CurPtr(nullptr), IsAtStartOfLine(true),
       IsAtStartOfStatement(true), IsParsingMSInlineAsm(false),
       IsPeeking(false) {
-  AllowAtInIdentifier = !StringRef(MAI.getCommentString()).startswith("@");
+
+  if (!StringRef(MAI.getCommentString()).startswith("@"))
+    setIdentifierCharSet(true, "", "@");
 }
 
 AsmLexer::~AsmLexer() {
 }
 
 void AsmLexer::setBuffer(StringRef Buf, const char *ptr) {
   CurBuf = Buf;
 
AsmToken AsmLexer::LexHexFloatLiteral(bool NoIntDigits) {
 
   if (CurPtr == ExpStart)
     return ReturnError(TokStart, "invalid hexadecimal floating-point constant: "
                                  "expected at least one exponent digit");
 
   return AsmToken(AsmToken::Real, StringRef(TokStart, CurPtr - TokStart));
 }
 
-/// LexIdentifier: [a-zA-Z_.][a-zA-Z0-9_$.@?]*
-static bool IsIdentifierChar(char c, bool AllowAt) {
-  return isalnum(c) || c == '_' || c == '$' || c == '.' ||
-         (c == '@' && AllowAt) || c == '?';
-}
 
 AsmToken AsmLexer::LexIdentifier() {
-  // Check for floating point literals.
-  if (CurPtr[-1] == '.' && isdigit(*CurPtr)) {
-    // Disambiguate a .1243foo identifier from a floating literal.
-    while (isdigit(*CurPtr))
-      ++CurPtr;
-    if (*CurPtr == 'e' || *CurPtr == 'E' ||
-        !IsIdentifierChar(*CurPtr, AllowAtInIdentifier))
-      return LexFloatLiteral();
-  }
-
-  while (IsIdentifierChar(*CurPtr, AllowAtInIdentifier))
+  while (IsAllowedIDBodyChar(*CurPtr))
     ++CurPtr;
-
-  // Handle . as a special case.
-  if (CurPtr == TokStart+1 && TokStart[0] == '.')
-    return AsmToken(AsmToken::Dot, StringRef(TokStart, 1));
-
   return AsmToken(AsmToken::Identifier, StringRef(TokStart, CurPtr - TokStart));
 }
 
 /// LexSlash: Slash: /
 ///           C-Style Comment: /* ... */
 AsmToken AsmLexer::LexSlash() {
   switch (*CurPtr) {
   case '*':

 bool AsmLexer::isAtStatementSeparator(const char *Ptr) {
   return strncmp(Ptr, MAI.getSeparatorString(),
                  strlen(MAI.getSeparatorString())) == 0;
 }
 
 AsmToken AsmLexer::LexToken() {
   TokStart = CurPtr;
   // This always consumes at least one character.
-  int CurChar = getNextChar();
+  const int CurChar = getNextChar();
 
   if (!IsPeeking && CurChar == '#' && IsAtStartOfStatement) {
     // If this starts with a '#', this may be a cpp
     // hash directive and otherwise a line comment.
     AsmToken TokenBuf[2];
     MutableArrayRef<AsmToken> Buf(TokenBuf, 2);
     size_t num = peekTokens(Buf, true);
     // There cannot be a space preceeding this
AsmToken AsmLexer::LexToken() {
   if (CurChar == EOF && !IsAtStartOfStatement) {
     IsAtStartOfLine = true;
     IsAtStartOfStatement = true;
     return AsmToken(AsmToken::EndOfStatement, StringRef(TokStart, 1));
   }
   IsAtStartOfLine = false;
   bool OldIsAtStartOfStatement = IsAtStartOfStatement;
   IsAtStartOfStatement = false;
+
+  if (CurChar == '.' && isdigit(*CurPtr)) {
+    if (!IsAllowedIDPrefixChar('.'))
+      return LexFloatLiteral();
+
+    const auto SavePos = CurPtr;
+    // Disambiguate a .1243foo identifier from a floating literal.
+    do { ++CurPtr; }
+    while (isdigit(*CurPtr));
+    if (*CurPtr == 'e' || *CurPtr == 'E' || !IsAllowedIDBodyChar(*CurPtr))
+      return LexFloatLiteral();
+    CurPtr = SavePos;
+  }
+
+  const bool IsIDPrefix = IsAllowedIDPrefixChar(CurChar);
+  if (IsIDPrefix && IsAllowedIDBodyChar(*CurPtr)) {
+    ++CurPtr;
+    return LexIdentifier();
+  }
+
   switch (CurChar) {
   default:
-    // Handle identifier: [a-zA-Z_.][a-zA-Z0-9_$.@]*
-    if (isalpha(CurChar) || CurChar == '_' || CurChar == '.')
-      return LexIdentifier();
+    if (IsIDPrefix)
+      return AsmToken(AsmToken::Identifier, StringRef(TokStart, 1));
 
     // Unknown character, emit an error.
     return ReturnError(TokStart, "invalid character in input");
   case EOF:
     IsAtStartOfLine = true;
     IsAtStartOfStatement = true;
     return AsmToken(AsmToken::Eof, StringRef(TokStart, 0));
   case 0:
   case ' ':
   case '\t':
     IsAtStartOfStatement = OldIsAtStartOfStatement;
     while (*CurPtr == ' ' || *CurPtr == '\t')
       CurPtr++;
     if (SkipSpace)
       return LexToken(); // Ignore whitespace.
     else
       return AsmToken(AsmToken::Space, StringRef(TokStart, CurPtr - TokStart));
   case '\n':
   case '\r':
     IsAtStartOfLine = true;
     IsAtStartOfStatement = true;
     return AsmToken(AsmToken::EndOfStatement, StringRef(TokStart, 1));
+  case '.': return AsmToken(AsmToken::Dot, StringRef(TokStart, 1));
   case ':': return AsmToken(AsmToken::Colon, StringRef(TokStart, 1));
   case '+': return AsmToken(AsmToken::Plus, StringRef(TokStart, 1));
   case '-': return AsmToken(AsmToken::Minus, StringRef(TokStart, 1));
   case '~': return AsmToken(AsmToken::Tilde, StringRef(TokStart, 1));
   case '(': return AsmToken(AsmToken::LParen, StringRef(TokStart, 1));
   case ')': return AsmToken(AsmToken::RParen, StringRef(TokStart, 1));
   case '[': return AsmToken(AsmToken::LBrac, StringRef(TokStart, 1));
   case ']': return AsmToken(AsmToken::RBrac, StringRef(TokStart, 1));

lib/MC/MCParser/MCAsmLexer.cpp

 #include "llvm/MC/MCParser/MCAsmLexer.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace llvm;
 
 MCAsmLexer::MCAsmLexer()
     : TokStart(nullptr), SkipSpace(true), IsAtStartOfStatement(true) {
   CurTok.emplace_back(AsmToken::Space, StringRef());
+  // Prefix char = [A-Za-z_.]
+  IdPrefixCharSet.resize(256);
+  IdPrefixCharSet.set((unsigned char)'a', (unsigned char)'z' + 1);
+  IdPrefixCharSet.set((unsigned char)'A', (unsigned char)'Z' + 1);
+  IdPrefixCharSet.set((unsigned char)'.');
+  IdPrefixCharSet.set((unsigned char)'_');
+
+  // Body char =  prefix + [0-9$?]
+  IdBodyCharSet = IdPrefixCharSet;
+  IdBodyCharSet.set((unsigned char)'0', (unsigned char)'9' + 1);
+  IdBodyCharSet.set((unsigned char)'$');
+  IdBodyCharSet.set((unsigned char)'?');
 }
 
 MCAsmLexer::~MCAsmLexer() {
 }
 
+void MCAsmLexer::setIdentifierCharSet(bool Value,
+                                      StringRef PfxCharSet,
+                                      StringRef BodyCharSet) {
+  if (Value) {
+    for (auto C : PfxCharSet)
+      IdPrefixCharSet.set((unsigned char)C);
+    for (auto C : BodyCharSet)
+      IdBodyCharSet.set((unsigned char)C);
+  } else {
+    for (auto C : PfxCharSet)
+      IdPrefixCharSet.reset((unsigned char)C);
+    for (auto C : BodyCharSet)
+      IdBodyCharSet.reset((unsigned char)C);
+  }
+}
+
 SMLoc MCAsmLexer::getLoc() const {
   return SMLoc::getFromPointer(TokStart);
 }
 
 SMLoc AsmToken::getLoc() const {
   return SMLoc::getFromPointer(Str.data());
 }
 
 SMLoc AsmToken::getEndLoc() const {
   return SMLoc::getFromPointer(Str.data() + Str.size());
 }
 
 SMRange AsmToken::getLocRange() const {
   return SMRange(getLoc(), getEndLoc());
 }

lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp

AMDGPUAsmParser(const MCSubtargetInfo &STI, MCAsmParser &_Parser,
         ForcedSDWA(false) {
     MCAsmParserExtension::Initialize(Parser);
 
     if (getSTI().getFeatureBits().none()) {
       // Set default features.
       copySTI().ToggleFeature("SOUTHERN_ISLANDS");
     }
 
+    getLexer().setIdentifierCharSet(true, "&", "");
+
     setAvailableFeatures(ComputeAvailableFeatures(getSTI().getFeatureBits()));
 
     {
       // TODO: make those pre-defined variables read-only.
       // Currently there is none suitable machinery in the core llvm-mc for this.
       // MCSymbol::isRedefinable is intended for another purpose, and
       // AsmParser::parseDirectiveSet() cannot be specialized for specific target.
       AMDGPU::IsaVersion Isa = AMDGPU::getIsaVersion(getSTI().getFeatureBits());

test/MC/AMDGPU/hsa.s

 // ELF: SHT_NOTE
 // ELF: 0000: 04000000 08000000 01000000 414D4400
 // ELF: 0010: 02000000 00000000 04000000 1B000000
 // ELF: 0020: 03000000 414D4400 04000700 07000000
 // ELF: 0030: 00000000 00000000 414D4400 414D4447
 // ELF: 0040: 50550000
 
 // ELF: Symbol {
-// ELF: Name: amd_kernel_code_t_minimal
+// ELF: Name: &amd_kernel_code_t_minimal
 // ELF: Type: AMDGPU_HSA_KERNEL (0xA)
 // ELF: Section: .text
 // ELF: }
 // ELF: Symbol {
-// ELF: Name: amd_kernel_code_t_test_all
+// ELF: Name: &amd_kernel_code_t_test_all
 // ELF: Type: AMDGPU_HSA_KERNEL (0xA)
 // ELF: Section: .text
 // ELF: }
 
 .text
 // ASM: .text
 
 .hsa_code_object_version 2,0
 // ASM: .hsa_code_object_version 2,0
 
 .hsa_code_object_isa 7,0,0,"AMD","AMDGPU"
 // ASM: .hsa_code_object_isa 7,0,0,"AMD","AMDGPU"
 
-.amdgpu_hsa_kernel amd_kernel_code_t_test_all
-.amdgpu_hsa_kernel amd_kernel_code_t_minimal
+.amdgpu_hsa_kernel &amd_kernel_code_t_test_all
+.amdgpu_hsa_kernel &amd_kernel_code_t_minimal
 
 
-amd_kernel_code_t_test_all:
+&amd_kernel_code_t_test_all:
 ; Test all amd_kernel_code_t members with non-default values.
 .amd_kernel_code_t
     kernel_code_version_major = 100
     kernel_code_version_minor = 100
     machine_kind = 0
     machine_version_major = 5
     machine_version_minor = 5
     machine_version_stepping = 5
.amd_kernel_code_t
     kernarg_segment_alignment = 5
     group_segment_alignment = 5
     private_segment_alignment = 5
     wavefront_size = 5
     call_convention = 1
     runtime_loader_kernel_symbol = 1
 .end_amd_kernel_code_t
 
-// ASM-LABEL: {{^}}amd_kernel_code_t_test_all:
+// ASM-LABEL: {{^\"\&}}amd_kernel_code_t_test_all{{\"}}:
 // ASM: .amd_kernel_code_t
 // ASM: amd_code_version_major = 100
 // ASM: amd_code_version_minor = 100
 // ASM: amd_machine_kind = 0
 // ASM: amd_machine_version_major = 5
 // ASM: amd_machine_version_minor = 5
 // ASM: amd_machine_version_stepping = 5
 // ASM: kernel_code_entry_byte_offset = 512

 // ASM: kernarg_segment_alignment = 5
 // ASM: group_segment_alignment = 5
 // ASM: private_segment_alignment = 5
 // ASM: wavefront_size = 5
 // ASM: call_convention = 1
 // ASM: runtime_loader_kernel_symbol = 1
 // ASM: .end_amd_kernel_code_t
 
-amd_kernel_code_t_minimal:
+&amd_kernel_code_t_minimal:
 .amd_kernel_code_t
   enable_sgpr_kernarg_segment_ptr = 1
   is_ptr64 = 1
   granulated_workitem_vgpr_count = 1
   granulated_wavefront_sgpr_count = 1
   user_sgpr_count = 2
   kernarg_segment_byte_size = 16
   wavefront_sgpr_count = 8
 //      wavefront_sgpr_count = 7
 ; wavefront_sgpr_count = 7
 // Make sure a blank line won't break anything:
 
 // Make sure a line with whitespace won't break anything:
 
   workitem_vgpr_count = 16
 .end_amd_kernel_code_t
 
-// ASM-LABEL: {{^}}amd_kernel_code_t_minimal:
+// ASM-LABEL: {{^\"&}}amd_kernel_code_t_minimal{{\"}}:
 // ASM: .amd_kernel_code_t
 // ASM: amd_code_version_major = 1
 // ASM: amd_code_version_minor = 0
 // ASM: amd_machine_kind = 1
 // ASM: amd_machine_version_major = 7
 // ASM: amd_machine_version_minor = 0
 // ASM: amd_machine_version_stepping = 0
 // ASM: kernel_code_entry_byte_offset = 256