[llvm] [TableGen] Detect invalid -D arguments and fail. (PR #102813)

[Rahul Joshi via llvm-commits]

https://github.com/jurahul updated https://github.com/llvm/llvm-project/pull/102813

>From ccfe9d3159d43704f8d58962f2eb15abf6a70aea Mon Sep 17 00:00:00 2001
From: Rahul Joshi <rjoshi at nvidia.com>
Date: Sun, 11 Aug 2024 06:05:56 -0700
Subject: [PATCH] [TableGen] Detect invalid -D arguments and fail.

- Detect invalid macro names specified on command line and fail if one
  found.
- Specifically, -DXYZ=1 for example, will fail instead is being silently
  accepted.
---
 llvm/lib/TableGen/TGLexer.cpp                 | 46 +++++++++++++------
 .../invalid-macro-name-command-line.td        |  9 ++++
 2 files changed, 42 insertions(+), 13 deletions(-)
 create mode 100644 llvm/test/TableGen/invalid-macro-name-command-line.td

diff --git a/llvm/lib/TableGen/TGLexer.cpp b/llvm/lib/TableGen/TGLexer.cpp
index 99d866a23a68ba..7e7ec77e9ecd18 100644
--- a/llvm/lib/TableGen/TGLexer.cpp
+++ b/llvm/lib/TableGen/TGLexer.cpp
@@ -44,6 +44,26 @@ struct {
 };
 } // end anonymous namespace
 
+// Returns a pointer past the end of a valid macro name beginning at `Str`.
+// Valid macro names match the regular expression [a-zA-Z_][0-9a-zA-Z_]*.
+static const char *lexMacroName(const char *Str) {
+  assert(Str);
+
+  // Macro names start with [a-zA-Z_].
+  char First = *Str;
+  if (First != '_' && !isalpha(First))
+    return Str;
+
+  // Eat the first character of the name.
+  ++Str;
+
+  // Match the rest of the identifier regex: [0-9a-zA-Z_]*
+  while (isalpha(*Str) || isdigit(*Str) || *Str == '_')
+    ++Str;
+
+  return Str;
+}
+
 TGLexer::TGLexer(SourceMgr &SM, ArrayRef<std::string> Macros) : SrcMgr(SM) {
   CurBuffer = SrcMgr.getMainFileID();
   CurBuf = SrcMgr.getMemoryBuffer(CurBuffer)->getBuffer();
@@ -54,9 +74,17 @@ TGLexer::TGLexer(SourceMgr &SM, ArrayRef<std::string> Macros) : SrcMgr(SM) {
   PrepIncludeStack.push_back(
       std::make_unique<std::vector<PreprocessorControlDesc>>());
 
-  // Put all macros defined in the command line into the DefinedMacros set.
-  for (const std::string &MacroName : Macros)
+  // Add all macros defined on the command line to the DefinedMacros set.
+  // Check invalid macro names and print fatal error if we find one.
+  for (StringRef MacroName : Macros) {
+    const char *Begin = MacroName.data();
+    const char *End = lexMacroName(Begin);
+    if (End != MacroName.end())
+      PrintFatalError(Twine("Invalid macro name `") + MacroName +
+                      Twine("` specified on command line."));
+
     DefinedMacros.insert(MacroName);
+  }
 }
 
 SMLoc TGLexer::getLoc() const {
@@ -710,9 +738,8 @@ bool TGLexer::prepEatPreprocessorDirective(tgtok::TokKind Kind) {
   return false;
 }
 
-tgtok::TokKind TGLexer::lexPreprocessor(
-    tgtok::TokKind Kind, bool ReturnNextLiveToken) {
-
+tgtok::TokKind TGLexer::lexPreprocessor(tgtok::TokKind Kind,
+                                        bool ReturnNextLiveToken) {
   // We must be looking at a preprocessing directive.  Eat it!
   if (!prepEatPreprocessorDirective(Kind))
     PrintFatalError("lexPreprocessor() called for unknown "
@@ -912,14 +939,7 @@ StringRef TGLexer::prepLexMacroName() {
     ++CurPtr;
 
   TokStart = CurPtr;
-  // Macro names start with [a-zA-Z_].
-  if (*CurPtr != '_' && !isalpha(*CurPtr))
-    return "";
-
-  // Match the rest of the identifier regex: [0-9a-zA-Z_]*
-  while (isalpha(*CurPtr) || isdigit(*CurPtr) || *CurPtr == '_')
-    ++CurPtr;
-
+  CurPtr = lexMacroName(CurPtr);
   return StringRef(TokStart, CurPtr - TokStart);
 }
 
diff --git a/llvm/test/TableGen/invalid-macro-name-command-line.td b/llvm/test/TableGen/invalid-macro-name-command-line.td
new file mode 100644
index 00000000000000..eafa96b0236889
--- /dev/null
+++ b/llvm/test/TableGen/invalid-macro-name-command-line.td
@@ -0,0 +1,9 @@
+// RUN: not llvm-tblgen %s -DMACRO=1 2>&1 | FileCheck %s --check-prefix=CHECK-TEST-1
+// RUN: not llvm-tblgen %s -D0MAC 2>&1 | FileCheck %s --check-prefix=CHECK-TEST-2
+// RUN: not llvm-tblgen %s -D_MAC# 2>&1 | FileCheck %s --check-prefix=CHECK-TEST-3
+// RUN: not llvm-tblgen %s -D 2>&1 | FileCheck %s --check-prefix=CHECK-TEST-4
+
+// CHECK-TEST-1: error: Invalid macro name `MACRO=1` specified on command line.
+// CHECK-TEST-2: error: Invalid macro name `0MAC` specified on command line.
+// CHECK-TEST-3: error: Invalid macro name `_MAC#` specified on command line.
+// CHECK-TEST-4: for the -D option: requires a value!