[llvm] [llvm]: tighten estimate of bits needed when parsing an integer (PR #205947)

[Jeremy Kun via llvm-commits]

https://github.com/j2kun created https://github.com/llvm/llvm-project/pull/205947

A colleague of mine noticed that `"12535824225335233"` parses in MLIR's integer attribute parser as a 68-bit integer, even though it is a 53-bit constant. I traced this back to `StringRef::consumeInteger`'s heuristic estimate of the bit size. This change tightens that estimate to at most a 1-bit overapproximation.

>From f60b433455657855f547b9d5a7785e8817f4caf0 Mon Sep 17 00:00:00 2001
From: Jeremy Kun <jkun at google.com>
Date: Thu, 25 Jun 2026 16:10:09 -0700
Subject: [PATCH] [llvm]: tighten estimate of bits needed when parsing an
 integer

---
 llvm/lib/Support/StringRef.cpp       | 17 +++++++++++++-
 llvm/unittests/ADT/StringRefTest.cpp | 35 ++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Support/StringRef.cpp b/llvm/lib/Support/StringRef.cpp
index 02270a6e223af..f62b6e602ccfd 100644
--- a/llvm/lib/Support/StringRef.cpp
+++ b/llvm/lib/Support/StringRef.cpp
@@ -527,7 +527,22 @@ bool StringRef::consumeInteger(unsigned Radix, APInt &Result) {
   while ((1U << Log2Radix) < Radix) Log2Radix++;
   bool IsPowerOf2Radix = ((1U << Log2Radix) == Radix);
 
-  unsigned BitWidth = Log2Radix * Str.size();
+  // (Over-)estimate the required number of bits tightly using a Q10 fixed-point
+  // approximation of log2(Radix) to avoid massive over-allocation for large
+  // non-power-of-two values (e.g. decimal).
+  unsigned BitWidth = 0;
+  if (Radix <= 36) {
+    static const uint16_t Log2RadixTable[] = {
+      0, 0,
+      1024, 1624, 2048, 2378, 2648, 2875, 3072, 3247, 3402, // 2..10
+      3543, 3672, 3790, 3899, 4001, 4096, 4186, 4271, 4350, // 11..19
+      4426, 4498, 4567, 4633, 4696, 4756, 4814, 4870, 4923, // 20..28
+      4975, 5025, 5074, 5120, 5166, 5210, 5253, 5295        // 29..36
+    };
+    BitWidth = (Str.size() * Log2RadixTable[Radix] + 1023) >> 10;
+  } else {
+    BitWidth = Log2Radix * Str.size();
+  }
   if (BitWidth < Result.getBitWidth())
     BitWidth = Result.getBitWidth(); // don't shrink the result
   else if (BitWidth > Result.getBitWidth())
diff --git a/llvm/unittests/ADT/StringRefTest.cpp b/llvm/unittests/ADT/StringRefTest.cpp
index cca0f3a2b46b0..db9c67c7f742f 100644
--- a/llvm/unittests/ADT/StringRefTest.cpp
+++ b/llvm/unittests/ADT/StringRefTest.cpp
@@ -968,6 +968,41 @@ TEST(StringRefTest, consumeIntegerSigned) {
   }
 }
 
+TEST(StringRefTest, consumeIntegerAPIntBitWidth) {
+  // Decimal large number (12535824225335233)
+  // 17 digits, (17 * 3402 + 1023) >> 10 = 57 bits.
+  {
+    APInt U;
+    StringRef Str = "12535824225335233";
+    bool Success = Str.consumeInteger(10, U);
+    ASSERT_FALSE(Success);
+    EXPECT_EQ(U.getZExtValue(), 12535824225335233ULL);
+    EXPECT_EQ(U.getBitWidth(), 57U);
+  }
+
+  // Hex version of same number (2c894405eaf7c1)
+  // 14 digits, (14 * 4096 + 1023) >> 10 = 56 bits.
+  {
+    APInt U;
+    StringRef Str = "2c894405eaf7c1";
+    bool Success = Str.consumeInteger(16, U);
+    ASSERT_FALSE(Success);
+    EXPECT_EQ(U.getZExtValue(), 12535824225335233ULL);
+    EXPECT_EQ(U.getBitWidth(), 56U);
+  }
+
+  // A very large decimal number (100 digits)
+  // (100 * 3402 + 1023) >> 10 = 333 bits.
+  {
+    APInt U;
+    std::string LargeDec(100, '9');
+    StringRef Str = LargeDec;
+    bool Success = Str.consumeInteger(10, U);
+    ASSERT_FALSE(Success);
+    EXPECT_EQ(U.getBitWidth(), 333U);
+  }
+}
+
 struct GetDoubleStrings {
   const char *Str;
   bool AllowInexact;