[llvm] [RISCV] Add partial validation of Z extension name to RISCVISAInfo::parseNormalizedArchString (PR #90895)

[Craig Topper via llvm-commits]

https://github.com/topperc created https://github.com/llvm/llvm-project/pull/90895

If 'z' is given as the complete extension name or with a digit after it, it will crash in the extension map compare function. Check for these cases and give an error.

>From 556e3009bb6044d983c0ec9aa4e594cfc6676e26 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Thu, 2 May 2024 12:13:42 -0700
Subject: [PATCH] [RISCV] Add partial validation of Z extension name to
 RISCVISAInfo::parseNormalizedArchString

If 'z' is given as the complete extension name or with a digit after it,
it will crash in the extension map compare function. Check for
these cases and give an error.
---
 llvm/lib/Support/RISCVISAUtils.cpp               |  3 ++-
 llvm/lib/TargetParser/RISCVISAInfo.cpp           |  5 +++++
 llvm/unittests/TargetParser/RISCVISAInfoTest.cpp | 10 +++++++++-
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Support/RISCVISAUtils.cpp b/llvm/lib/Support/RISCVISAUtils.cpp
index ca7518f71907b5..46efe93695074f 100644
--- a/llvm/lib/Support/RISCVISAUtils.cpp
+++ b/llvm/lib/Support/RISCVISAUtils.cpp
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/RISCVISAUtils.h"
+#include "llvm/ADT/StringExtras.h"
 #include <cassert>
 
 using namespace llvm;
@@ -35,7 +36,7 @@ enum RankFlags {
 // Get the rank for single-letter extension, lower value meaning higher
 // priority.
 static unsigned singleLetterExtensionRank(char Ext) {
-  assert(Ext >= 'a' && Ext <= 'z');
+  assert(isLower(Ext));
   switch (Ext) {
   case 'i':
     return 0;
diff --git a/llvm/lib/TargetParser/RISCVISAInfo.cpp b/llvm/lib/TargetParser/RISCVISAInfo.cpp
index c1d50afee09b08..d244326537faff 100644
--- a/llvm/lib/TargetParser/RISCVISAInfo.cpp
+++ b/llvm/lib/TargetParser/RISCVISAInfo.cpp
@@ -485,6 +485,11 @@ RISCVISAInfo::parseNormalizedArchString(StringRef Arch) {
     if (MajorVersionStr.getAsInteger(10, MajorVersion))
       return createStringError(errc::invalid_argument,
                                "failed to parse major version number");
+
+    if (ExtName[0] == 'z' && (ExtName.size() == 1 || isDigit(ExtName[1])))
+      return createStringError(errc::invalid_argument,
+                               "'z' must be followed by a letter");
+
     ISAInfo->addExtension(ExtName, {MajorVersion, MinorVersion});
   }
   ISAInfo->updateImpliedLengths();
diff --git a/llvm/unittests/TargetParser/RISCVISAInfoTest.cpp b/llvm/unittests/TargetParser/RISCVISAInfoTest.cpp
index ec886bad4f67f7..eb8eab73686931 100644
--- a/llvm/unittests/TargetParser/RISCVISAInfoTest.cpp
+++ b/llvm/unittests/TargetParser/RISCVISAInfoTest.cpp
@@ -46,7 +46,7 @@ TEST(ParseNormalizedArchString, RejectsMalformedInputs) {
   }
 }
 
-TEST(ParseNormalizedArchString, OnlyVersion) {
+TEST(ParseNormalizedArchString, RejectsOnlyVersion) {
   for (StringRef Input : {"rv64i2p0_1p0", "rv32i2p0_1p0"}) {
     EXPECT_EQ(
         toString(RISCVISAInfo::parseNormalizedArchString(Input).takeError()),
@@ -54,6 +54,14 @@ TEST(ParseNormalizedArchString, OnlyVersion) {
   }
 }
 
+TEST(ParseNormalizedArchString, RejectsBadZ) {
+  for (StringRef Input : {"rv64i2p0_z1p0", "rv32i2p0_z2a1p0"}) {
+    EXPECT_EQ(
+        toString(RISCVISAInfo::parseNormalizedArchString(Input).takeError()),
+        "'z' must be followed by a letter");
+  }
+}
+
 TEST(ParseNormalizedArchString, AcceptsValidBaseISAsAndSetsXLen) {
   auto MaybeRV32I = RISCVISAInfo::parseNormalizedArchString("rv32i2p0");
   ASSERT_THAT_EXPECTED(MaybeRV32I, Succeeded());