[flang] Fold PACK()

[klausler]

Summary

Implement compile-time constant folding for the transformational
intrinsic function PACK.

[PeteSteinfeld]

Very clever test!

All builds, tests, and looks good.

flang/lib/Evaluate/fold-implementation.h

public:
   std::optional<Constant<T>> GetConstantComponent(
       Component &, const std::vector<Constant<SubscriptInteger>> * = nullptr);
   std::optional<Constant<T>> Folding(ArrayRef &);
   Expr<T> Folding(Designator<T> &&);
   Constant<T> *Folding(std::optional<ActualArgument> &);
 
   Expr<T> CSHIFT(FunctionRef<T> &&);
   Expr<T> EOSHIFT(FunctionRef<T> &&);
+  Expr<T> PACK(FunctionRef<T> &&);
   Expr<T> RESHAPE(FunctionRef<T> &&);
 
 private:
   FoldingContext &context_;
 };
 
 std::optional<Constant<SubscriptInteger>> GetConstantSubscript(
     FoldingContext &, Subscript &, const NamedEntity &, int dim);
if (*dim < 1 || *dim > array->Rank()) {
     int zbDim{static_cast<int>(*dim) - 1};
     bool ok{true};
     if (shift->Rank() > 0) {
       int k{0};
       for (int j{0}; j < rank; ++j) {
         if (j != zbDim) {
           if (array->shape()[j] != shift->shape()[k]) {
             context_.messages().Say(
-                "Invalid 'shift=' argument in CSHIFT; extent on dimension %d is %jd but must be %jd"_err_en_US,
+                "Invalid 'shift=' argument in CSHIFT: extent on dimension %d is %jd but must be %jd"_err_en_US,
                 k + 1, static_cast<std::intmax_t>(shift->shape()[k]),
                 static_cast<std::intmax_t>(array->shape()[j]));
             ok = false;
           }
           ++k;
         }
       }
     }
template <typename T> Expr<T> Folder<T>::EOSHIFT(FunctionRef<T> &&funcRef) {
   }
   // Arguments are constant
   if (*dim < 1 || *dim > array->Rank()) {
     context_.messages().Say(
         "Invalid 'dim=' argument (%jd) in EOSHIFT"_err_en_US,
         static_cast<std::intmax_t>(*dim));
   } else if (shift->Rank() > 0 && shift->Rank() != array->Rank() - 1) {
     // message already emitted from intrinsic look-up
+  } else if (boundary && boundary->Rank() > 0 &&
+      boundary->Rank() != array->Rank() - 1) {
+    // ditto
   } else {
     int rank{array->Rank()};
     int zbDim{static_cast<int>(*dim) - 1};
     bool ok{true};
     if (shift->Rank() > 0) {
       int k{0};
       for (int j{0}; j < rank; ++j) {
         if (j != zbDim) {
           if (array->shape()[j] != shift->shape()[k]) {
             context_.messages().Say(
-                "Invalid 'shift=' argument in EOSHIFT; extent on dimension %d is %jd but must be %jd"_err_en_US,
+                "Invalid 'shift=' argument in EOSHIFT: extent on dimension %d is %jd but must be %jd"_err_en_US,
                 k + 1, static_cast<std::intmax_t>(shift->shape()[k]),
                 static_cast<std::intmax_t>(array->shape()[j]));
             ok = false;
           }
-          if (boundary && array->shape()[j] != boundary->shape()[k]) {
+          ++k;
+        }
+      }
+    }
+    if (boundary && boundary->Rank() > 0) {
+      int k{0};
+      for (int j{0}; j < rank; ++j) {
+        if (j != zbDim) {
+          if (array->shape()[j] != boundary->shape()[k]) {
             context_.messages().Say(
-                "Invalid 'boundary=' argument in EOSHIFT; extent on dimension %d is %jd but must be %jd"_err_en_US,
-                k + 1, static_cast<std::intmax_t>(shift->shape()[k]),
+                "Invalid 'boundary=' argument in EOSHIFT: extent on dimension %d is %jd but must be %jd"_err_en_US,
+                k + 1, static_cast<std::intmax_t>(boundary->shape()[k]),
                 static_cast<std::intmax_t>(array->shape()[j]));
             ok = false;
           }
           ++k;
         }
       }
     }
     if (ok) {
if (ok) {
       return Expr<T>{PackageConstant<T>(
           std::move(resultElements), *array, array->shape())};
     }
   }
   // Invalid, prevent re-folding
   return MakeInvalidIntrinsic(std::move(funcRef));
 }
 
+template <typename T> Expr<T> Folder<T>::PACK(FunctionRef<T> &&funcRef) {
+  auto args{funcRef.arguments()};
+  CHECK(args.size() == 3);
+  const auto *array{UnwrapConstantValue<T>(args[0])};
+  const auto *vector{UnwrapConstantValue<T>(args[2])};
+  auto convertedMask{Fold(context_,
+      ConvertToType<LogicalResult>(
+          Expr<SomeLogical>{DEREF(UnwrapExpr<Expr<SomeLogical>>(args[1]))}))};
+  const auto *mask{UnwrapConstantValue<LogicalResult>(convertedMask)};
+  if (!array || !mask || (args[2] && !vector)) {
+    return Expr<T>{std::move(funcRef)};
+  }
+  // Arguments are constant.
+  ConstantSubscript arrayElements{GetSize(array->shape())};
+  ConstantSubscript truths{0};
+  ConstantSubscripts maskAt{mask->lbounds()};
+  if (mask->Rank() == 0) {
+    if (mask->At(maskAt).IsTrue()) {
+      truths = arrayElements;
+    }
+  } else if (array->shape() != mask->shape()) {
+    // Error already emitted from intrinsic processing
+    return MakeInvalidIntrinsic(std::move(funcRef));
+  } else {
+    for (ConstantSubscript j{0}; j < arrayElements;
+         ++j, mask->IncrementSubscripts(maskAt)) {
+      if (mask->At(maskAt).IsTrue()) {
+        ++truths;
+      }
+    }
+  }
+  std::vector<Scalar<T>> resultElements;
+  ConstantSubscripts arrayAt{array->lbounds()};
+  ConstantSubscript resultSize{truths};
+  if (vector) {
+    resultSize = vector->shape().at(0);
+    if (resultSize < truths) {
+      context_.messages().Say(
+          "Invalid 'vector=' argument in PACK: the 'mask=' argument has %jd true elements, but the vector has only %jd elements"_err_en_US,
+          static_cast<std::intmax_t>(truths),
+          static_cast<std::intmax_t>(resultSize));
+      return MakeInvalidIntrinsic(std::move(funcRef));
+    }
+  }
+  for (ConstantSubscript j{0}; j < truths;) {
+    if (mask->At(maskAt).IsTrue()) {
+      resultElements.push_back(array->At(arrayAt));
+      ++j;
+    }
+    array->IncrementSubscripts(arrayAt);
+    mask->IncrementSubscripts(maskAt);
+  }
+  if (vector) {
+    ConstantSubscripts vectorAt{vector->lbounds()};
+    vectorAt.at(0) += truths;
+    for (ConstantSubscript j{truths}; j < resultSize; ++j) {
+      resultElements.push_back(vector->At(vectorAt));
+      ++vectorAt[0];
+    }
+  }
+  return Expr<T>{PackageConstant<T>(std::move(resultElements), *array,
+      ConstantSubscripts{static_cast<ConstantSubscript>(resultSize)})};
+}
+
 template <typename T> Expr<T> Folder<T>::RESHAPE(FunctionRef<T> &&funcRef) {
   auto args{funcRef.arguments()};
   CHECK(args.size() == 4);
   const auto *source{UnwrapConstantValue<T>(args[0])};
   const auto *pad{UnwrapConstantValue<T>(args[2])};
   std::optional<std::vector<ConstantSubscript>> shape{
       GetIntegerVector<ConstantSubscript>(args[1])};
   std::optional<std::vector<int>> order{GetIntegerVector<int>(args[3])};
for (std::optional<ActualArgument> &arg : args) {
     }
   }
   if (auto *intrinsic{std::get_if<SpecificIntrinsic>(&funcRef.proc().u)}) {
     const std::string name{intrinsic->name};
     if (name == "cshift") {
       return Folder<T>{context}.CSHIFT(std::move(funcRef));
     } else if (name == "eoshift") {
       return Folder<T>{context}.EOSHIFT(std::move(funcRef));
+    } else if (name == "pack") {
+      return Folder<T>{context}.PACK(std::move(funcRef));
     } else if (name == "reshape") {
       return Folder<T>{context}.RESHAPE(std::move(funcRef));
     }
-    // TODO: eoshift, pack, spread, unpack, transpose
+    // TODO: spread, unpack, transpose
     // TODO: extends_type_of, same_type_as
     if constexpr (!std::is_same_v<T, SomeDerived>) {
       return FoldIntrinsicFunction(context, std::move(funcRef));
     }
   }
   return Expr<T>{std::move(funcRef)};
 }
 

flang/test/Evaluate/folding19.f90

subroutine s1(a,b)
     integer :: ub2(ubound(a,0))
     !CHECK: error: DIM=2 dimension is out of range for rank-1 array
     integer :: ub3(ubound(a,2))
     !CHECK: error: DIM=0 dimension is out of range for rank-1 array
     integer :: lb2(lbound(b,0))
     !CHECK: error: DIM=2 dimension is out of range for rank-1 array
     integer :: lb3(lbound(b,2))
   end subroutine
+  subroutine s2
+    integer, parameter :: array(2,3) = reshape([(j, j=1, 6)], shape(array))
+    integer :: x(2, 3)
+    !CHECK: error: Invalid 'dim=' argument (0) in CSHIFT
+    x = cshift(array, [1, 2], dim=0)
+    !CHECK: error: Invalid 'shift=' argument in CSHIFT: extent on dimension 1 is 2 but must be 3
+    x = cshift(array, [1, 2], dim=1)
+  end subroutine
+  subroutine s3
+    integer, parameter :: array(2,3) = reshape([(j, j=1, 6)], shape(array))
+    integer :: x(2, 3)
+    !CHECK: error: Invalid 'dim=' argument (0) in EOSHIFT
+    x = eoshift(array, [1, 2], dim=0)
+    !CHECK: error: Invalid 'shift=' argument in EOSHIFT: extent on dimension 1 is 2 but must be 3
+    x = eoshift(array, [1, 2], dim=1)
+    !CHECK: error: Invalid 'boundary=' argument in EOSHIFT: extent on dimension 1 is 3 but must be 2
+    x = eoshift(array, 1, [0, 0, 0], 2)
+  end subroutine
+  subroutine s4
+    integer, parameter :: array(2,3) = reshape([(j, j=1, 6)], shape(array))
+    logical, parameter :: mask(*,*) = reshape([(.true., j=1,3),(.false., j=1,3)], shape(array))
+    integer :: x(3)
+    !CHECK: error: Invalid 'vector=' argument in PACK: the 'mask=' argument has 3 true elements, but the vector has only 2 elements
+    x = pack(array, mask, [0,0])
+  end subroutine
 end module
 

flang/test/Evaluate/folding24.f90

+! RUN: %S/test_folding.sh %s %t %flang_fc1
+! REQUIRES: shell
+! Tests folding of PACK (valid cases)
+module m
+  integer, parameter :: arr(2,3) = reshape([1, 2, 3, 4, 5, 6], shape(arr))
+  logical, parameter :: odds(*,*) = modulo(arr, 2) /= 0
+  integer, parameter :: vect(*) = [(j, j=-10, -1)]
+  logical, parameter :: test_pack_1 = all(pack(arr, .true.) == [arr])
+  logical, parameter :: test_pack_2 = all(pack(arr, .false.) == [integer::])
+  logical, parameter :: test_pack_3 = all(pack(arr, odds) == [1, 3, 5])
+  logical, parameter :: test_pack_4 = all(pack(arr, .not. odds) == [2, 4, 6])
+  logical, parameter :: test_pack_5 = all(pack(arr, .true., vect) == [1, 2, 3, 4, 5, 6, -4, -3, -2, -1])
+  logical, parameter :: test_pack_6 = all(pack(arr, .false., vect) == vect)
+  logical, parameter :: test_pack_7 = all(pack(arr, odds, vect) == [1, 3, 5, -7, -6, -5, -4, -3, -2, -1])
+  logical, parameter :: test_pack_8 = all(pack(arr, .not. odds, vect) == [2, 4, 6, -7, -6, -5, -4, -3, -2, -1])
+end module