[flang-commits] [flang] bd5371e - [flang] Fix LBOUND() folding for constant arrays

[Mike Kashkarov via flang-commits]

Author: Mike Kashkarov
Date: 2022-04-20T20:04:31+03:00
New Revision: bd5371e4fc607aae1a77f29273fd130b59e72409

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

LOG: [flang] Fix LBOUND() folding for constant arrays

Previously constant folding uses 'dim' without checks which leads to ICE if we
do not have DIM= parameter. And for inputs without DIM= we need to form an
array of rank size with computed bounds instead of single value.

Add additional PackageConstant function to simplify 'if (dim)' handling since we
need to distinguish between scalar initialization in case of DIM= argument and
rank=1 array.

Also add a few more tests with 'parameter' type to verify folding for constant
arrays.

Reviewed By: jeanPerier

Differential Revision: https://reviews.llvm.org/D123237

Added: 
    

Modified: 
    flang/lib/Evaluate/fold-integer.cpp
    flang/test/Evaluate/folding08.f90

Removed: 
    


################################################################################
diff  --git a/flang/lib/Evaluate/fold-integer.cpp b/flang/lib/Evaluate/fold-integer.cpp
index d07af216d230a..cbecda0ecebda 100644
--- a/flang/lib/Evaluate/fold-integer.cpp
+++ b/flang/lib/Evaluate/fold-integer.cpp
@@ -12,30 +12,52 @@
 
 namespace Fortran::evaluate {
 
+// Given a collection of ConstantSubscripts values, package them as a Constant.
+// Return scalar value if asScalar == true and shape-dim array otherwise.
+template <typename T>
+Expr<T> PackageConstantBounds(
+    const ConstantSubscripts &&bounds, bool asScalar = false) {
+  if (asScalar) {
+    return Expr<T>{Constant<T>{bounds.at(0)}};
+  } else {
+    // As rank-dim array
+    const int rank{GetRank(bounds)};
+    std::vector<Scalar<T>> packed(rank);
+    std::transform(bounds.begin(), bounds.end(), packed.begin(),
+        [](ConstantSubscript x) { return Scalar<T>(x); });
+    return Expr<T>{Constant<T>{std::move(packed), ConstantSubscripts{rank}}};
+  }
+}
+
 // Class to retrieve the constant lower bound of an expression which is an
 // array that devolves to a type of Constant<T>
 class GetConstantArrayLboundHelper {
 public:
-  GetConstantArrayLboundHelper(ConstantSubscript dim) : dim_{dim} {}
+  GetConstantArrayLboundHelper(std::optional<ConstantSubscript> dim)
+      : dim_{dim} {}
 
-  template <typename T> ConstantSubscript GetLbound(const T &) {
+  template <typename T> ConstantSubscripts GetLbound(const T &) {
     // The method is needed for template expansion, but we should never get
     // here in practice.
     CHECK(false);
-    return 0;
+    return {0};
   }
 
-  template <typename T> ConstantSubscript GetLbound(const Constant<T> &x) {
+  template <typename T> ConstantSubscripts GetLbound(const Constant<T> &x) {
     // Return the lower bound
-    return x.lbounds()[dim_];
+    if (dim_) {
+      return {x.lbounds().at(*dim_)};
+    } else {
+      return x.lbounds();
+    }
   }
 
-  template <typename T> ConstantSubscript GetLbound(const Parentheses<T> &x) {
+  template <typename T> ConstantSubscripts GetLbound(const Parentheses<T> &x) {
     // Strip off the parentheses
     return GetLbound(x.left());
   }
 
-  template <typename T> ConstantSubscript GetLbound(const Expr<T> &x) {
+  template <typename T> ConstantSubscripts GetLbound(const Expr<T> &x) {
     // recurse through Expr<T>'a until we hit a constant
     return common::visit([&](const auto &inner) { return GetLbound(inner); },
         //      [&](const auto &) { return 0; },
@@ -43,7 +65,7 @@ class GetConstantArrayLboundHelper {
   }
 
 private:
-  ConstantSubscript dim_;
+  std::optional<ConstantSubscript> dim_;
 };
 
 template <int KIND>
@@ -89,16 +111,13 @@ Expr<Type<TypeCategory::Integer, KIND>> LBOUND(FoldingContext &context,
         }
       }
       if (IsActuallyConstant(*array)) {
-        return Expr<T>{GetConstantArrayLboundHelper{*dim}.GetLbound(*array)};
+        const ConstantSubscripts bounds{
+            GetConstantArrayLboundHelper{dim}.GetLbound(*array)};
+        return PackageConstantBounds<T>(std::move(bounds), dim.has_value());
       }
       if (lowerBoundsAreOne) {
-        if (dim) {
-          return Expr<T>{1};
-        } else {
-          std::vector<Scalar<T>> ones(rank, Scalar<T>{1});
-          return Expr<T>{
-              Constant<T>{std::move(ones), ConstantSubscripts{rank}}};
-        }
+        ConstantSubscripts ones(rank, ConstantSubscript{1});
+        return PackageConstantBounds<T>(std::move(ones), dim.has_value());
       }
     }
   }

diff  --git a/flang/test/Evaluate/folding08.f90 b/flang/test/Evaluate/folding08.f90
index 7432cbd86fb08..a5b65e0246efb 100644
--- a/flang/test/Evaluate/folding08.f90
+++ b/flang/test/Evaluate/folding08.f90
@@ -77,4 +77,22 @@ subroutine test2
       end block
     end associate
   end subroutine
+  subroutine test3_lbound_parameter
+    ! Test lbound with constant arrays
+    integer, parameter :: a1(1) = 0
+    integer, parameter :: lba1(*) = lbound(a1)
+    logical, parameter :: test_lba1 = all(lba1 == [1])
+
+    integer, parameter :: a2(0:1) = 0
+    integer, parameter :: lba2(*) = lbound(a2)
+    logical, parameter :: test_lba2 = all(lba2 == [0])
+
+    integer, parameter :: a3(-10:-5,1,4:6) = 0
+    integer, parameter :: lba3(*) = lbound(a3)
+    logical, parameter :: test_lba3 = all(lba3 == [-10, 1, 4])
+    ! Exercise with DIM=
+    logical, parameter :: test_lba3_dim = lbound(a3, 1) == -10 .and. &
+         lbound(a3, 2) == 1 .and. &
+         lbound(a3, 3) == 4
+  end subroutine
 end