[DataLayout] Refactor parsing of "ni" specification (#104546)

Split off of #104545 to reduce patch size.
This introduces `parseAddrSpace` function, intended as a replacement for
`getAddrSpace`, which doesn't check for trailing characters after the
address space number. `getAddrSpace` will be removed after switching all
uses to `parseAddrSpace`.

Pull Request: #104546

Author: s-barannikov

Diff for: llvm/lib/IR/DataLayout.cpp

@@ -272,6 +272,22 @@ static Error reportError(const Twine &Message) {
   return createStringError(inconvertibleErrorCode(), Message);
 }
 
+static Error createSpecFormatError(Twine Format) {
+  return createStringError("malformed specification, must be of the form \"" +
+                           Format + "\"");
+}
+
+/// Attempts to parse an address space component of a specification.
+static Error parseAddrSpace(StringRef Str, unsigned &AddrSpace) {
+  if (Str.empty())
+    return createStringError("address space component cannot be empty");
+
+  if (!to_integer(Str, AddrSpace, 10) || !isUInt<24>(AddrSpace))
+    return createStringError("address space must be a 24-bit integer");
+
+  return Error::success();
+}
+
 /// Checked version of split, to ensure mandatory subparts.
 static Error split(StringRef Str, char Separator,
                    std::pair<StringRef, StringRef> &Split) {
@@ -313,6 +329,26 @@ static Error getAddrSpace(StringRef R, unsigned &AddrSpace) {
 }
 
 Error DataLayout::parseSpecification(StringRef Spec) {
+  // The "ni" specifier is the only two-character specifier. Handle it first.
+  if (Spec.starts_with("ni")) {
+    // ni:<address space>[:<address space>]...
+    StringRef Rest = Spec.drop_front(2);
+
+    // Drop the first ':', then split the rest of the string the usual way.
+    if (!Rest.consume_front(":"))
+      return createSpecFormatError("ni:<address space>[:<address space>]...");
+
+    for (StringRef Str : split(Rest, ':')) {
+      unsigned AddrSpace;
+      if (Error Err = parseAddrSpace(Str, AddrSpace))
+        return Err;
+      if (AddrSpace == 0)
+        return createStringError("address space 0 cannot be non-integral");
+      NonIntegralAddressSpaces.push_back(AddrSpace);
+    }
+    return Error::success();
+  }
+
   // Split at ':'.
   std::pair<StringRef, StringRef> Split;
   if (Error Err = ::split(Spec, ':', Split))
@@ -322,22 +358,6 @@ Error DataLayout::parseSpecification(StringRef Spec) {
   StringRef &Tok = Split.first;   // Current token.
   StringRef &Rest = Split.second; // The rest of the string.
 
-  if (Tok == "ni") {
-    do {
-      if (Error Err = ::split(Rest, ':', Split))
-        return Err;
-      Rest = Split.second;
-      unsigned AS;
-      if (Error Err = getInt(Split.first, AS))
-        return Err;
-      if (AS == 0)
-        return reportError("Address space 0 can never be non-integral");
-      NonIntegralAddressSpaces.push_back(AS);
-    } while (!Rest.empty());
-
-    return Error::success();
-  }
-
   char SpecifierChar = Tok.front();
   Tok = Tok.substr(1);
 

Diff for: llvm/unittests/IR/DataLayoutTest.cpp

@@ -134,6 +134,47 @@ TEST(DataLayout, LayoutStringFormat) {
         FailedWithMessage("empty specification is not allowed"));
 }
 
+TEST(DataLayout, ParseNonIntegralAddrSpace) {
+  for (StringRef Str : {"ni:1", "ni:16777215", "ni:1:16777215"})
+    EXPECT_THAT_EXPECTED(DataLayout::parse(Str), Succeeded());
+
+  for (StringRef Str : {"ni", "ni42", "nix"})
+    EXPECT_THAT_EXPECTED(
+        DataLayout::parse(Str),
+        FailedWithMessage("malformed specification, must be of the form "
+                          "\"ni:<address space>[:<address space>]...\""));
+
+  for (StringRef Str : {"ni:", "ni::42", "ni:42:"})
+    EXPECT_THAT_EXPECTED(
+        DataLayout::parse(Str),
+        FailedWithMessage("address space component cannot be empty"));
+
+  for (StringRef Str : {"ni:x", "ni:42:0x1", "ni:16777216", "ni:42:16777216"})
+    EXPECT_THAT_EXPECTED(
+        DataLayout::parse(Str),
+        FailedWithMessage("address space must be a 24-bit integer"));
+
+  for (StringRef Str : {"ni:0", "ni:42:0"})
+    EXPECT_THAT_EXPECTED(
+        DataLayout::parse(Str),
+        FailedWithMessage("address space 0 cannot be non-integral"));
+}
+
+TEST(DataLayout, IsNonIntegralAddressSpace) {
+  DataLayout Default;
+  EXPECT_THAT(Default.getNonIntegralAddressSpaces(), ::testing::SizeIs(0));
+  EXPECT_FALSE(Default.isNonIntegralAddressSpace(0));
+  EXPECT_FALSE(Default.isNonIntegralAddressSpace(1));
+
+  DataLayout Custom = cantFail(DataLayout::parse("ni:2:16777215"));
+  EXPECT_THAT(Custom.getNonIntegralAddressSpaces(),
+              ::testing::ElementsAreArray({2U, 16777215U}));
+  EXPECT_FALSE(Custom.isNonIntegralAddressSpace(0));
+  EXPECT_FALSE(Custom.isNonIntegralAddressSpace(1));
+  EXPECT_TRUE(Custom.isNonIntegralAddressSpace(2));
+  EXPECT_TRUE(Custom.isNonIntegralAddressSpace(16777215));
+}
+
 TEST(DataLayoutTest, CopyAssignmentInvalidatesStructLayout) {
   DataLayout DL1 = cantFail(DataLayout::parse("p:32:32"));
   DataLayout DL2 = cantFail(DataLayout::parse("p:64:64"));