[DataLayout] Refactor parsing of "ni" specification

llvm/lib/IR/DataLayout.cpp

@@ -272,6 +272,27 @@ 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,
+/// commonly referred to as <n> or <address space>.
+///
+/// \p Name should be the name of the component as specified by LangRef,
+/// it is used in diagnostic messages.
+static Error parseAddrSpace(StringRef Str, unsigned &AddrSpace,
+                            StringRef Name) {
+  if (Str.empty())
+    return createStringError(Name + " is required");
+
+  if (!to_integer(Str, AddrSpace, 10) || !isUInt<24>(AddrSpace))
+    return createStringError(Name + " 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 +334,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, "<address space>"))
+        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 +363,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);
 

llvm/unittests/IR/DataLayoutTest.cpp

@@ -134,6 +134,46 @@ 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> is required"));
+
+  for (StringRef Str : {"ni:x", "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"));