[flang-commits] [flang] 4133a85 - [flang] Accommodate arrays with a zero-extent dimension in location folding

[Peter Klausler via flang-commits]

Author: Peter Klausler
Date: 2022-03-07T10:44:18-08:00
New Revision: 4133a85c867fc0b5d6d53de85063f66562d3cb46

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

LOG: [flang] Accommodate arrays with a zero-extent dimension in location folding

The index incrementation code used for FINDLOC, MAXLOC, and MINLOC folding
would crash if the array had a zero extent on the dimension selected with
a DIM= argument since the subscript passed to IncrementSubscripts would
have a value less than the lower bound.  Fix, and add tests.

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

Added: 
    

Modified: 
    flang/lib/Evaluate/constant.cpp
    flang/lib/Evaluate/fold-implementation.h
    flang/lib/Evaluate/fold-integer.cpp
    flang/test/Evaluate/fold-findloc.f90

Removed: 
    


################################################################################
diff  --git a/flang/lib/Evaluate/constant.cpp b/flang/lib/Evaluate/constant.cpp
index a4f5f72a35c7b..9f2040ce8575f 100644
--- a/flang/lib/Evaluate/constant.cpp
+++ b/flang/lib/Evaluate/constant.cpp
@@ -68,7 +68,7 @@ bool ConstantBounds::IncrementSubscripts(
     if (++indices[k] < lb + shape_[k]) {
       return true;
     } else {
-      CHECK(indices[k] == lb + shape_[k]);
+      CHECK(indices[k] == lb + std::max<ConstantSubscript>(shape_[k], 1));
       indices[k] = lb;
     }
   }

diff  --git a/flang/lib/Evaluate/fold-implementation.h b/flang/lib/Evaluate/fold-implementation.h
index 43b7f931826ab..14d9c981d0f41 100644
--- a/flang/lib/Evaluate/fold-implementation.h
+++ b/flang/lib/Evaluate/fold-implementation.h
@@ -612,7 +612,7 @@ template <typename T> Expr<T> Folder<T>::CSHIFT(FunctionRef<T> &&funcRef) {
             zbDimIndex = 0;
           }
         }
-        arrayAt[zbDim] = dimLB + dimExtent - 1;
+        arrayAt[zbDim] = dimLB + std::max<ConstantSubscript>(dimExtent, 1) - 1;
         array->IncrementSubscripts(arrayAt);
         shift->IncrementSubscripts(shiftAt);
       }
@@ -726,7 +726,7 @@ template <typename T> Expr<T> Folder<T>::EOSHIFT(FunctionRef<T> &&funcRef) {
             DIE("no derived type boundary");
           }
         }
-        arrayAt[zbDim] = dimLB + dimExtent - 1;
+        arrayAt[zbDim] = dimLB + std::max<ConstantSubscript>(dimExtent, 1) - 1;
         array->IncrementSubscripts(arrayAt);
         shift->IncrementSubscripts(shiftAt);
         if (boundary) {

diff  --git a/flang/lib/Evaluate/fold-integer.cpp b/flang/lib/Evaluate/fold-integer.cpp
index 6a110cbe7e8ce..df5fa0bd96123 100644
--- a/flang/lib/Evaluate/fold-integer.cpp
+++ b/flang/lib/Evaluate/fold-integer.cpp
@@ -284,11 +284,12 @@ template <WhichLocation WHICH> class LocationHelper {
           }
         }
         resultIndices.emplace_back(hit);
-        at[zbDim] = dimLength;
+        at[zbDim] = std::max<ConstantSubscript>(dimLength, 1);
         array->IncrementSubscripts(at);
         at[zbDim] = 1;
         if (mask) {
-          maskAt[zbDim] = mask->lbounds()[zbDim] + dimLength - 1;
+          maskAt[zbDim] = mask->lbounds()[zbDim] +
+              std::max<ConstantSubscript>(dimLength, 1) - 1;
           mask->IncrementSubscripts(maskAt);
           maskAt[zbDim] = mask->lbounds()[zbDim];
         }

diff  --git a/flang/test/Evaluate/fold-findloc.f90 b/flang/test/Evaluate/fold-findloc.f90
index 0c0006657ebcb..399afca97de55 100644
--- a/flang/test/Evaluate/fold-findloc.f90
+++ b/flang/test/Evaluate/fold-findloc.f90
@@ -3,6 +3,7 @@
 module m1
   integer, parameter :: ia1(2:6) = [1, 2, 3, 2, 1]
   integer, parameter :: ia2(2:3,2:4) = reshape([1, 2, 3, 3, 2, 1], shape(ia2))
+  integer, parameter :: ia3(2,0,2) = 0 ! middle dimension has zero extent
 
   logical, parameter :: test_fi1a = all(findloc(ia1, 1) == 1)
   logical, parameter :: test_fi1ar = rank(findloc(ia1, 1)) == 1
@@ -54,4 +55,9 @@ module m1
   logical, parameter :: test_ni2e = all(minloc(ia2, dim=1) == [1, 1, 2])
   logical, parameter :: test_ni2f = all(minloc(ia2, dim=1, back=.true.) == [1, 2, 2])
   logical, parameter :: test_ni2g = all(minloc(ia2, dim=2) == [1, 3])
+
+  logical, parameter :: test_xi3a = all(maxloc(ia3) == [0,0,0])
+  logical, parameter :: test_xi3b = all(maxloc(ia3, back=.true.) == [0,0,0])
+  logical, parameter :: test_xi3c = all(maxloc(ia3, dim=2) == reshape([0,0,0,0],shape=[2,2]))
+  logical, parameter :: test_xi3d = all(maxloc(ia3, dim=2, back=.true.) == reshape([0,0,0,0],shape=[2,2]))
 end module