[flang] Fold COUNT()

[klausler]

Summary

Complete folding of the intrinsic reduction function COUNT() for all
cases, including partial reductions with DIM= arguments.

[PeteSteinfeld]

Looks great!

flang/lib/Evaluate/CMakeLists.txt

   fold-integer.cpp
   fold-logical.cpp
   fold-real.cpp
+  fold-reduction.cpp
   formatting.cpp
   host.cpp
   initial-image.cpp

flang/lib/Evaluate/fold-implementation.h

     // Build and return constant result
     if constexpr (TR::category == TypeCategory::Character) {
       auto len{static_cast<ConstantSubscript>(
-          results.size() ? results[0].length() : 0)};
+          results.empty() ? 0 : results[0].length())};
       return Expr<TR>{Constant<TR>{len, std::move(results), std::move(shape)}};
     } else {
       return Expr<TR>{Constant<TR>{std::move(results), std::move(shape)}};

   if (constantArgs.size() != funcRef.arguments().size()) {
     return Expr<T>(std::move(funcRef));
   }
-  CHECK(constantArgs.size() > 0);
+  CHECK(!constantArgs.empty());
   Expr<T> result{std::move(*constantArgs[0])};
   for (std::size_t i{1}; i < constantArgs.size(); ++i) {
     Extremum<T> extremum{order, result, Expr<T>{std::move(*constantArgs[i])}};

     Expr<T> folded{Fold(context_, common::Clone(expr.value()))};
     if (const auto *c{UnwrapConstantValue<T>(folded)}) {
       // Copy elements in Fortran array element order
-      if (c->size() > 0) {
+      if (!c->empty()) {
         ConstantSubscripts index{c->lbounds()};
         do {
           elements_.emplace_back(c->At(index));

 std::optional<Expr<T>> AsFlatArrayConstructor(const Expr<T> &expr) {
   if (const auto *c{UnwrapConstantValue<T>(expr)}) {
     ArrayConstructor<T> result{expr};
-    if (c->size() > 0) {
+    if (!c->empty()) {
       ConstantSubscripts at{c->lbounds()};
       do {
         result.Push(Expr<T>{Constant<T>{c->At(at)}});

flang/lib/Evaluate/fold-integer.cpp

   return Expr<T>{std::move(funcRef)};
 }
 
+// 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) {
+      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)};
+}
+
 // 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);
-  using Element = Scalar<T>;
   std::optional<ConstantSubscript> 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) {
+    auto accumulator{[&](Scalar<T> &element, const ConstantSubscripts &at) {
       element = (element.*operation)(array->At(at));
     }};
-    return Expr<T>{DoReduction(*array, dim, identity, accumulator)};
+    return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
   }
   return Expr<T>{std::move(ref)};
 }

           cx->u);
     }
   } else if (name == "count") {
-    if (!args[1]) { // TODO: COUNT(x,DIM=d)
-      if (const auto *constant{UnwrapConstantValue<LogicalResult>(args[0])}) {
-        std::int64_t result{0};
-        for (const auto &element : constant->values()) {
-          if (element.IsTrue()) {
-            ++result;
-          }
-        }
-        return Expr<T>{result};
-      }
-    }
+    return FoldCount<T>(context, std::move(funcRef));
   } else if (name == "digits") {
     if (const auto *cx{UnwrapExpr<Expr<SomeInteger>>(args[0])}) {
       return Expr<T>{std::visit(

flang/lib/Evaluate/fold-logical.cpp

     auto accumulator{[&](Element &element, const ConstantSubscripts &at) {
       element = (element.*operation)(array->At(at));
     }};
-    return Expr<T>{DoReduction(*array, dim, identity, accumulator)};
+    return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
   }
   return Expr<T>{std::move(ref)};
 }

flang/lib/Evaluate/fold-reduction.h

 //
 //===----------------------------------------------------------------------===//
 
-// TODO: ALL, ANY, COUNT, DOT_PRODUCT, FINDLOC, IALL, IANY, IPARITY,
-// NORM2, MAXLOC, MINLOC, PARITY, PRODUCT, SUM
+// TODO: DOT_PRODUCT, FINDLOC, NORM2, MAXLOC, MINLOC, PARITY
 
 #ifndef FORTRAN_EVALUATE_FOLD_REDUCTION_H_
 #define FORTRAN_EVALUATE_FOLD_REDUCTION_H_

 
 namespace Fortran::evaluate {
 
+// Folds & validates a DIM= actual argument.
+std::optional<ConstantSubscript> CheckDIM(
+    FoldingContext &, std::optional<ActualArgument> &, int rank);
+
 // 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

     return std::nullopt;
   }
   if (dimIndex && arg.size() >= *dimIndex + 1 && arg[*dimIndex]) {
-    if (auto *dimConst{
-            Folder<SubscriptInteger>{context}.Folding(arg[*dimIndex])}) {
-      if (auto dimScalar{dimConst->GetScalarValue()}) {
-        dim.emplace(dimScalar->ToInt64());
-        if (*dim < 1 || *dim > folded->Rank()) {
-          context.messages().Say(
-              "DIM=%jd is not valid for an array of rank %d"_err_en_US,
-              static_cast<std::intmax_t>(*dim), folded->Rank());
-          dim.reset();
-        }
-      }
-    }
+    dim = CheckDIM(context, arg[*dimIndex], folded->Rank());
     if (!dim) {
       return std::nullopt;
     }

 
 // Generalized reduction to an array of one dimension fewer (w/ DIM=)
 // or to a scalar (w/o DIM=).
-template <typename T, typename ACCUMULATOR>
-static Constant<T> DoReduction(const Constant<T> &array,
+template <typename T, typename ACCUMULATOR, typename ARRAY>
+static Constant<T> DoReduction(const Constant<ARRAY> &array,
     std::optional<ConstantSubscript> &dim, const Scalar<T> &identity,
     ACCUMULATOR &accumulator) {
   ConstantSubscripts at{array.lbounds()};

         element = array->At(at);
       }
     }};
-    return Expr<T>{DoReduction(*array, dim, identity, accumulator)};
+    return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
   }
   return Expr<T>{std::move(ref)};
 }

       context.messages().Say(
           "PRODUCT() of %s data overflowed"_en_US, T::AsFortran());
     } else {
-      return Expr<T>{DoReduction(*array, dim, identity, accumulator)};
+      return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
     }
   }
   return Expr<T>{std::move(ref)};

       context.messages().Say(
           "SUM() of %s data overflowed"_en_US, T::AsFortran());
     } else {
-      return Expr<T>{DoReduction(*array, dim, identity, accumulator)};
+      return Expr<T>{DoReduction<T>(*array, dim, identity, accumulator)};
     }
   }
   return Expr<T>{std::move(ref)};

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)}) {
+      if (auto dimScalar{dimConst->GetScalarValue()}) {
+        auto dim{dimScalar->ToInt64()};
+        if (dim >= 1 && dim <= rank) {
+          return {dim};
+        } 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);
+        }
+      }
+    }
+  }
+  return std::nullopt;
+}
+
+} // namespace Fortran::evaluate

flang/test/Evaluate/folding29.f90

+! RUN: %python %S/test_folding.py %s %flang_fc1
+! Tests folding of COUNT()
+module m
+  logical, parameter :: arr(3,4) = reshape([(modulo(j, 2) == 1, j = 1, size(arr))], shape(arr))
+  logical, parameter :: test_1 = count([1, 2, 3, 2, 1] < [(j, j=1, 5)]) == 2
+  logical, parameter :: test_2 = count(arr) == 6
+  logical, parameter :: test_3 = all(count(arr, dim=1) == [2, 1, 2, 1])
+  logical, parameter :: test_4 = all(count(arr, dim=2) == [2, 2, 2])
+  logical, parameter :: test_5 = count(logical(arr, kind=1)) == 6
+  logical, parameter :: test_6 = count(logical(arr, kind=2)) == 6
+end module