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

[klausler]

Summary

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.

[PeteSteinfeld]

All builds and tests correctly and looks good.

flang/lib/Evaluate/constant.cpp

bool ConstantBounds::IncrementSubscripts(
   CHECK(!dimOrder || static_cast<int>(dimOrder->size()) == rank);
   for (int j{0}; j < rank; ++j) {
     ConstantSubscript k{dimOrder ? (*dimOrder)[j] : j};
     auto lb{lbounds_[k]};
     CHECK(indices[k] >= lb);
     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;
     }
   }
   return false; // all done
 }
 
 std::optional<std::vector<int>> ValidateDimensionOrder(
     int rank, const std::vector<int> &order) {

flang/lib/Evaluate/fold-implementation.h

if (ok) {
         }
         for (ConstantSubscript j{0}; j < dimExtent; ++j) {
           arrayAt[zbDim] = dimLB + zbDimIndex;
           resultElements.push_back(array->At(arrayAt));
           if (++zbDimIndex == dimExtent) {
             zbDimIndex = 0;
           }
         }
-        arrayAt[zbDim] = dimLB + dimExtent - 1;
+        arrayAt[zbDim] = dimLB + std::max<ConstantSubscript>(dimExtent, 1) - 1;
         array->IncrementSubscripts(arrayAt);
         shift->IncrementSubscripts(shiftAt);
       }
       return Expr<T>{PackageConstant<T>(
           std::move(resultElements), *array, array->shape())};
     }
   }
   // Invalid, prevent re-folding
if (ok) {
           } else if constexpr (T::category == TypeCategory::Character) {
             auto len{static_cast<std::size_t>(array->LEN())};
             typename Scalar<T>::value_type space{' '};
             resultElements.emplace_back(len, space);
           } else {
             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) {
           boundary->IncrementSubscripts(boundaryAt);
         }
       }
       return Expr<T>{PackageConstant<T>(
           std::move(resultElements), *array, array->shape())};

flang/lib/Evaluate/fold-integer.cpp

if (dim) { // DIM=
             if constexpr (WHICH == WhichLocation::Findloc) {
               if (!back) {
                 break;
               }
             }
           }
         }
         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];
         }
       }
     } else { // no DIM=
       resultShape = ConstantSubscripts{array->Rank()}; // always a vector
       ConstantSubscript n{GetSize(array->shape())};
       resultIndices = ConstantSubscripts(array->Rank(), 0);

flang/test/Evaluate/fold-findloc.f90

 ! RUN: %python %S/test_folding.py %s %flang_fc1
 ! Tests folding of FINDLOC, MAXLOC, & MINLOC
 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
   logical, parameter :: test_fi1ak = kind(findloc(ia1, 1, kind=2)) == 2
   logical, parameter :: test_fi1ad = findloc(ia1, 1, dim=1) == 1
   logical, parameter :: test_fi1adr = rank(findloc(ia1, 1, dim=1)) == 0
   logical, parameter :: test_fi1b = all(findloc(ia1, 1, back=.true.) == 5)
   logical, parameter :: test_fi1c = all(findloc(ia1, 2, mask=[.false., .false., .true., .true., .true.]) == 4)
module m1
   logical, parameter :: test_ni2a = all(minloc(ia2) == [1, 1])
   logical, parameter :: test_ni2ar = rank(minloc(ia2)) == 1
   logical, parameter :: test_ni2b = all(minloc(ia2, back=.true.) == [2, 3])
   logical, parameter :: test_ni2c = all(minloc(ia2, mask=reshape([.false., .true., .true., .false., .true., .false.], shape(ia2))) == [2, 1])
   logical, parameter :: test_ni2d = all(minloc(ia2, mask=reshape([.false., .true., .true., .false., .true., .false.], shape(ia2)), back=.true.) == [1, 3])
   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