From b3c8d214953173e0f4aa798dbf855f920a6690fc Mon Sep 17 00:00:00 2001
From: Anirudh Prasad <anirudh_prasad@hotmail.com>
Date: Thu, 29 Apr 2021 11:27:56 -0400
Subject: [PATCH] [AsmParser][SystemZ][z/OS] Reject "Dot" as current PC on z/OS

- Currently, the "." (Dot) character, when not identifying an Identifier or a Constant, refers to the current PC (Program Counter)
- However, in z/OS, for the HLASM dialect, it strictly accepts only the "*" as the current PC (Support for this will be put up in a follow-up patch)
- The changes in this patch allow individual platforms to choose whether they would like to use the "." (Dot) character as a marker for the current PC or not.
- It is achieved by introducing a new field in MCAsmInfo.h called `DotIsPC` (similar to `DollarIsPC`)

Reviewed By: abhina.sreeskantharajan

Differential Revision: https://reviews.llvm.org/D100975
---
 include/llvm/MC/MCAsmInfo.h                   |  5 ++++
 lib/MC/MCParser/AsmParser.cpp                 |  3 ++
 .../SystemZ/MCTargetDesc/SystemZMCAsmInfo.cpp |  1 +
 unittests/MC/SystemZ/SystemZAsmLexerTest.cpp  | 30 +++++++++++++++++++
 4 files changed, 39 insertions(+)

diff --git a/include/llvm/MC/MCAsmInfo.h b/include/llvm/MC/MCAsmInfo.h
index f72cc709389..a0e7ba6e2a5 100644
--- a/include/llvm/MC/MCAsmInfo.h
+++ b/include/llvm/MC/MCAsmInfo.h
@@ -118,6 +118,10 @@ protected:
   /// the current PC.  Defaults to false.
   bool DollarIsPC = false;
 
+  /// Allow '.' token, when not referencing an identifier or constant, to refer
+  /// to the current PC. Defaults to true.
+  bool DotIsPC = true;
+
   /// This string, if specified, is used to separate instructions from each
   /// other when on the same line.  Defaults to ';'
   const char *SeparatorString;
@@ -593,6 +597,7 @@ public:
 
   unsigned getMinInstAlignment() const { return MinInstAlignment; }
   bool getDollarIsPC() const { return DollarIsPC; }
+  bool getDotIsPC() const { return DotIsPC; }
   const char *getSeparatorString() const { return SeparatorString; }
 
   /// This indicates the column (zero-based) at which asm comments should be
diff --git a/lib/MC/MCParser/AsmParser.cpp b/lib/MC/MCParser/AsmParser.cpp
index f2fb5714a6b..5d037931bb0 100644
--- a/lib/MC/MCParser/AsmParser.cpp
+++ b/lib/MC/MCParser/AsmParser.cpp
@@ -1245,6 +1245,9 @@ bool AsmParser::parsePrimaryExpr(const MCExpr *&Res, SMLoc &EndLoc,
     return false;
   }
   case AsmToken::Dot: {
+    if (!MAI.getDotIsPC())
+      return TokError("cannot use . as current PC");
+
     // This is a '.' reference, which references the current PC.  Emit a
     // temporary label to the streamer and refer to it.
     MCSymbol *Sym = Ctx.createTempSymbol();
diff --git a/lib/Target/SystemZ/MCTargetDesc/SystemZMCAsmInfo.cpp b/lib/Target/SystemZ/MCTargetDesc/SystemZMCAsmInfo.cpp
index 4bb9e15d4c4..6422ca93b6d 100644
--- a/lib/Target/SystemZ/MCTargetDesc/SystemZMCAsmInfo.cpp
+++ b/lib/Target/SystemZ/MCTargetDesc/SystemZMCAsmInfo.cpp
@@ -27,6 +27,7 @@ SystemZMCAsmInfo::SystemZMCAsmInfo(const Triple &TT) {
   AllowAtAtStartOfIdentifier = (AssemblerDialect == AD_HLASM);
   AllowDollarAtStartOfIdentifier = (AssemblerDialect == AD_HLASM);
   AllowHashAtStartOfIdentifier = (AssemblerDialect == AD_HLASM);
+  DotIsPC = (AssemblerDialect == AD_ATT);
 
   ZeroDirective = "\t.space\t";
   Data64bitsDirective = "\t.quad\t";
diff --git a/unittests/MC/SystemZ/SystemZAsmLexerTest.cpp b/unittests/MC/SystemZ/SystemZAsmLexerTest.cpp
index 81180966708..6c868ce8616 100644
--- a/unittests/MC/SystemZ/SystemZAsmLexerTest.cpp
+++ b/unittests/MC/SystemZ/SystemZAsmLexerTest.cpp
@@ -47,6 +47,7 @@ public:
   void setAllowHashAtStartOfIdentifier(bool Value) {
     AllowHashAtStartOfIdentifier = Value;
   }
+  void setAllowDotIsPC(bool Value) { DotIsPC = Value; }
 };
 
 // Setup a testing class that the GTest framework can call.
@@ -55,6 +56,7 @@ protected:
   static void SetUpTestCase() {
     LLVMInitializeSystemZTargetInfo();
     LLVMInitializeSystemZTargetMC();
+    LLVMInitializeSystemZAsmParser();
   }
 
   std::unique_ptr<MCRegisterInfo> MRI;
@@ -63,6 +65,8 @@ protected:
   std::unique_ptr<MCStreamer> Str;
   std::unique_ptr<MCAsmParser> Parser;
   std::unique_ptr<MCContext> Ctx;
+  std::unique_ptr<MCSubtargetInfo> STI;
+  std::unique_ptr<MCTargetAsmParser> TargetAsmParser;
 
   SourceMgr SrcMgr;
   std::string TripleName;
@@ -85,6 +89,12 @@ protected:
     MRI.reset(TheTarget->createMCRegInfo(TripleName));
     EXPECT_NE(MRI, nullptr);
 
+    MII.reset(TheTarget->createMCInstrInfo());
+    EXPECT_NE(MII, nullptr);
+
+    STI.reset(TheTarget->createMCSubtargetInfo(TripleName, "z10", ""));
+    EXPECT_NE(STI, nullptr);
+
     std::unique_ptr<MCAsmInfo> MAI;
     MAI.reset(TheTarget->createMCAsmInfo(*MRI, TripleName, MCOptions));
     EXPECT_NE(MAI, nullptr);
@@ -109,6 +119,10 @@ protected:
     Str.reset(TheTarget->createNullStreamer(*Ctx));
 
     Parser.reset(createMCAsmParser(SrcMgr, *Ctx, *Str, *MUPMAI));
+
+    TargetAsmParser.reset(
+        TheTarget->createMCAsmParser(*STI, *Parser, *MII, MCOptions));
+    Parser->setTargetParser(*TargetAsmParser);
   }
 
   void lexAndCheckTokens(StringRef AsmStr,
@@ -655,4 +669,20 @@ TEST_F(SystemZAsmLexerTest, CheckAcceptHashAtStartOfIdentifier4) {
       {AsmToken::Identifier, AsmToken::EndOfStatement, AsmToken::Eof});
   lexAndCheckTokens(AsmStr, ExpectedTokens);
 }
+
+TEST_F(SystemZAsmLexerTest, CheckRejectDotAsCurrentPC) {
+  StringRef AsmStr = ".-4";
+
+  // Setup.
+  MUPMAI->setAllowDotIsPC(false);
+  setupCallToAsmParser(AsmStr);
+
+  // Lex initially to get the string.
+  Parser->getLexer().Lex();
+
+  const MCExpr *Expr;
+  bool ParsePrimaryExpr = Parser->parseExpression(Expr);
+  EXPECT_EQ(ParsePrimaryExpr, true);
+  EXPECT_EQ(Parser->hasPendingError(), true);
+}
 } // end anonymous namespace