[flang] Fix combined folding of FINDLOC/MAXLOC/MINLOC

[klausler]

Summary

The tests for folding these intrinsics neglected to name the
logical scalars with a leading "test_", so test failures caused
by recent work to implement a combined constant folding facility
for these intrinsics wasn't catching some bugs. This patch fixes
the tests and the bugs.

[jeanPerier]

LGTM

flang/lib/Evaluate/fold-integer.cpp

template <typename T> Result Test() const {
     array->SetLowerBoundsToOne();
     ConstantSubscripts at{array->lbounds()}, maskAt, resultIndices, resultShape;
     if (mask) {
       mask->SetLowerBoundsToOne();
       maskAt = mask->lbounds();
     }
     if (dim) { // DIM=
       if (*dim < 1 || *dim > array->Rank()) {
-        context_.messages().Say(
-            "FINDLOC(DIM=%d) is out of range"_err_en_US, *dim);
+        context_.messages().Say("DIM=%d is out of range"_err_en_US, *dim);
         return std::nullopt;
       }
       int zbDim{*dim - 1};
       resultShape = array->shape();
       resultShape.erase(
           resultShape.begin() + zbDim); // scalar if array is vector
       ConstantSubscript dimLength{array->shape()[zbDim]};
       ConstantSubscript n{GetSize(resultShape)};
       for (ConstantSubscript j{0}; j < n; ++j) {
-        ConstantSubscript hit{array->lbounds()[zbDim] - 1};
+        ConstantSubscript hit{0};
+        if constexpr (WHICH == WhichLocation::Maxloc ||
+            WHICH == WhichLocation::Minloc) {
           value.reset();
+        }
         for (ConstantSubscript k{0}; k < dimLength;
              ++k, ++at[zbDim], mask && ++maskAt[zbDim]) {
           if ((!mask || mask->At(maskAt).IsTrue()) &&
               IsHit(array->At(at), value, relation)) {
             hit = at[zbDim];
+            if constexpr (WHICH == WhichLocation::Findloc) {
               if (!back) {
                 break;
               }
             }
           }
+        }
         resultIndices.emplace_back(hit);
-        at[zbDim] = array->lbounds()[zbDim] + dimLength - 1;
+        at[zbDim] = dimLength;
         array->IncrementSubscripts(at);
-        at[zbDim] = array->lbounds()[zbDim];
+        at[zbDim] = 1;
         if (mask) {
           maskAt[zbDim] = mask->lbounds()[zbDim] + dimLength - 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);
       for (ConstantSubscript j{0}; j < n; ++j, array->IncrementSubscripts(at),
            mask && mask->IncrementSubscripts(maskAt)) {
         if ((!mask || mask->At(maskAt).IsTrue()) &&
             IsHit(array->At(at), value, relation)) {
           resultIndices = at;
+          if constexpr (WHICH == WhichLocation::Findloc) {
             if (!back) {
               break;
             }
           }
         }
       }
+    }
     std::vector<Scalar<SubscriptInteger>> resultElements;
     for (ConstantSubscript j : resultIndices) {
       resultElements.emplace_back(j);
     }
     return Constant<SubscriptInteger>{
         std::move(resultElements), std::move(resultShape)};
   }
 
 private:
   template <typename T>
   bool IsHit(typename Constant<T>::Element element,
       std::optional<Constant<T>> &value,
       [[maybe_unused]] RelationalOperator relation) const {
     std::optional<Expr<LogicalResult>> cmp;
+    bool result{true};
     if (value) {
       if constexpr (T::category == TypeCategory::Logical) {
         // array(at) .EQV. value?
         static_assert(WHICH == WhichLocation::Findloc);
         cmp.emplace(
             ConvertToType<LogicalResult>(Expr<T>{LogicalOperation<T::kind>{
                 LogicalOperator::Eqv, Expr<T>{Constant<T>{std::move(element)}},
                 Expr<T>{Constant<T>{*value}}}}));
       } else { // compare array(at) to value
         cmp.emplace(
             PackageRelation(relation, Expr<T>{Constant<T>{std::move(element)}},
                 Expr<T>{Constant<T>{*value}}));
       }
       Expr<LogicalResult> folded{Fold(context_, std::move(*cmp))};
-      return GetScalarConstantValue<LogicalResult>(folded).value().IsTrue();
-    } else { // first unmasked element seen for MAXLOC/MINLOC
+      result = GetScalarConstantValue<LogicalResult>(folded).value().IsTrue();
+    } else {
+      // first unmasked element for MAXLOC/MINLOC - always take it
+    }
+    if constexpr (WHICH == WhichLocation::Maxloc ||
+        WHICH == WhichLocation::Minloc) {
+      if (result) {
         value.emplace(std::move(element));
-      return true;
       }
     }
+    return result;
+  }
 
   static constexpr int dimArg{WHICH == WhichLocation::Findloc ? 2 : 1};
   static constexpr int maskArg{dimArg + 1};
   static constexpr int backArg{maskArg + 2};
 
   DynamicType type_;
   ActualArguments &arg_;
   FoldingContext &context_;

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))
+
+  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)
+
+  logical, parameter :: test_fi2a = all(findloc(ia2, 1) == [1, 1])
+  logical, parameter :: test_fi2ar = rank(findloc(ia2, 1)) == 1
+  logical, parameter :: test_fi2b = all(findloc(ia2, 1, back=.true.) == [2, 3])
+  logical, parameter :: test_fi2c = all(findloc(ia2, 2, mask=reshape([.false., .false., .true., .true., .true., .false.], shape(ia2))) == [1, 3])
+  logical, parameter :: test_fi2d = all(findloc(ia2, 1, dim=1) == [1, 0, 2])
+  logical, parameter :: test_fi2e = all(findloc(ia2, 1, dim=2) == [1, 3])
+
+  logical, parameter :: test_xi1a = all(maxloc(ia1) == 3)
+  logical, parameter :: test_xi1ar = rank(maxloc(ia1)) == 1
+  logical, parameter :: test_xi1ak = kind(maxloc(ia1, kind=2)) == 2
+  logical, parameter :: test_xi1ad = maxloc(ia1, dim=1) == 3
+  logical, parameter :: test_xi1adr = rank(maxloc(ia1, dim=1)) == 0
+  logical, parameter :: test_xi1b = all(maxloc(ia1, back=.true.) == 3)
+  logical, parameter :: test_xi1c = all(maxloc(ia1, mask=[.false., .true., .false., .true., .true.]) == 2)
+  logical, parameter :: test_xi1d = all(maxloc(ia1, mask=[.false., .true., .false., .true., .true.], back=.true.) == 4)
+
+  logical, parameter :: test_xi2a = all(maxloc(ia2) == [1, 2])
+  logical, parameter :: test_xi2ar = rank(maxloc(ia2)) == 1
+  logical, parameter :: test_xi2b = all(maxloc(ia2, back=.true.) == [2, 2])
+  logical, parameter :: test_xi2c = all(maxloc(ia2, mask=reshape([.false., .true., .true., .true., .true., .true.], shape(ia2))) == [1, 2])
+  logical, parameter :: test_xi2d = all(maxloc(ia2, mask=reshape([.false., .true., .true., .true., .true., .true.], shape(ia2)), back=.true.) == [2, 2])
+  logical, parameter :: test_xi2e = all(maxloc(ia2, dim=1) == [2, 1, 1])
+  logical, parameter :: test_xi2f = all(maxloc(ia2, dim=1, back=.true.) == [2, 2, 1])
+  logical, parameter :: test_xi2g = all(maxloc(ia2, dim=2) == [2, 2])
+
+  logical, parameter :: test_ni1a = all(minloc(ia1) == 1)
+  logical, parameter :: test_ni1ar = rank(minloc(ia1)) == 1
+  logical, parameter :: test_ni1ak = kind(minloc(ia1, kind=2)) == 2
+  logical, parameter :: test_ni1ad = minloc(ia1, dim=1) == 1
+  logical, parameter :: test_ni1adr = rank(minloc(ia1, dim=1)) == 0
+  logical, parameter :: test_ni1b = all(minloc(ia1, back=.true.) == 5)
+  logical, parameter :: test_ni1c = all(minloc(ia1, mask=[.false., .true., .true., .true., .false.]) == 2)
+  logical, parameter :: test_ni1d = all(minloc(ia1, mask=[.false., .true., .true., .true., .false.], back=.true.) == 4)
+
+  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])
+end module

flang/test/Evaluate/folding30.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))
-
-  logical, parameter :: fi1a = all(findloc(ia1, 1) == 1)
-  logical, parameter :: fi1ar = rank(findloc(ia1, 1)) == 1
-  logical, parameter :: fi1ak = kind(findloc(ia1, 1, kind=2)) == 2
-  logical, parameter :: fi1ad = findloc(ia1, 1, dim=1) == 1
-  logical, parameter :: fi1adr = rank(findloc(ia1, 1, dim=1)) == 0
-  logical, parameter :: fi1b = all(findloc(ia1, 1, back=.true.) == 5)
-  logical, parameter :: fi1c = all(findloc(ia1, 2, mask=[.false., .false., .true., .true., .true.]) == 4)
-
-  logical, parameter :: fi2a = all(findloc(ia2, 1) == [1, 1])
-  logical, parameter :: fi2ar = rank(findloc(ia2, 1)) == 1
-  logical, parameter :: fi2b = all(findloc(ia2, 1, back=.true.) == [2, 3])
-  logical, parameter :: fi2c = all(findloc(ia2, 2, mask=reshape([.false., .false., .true., .true., .true., .false.], shape(ia2))) == [1, 3])
-  logical, parameter :: fi2d = all(findloc(ia2, 1, dim=1) == [1, 0, 2])
-  logical, parameter :: fi2e = all(findloc(ia2, 1, dim=2) == [1, 3])
-
-  logical, parameter :: xi1a = all(maxloc(ia1) == 3)
-  logical, parameter :: xi1ar = rank(maxloc(ia1)) == 1
-  logical, parameter :: xi1ak = kind(maxloc(ia1, kind=2)) == 2
-  logical, parameter :: xi1ad = maxloc(ia1, dim=1) == 3
-  logical, parameter :: xi1adr = rank(maxloc(ia1, dim=1)) == 0
-  logical, parameter :: xi1b = all(maxloc(ia1, back=.true.) == 3)
-  logical, parameter :: xi1c = all(maxloc(ia1, mask=[.false., .true., .false., .true., .true.]) == 2)
-  logical, parameter :: xi1d = all(maxloc(ia1, mask=[.false., .true., .false., .true., .true.], back=.true.) == 4)
-
-  logical, parameter :: xi2a = all(maxloc(ia2) == [1, 2])
-  logical, parameter :: xi2ar = rank(maxloc(ia2)) == 1
-  logical, parameter :: xi2b = all(maxloc(ia2, back=.true.) == [2, 2])
-  logical, parameter :: xi2c = all(maxloc(ia2, mask=reshape([.false., .true., .true., .false., .true., .true.], shape(ia2))) == [2, 1])
-  logical, parameter :: xi2d = all(maxloc(ia2, mask=reshape([.false., .true., .true., .false., .true., .true.], shape(ia2)), back=.true.) == [1, 3])
-  logical, parameter :: xi2e = all(maxloc(ia2, dim=1) == [2, 1, 1])
-  logical, parameter :: xi2f = all(maxloc(ia2, dim=1, back=.true.) == [2, 2, 1])
-  logical, parameter :: xi2g = all(maxloc(ia2, dim=2) == [2, 2])
-
-  logical, parameter :: ni1a = all(minloc(ia1) == 1)
-  logical, parameter :: ni1ar = rank(minloc(ia1)) == 1
-  logical, parameter :: ni1ak = kind(minloc(ia1, kind=2)) == 2
-  logical, parameter :: ni1ad = minloc(ia1, dim=1) == 1
-  logical, parameter :: ni1adr = rank(minloc(ia1, dim=1)) == 0
-  logical, parameter :: ni1b = all(minloc(ia1, back=.true.) == 5)
-  logical, parameter :: ni1c = all(minloc(ia1, mask=[.false., .true., .true., .true., .false.]) == 2)
-  logical, parameter :: ni1d = all(minloc(ia1, mask=[.false., .true., .true., .true., .false.], back=.true.) == 4)
-
-  logical, parameter :: ni2a = all(minloc(ia2) == [1, 1])
-  logical, parameter :: ni2ar = rank(minloc(ia2)) == 1
-  logical, parameter :: ni2b = all(minloc(ia2, back=.true.) == [2, 3])
-  logical, parameter :: ni2c = all(minloc(ia2, mask=reshape([.false., .true., .true., .false., .true., .false.], shape(ia2))) == [2, 1])
-  logical, parameter :: ni2d = all(minloc(ia2, mask=reshape([.false., .true., .true., .false., .true., .false.], shape(ia2)), back=.true.) == [1, 3])
-  logical, parameter :: ni2e = all(minloc(ia2, dim=1) == [1, 1, 2])
-  logical, parameter :: ni2f = all(minloc(ia2, dim=1, back=.true.) == [1, 2, 2])
-  logical, parameter :: ni2g = all(minloc(ia2, dim=2) == [1, 3])
-end module