MIR: Validate LLT types when parsing

[arsenm]

Summary

Fixes asserting on anything invalid looking

[paquette]

This seems like a pretty good idea... LGTM, aside from the one nit.

lib/CodeGen/MIRParser/MIParser.cpp

1344

Where are these values documented? Can we have a comment that points out where these are documented? I trust that the values are correct, but it'd be nice to know.

[arsenm]

lib/CodeGen/MIRParser/MIParser.cpp

1344

In LowLevelTypeImpl.h

[arsenm]

r353107

lib/CodeGen/MIRParser/MIParser.cpp

 
 bool MIParser::parseIRConstant(StringRef::iterator Loc, const Constant *&C) {
   if (parseIRConstant(Loc, StringRef(Loc, Token.range().end() - Loc), C))
     return true;
   lex();
   return false;
 }
 
+static bool verifyScalarSize(uint64_t Size) {
+  return Size != 0 && isUInt<16>(Size);
+}
+
+static bool verifyVectorElementCount(uint64_t NumElts) {
+  return NumElts != 0 && isUInt<16>(NumElts);
+}
+
+static bool verifyAddrSpace(uint64_t AddrSpace) {
+  return isUInt<24>(AddrSpace);
+}
+
 bool MIParser::parseLowLevelType(StringRef::iterator Loc, LLT &Ty) {
   if (Token.range().front() == 's' || Token.range().front() == 'p') {
     StringRef SizeStr = Token.range().drop_front();
     if (SizeStr.size() == 0 || !llvm::all_of(SizeStr, isdigit))
       return error("expected integers after 's'/'p' type character");
   }
 
   if (Token.range().front() == 's') {
-    Ty = LLT::scalar(APSInt(Token.range().drop_front()).getZExtValue());
+    auto ScalarSize = APSInt(Token.range().drop_front()).getZExtValue();
+    if (!verifyScalarSize(ScalarSize))
+      return error("invalid size for scalar type");
+
+    Ty = LLT::scalar(ScalarSize);
     lex();
     return false;
   } else if (Token.range().front() == 'p') {
     const DataLayout &DL = MF.getDataLayout();
-    unsigned AS = APSInt(Token.range().drop_front()).getZExtValue();
+    uint64_t AS = APSInt(Token.range().drop_front()).getZExtValue();
+    if (!verifyAddrSpace(AS))
+      return error("invalid address space number");
+
     Ty = LLT::pointer(AS, DL.getPointerSizeInBits(AS));
     lex();
     return false;
   }
 
   // Now we're looking for a vector.
   if (Token.isNot(MIToken::less))
     return error(Loc,
                  "expected sN, pA, <M x sN>, or <M x pA> for GlobalISel type");
   lex();
 
   if (Token.isNot(MIToken::IntegerLiteral))
     return error(Loc, "expected <M x sN> or <M x pA> for vector type");
   uint64_t NumElements = Token.integerValue().getZExtValue();
+  if (!verifyVectorElementCount(NumElements))
+    return error("invalid number of vector elements");
+
   lex();
 
   if (Token.isNot(MIToken::Identifier) || Token.stringValue() != "x")
     return error(Loc, "expected <M x sN> or <M x pA> for vector type");
   lex();
 
   if (Token.range().front() != 's' && Token.range().front() != 'p')
     return error(Loc, "expected <M x sN> or <M x pA> for vector type");
   StringRef SizeStr = Token.range().drop_front();
   if (SizeStr.size() == 0 || !llvm::all_of(SizeStr, isdigit))
     return error("expected integers after 's'/'p' type character");
 
-  if (Token.range().front() == 's')
-    Ty = LLT::scalar(APSInt(Token.range().drop_front()).getZExtValue());
-  else if (Token.range().front() == 'p') {
+  if (Token.range().front() == 's') {
+    auto ScalarSize = APSInt(Token.range().drop_front()).getZExtValue();
+    if (!verifyScalarSize(ScalarSize))
+      return error("invalid size for scalar type");
+    Ty = LLT::scalar(ScalarSize);
+  } else if (Token.range().front() == 'p') {
     const DataLayout &DL = MF.getDataLayout();
-    unsigned AS = APSInt(Token.range().drop_front()).getZExtValue();
+    uint64_t AS = APSInt(Token.range().drop_front()).getZExtValue();
+    if (!verifyAddrSpace(AS))
+      return error("invalid address space number");
+
     Ty = LLT::pointer(AS, DL.getPointerSizeInBits(AS));
   } else
     return error(Loc, "expected <M x sN> or <M x pA> for vector type");
   lex();
 
   if (Token.isNot(MIToken::greater))
     return error(Loc, "expected <M x sN> or <M x pA> for vector type");
   lex();

test/CodeGen/MIR/AArch64/parse-low-level-type-invalid10.mir

+# RUN: not llc -mtriple=aarch64-- -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
+
+# When a low-level type pointer has an address space greater than supported, and make sure an implicit truncate to 32-bits doesn't happen.
+
+---
+name: test_address_space_number_too_big64
+body: |
+  bb.0:
+    liveins: $x0
+    ; CHECK: [[@LINE+1]]:10: invalid address space number
+    %0:_(p17179869185) = G_IMPLICIT_DEF
+...

test/CodeGen/MIR/AArch64/parse-low-level-type-invalid4.mir

+# RUN: not llc -mtriple=aarch64-- -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
+# When a low-level type is 0 bits
+---
+name: test_scalar_size_0
+body: |
+  bb.0:
+    liveins: $x0
+    ; CHECK: [[@LINE+1]]:10: invalid size for scalar type
+    %0:_(s0) = G_IMPLICIT_DEF
+...

test/CodeGen/MIR/AArch64/parse-low-level-type-invalid5.mir

+# RUN: not llc -mtriple=aarch64-- -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
+# When a low-level type is larger than supported
+---
+name: test_scalar_size_65536
+body: |
+  bb.0:
+    liveins: $x0
+    ; CHECK: [[@LINE+1]]:10: invalid size for scalar type
+    %0:_(s65536) = G_IMPLICIT_DEF
+...

test/CodeGen/MIR/AArch64/parse-low-level-type-invalid6.mir

+# RUN: not llc -mtriple=aarch64-- -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
+# When a low-level type vector has a 0-bit element
+---
+name: test_vector_element_size_0
+body: |
+  bb.0:
+    liveins: $x0
+    ; CHECK: [[@LINE+1]]:15: invalid size for scalar type
+    %0:_(<2 x s0>) = G_IMPLICIT_DEF
+...

test/CodeGen/MIR/AArch64/parse-low-level-type-invalid7.mir

+# RUN: not llc -mtriple=aarch64-- -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
+# When a low-level type vector has 0 elements
+---
+name: test_vector_0_elements
+body: |
+  bb.0:
+    liveins: $x0
+    ; CHECK: [[@LINE+1]]:11: invalid number of vector elements
+    %0:_(<0 x s1>) = G_IMPLICIT_DEF
+...

test/CodeGen/MIR/AArch64/parse-low-level-type-invalid8.mir

+# RUN: not llc -mtriple=aarch64-- -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
+# When a low-level type vector has more elements than supported
+---
+name: test_vector_too_many_elements
+body: |
+  bb.0:
+    liveins: $x0
+    ; CHECK: [[@LINE+1]]:11: invalid number of vector elements
+    %0:_(<65536 x s1>) = G_IMPLICIT_DEF
+...

test/CodeGen/MIR/AArch64/parse-low-level-type-invalid9.mir

+# RUN: not llc -mtriple=aarch64-- -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
+# When a low-level type pointer has an address space greater than supported.
+---
+name: test_address_space_number_too_big
+body: |
+  bb.0:
+    liveins: $x0
+    ; CHECK: [[@LINE+1]]:10: invalid address space number
+    %0:_(p16777216) = G_IMPLICIT_DEF
+...