[flang-commits] [flang] [flang] Check for overflows in RESHAPE folding (PR #68342)

[Leandro Lupori via flang-commits]

https://github.com/luporl created https://github.com/llvm/llvm-project/pull/68342

None

>From daa8de4653bc9e4fd7f8b5827b3b0e0cc3be9213 Mon Sep 17 00:00:00 2001
From: Leandro Lupori <leandro.lupori at linaro.org>
Date: Wed, 4 Oct 2023 18:15:38 +0000
Subject: [PATCH] [flang] Check for overflows in RESHAPE folding

---
 flang/include/flang/Evaluate/constant.h  |  1 +
 flang/lib/Evaluate/constant.cpp          | 14 ++++++++++++++
 flang/lib/Evaluate/fold-implementation.h |  3 +++
 flang/test/Semantics/reshape.f90         |  5 +++++
 4 files changed, 23 insertions(+)

diff --git a/flang/include/flang/Evaluate/constant.h b/flang/include/flang/Evaluate/constant.h
index 04474e2f49a0f88..c5e08d38f793234 100644
--- a/flang/include/flang/Evaluate/constant.h
+++ b/flang/include/flang/Evaluate/constant.h
@@ -47,6 +47,7 @@ inline int GetRank(const ConstantSubscripts &s) {
 }
 
 std::size_t TotalElementCount(const ConstantSubscripts &);
+bool TotalElementCountOverflows(const ConstantSubscripts &);
 
 // Validate dimension re-ordering like ORDER in RESHAPE.
 // On success, return a vector that can be used as dimOrder in
diff --git a/flang/lib/Evaluate/constant.cpp b/flang/lib/Evaluate/constant.cpp
index 084836b4ec36773..76403364ce503e3 100644
--- a/flang/lib/Evaluate/constant.cpp
+++ b/flang/lib/Evaluate/constant.cpp
@@ -84,6 +84,20 @@ std::size_t TotalElementCount(const ConstantSubscripts &shape) {
   return static_cast<std::size_t>(GetSize(shape));
 }
 
+bool TotalElementCountOverflows(const ConstantSubscripts &shape) {
+  uint64_t size{1};
+  for (auto dim : shape) {
+    CHECK(dim >= 0);
+    uint64_t osize{size};
+    size = osize * dim;
+    if (size > std::numeric_limits<decltype(dim)>::max() ||
+        (dim != 0 && size / dim != osize)) {
+      return true;
+    }
+  }
+  return false;
+}
+
 bool ConstantBounds::IncrementSubscripts(
     ConstantSubscripts &indices, const std::vector<int> *dimOrder) const {
   int rank{GetRank(shape_)};
diff --git a/flang/lib/Evaluate/fold-implementation.h b/flang/lib/Evaluate/fold-implementation.h
index 2a40018cd5a3865..0748a27268aee26 100644
--- a/flang/lib/Evaluate/fold-implementation.h
+++ b/flang/lib/Evaluate/fold-implementation.h
@@ -878,6 +878,9 @@ template <typename T> Expr<T> Folder<T>::RESHAPE(FunctionRef<T> &&funcRef) {
   } else if (HasNegativeExtent(shape.value())) {
     context_.messages().Say(
         "'shape=' argument must not have a negative extent"_err_en_US);
+  } else if (TotalElementCountOverflows(shape.value())) {
+    context_.messages().Say(
+        "'shape=' argument has too many elements"_err_en_US);
   } else {
     int rank{GetRank(shape.value())};
     std::size_t resultElements{TotalElementCount(shape.value())};
diff --git a/flang/test/Semantics/reshape.f90 b/flang/test/Semantics/reshape.f90
index 2e9b5adf3ff0e50..4f9be6510e5dcd4 100644
--- a/flang/test/Semantics/reshape.f90
+++ b/flang/test/Semantics/reshape.f90
@@ -44,6 +44,11 @@ program reshaper
   type(dType), parameter :: array19(*) = [dType::dType(field=[1,2])]
   logical, parameter :: lVar = all(array19(:)%field(1) == [2])
 
+  integer(8), parameter :: I64_MAX = INT(z'7fffffffffffffff', kind=8)
+  integer(8), parameter :: huge_shape(2) = [I64_MAX, I64_MAX]
+  !ERROR: 'shape=' argument has too many elements
+  integer :: array21(I64_MAX, I64_MAX) = RESHAPE([1, 2, 3], huge_shape)
+
   !ERROR: Size of 'shape=' argument must not be greater than 15
   CALL ext_sub(RESHAPE([(n, n=1,20)], &
     [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]))