[flang] Accommodate arrays with a zero-extent dimension in location f…

…olding 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
llvm · Mar 7, 2022 · 4133a85 · 4133a85
1 parent 30b6b6a
commit 4133a85
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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