[ms] [llvm-ml] Disambiguate size directives and variable declarations

[epastor]

Summary

MASM allows statements of the form:
<VAR> DWORD 5
to declare a variable with name <VAR>, while:
call dword ptr [<value>]
is a valid instruction. To disambiguate, we recognize size directives by the trailing "ptr" token.

As discussed in https://lists.llvm.org/pipermail/llvm-dev/2021-May/150774.html

[epastor]

Fix a copy/paste error

[epastor]

Updating use of peekTok to account for macros & included files

[epastor]

Add test (confirmed to fail before this patch)

[rnk]

lgtm

llvm/lib/MC/MCParser/MasmParser.cpp

public:
   bool parseIdentifier(StringRef &Res) override;
   void eatToEndOfStatement() override;
 
   bool checkForValidSection() override;
 
   /// }
 
 private:
+  const AsmToken peekTok(bool ShouldSkipSpace = true);
+
   bool parseStatement(ParseStatementInfo &Info,
                       MCAsmParserSemaCallback *SI);
   bool parseCurlyBlockScope(SmallVectorImpl<AsmRewrite>& AsmStrRewrites);
   bool parseCppHashLineFilenameComment(SMLoc L);
 
   bool expandMacro(raw_svector_ostream &OS, StringRef Body,
                    ArrayRef<MCAsmMacroParameter> Parameters,
                    ArrayRef<MCAsmMacroArgument> A,
if (it != Variables.end() && it->second.IsText) {
 
       // Jump to the macro instantiation and prime the lexer.
       CurBuffer = SrcMgr.AddNewSourceBuffer(std::move(Instantiation),
                                             getTok().getEndLoc());
       Lexer.setBuffer(SrcMgr.getMemoryBuffer(CurBuffer)->getBuffer(), nullptr,
                       /*EndStatementAtEOF=*/false);
       EndStatementAtEOFStack.push_back(false);
       tok = &Lexer.Lex();
-    } else if (M && M->IsFunction && Lexer.peekTok().is(AsmToken::LParen)) {
+    } else if (M && M->IsFunction && peekTok().is(AsmToken::LParen)) {
       // This is a macro function invocation; expand it in place.
       const AsmToken MacroTok = *tok;
       tok = &Lexer.Lex();
       if (handleMacroInvocation(M, MacroTok.getLoc())) {
         Lexer.UnLex(AsmToken(AsmToken::Error, MacroTok.getIdentifier()));
         tok = &Lexer.Lex();
       }
       continue;
     } else {
       break;
     }
   }
 
   // Parse comments here to be deferred until end of next statement.
   while (tok->is(AsmToken::Comment)) {
     if (MAI.preserveAsmComments())
       Out.addExplicitComment(Twine(tok->getString()));
     tok = &Lexer.Lex();
   }
 
   // Recognize and bypass line continuations.
   while (tok->is(AsmToken::BackSlash) &&
-         Lexer.peekTok().is(AsmToken::EndOfStatement)) {
+         peekTok().is(AsmToken::EndOfStatement)) {
     // Eat both the backslash and the end of statement.
     Lexer.Lex();
     tok = &Lexer.Lex();
   }
 
   if (tok->is(AsmToken::Eof)) {
     // If this is the end of an included file, pop the parent file off the
     // include stack.
     SMLoc ParentIncludeLoc = SrcMgr.getParentIncludeLoc(CurBuffer);
     if (ParentIncludeLoc != SMLoc()) {
       EndStatementAtEOFStack.pop_back();
       jumpToLoc(ParentIncludeLoc, 0, EndStatementAtEOFStack.back());
       return Lex();
     }
     EndStatementAtEOFStack.pop_back();
     assert(EndStatementAtEOFStack.empty());
   }
 
   return *tok;
 }
 
+const AsmToken MasmParser::peekTok(bool ShouldSkipSpace) {
+  AsmToken Tok;
+
+  MutableArrayRef<AsmToken> Buf(Tok);
+  size_t ReadCount = Lexer.peekTokens(Buf, ShouldSkipSpace);
+
+  if (ReadCount == 0) {
+    // If this is the end of an included file, pop the parent file off the
+    // include stack.
+    SMLoc ParentIncludeLoc = SrcMgr.getParentIncludeLoc(CurBuffer);
+    if (ParentIncludeLoc != SMLoc()) {
+      EndStatementAtEOFStack.pop_back();
+      jumpToLoc(ParentIncludeLoc, 0, EndStatementAtEOFStack.back());
+      return peekTok(ShouldSkipSpace);
+    }
+    EndStatementAtEOFStack.pop_back();
+    assert(EndStatementAtEOFStack.empty());
+  }
+
+  assert(ReadCount == 1);
+  return Tok;
+}
+
 bool MasmParser::enabledGenDwarfForAssembly() {
   // Check whether the user specified -g.
   if (!getContext().getGenDwarfForAssembly())
     return false;
   // If we haven't encountered any .file directives (which would imply that
   // the assembler source was produced with debug info already) then emit one
   // describing the assembler source file itself.
   if (getContext().getGenDwarfFileNumber() == 0) {
if (IDIt != Structs.end())
     return parseDirectiveStructValue(/*Structure=*/IDIt->getValue(), IDVal,
                                      IDLoc);
 
   // Non-conditional Microsoft directives sometimes follow their first argument.
   const AsmToken nextTok = getTok();
   const StringRef nextVal = nextTok.getString();
   const SMLoc nextLoc = nextTok.getLoc();
 
+  const AsmToken afterNextTok = peekTok();
+
   // There are several entities interested in parsing infix directives:
   //
   // 1. Asm parser extensions. For example, platform-specific parsers
   //    (like the ELF parser) register themselves as extensions.
   // 2. The generic directive parser implemented by this class. These are
   //    all the directives that behave in a target and platform independent
   //    manner, or at least have a default behavior that's shared between
   //    all targets and platforms.
bool MasmParser::parseStatement(ParseStatementInfo &Info,
   default:
     break;
   case DK_ASSIGN:
   case DK_EQU:
   case DK_TEXTEQU:
     Lex();
     return parseDirectiveEquate(nextVal, IDVal, DirKind);
   case DK_BYTE:
+    if (afterNextTok.is(AsmToken::Identifier) &&
+        afterNextTok.getString().equals_lower("ptr")) {
+      // Size directive; part of an instruction.
+      break;
+    }
+    LLVM_FALLTHROUGH;
   case DK_SBYTE:
   case DK_DB:
     Lex();
     return parseDirectiveNamedValue(nextVal, 1, IDVal, IDLoc);
   case DK_WORD:
+    if (afterNextTok.is(AsmToken::Identifier) &&
+        afterNextTok.getString().equals_lower("ptr")) {
+      // Size directive; part of an instruction.
+      break;
+    }
+    LLVM_FALLTHROUGH;
   case DK_SWORD:
   case DK_DW:
     Lex();
     return parseDirectiveNamedValue(nextVal, 2, IDVal, IDLoc);
   case DK_DWORD:
+    if (afterNextTok.is(AsmToken::Identifier) &&
+        afterNextTok.getString().equals_lower("ptr")) {
+      // Size directive; part of an instruction.
+      break;
+    }
+    LLVM_FALLTHROUGH;
   case DK_SDWORD:
   case DK_DD:
     Lex();
     return parseDirectiveNamedValue(nextVal, 4, IDVal, IDLoc);
   case DK_FWORD:
+    if (afterNextTok.is(AsmToken::Identifier) &&
+        afterNextTok.getString().equals_lower("ptr")) {
+      // Size directive; part of an instruction.
+      break;
+    }
+    LLVM_FALLTHROUGH;
   case DK_DF:
     Lex();
     return parseDirectiveNamedValue(nextVal, 6, IDVal, IDLoc);
   case DK_QWORD:
+    if (afterNextTok.is(AsmToken::Identifier) &&
+        afterNextTok.getString().equals_lower("ptr")) {
+      // Size directive; part of an instruction.
+      break;
+    }
+    LLVM_FALLTHROUGH;
   case DK_SQWORD:
   case DK_DQ:
     Lex();
     return parseDirectiveNamedValue(nextVal, 8, IDVal, IDLoc);
   case DK_REAL4:
     Lex();
     return parseDirectiveNamedRealValue(nextVal, APFloat::IEEEsingle(), 4,
                                         IDVal, IDLoc);
bool MasmParser::parseMacroArguments(const MCAsmMacro *M,
   // Parse two kinds of macro invocations:
   // - macros defined without any parameters accept an arbitrary number of them
   // - macros defined with parameters accept at most that many of them
   for (unsigned Parameter = 0; !NParameters || Parameter < NParameters;
        ++Parameter) {
     SMLoc IDLoc = Lexer.getLoc();
     MCAsmMacroParameter FA;
 
-    if (Lexer.is(AsmToken::Identifier) && Lexer.peekTok().is(AsmToken::Equal)) {
+    if (Lexer.is(AsmToken::Identifier) && peekTok().is(AsmToken::Equal)) {
       if (parseIdentifier(FA.Name))
         return Error(IDLoc, "invalid argument identifier for formal argument");
 
       if (Lexer.isNot(AsmToken::Equal))
         return TokError("expected '=' after formal parameter identifier");
 
       Lex();
 
bool MasmParser::parseIdentifier(StringRef &Res) {
   // be separate tokens. At this level, we have already lexed so we cannot
   // (currently) handle this as a context dependent token, instead we detect
   // adjacent tokens and return the combined identifier.
   if (Lexer.is(AsmToken::Dollar) || Lexer.is(AsmToken::At)) {
     SMLoc PrefixLoc = getLexer().getLoc();
 
     // Consume the prefix character, and check for a following identifier.
 
-    AsmToken Buf[1];
-    Lexer.peekTokens(Buf, false);
+    AsmToken nextTok = peekTok(false);
 
-    if (Buf[0].isNot(AsmToken::Identifier))
+    if (nextTok.isNot(AsmToken::Identifier))
       return true;
 
     // We have a '$' or '@' followed by an identifier, make sure they are adjacent.
-    if (PrefixLoc.getPointer() + 1 != Buf[0].getLoc().getPointer())
+    if (PrefixLoc.getPointer() + 1 != nextTok.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);
     Lex(); // Parser Lex to maintain invariants.

 }
 
 bool MasmParser::parseRealInstList(const fltSemantics &Semantics,
                                    SmallVectorImpl<APInt> &ValuesAsInt,
                                    const AsmToken::TokenKind EndToken) {
   while (getTok().isNot(EndToken) ||
          (EndToken == AsmToken::Greater &&
           getTok().isNot(AsmToken::GreaterGreater))) {
-    const AsmToken NextTok = Lexer.peekTok();
+    const AsmToken NextTok = peekTok();
     if (NextTok.is(AsmToken::Identifier) &&
         NextTok.getString().equals_lower("dup")) {
       const MCExpr *Value;
       if (parseExpression(Value) || parseToken(AsmToken::Identifier))
         return true;
       const MCConstantExpr *MCE = dyn_cast<MCConstantExpr>(Value);
       if (!MCE)
         return Error(Value->getLoc(),

 }
 
 bool MasmParser::parseStructInstList(
     const StructInfo &Structure, std::vector<StructInitializer> &Initializers,
     const AsmToken::TokenKind EndToken) {
   while (getTok().isNot(EndToken) ||
          (EndToken == AsmToken::Greater &&
           getTok().isNot(AsmToken::GreaterGreater))) {
-    const AsmToken NextTok = Lexer.peekTok();
+    const AsmToken NextTok = peekTok();
     if (NextTok.is(AsmToken::Identifier) &&
         NextTok.getString().equals_lower("dup")) {
       const MCExpr *Value;
       if (parseExpression(Value) || parseToken(AsmToken::Identifier))
         return true;
       const MCConstantExpr *MCE = dyn_cast<MCConstantExpr>(Value);
       if (!MCE)
         return Error(Value->getLoc(),
if (getLexer().is(AsmToken::Identifier)) {
             return TokError("unexpected token in '" + EndToken.getIdentifier() +
                             "' directive");
           break;
         } else {
           // Otherwise we just found the end of an inner macro.
           --MacroDepth;
         }
       } else if (getTok().getIdentifier().equals_lower("exitm")) {
-        if (MacroDepth == 0 &&
-            getLexer().peekTok().isNot(AsmToken::EndOfStatement)) {
+        if (MacroDepth == 0 && peekTok().isNot(AsmToken::EndOfStatement)) {
           IsMacroFunction = true;
         }
       } else if (isMacroLikeDirective()) {
         // We allow nested macros. Those aren't instantiated until the
         // outermost macro is expanded so just ignore them for now.
         ++MacroDepth;
       }
     }
bool IsMacroLike = StringSwitch<bool>(getTok().getIdentifier())
                            .CasesLower("repeat", "rept", true)
                            .CaseLower("while", true)
                            .CasesLower("for", "irp", true)
                            .CasesLower("forc", "irpc", true)
                            .Default(false);
     if (IsMacroLike)
       return true;
   }
-  if (getLexer().peekTok().is(AsmToken::Identifier) &&
-      getLexer().peekTok().getIdentifier().equals_lower("macro"))
+  if (peekTok().is(AsmToken::Identifier) &&
+      peekTok().getIdentifier().equals_lower("macro"))
     return true;
 
   return false;
 }
 
 MCAsmMacro *MasmParser::parseMacroLikeBody(SMLoc DirectiveLoc) {
   AsmToken EndToken, StartToken = getTok();