[llvm] 298e292 - [IR] Add overflow check in AllocaInst::getAllocationSize (#97170)

[via llvm-commits]

Author: Tsz Chan
Date: 2024-07-03T12:48:04+02:00
New Revision: 298e292a76289d93f7c1a80c26e354830c6080e4

URL: https://github.com/llvm/llvm-project/commit/298e292a76289d93f7c1a80c26e354830c6080e4
DIFF: https://github.com/llvm/llvm-project/commit/298e292a76289d93f7c1a80c26e354830c6080e4.diff

LOG: [IR] Add overflow check in AllocaInst::getAllocationSize (#97170)

Fixes #91380.

Added: 
    

Modified: 
    llvm/lib/IR/Instructions.cpp
    llvm/unittests/IR/InstructionsTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/IR/Instructions.cpp b/llvm/lib/IR/Instructions.cpp
index 494306815d90d..7a8cf8c230498 100644
--- a/llvm/lib/IR/Instructions.cpp
+++ b/llvm/lib/IR/Instructions.cpp
@@ -37,6 +37,7 @@
 #include "llvm/IR/Value.h"
 #include "llvm/Support/AtomicOrdering.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/CheckedArithmetic.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/ModRef.h"
@@ -65,7 +66,11 @@ AllocaInst::getAllocationSize(const DataLayout &DL) const {
     if (!C)
       return std::nullopt;
     assert(!Size.isScalable() && "Array elements cannot have a scalable size");
-    Size *= C->getZExtValue();
+    auto CheckedProd =
+        checkedMulUnsigned(Size.getKnownMinValue(), C->getZExtValue());
+    if (!CheckedProd)
+      return std::nullopt;
+    return TypeSize::getFixed(*CheckedProd);
   }
   return Size;
 }
@@ -73,9 +78,13 @@ AllocaInst::getAllocationSize(const DataLayout &DL) const {
 std::optional<TypeSize>
 AllocaInst::getAllocationSizeInBits(const DataLayout &DL) const {
   std::optional<TypeSize> Size = getAllocationSize(DL);
-  if (Size)
-    return *Size * 8;
-  return std::nullopt;
+  if (!Size)
+    return std::nullopt;
+  auto CheckedProd = checkedMulUnsigned(Size->getKnownMinValue(),
+                                        static_cast<TypeSize::ScalarTy>(8));
+  if (!CheckedProd)
+    return std::nullopt;
+  return TypeSize::get(*CheckedProd, Size->isScalable());
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/llvm/unittests/IR/InstructionsTest.cpp b/llvm/unittests/IR/InstructionsTest.cpp
index b6044b2862920..4c1e9a9acb29a 100644
--- a/llvm/unittests/IR/InstructionsTest.cpp
+++ b/llvm/unittests/IR/InstructionsTest.cpp
@@ -1750,6 +1750,7 @@ TEST(InstructionsTest, AllocaInst) {
         %F = alloca [2 x half]
         %G = alloca [2 x [3 x i128]]
         %H = alloca %T
+        %I = alloca i32, i64 9223372036854775807
         ret void
       }
     )");
@@ -1766,6 +1767,7 @@ TEST(InstructionsTest, AllocaInst) {
   AllocaInst &F = cast<AllocaInst>(*It++);
   AllocaInst &G = cast<AllocaInst>(*It++);
   AllocaInst &H = cast<AllocaInst>(*It++);
+  AllocaInst &I = cast<AllocaInst>(*It++);
   EXPECT_EQ(A.getAllocationSizeInBits(DL), TypeSize::getFixed(32));
   EXPECT_EQ(B.getAllocationSizeInBits(DL), TypeSize::getFixed(128));
   EXPECT_FALSE(C.getAllocationSizeInBits(DL));
@@ -1774,6 +1776,7 @@ TEST(InstructionsTest, AllocaInst) {
   EXPECT_EQ(F.getAllocationSizeInBits(DL), TypeSize::getFixed(32));
   EXPECT_EQ(G.getAllocationSizeInBits(DL), TypeSize::getFixed(768));
   EXPECT_EQ(H.getAllocationSizeInBits(DL), TypeSize::getFixed(160));
+  EXPECT_FALSE(I.getAllocationSizeInBits(DL));
 }
 
 TEST(InstructionsTest, InsertAtBegin) {