[flang] Support FINDLOC/MAXLOC/MINLOC with scalar mask

[FruitClover]

Summary

Previously MASK= elements were accessed in assumption that mask is an array of
input argument rank (and in combination with explicit DIM= argument we had
out-of-bounds access), but for MAXLOC/MINLOC/FINDLOC mask should be be
conformable and could be scalar.

Add new regression tests with scalar mask for verification.

[FruitClover]

Update comments

[FruitClover]

clang-format

flang/lib/Evaluate/fold-integer.cpp

const RelationalOperator relation{WHICH == WhichLocation::Findloc
             : WHICH == WhichLocation::Maxloc
             ? (back ? RelationalOperator::GE : RelationalOperator::GT)
             : back ? RelationalOperator::LE
                    : RelationalOperator::LT};
     // Use lower bounds of 1 exclusively.
     array->SetLowerBoundsToOne();
     ConstantSubscripts at{array->lbounds()}, maskAt, resultIndices, resultShape;
     if (mask) {
+      if (auto scalarMask{mask->GetScalarValue()}) {
+        // Convert into array in case of scalar MASK= (for
+        // MAXLOC/MINLOC/FINDLOC mask should be be conformable)
+        ConstantSubscript n{GetSize(array->shape())};
+        std::vector<Scalar<LogicalResult>> mask_elements(
+            n, Scalar<LogicalResult>{scalarMask.value()});
+        *mask = Constant<LogicalResult>{
+            std::move(mask_elements), ConstantSubscripts{n}};
+      }
       mask->SetLowerBoundsToOne();
       maskAt = mask->lbounds();
     }
     if (dim) { // DIM=
       if (*dim < 1 || *dim > array->Rank()) {
         context_.messages().Say("DIM=%d is out of range"_err_en_US, *dim);
         return std::nullopt;
       }

flang/test/Evaluate/fold-findloc.f90

module m1
   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]))
 
   character(1), parameter :: a(4) = ['a', 'b', 'a', 'b']
   logical, parameter :: test_char1 = all(maxloc(a).eq.[2])
   logical, parameter :: test_char2 = all(minloc(a).eq.[1])
   logical, parameter :: test_char3 = all(maxloc(a, back=.true.).eq.[4])
   logical, parameter :: test_char4 = all(minloc(a, back=.true.).eq.[3])
+
+  ! Check with scalar MASK=
+  logical, parameter:: test_mia1_mt  = all(minloc(ia1, mask=.true.) == 1)
+  logical, parameter:: test_mia1_mtd = all(minloc(ia1, mask=.true., dim=1) == [1])
+  logical, parameter:: test_xia1_mt  = all(maxloc(ia1, mask=.true.) == 3)
+  logical, parameter:: test_xia1_mtd = all(maxloc(ia1, mask=.true., dim=1) == [3])
+  logical, parameter:: test_fia1_mt  = all(findloc(ia1, 1, mask=.true.) == 1)
+  logical, parameter:: test_fia1_mtd = all(findloc(ia1, 1, mask=.true., dim=1) == [1])
+
+  logical, parameter:: test_mia1_mf  = all(minloc(ia1, mask=.false.) == 0)
+  logical, parameter:: test_mia1_mfd = all(minloc(ia1, mask=.false., dim=1) == [0])
+  logical, parameter:: test_xia1_mf  = all(maxloc(ia1, mask=.false.) == 0)
+  logical, parameter:: test_xia1_mfd = all(maxloc(ia1, mask=.false., dim=1) == [0])
+  logical, parameter:: test_fia1_mf  = all(findloc(ia1, 1, mask=.false.) == 0)
+  logical, parameter:: test_fia1_mfd = all(findloc(ia1, 1, mask=.false., dim=1) == [0])
 end module