From d046332679ada51644bff7f7a06352ae97d7b97e Mon Sep 17 00:00:00 2001
From: Nirav Dave <niravd@google.com>
Date: Wed, 12 Oct 2016 13:58:07 +0000
Subject: [PATCH] [MC] Fix Error Location for ParseIdentifier

Prevent partial parsing of '$' or '@' of invalid identifiers and fixup
workaround points. NFC Intended.

llvm-svn: 284017
---
 llvm/lib/MC/MCParser/AsmParser.cpp              | 15 +++++++++++----
 llvm/test/MC/AsmParser/at-pseudo-variable-bad.s |  2 +-
 llvm/test/MC/Mips/cpsetup-bad.s                 |  2 +-
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/MC/MCParser/AsmParser.cpp b/llvm/lib/MC/MCParser/AsmParser.cpp
index f4de634..e817a8f 100644
--- a/llvm/lib/MC/MCParser/AsmParser.cpp
+++ b/llvm/lib/MC/MCParser/AsmParser.cpp
@@ -924,8 +924,10 @@ bool AsmParser::parsePrimaryExpr(const MCExpr *&Res, SMLoc &EndLoc) {
   case AsmToken::Identifier: {
     StringRef Identifier;
     if (parseIdentifier(Identifier)) {
-      if (FirstTokenKind == AsmToken::Dollar) {
+      // We may have failed but $ may be a valid token.
+      if (getTok().is(AsmToken::Dollar)) {
         if (Lexer.getMAI().getDollarIsPC()) {
+          Lex();
           // This is a '$' reference, which references the current PC.  Emit a
           // temporary label to the streamer and refer to it.
           MCSymbol *Sym = Ctx.createTempSymbol();
@@ -2607,14 +2609,19 @@ bool AsmParser::parseIdentifier(StringRef &Res) {
     SMLoc PrefixLoc = getLexer().getLoc();
 
     // Consume the prefix character, and check for a following identifier.
-    Lexer.Lex(); // Lexer's Lex guarantees consecutive token.
-    if (Lexer.isNot(AsmToken::Identifier))
+
+    AsmToken Buf[1];
+    Lexer.peekTokens(Buf, false);
+
+    if (Buf[0].isNot(AsmToken::Identifier))
       return true;
 
     // We have a '$' or '@' followed by an identifier, make sure they are adjacent.
-    if (PrefixLoc.getPointer() + 1 != getTok().getLoc().getPointer())
+    if (PrefixLoc.getPointer() + 1 != Buf[0].getLoc().getPointer())
       return true;
 
+    // eat $ or @
+    Lexer.Lex(); // Lexer's Lex guarantees consecutive token.
     // Construct the joined identifier and consume the token.
     Res =
         StringRef(PrefixLoc.getPointer(), getTok().getIdentifier().size() + 1);
diff --git a/llvm/test/MC/AsmParser/at-pseudo-variable-bad.s b/llvm/test/MC/AsmParser/at-pseudo-variable-bad.s
index 885f18b..0859cab 100644
--- a/llvm/test/MC/AsmParser/at-pseudo-variable-bad.s
+++ b/llvm/test/MC/AsmParser/at-pseudo-variable-bad.s
@@ -6,7 +6,7 @@ add  $1\@, %eax
 .macro A @
   mov  %eax, %eax
 .endm
-# CHECK: :[[@LINE-3]]:11: error: expected identifier in '.macro' directive
+# CHECK: :[[@LINE-3]]:10: error: expected identifier in '.macro' directive
 
 .rept 2
   addi $8, $8, \@
diff --git a/llvm/test/MC/Mips/cpsetup-bad.s b/llvm/test/MC/Mips/cpsetup-bad.s
index ec6525a..4b21486 100644
--- a/llvm/test/MC/Mips/cpsetup-bad.s
+++ b/llvm/test/MC/Mips/cpsetup-bad.s
@@ -13,7 +13,7 @@ t1:
         .cpsetup $31, $32, __cerror
 # ASM: :[[@LINE-1]]:23: error: invalid register
         .cpsetup $25, $2, $3
-# ASM: :[[@LINE-1]]:28: error: expected expression
+# ASM: :[[@LINE-1]]:27: error: expected expression
         .cpsetup $25, $2, 4
 # ASM: :[[@LINE-1]]:28: error: expected symbol
         .cpsetup $25, $2, 4+65
-- 
2.7.4