[flang] Fold FINDLOC()

[klausler]

Summary

Fold the transformational intrinsic function FINDLOC() for
all combinations of optional arguments and data types.

flang/include/flang/Evaluate/constant.h

 public:
   ConstantBounds() = default;
   explicit ConstantBounds(const ConstantSubscripts &shape);
   explicit ConstantBounds(ConstantSubscripts &&shape);
   ~ConstantBounds();
   const ConstantSubscripts &shape() const { return shape_; }
   const ConstantSubscripts &lbounds() const { return lbounds_; }
   void set_lbounds(ConstantSubscripts &&);
+  void SetLowerBoundsToOne();
   int Rank() const { return GetRank(shape_); }
   Constant<SubscriptInteger> SHAPE() const;
 
   // If no optional dimension order argument is passed, increments a vector of
   // subscripts in Fortran array order (first dimension varying most quickly).
   // Otherwise, increments the vector of subscripts according to the given
   // dimension order (dimension dimOrder[0] varying most quickly; dimension
   // indexing is zero based here). Returns false when last element was visited.
public:
   std::optional<Scalar<T>> GetScalarValue() const {
     if (ConstantBounds::Rank() == 0) {
       return Base::values_.at(0);
     } else {
       return std::nullopt;
     }
   }
 
-  // Apply subscripts.  An empty subscript list is allowed for
-  // a scalar constant.
+  // Apply subscripts.  Excess subscripts are ignored, including the
+  // case of a scalar.
   Element At(const ConstantSubscripts &) const;
 
   Constant Reshape(ConstantSubscripts &&) const;
   std::size_t CopyFrom(const Constant &source, std::size_t count,
       ConstantSubscripts &resultSubscripts, const std::vector<int> *dimOrder);
 };
 
 template <int KIND>

flang/include/flang/Parser/provenance.h

 #include <utility>
 #include <variant>
 #include <vector>
 
 namespace Fortran::parser {
 
 // Each character in the contiguous source stream built by the
 // prescanner corresponds to a particular character in a source file,
-// include file, macro expansion, or compiler-inserted padding.
+// include file, macro expansion, or compiler-inserted text.
 // The location of this original character to which a parsable character
 // corresponds is its provenance.
 //
 // Provenances are offsets into an (unmaterialized) marshaling of the
 // entire contents of all the original source files, include files, macro
 // expansions, &c. for each visit to each source.  These origins of the
 // original source characters constitute a forest whose roots are
 // the original source files named on the compiler's command line.

flang/lib/Evaluate/constant.cpp

 //===-- lib/Evaluate/constant.cpp -----------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "flang/Evaluate/constant.h"
 #include "flang/Evaluate/expression.h"
 #include "flang/Evaluate/shape.h"
 #include "flang/Evaluate/type.h"
 #include <string>
 
 namespace Fortran::evaluate {
 
-std::size_t TotalElementCount(const ConstantSubscripts &shape) {
-  std::size_t size{1};
-  for (auto dim : shape) {
-    CHECK(dim >= 0);
-    size *= dim;
-  }
-  return size;
-}
-
 ConstantBounds::ConstantBounds(const ConstantSubscripts &shape)
     : shape_(shape), lbounds_(shape_.size(), 1) {}
 
 ConstantBounds::ConstantBounds(ConstantSubscripts &&shape)
     : shape_(std::move(shape)), lbounds_(shape_.size(), 1) {}
 
 ConstantBounds::~ConstantBounds() = default;
 
 void ConstantBounds::set_lbounds(ConstantSubscripts &&lb) {
   CHECK(lb.size() == shape_.size());
   lbounds_ = std::move(lb);
 }
 
+void ConstantBounds::SetLowerBoundsToOne() {
+  for (auto &n : lbounds_) {
+    n = 1;
+  }
+}
+
 Constant<SubscriptInteger> ConstantBounds::SHAPE() const {
   return AsConstantShape(shape_);
 }
 
 ConstantSubscript ConstantBounds::SubscriptsToOffset(
     const ConstantSubscripts &index) const {
   CHECK(GetRank(index) == GetRank(shape_));
   ConstantSubscript stride{1}, offset{0};
   int dim{0};
   for (auto j : index) {
     auto lb{lbounds_[dim]};
     auto extent{shape_[dim++]};
     CHECK(j >= lb && j < lb + extent);
     offset += stride * (j - lb);
     stride *= extent;
   }
   return offset;
 }
 
+std::size_t TotalElementCount(const ConstantSubscripts &shape) {
+  return static_cast<std::size_t>(GetSize(shape));
+}
+
 bool ConstantBounds::IncrementSubscripts(
     ConstantSubscripts &indices, const std::vector<int> *dimOrder) const {
   int rank{GetRank(shape_)};
   CHECK(GetRank(indices) == rank);
   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]};

flang/lib/Evaluate/fold-character.cpp

if (name == "achar" || name == "char") {
     }
   } else if (name == "trim") { // not elemental
     if (auto scalar{
             GetScalarConstantArguments<T>(context, funcRef.arguments())}) {
       return Expr<T>{Constant<T>{
           CharacterUtils<KIND>::TRIM(std::get<Scalar<T>>(*scalar))}};
     }
   }
-  // TODO: findloc, maxloc, minloc, transfer
+  // TODO: maxloc, minloc, transfer
   return Expr<T>{std::move(funcRef)};
 }
 
 template <int KIND>
 Expr<Type<TypeCategory::Character, KIND>> FoldOperation(
     FoldingContext &context, Concat<KIND> &&x) {
   if (auto array{ApplyElementwise(context, x)}) {
     return *array;

flang/lib/Evaluate/fold-integer.cpp

 // COUNT()
 template <typename T>
 static Expr<T> FoldCount(FoldingContext &context, FunctionRef<T> &&ref) {
   static_assert(T::category == TypeCategory::Integer);
   ActualArguments &arg{ref.arguments()};
   if (const Constant<LogicalResult> *mask{arg.empty()
               ? nullptr
               : Folder<LogicalResult>{context}.Folding(arg[0])}) {
-    std::optional<ConstantSubscript> dim;
-    if (arg.size() > 1 && arg[1]) {
-      dim = CheckDIM(context, arg[1], mask->Rank());
-      if (!dim) {
-        mask = nullptr;
-      }
-    }
-    if (mask) {
+    std::optional<int> dim;
+    if (CheckReductionDIM(dim, context, arg, 1, mask->Rank())) {
       auto accumulator{[&](Scalar<T> &element, const ConstantSubscripts &at) {
         if (mask->At(at).IsTrue()) {
           element = element.AddSigned(Scalar<T>{1}).value;
         }
       }};
       return Expr<T>{DoReduction<T>(*mask, dim, Scalar<T>{}, accumulator)};
     }
   }
   return Expr<T>{std::move(ref)};
 }
 
+// FINDLOC()
+class FindlocHelper {
+public:
+  FindlocHelper(
+      DynamicType &&type, ActualArguments &arg, FoldingContext &context)
+      : type_{type}, arg_{arg}, context_{context} {}
+  using Result = std::optional<Constant<SubscriptInteger>>;
+  using Types = AllIntrinsicTypes;
+
+  template <typename T> Result Test() const {
+    if (T::category != type_.category() || T::kind != type_.kind()) {
+      return std::nullopt;
+    }
+    CHECK(arg_.size() == 6);
+    Folder<T> folder{context_};
+    Constant<T> *array{folder.Folding(arg_[0])};
+    Constant<T> *value{folder.Folding(arg_[1])};
+    if (!array || !value) {
+      return std::nullopt;
+    }
+    std::optional<int> dim;
+    Constant<LogicalResult> *mask{
+        GetReductionMASK(arg_[3], array->shape(), context_)};
+    if ((!mask && arg_[3]) ||
+        !CheckReductionDIM(dim, context_, arg_, 2, array->Rank())) {
+      return std::nullopt;
+    }
+    bool back{false};
+    if (arg_[5]) {
+      const auto *backConst{Folder<LogicalResult>{context_}.Folding(arg_[5])};
+      if (backConst) {
+        back = backConst->GetScalarValue().value().IsTrue();
+      } else {
+        return std::nullopt;
+      }
+    }
+    // Use lower bounds of 1 exclusively.
+    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);
+        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};
+        for (ConstantSubscript k{0}; k < dimLength;
+             ++k, ++at[zbDim], mask && ++maskAt[zbDim]) {
+          if ((!mask || mask->At(maskAt).IsTrue()) &&
+              IsHit(array->At(at), *value)) {
+            hit = at[zbDim];
+            if (!back) {
+              break;
+            }
+          }
+        }
+        resultIndices.emplace_back(hit);
+        at[zbDim] = array->lbounds()[zbDim] + dimLength - 1;
+        array->IncrementSubscripts(at);
+        at[zbDim] = array->lbounds()[zbDim];
+        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)) {
+          resultIndices = at;
+          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, Constant<T> value) const {
+    std::optional<Expr<LogicalResult>> cmp;
+    if constexpr (T::category == TypeCategory::Logical) {
+      // array(at) .EQV. value?
+      cmp.emplace(
+          ConvertToType<LogicalResult>(Expr<T>{LogicalOperation<T::kind>{
+              LogicalOperator::Eqv, Expr<T>{Constant<T>{std::move(element)}},
+              Expr<T>{std::move(value)}}}));
+    } else { // array(at) .EQ. value?
+      cmp.emplace(PackageRelation(RelationalOperator::EQ,
+          Expr<T>{Constant<T>{std::move(element)}}, Expr<T>{std::move(value)}));
+    }
+    Expr<LogicalResult> folded{Fold(context_, std::move(*cmp))};
+    return GetScalarConstantValue<LogicalResult>(folded).value().IsTrue();
+  }
+
+  DynamicType type_;
+  ActualArguments &arg_;
+  FoldingContext &context_;
+};
+
+static std::optional<Constant<SubscriptInteger>> FoldFindlocCall(
+    ActualArguments &arg, FoldingContext &context) {
+  CHECK(arg.size() == 6);
+  if (arg[0]) {
+    if (auto type{arg[0]->GetType()}) {
+      return common::SearchTypes(FindlocHelper{std::move(*type), arg, context});
+    }
+  }
+  return std::nullopt;
+}
+
+template <typename T>
+static Expr<T> FoldFindloc(FoldingContext &context, FunctionRef<T> &&ref) {
+  static_assert(T::category == TypeCategory::Integer);
+  if (std::optional<Constant<SubscriptInteger>> found{
+          FoldFindlocCall(ref.arguments(), context)}) {
+    return Expr<T>{Fold(
+        context, ConvertToType<T>(Expr<SubscriptInteger>{std::move(*found)}))};
+  } else {
+    return Expr<T>{std::move(ref)};
+  }
+}
+
 // for IALL, IANY, & IPARITY
 template <typename T>
 static Expr<T> FoldBitReduction(FoldingContext &context, FunctionRef<T> &&ref,
     Scalar<T> (Scalar<T>::*operation)(const Scalar<T> &) const,
     Scalar<T> identity) {
   static_assert(T::category == TypeCategory::Integer);
-  std::optional<ConstantSubscript> dim;
+  std::optional<int> dim;
   if (std::optional<Constant<T>> array{
           ProcessReductionArgs<T>(context, ref.arguments(), dim, identity,
               /*ARRAY=*/0, /*DIM=*/1, /*MASK=*/2)}) {
     auto accumulator{[&](Scalar<T> &element, const ConstantSubscripts &at) {
       element = (element.*operation)(array->At(at));
     }};
     return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
   }
if (auto *sx{UnwrapExpr<Expr<SomeReal>>(args[0])}) {
             using TR = typename std::decay_t<decltype(x)>::Result;
             return FoldElementalIntrinsic<T, TR>(context, std::move(funcRef),
                 &Scalar<TR>::template EXPONENT<Scalar<T>>);
           },
           sx->u);
     } else {
       DIE("exponent argument must be real");
     }
+  } else if (name == "findloc") {
+    return FoldFindloc<T>(context, std::move(funcRef));
   } else if (name == "huge") {
     return Expr<T>{Scalar<T>::HUGE()};
   } else if (name == "iachar" || name == "ichar") {
     auto *someChar{UnwrapExpr<Expr<SomeCharacter>>(args[0])};
     CHECK(someChar);
     if (auto len{ToInt64(someChar->LEN())}) {
       if (len.value() != 1) {
         // Do not die, this was not checked before
if (auto info{
             Fold(context, Expr<T>{8} * ConvertToType<T>(std::move(*bytes)))};
       }
     }
   } else if (name == "sum") {
     return FoldSum<T>(context, std::move(funcRef));
   } else if (name == "ubound") {
     return UBOUND(context, std::move(funcRef));
   }
-  // TODO: dot_product, findloc, ibits, image_status, ishftc,
+  // TODO: dot_product, ibits, image_status, ishftc,
   // matmul, maxloc, minloc, sign, transfer
   return Expr<T>{std::move(funcRef)};
 }
 
 // Substitutes a bare type parameter reference with its value if it has one now
 // in an instantiation.  Bare LEN type parameters are substituted only when
 // the known value is constant.
 Expr<TypeParamInquiry::Result> FoldOperation(

flang/lib/Evaluate/fold-logical.cpp

 
 // for ALL & ANY
 template <typename T>
 static Expr<T> FoldAllAny(FoldingContext &context, FunctionRef<T> &&ref,
     Scalar<T> (Scalar<T>::*operation)(const Scalar<T> &) const,
     Scalar<T> identity) {
   static_assert(T::category == TypeCategory::Logical);
   using Element = Scalar<T>;
-  std::optional<ConstantSubscript> dim;
+  std::optional<int> dim;
   if (std::optional<Constant<T>> array{
           ProcessReductionArgs<T>(context, ref.arguments(), dim, identity,
               /*ARRAY(MASK)=*/0, /*DIM=*/1)}) {
     auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
       element = (element.*operation)(array->At(at));
     }};
     return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
   }

flang/lib/Evaluate/fold-reduction.h

 //===-- lib/Evaluate/fold-reduction.h -------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-// TODO: DOT_PRODUCT, FINDLOC, NORM2, MAXLOC, MINLOC, PARITY
+// TODO: DOT_PRODUCT, NORM2, MAXLOC, MINLOC, PARITY
 
 #ifndef FORTRAN_EVALUATE_FOLD_REDUCTION_H_
 #define FORTRAN_EVALUATE_FOLD_REDUCTION_H_
 
 #include "fold-implementation.h"
 
 namespace Fortran::evaluate {
 
-// Folds & validates a DIM= actual argument.
-std::optional<ConstantSubscript> CheckDIM(
-    FoldingContext &, std::optional<ActualArgument> &, int rank);
+// Fold and validate a DIM= argument.  Returns false on error.
+bool CheckReductionDIM(std::optional<int> &dim, FoldingContext &,
+    ActualArguments &, std::optional<int> dimIndex, int rank);
+
+// Fold and validate a MASK= argument.  Return null on error, absent MASK=, or
+// non-constant MASK=.
+Constant<LogicalResult> *GetReductionMASK(
+    std::optional<ActualArgument> &maskArg, const ConstantSubscripts &shape,
+    FoldingContext &);
 
 // Common preprocessing for reduction transformational intrinsic function
 // folding.  If the intrinsic can have DIM= &/or MASK= arguments, extract
 // and check them.  If a MASK= is present, apply it to the array data and
 // substitute identity values for elements corresponding to .FALSE. in
 // the mask.  If the result is present, the intrinsic call can be folded.
 template <typename T>
 static std::optional<Constant<T>> ProcessReductionArgs(FoldingContext &context,
-    ActualArguments &arg, std::optional<ConstantSubscript> &dim,
-    const Scalar<T> &identity, int arrayIndex,
-    std::optional<std::size_t> dimIndex = std::nullopt,
-    std::optional<std::size_t> maskIndex = std::nullopt) {
+    ActualArguments &arg, std::optional<int> &dim, const Scalar<T> &identity,
+    int arrayIndex, std::optional<int> dimIndex = std::nullopt,
+    std::optional<int> maskIndex = std::nullopt) {
   if (arg.empty()) {
     return std::nullopt;
   }
   Constant<T> *folded{Folder<T>{context}.Folding(arg[arrayIndex])};
   if (!folded || folded->Rank() < 1) {
     return std::nullopt;
   }
-  if (dimIndex && arg.size() >= *dimIndex + 1 && arg[*dimIndex]) {
-    dim = CheckDIM(context, arg[*dimIndex], folded->Rank());
-    if (!dim) {
-      return std::nullopt;
-    }
-  }
-  if (maskIndex && arg.size() >= *maskIndex + 1 && arg[*maskIndex]) {
-    if (Constant<LogicalResult> *
-        mask{Folder<LogicalResult>{context}.Folding(arg[*maskIndex])}) {
-      if (CheckConformance(context.messages(), AsShape(folded->shape()),
-              AsShape(mask->shape()),
-              CheckConformanceFlags::RightScalarExpandable, "ARRAY=", "MASK=")
-              .value_or(false)) {
-        // Apply the mask in place to the array
-        std::size_t n{folded->size()};
-        std::vector<typename Constant<T>::Element> elements;
-        if (auto scalarMask{mask->GetScalarValue()}) {
-          if (scalarMask->IsTrue()) {
-            return Constant<T>{*folded};
-          } else { // MASK=.FALSE.
-            elements = std::vector<typename Constant<T>::Element>(n, identity);
-          }
-        } else { // mask is an array; test its elements
-          elements = std::vector<typename Constant<T>::Element>(n, identity);
-          ConstantSubscripts at{folded->lbounds()};
-          for (std::size_t j{0}; j < n; ++j, folded->IncrementSubscripts(at)) {
-            if (mask->values()[j].IsTrue()) {
-              elements[j] = folded->At(at);
-            }
-          }
-        }
-        if constexpr (T::category == TypeCategory::Character) {
-          return Constant<T>{static_cast<ConstantSubscript>(identity.size()),
-              std::move(elements), ConstantSubscripts{folded->shape()}};
-        } else {
-          return Constant<T>{
-              std::move(elements), ConstantSubscripts{folded->shape()}};
-        }
-      } else {
-        return std::nullopt;
-      }
-    } else {
-      return std::nullopt;
-    }
-  } else {
+  if (!CheckReductionDIM(dim, context, arg, dimIndex, folded->Rank())) {
+  if (maskIndex && static_cast<std::size_t>(*maskIndex) < arg.size() &&
+      arg[*maskIndex]) {
+    if (const Constant<LogicalResult> *mask{
+            GetReductionMASK(arg[*maskIndex], folded->shape(), context)}) {
     return Constant<T>{*folded};
   }
 }
 
 // Generalized reduction to an array of one dimension fewer (w/ DIM=)
 // or to a scalar (w/o DIM=).
 template <typename T, typename ACCUMULATOR, typename ARRAY>
 static Constant<T> DoReduction(const Constant<ARRAY> &array,
-    std::optional<ConstantSubscript> &dim, const Scalar<T> &identity,
+    std::optional<int> &dim, const Scalar<T> &identity,
     ACCUMULATOR &accumulator) {
   ConstantSubscripts at{array.lbounds()};
   std::vector<typename Constant<T>::Element> elements;
   ConstantSubscripts resultShape; // empty -> scalar
   if (dim) { // DIM= is present, so result is an array
     resultShape = array.shape();
     resultShape.erase(resultShape.begin() + (*dim - 1));
     ConstantSubscript dimExtent{array.shape().at(*dim - 1)};

 // MAXVAL & MINVAL
 template <typename T>
 static Expr<T> FoldMaxvalMinval(FoldingContext &context, FunctionRef<T> &&ref,
     RelationalOperator opr, const Scalar<T> &identity) {
   static_assert(T::category == TypeCategory::Integer ||
       T::category == TypeCategory::Real ||
       T::category == TypeCategory::Character);
   using Element = Scalar<T>;
-  std::optional<ConstantSubscript> dim;
+  std::optional<int> dim;
   if (std::optional<Constant<T>> array{
           ProcessReductionArgs<T>(context, ref.arguments(), dim, identity,
               /*ARRAY=*/0, /*DIM=*/1, /*MASK=*/2)}) {
     auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
       Expr<LogicalResult> test{PackageRelation(opr,
           Expr<T>{Constant<T>{array->At(at)}}, Expr<T>{Constant<T>{element}})};
       auto folded{GetScalarConstantValue<LogicalResult>(
           test.Rewrite(context, std::move(test)))};

 // PRODUCT
 template <typename T>
 static Expr<T> FoldProduct(
     FoldingContext &context, FunctionRef<T> &&ref, Scalar<T> identity) {
   static_assert(T::category == TypeCategory::Integer ||
       T::category == TypeCategory::Real ||
       T::category == TypeCategory::Complex);
   using Element = typename Constant<T>::Element;
-  std::optional<ConstantSubscript> dim;
+  std::optional<int> dim;
   if (std::optional<Constant<T>> array{
           ProcessReductionArgs<T>(context, ref.arguments(), dim, identity,
               /*ARRAY=*/0, /*DIM=*/1, /*MASK=*/2)}) {
     bool overflow{false};
     auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
       if constexpr (T::category == TypeCategory::Integer) {
         auto prod{element.MultiplySigned(array->At(at))};
         overflow |= prod.SignedMultiplicationOverflowed();

 
 // SUM
 template <typename T>
 static Expr<T> FoldSum(FoldingContext &context, FunctionRef<T> &&ref) {
   static_assert(T::category == TypeCategory::Integer ||
       T::category == TypeCategory::Real ||
       T::category == TypeCategory::Complex);
   using Element = typename Constant<T>::Element;
-  std::optional<ConstantSubscript> dim;
+  std::optional<int> dim;
   Element identity{}, correction{};
   if (std::optional<Constant<T>> array{
           ProcessReductionArgs<T>(context, ref.arguments(), dim, identity,
               /*ARRAY=*/0, /*DIM=*/1, /*MASK=*/2)}) {
     bool overflow{false};
     auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
       if constexpr (T::category == TypeCategory::Integer) {
         auto sum{element.AddSigned(array->At(at))};

flang/lib/Evaluate/fold-reduction.cpp

 //===-- lib/Evaluate/fold-reduction.cpp -----------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "fold-reduction.h"
 
 namespace Fortran::evaluate {
-
-std::optional<ConstantSubscript> CheckDIM(
-    FoldingContext &context, std::optional<ActualArgument> &arg, int rank) {
-  if (arg) {
-    if (auto *dimConst{Folder<SubscriptInteger>{context}.Folding(arg)}) {
+bool CheckReductionDIM(std::optional<int> &dim, FoldingContext &context,
+    ActualArguments &arg, std::optional<int> dimIndex, int rank) {
+  if (dimIndex && static_cast<std::size_t>(*dimIndex) < arg.size()) {
+    if (auto *dimConst{
+            Folder<SubscriptInteger>{context}.Folding(arg[*dimIndex])}) {
       if (auto dimScalar{dimConst->GetScalarValue()}) {
-        auto dim{dimScalar->ToInt64()};
-        if (dim >= 1 && dim <= rank) {
-          return {dim};
+        auto dimVal{dimScalar->ToInt64()};
+        if (dimVal >= 1 && dimVal <= rank) {
+          dim = dimVal;
         } else {
           context.messages().Say(
               "DIM=%jd is not valid for an array of rank %d"_err_en_US,
-              static_cast<std::intmax_t>(dim), rank);
+              static_cast<std::intmax_t>(dimVal), rank);
+          return false;
         }
       }
     }
   }
-  return std::nullopt;
+  return true;
 }
 
+Constant<LogicalResult> *GetReductionMASK(
+    std::optional<ActualArgument> &maskArg, const ConstantSubscripts &shape,
+    FoldingContext &context) {
+  Constant<LogicalResult> *mask{
+      Folder<LogicalResult>{context}.Folding(maskArg)};
+  if (mask &&
+      !CheckConformance(context.messages(), AsShape(shape),
+          AsShape(mask->shape()), CheckConformanceFlags::RightScalarExpandable,
+          "ARRAY=", "MASK=")
+           .value_or(false)) {
+    mask = nullptr;
+  }
+  return mask;
+}
 } // namespace Fortran::evaluate

flang/lib/Evaluate/shape.cpp

if (dim) {
       return std::nullopt;
     }
   }
   return extent;
 }
 
 ConstantSubscript GetSize(const ConstantSubscripts &shape) {
   ConstantSubscript size{1};
-  for (auto dim : std::move(shape)) {
+  for (auto dim : shape) {
+    CHECK(dim >= 0);
     size *= dim;
   }
   return size;
 }
 
 bool ContainsAnyImpliedDoIndex(const ExtentExpr &expr) {
   struct MyVisitor : public AnyTraverse<MyVisitor> {
     using Base = AnyTraverse<MyVisitor>;
if (intrinsic->name == "shape" || intrinsic->name == "lbound" ||
             if (*dim >= 1 &&
                 static_cast<std::size_t>(*dim) <= arrayShape->size()) {
               arrayShape->erase(arrayShape->begin() + (*dim - 1));
               return std::move(*arrayShape);
             }
           }
         }
       }
-    } else if (intrinsic->name == "maxloc" || intrinsic->name == "minloc") {
-      // TODO: FINDLOC
-      if (call.arguments().size() >= 2) {
+    } else if (intrinsic->name == "findloc" || intrinsic->name == "maxloc" ||
+        intrinsic->name == "minloc") {
+      std::size_t dimIndex{intrinsic->name == "findloc" ? 2u : 1u};
+      if (call.arguments().size() > dimIndex) {
         if (auto arrayShape{
                 (*this)(UnwrapExpr<Expr<SomeType>>(call.arguments().at(0)))}) {
           auto rank{static_cast<int>(arrayShape->size())};
           if (const auto *dimArg{
-                  UnwrapExpr<Expr<SomeType>>(call.arguments()[1])}) {
+                  UnwrapExpr<Expr<SomeType>>(call.arguments()[dimIndex])}) {
             auto dim{ToInt64(*dimArg)};
             if (dim && *dim >= 1 && *dim <= rank) {
               arrayShape->erase(arrayShape->begin() + (*dim - 1));
               return std::move(*arrayShape);
             }
           } else {
             // xxxLOC(no DIM=) result is vector(1:RANK(ARRAY=))
             return Shape{ExtentExpr{rank}};

flang/test/Evaluate/folding30.f90

+! RUN: %python %S/test_folding.py %s %flang_fc1
+! Tests folding of FINDLOC
+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 :: ti1a = all(findloc(ia1, 1) == 1)
+  logical, parameter :: ti1ar = rank(findloc(ia1, 1)) == 1
+  logical, parameter :: ti1ak = kind(findloc(ia1, 1, kind=2)) == 2
+  logical, parameter :: ti1ad = findloc(ia1, 1, dim=1) == 1
+  logical, parameter :: ti1adr = rank(findloc(ia1, 1, dim=1)) == 0
+  logical, parameter :: ti1b = all(findloc(ia1, 1, back=.true.) == 5)
+  logical, parameter :: ti1c = all(findloc(ia1, 2, mask=[.false., .false., .true., .true., .true.]) == 4)
+
+  logical, parameter :: ti2a = all(findloc(ia2, 1) == [1, 1])
+  logical, parameter :: ti2ar = rank(findloc(ia2, 1)) == 1
+  logical, parameter :: ti2b = all(findloc(ia2, 1, back=.true.) == [2, 3])
+  logical, parameter :: ti2c = all(findloc(ia2, 2, mask=reshape([.false., .false., .true., .true., .true., .false.], shape(ia2))) == [1, 3])
+  logical, parameter :: ti2d = all(findloc(ia2, 1, dim=1) == [1, 0, 2])
+  logical, parameter :: ti2e = all(findloc(ia2, 1, dim=2) == [1, 3])
+end module