Enhance the type system to support arbitrary precision integers, which are
important for low-bitwidth inference cases and hardware synthesis targets.

Rename 'int' to 'affineint' to avoid confusion between "the integers" and "the int
type".

PiperOrigin-RevId: 202751508
diff --git a/lib/Parser/Lexer.cpp b/lib/Parser/Lexer.cpp
index 17755e0..8943200 100644
--- a/lib/Parser/Lexer.cpp
+++ b/lib/Parser/Lexer.cpp
@@ -140,6 +140,7 @@
 /// Lex a bare identifier or keyword that starts with a letter.
 ///
 ///   bare-id ::= letter (letter|digit|[_])*
+///   integer-type ::= `i[1-9][0-9]*`
 ///
 Token Lexer::lexBareIdentifierOrKeyword(const char *tokStart) {
   // Match the rest of the identifier regex: [0-9a-zA-Z_]*
@@ -149,6 +150,15 @@
   // Check to see if this identifier is a keyword.
   StringRef spelling(tokStart, curPtr-tokStart);
 
+  // Check for i123.
+  if (tokStart[0] == 'i') {
+    bool allDigits = true;
+    for (auto c : spelling.drop_front())
+      allDigits &= isdigit(c) != 0;
+    if (allDigits && spelling.size() != 1)
+      return Token(Token::inttype, spelling);
+  }
+
   Token::Kind kind = llvm::StringSwitch<Token::Kind>(spelling)
 #define TOK_KEYWORD(SPELLING) \
     .Case(#SPELLING, Token::kw_##SPELLING)
diff --git a/lib/Parser/Parser.cpp b/lib/Parser/Parser.cpp
index 6927050..1bfa331 100644
--- a/lib/Parser/Parser.cpp
+++ b/lib/Parser/Parser.cpp
@@ -121,7 +121,7 @@
   // as the results of their action.
 
   // Type parsing.
-  PrimitiveType *parsePrimitiveType();
+  Type *parsePrimitiveType();
   Type *parseElementType();
   VectorType *parseVectorType();
   ParseResult parseDimensionListRanked(SmallVectorImpl<int> &dimensions);
@@ -218,12 +218,11 @@
 
 /// Parse the low-level fixed dtypes in the system.
 ///
-///   primitive-type
-///      ::= `f16` | `bf16` | `f32` | `f64`      // Floating point
-///        | `i1` | `i8` | `i16` | `i32` | `i64` // Sized integers
-///        | `int`
+///   primitive-type ::= `f16` | `bf16` | `f32` | `f64`
+///   primitive-type ::= integer-type
+///   primitive-type ::= `affineint`
 ///
-PrimitiveType *Parser::parsePrimitiveType() {
+Type *Parser::parsePrimitiveType() {
   switch (curToken.getKind()) {
   default:
     return (emitError("expected type"), nullptr);
@@ -239,24 +238,16 @@
   case Token::kw_f64:
     consumeToken(Token::kw_f64);
     return Type::getF64(context);
-  case Token::kw_i1:
-    consumeToken(Token::kw_i1);
-    return Type::getI1(context);
-  case Token::kw_i8:
-    consumeToken(Token::kw_i8);
-    return Type::getI8(context);
-  case Token::kw_i16:
-    consumeToken(Token::kw_i16);
-    return Type::getI16(context);
-  case Token::kw_i32:
-    consumeToken(Token::kw_i32);
-    return Type::getI32(context);
-  case Token::kw_i64:
-    consumeToken(Token::kw_i64);
-    return Type::getI64(context);
-  case Token::kw_int:
-    consumeToken(Token::kw_int);
-    return Type::getInt(context);
+  case Token::kw_affineint:
+    consumeToken(Token::kw_affineint);
+    return Type::getAffineInt(context);
+  case Token::inttype: {
+    auto width = curToken.getIntTypeBitwidth();
+    if (!width.hasValue())
+      return (emitError("invalid integer width"), nullptr);
+    consumeToken(Token::inttype);
+    return Type::getInt(width.getValue(), context);
+  }
   }
 }
 
@@ -419,11 +410,9 @@
     return (emitError("expected '>' in memref type"), nullptr);
 
   // FIXME: Add an IR representation for memref types.
-  return Type::getI1(context);
+  return Type::getInt(1, context);
 }
 
-
-
 /// Parse a function type.
 ///
 ///   function-type ::= type-list-parens `->` type-list
@@ -445,7 +434,6 @@
   return FunctionType::get(arguments, results, context);
 }
 
-
 /// Parse an arbitrary type.
 ///
 ///   type ::= primitive-type
diff --git a/lib/Parser/Token.cpp b/lib/Parser/Token.cpp
index 5563255..e1e4bed 100644
--- a/lib/Parser/Token.cpp
+++ b/lib/Parser/Token.cpp
@@ -48,6 +48,18 @@
   return result;
 }
 
+/// For an inttype token, return its bitwidth.
+Optional<unsigned> Token::getIntTypeBitwidth() const {
+ unsigned result = 0;
+  if (spelling[1] == '0' ||
+      spelling.drop_front().getAsInteger(10, result) ||
+      // Arbitrary but large limit on bitwidth.
+      result > 4096 || result == 0)
+    return None;
+  return result;
+}
+
+
 /// Given a 'string' token, return its value, including removing the quote
 /// characters and unescaping the contents of the string.
 std::string Token::getStringValue() const {
diff --git a/lib/Parser/Token.h b/lib/Parser/Token.h
index e5e4fc4..bc9e8e4 100644
--- a/lib/Parser/Token.h
+++ b/lib/Parser/Token.h
@@ -73,6 +73,9 @@
   /// return None.
   Optional<unsigned> getUnsignedIntegerValue() const;
 
+  /// For an inttype token, return its bitwidth.
+  Optional<unsigned> getIntTypeBitwidth() const;
+
   /// Given a 'string' token, return its value, including removing the quote
   /// characters and unescaping the contents of the string.
   std::string getStringValue() const;
diff --git a/lib/Parser/TokenKinds.def b/lib/Parser/TokenKinds.def
index 72d769a..73a30df 100644
--- a/lib/Parser/TokenKinds.def
+++ b/lib/Parser/TokenKinds.def
@@ -58,6 +58,7 @@
 // Literals
 TOK_LITERAL(integer)                    // 42
 TOK_LITERAL(string)                     // "foo"
+TOK_LITERAL(inttype)                    // i421
 
 // Punctuation.
 TOK_PUNCTUATION(arrow,            "->")
@@ -84,6 +85,7 @@
 // TODO: More operator tokens
 
 // Keywords.  These turn "foo" into Token::kw_foo enums.
+TOK_KEYWORD(affineint)
 TOK_KEYWORD(bf16)
 TOK_KEYWORD(br)
 TOK_KEYWORD(cfgfunc)
@@ -91,12 +93,6 @@
 TOK_KEYWORD(f16)
 TOK_KEYWORD(f32)
 TOK_KEYWORD(f64)
-TOK_KEYWORD(i1)
-TOK_KEYWORD(i16)
-TOK_KEYWORD(i32)
-TOK_KEYWORD(i64)
-TOK_KEYWORD(i8)
-TOK_KEYWORD(int)
 TOK_KEYWORD(memref)
 TOK_KEYWORD(mlfunc)
 TOK_KEYWORD(return)