[llvm] [llvm::transforms] Add overflow check in AllocaInst::getAllocationSize (PR #96668)

via llvm-commits llvm-commits at lists.llvm.org
Tue Jun 25 10:03:08 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-ir

Author: Tsz Chan (tszhin-swe)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/96668.diff


2 Files Affected:

- (modified) llvm/lib/IR/Instructions.cpp (+11-10) 
- (modified) llvm/unittests/IR/InstructionsTest.cpp (+3) 


``````````diff
diff --git a/llvm/lib/IR/Instructions.cpp b/llvm/lib/IR/Instructions.cpp
index 2e1fd2a96ffff..dee9c0acfcab8 100644
--- a/llvm/lib/IR/Instructions.cpp
+++ b/llvm/lib/IR/Instructions.cpp
@@ -41,6 +41,7 @@
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/ModRef.h"
 #include "llvm/Support/TypeSize.h"
+#include "llvm/Support/CheckedArithmetic.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
@@ -60,21 +61,21 @@ static cl::opt<bool> DisableI2pP2iOpt(
 std::optional<TypeSize>
 AllocaInst::getAllocationSize(const DataLayout &DL) const {
   TypeSize Size = DL.getTypeAllocSize(getAllocatedType());
-  if (isArrayAllocation()) {
-    auto *C = dyn_cast<ConstantInt>(getArraySize());
-    if (!C)
-      return std::nullopt;
-    assert(!Size.isScalable() && "Array elements cannot have a scalable size");
-    Size *= C->getZExtValue();
+  if (!isArrayAllocation()) {
+    return Size;
   }
-  return Size;
+  auto *C = dyn_cast<ConstantInt>(getArraySize());
+  if (!C)
+    return std::nullopt;
+  assert(!Size.isScalable() && "Array elements cannot have a scalable size");
+  return checkedMul(Size, TypeSize::getFixed(C->getZExtValue()));
 }
 
 std::optional<TypeSize>
 AllocaInst::getAllocationSizeInBits(const DataLayout &DL) const {
-  std::optional<TypeSize> Size = getAllocationSize(DL);
-  if (Size)
-    return *Size * 8;
+  std::optional<TypeSize> optSize = getAllocationSize(DL);
+  if (optSize)
+    return checkedMul(*optSize, TypeSize::getFixed(8));
   return std::nullopt;
 }
 
diff --git a/llvm/unittests/IR/InstructionsTest.cpp b/llvm/unittests/IR/InstructionsTest.cpp
index b6044b2862920..7699fd3ee6681 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, i32 1000000000 ; Allocation causing overflow
         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_EQ(I.getAllocationSizeInBits(DL), TypeSize::getFixed(160));
 }
 
 TEST(InstructionsTest, InsertAtBegin) {

``````````

</details>


https://github.com/llvm/llvm-project/pull/96668


More information about the llvm-commits mailing list