[flang] Fold LCOBOUND & UCOBOUND

[klausler]

Implement constant folding for LCOBOUND and UCOBOUND intrinsic functions. Moves some error detection code from intrinsics.cpp to fold-integer.cpp so that erroneous calls get properly flagged and converted into known errors.

[llvmbot]

@llvm/pr-subscribers-flang-semantics

Author: Peter Klausler (klausler)

Changes

Implement constant folding for LCOBOUND and UCOBOUND intrinsic functions. Moves some error detection code from intrinsics.cpp to fold-integer.cpp so that erroneous calls get properly flagged and converted into known errors.

---

Full diff: https://github.com/llvm/llvm-project/pull/121411.diff

6 Files Affected:

• (modified) flang/include/flang/Evaluate/shape.h (+8)
• (modified) flang/lib/Evaluate/fold-integer.cpp (+57)
• (modified) flang/lib/Evaluate/intrinsics.cpp (-25)
• (modified) flang/lib/Evaluate/shape.cpp (+52)
• (modified) flang/test/Semantics/lcobound.f90 (+11-8)
• (modified) flang/test/Semantics/ucobound.f90 (+11-8)

diff --git a/flang/include/flang/Evaluate/shape.h b/flang/include/flang/Evaluate/shape.h
index e33044c0d34e56..e679a001235490 100644
--- a/flang/include/flang/Evaluate/shape.h
+++ b/flang/include/flang/Evaluate/shape.h
@@ -117,6 +117,14 @@ MaybeExtentExpr GetExtent(const Subscript &, const NamedEntity &, int dimension,
 MaybeExtentExpr GetExtent(FoldingContext &, const Subscript &,
     const NamedEntity &, int dimension, bool invariantOnly = true);
 
+// Similar analyses for coarrays
+MaybeExtentExpr GetLCOBOUND(
+    const Symbol &, int dimension, bool invariantOnly = true);
+MaybeExtentExpr GetUCOBOUND(
+    const Symbol &, int dimension, bool invariantOnly = true);
+Shape GetLCOBOUNDs(const Symbol &, bool invariantOnly = true);
+Shape GetUCOBOUNDs(const Symbol &, bool invariantOnly = true);
+
 // Compute an element count for a triplet or trip count for a DO.
 ExtentExpr CountTrips(
     ExtentExpr &&lower, ExtentExpr &&upper, ExtentExpr &&stride);
diff --git a/flang/lib/Evaluate/fold-integer.cpp b/flang/lib/Evaluate/fold-integer.cpp
index 26ae33faffe1e2..352dec4bb5ee26 100644
--- a/flang/lib/Evaluate/fold-integer.cpp
+++ b/flang/lib/Evaluate/fold-integer.cpp
@@ -71,6 +71,28 @@ static bool CheckDimArg(const std::optional<ActualArgument> &dimArg,
   return true;
 }
 
+static bool CheckCoDimArg(const std::optional<ActualArgument> &dimArg,
+    const Symbol &symbol, parser::ContextualMessages &messages,
+    std::optional<int> &dimVal) {
+  dimVal.reset();
+  if (int corank{symbol.Corank()}; corank > 0) {
+    if (auto dim64{ToInt64(dimArg)}) {
+      if (*dim64 < 1) {
+        messages.Say("DIM=%jd dimension must be positive"_err_en_US, *dim64);
+        return false;
+      } else if (*dim64 > corank) {
+        messages.Say(
+            "DIM=%jd dimension is out of range for corank-%d coarray"_err_en_US,
+            *dim64, corank);
+        return false;
+      } else {
+        dimVal = static_cast<int>(*dim64 - 1); // 1-based to 0-based
+      }
+    }
+  }
+  return true;
+}
+
 // Class to retrieve the constant bound of an expression which is an
 // array that devolves to a type of Constant<T>
 class GetConstantArrayBoundHelper {
@@ -264,6 +286,37 @@ Expr<Type<TypeCategory::Integer, KIND>> UBOUND(FoldingContext &context,
   return Expr<T>{std::move(funcRef)};
 }
 
+// LCOBOUND() & UCOBOUND()
+template <int KIND>
+Expr<Type<TypeCategory::Integer, KIND>> COBOUND(FoldingContext &context,
+    FunctionRef<Type<TypeCategory::Integer, KIND>> &&funcRef, bool isUCOBOUND) {
+  using T = Type<TypeCategory::Integer, KIND>;
+  ActualArguments &args{funcRef.arguments()};
+  if (const Symbol * coarray{UnwrapWholeSymbolOrComponentDataRef(args[0])}) {
+    std::optional<int> dim;
+    if (funcRef.Rank() == 0) {
+      // Optional DIM= argument is present: result is scalar.
+      if (!CheckCoDimArg(args[1], *coarray, context.messages(), dim)) {
+        return MakeInvalidIntrinsic<T>(std::move(funcRef));
+      } else if (!dim) {
+        // DIM= is present but not constant, or error
+        return Expr<T>{std::move(funcRef)};
+      }
+    }
+    if (dim) {
+      if (auto cb{isUCOBOUND ? GetUCOBOUND(*coarray, *dim)
+                             : GetLCOBOUND(*coarray, *dim)}) {
+        return Fold(context, ConvertToType<T>(std::move(*cb)));
+      }
+    } else if (auto cbs{
+                   AsExtentArrayExpr(isUCOBOUND ? GetUCOBOUNDs(*coarray)
+                                                : GetLCOBOUNDs(*coarray))}) {
+      return Fold(context, ConvertToType<T>(Expr<ExtentType>{std::move(*cbs)}));
+    }
+  }
+  return Expr<T>{std::move(funcRef)};
+}
+
 // COUNT()
 template <typename T, int MASK_KIND> class CountAccumulator {
   using MaskT = Type<TypeCategory::Logical, MASK_KIND>;
@@ -1105,6 +1158,8 @@ Expr<Type<TypeCategory::Integer, KIND>> FoldIntrinsicFunction(
     }
   } else if (name == "lbound") {
     return LBOUND(context, std::move(funcRef));
+  } else if (name == "lcobound") {
+    return COBOUND(context, std::move(funcRef), /*isUCOBOUND=*/false);
   } else if (name == "leadz" || name == "trailz" || name == "poppar" ||
       name == "popcnt") {
     if (auto *sn{UnwrapExpr<Expr<SomeKind<T::category>>>(args[0])}) {
@@ -1396,6 +1451,8 @@ Expr<Type<TypeCategory::Integer, KIND>> FoldIntrinsicFunction(
     }
   } else if (name == "ubound") {
     return UBOUND(context, std::move(funcRef));
+  } else if (name == "ucobound") {
+    return COBOUND(context, std::move(funcRef), /*isUCOBOUND=*/true);
   } else if (name == "__builtin_numeric_storage_size") {
     if (!context.moduleFileName()) {
       // Don't fold this reference until it appears in the module file
diff --git a/flang/lib/Evaluate/intrinsics.cpp b/flang/lib/Evaluate/intrinsics.cpp
index 28805efb177ee2..f85ebe60336e5a 100644
--- a/flang/lib/Evaluate/intrinsics.cpp
+++ b/flang/lib/Evaluate/intrinsics.cpp
@@ -3119,27 +3119,6 @@ static bool CheckForNonPositiveValues(FoldingContext &context,
   return ok;
 }
 
-static bool CheckDimAgainstCorank(SpecificCall &call, FoldingContext &context) {
-  bool ok{true};
-  if (const auto &coarrayArg{call.arguments[0]}) {
-    if (const auto &dimArg{call.arguments[1]}) {
-      if (const auto *symbol{
-              UnwrapWholeSymbolDataRef(coarrayArg->UnwrapExpr())}) {
-        const auto corank = symbol->Corank();
-        if (const auto dimNum{ToInt64(dimArg->UnwrapExpr())}) {
-          if (dimNum < 1 || dimNum > corank) {
-            ok = false;
-            context.messages().Say(dimArg->sourceLocation(),
-                "DIM=%jd dimension is out of range for coarray with corank %d"_err_en_US,
-                static_cast<std::intmax_t>(*dimNum), corank);
-          }
-        }
-      }
-    }
-  }
-  return ok;
-}
-
 static bool CheckAtomicDefineAndRef(FoldingContext &context,
     const std::optional<ActualArgument> &atomArg,
     const std::optional<ActualArgument> &valueArg,
@@ -3207,8 +3186,6 @@ static bool ApplySpecificChecks(SpecificCall &call, FoldingContext &context) {
     if (const auto &arg{call.arguments[0]}) {
       ok = CheckForNonPositiveValues(context, *arg, name, "image");
     }
-  } else if (name == "lcobound") {
-    return CheckDimAgainstCorank(call, context);
   } else if (name == "loc") {
     const auto &arg{call.arguments[0]};
     ok =
@@ -3218,8 +3195,6 @@ static bool ApplySpecificChecks(SpecificCall &call, FoldingContext &context) {
           arg ? arg->sourceLocation() : context.messages().at(),
           "Argument of LOC() must be an object or procedure"_err_en_US);
     }
-  } else if (name == "ucobound") {
-    return CheckDimAgainstCorank(call, context);
   }
   return ok;
 }
diff --git a/flang/lib/Evaluate/shape.cpp b/flang/lib/Evaluate/shape.cpp
index c62d0cb0ff29dd..f006fe598c4224 100644
--- a/flang/lib/Evaluate/shape.cpp
+++ b/flang/lib/Evaluate/shape.cpp
@@ -723,6 +723,58 @@ Shape GetUBOUNDs(const NamedEntity &base, bool invariantOnly) {
   return GetUBOUNDs(nullptr, base, invariantOnly);
 }
 
+MaybeExtentExpr GetLCOBOUND(
+    const Symbol &symbol0, int dimension, bool invariantOnly) {
+  const Symbol &symbol{ResolveAssociations(symbol0)};
+  if (const auto *object{symbol.detailsIf<semantics::ObjectEntityDetails>()}) {
+    int corank{object->coshape().Rank()};
+    if (dimension < corank) {
+      const semantics::ShapeSpec &shapeSpec{object->coshape()[dimension]};
+      if (const auto &lcobound{shapeSpec.lbound().GetExplicit()}) {
+        if (!invariantOnly || IsScopeInvariantExpr(*lcobound)) {
+          return *lcobound;
+        }
+      }
+    }
+  }
+  return std::nullopt;
+}
+
+MaybeExtentExpr GetUCOBOUND(
+    const Symbol &symbol0, int dimension, bool invariantOnly) {
+  const Symbol &symbol{ResolveAssociations(symbol0)};
+  if (const auto *object{symbol.detailsIf<semantics::ObjectEntityDetails>()}) {
+    int corank{object->coshape().Rank()};
+    if (dimension < corank - 1) {
+      const semantics::ShapeSpec &shapeSpec{object->coshape()[dimension]};
+      if (const auto ucobound{shapeSpec.ubound().GetExplicit()}) {
+        if (!invariantOnly || IsScopeInvariantExpr(*ucobound)) {
+          return *ucobound;
+        }
+      }
+    }
+  }
+  return std::nullopt;
+}
+
+Shape GetLCOBOUNDs(const Symbol &symbol, bool invariantOnly) {
+  Shape result;
+  int corank{symbol.Corank()};
+  for (int dim{0}; dim < corank; ++dim) {
+    result.emplace_back(GetLCOBOUND(symbol, dim, invariantOnly));
+  }
+  return result;
+}
+
+Shape GetUCOBOUNDs(const Symbol &symbol, bool invariantOnly) {
+  Shape result;
+  int corank{symbol.Corank()};
+  for (int dim{0}; dim < corank; ++dim) {
+    result.emplace_back(GetUCOBOUND(symbol, dim, invariantOnly));
+  }
+  return result;
+}
+
 auto GetShapeHelper::operator()(const Symbol &symbol) const -> Result {
   return common::visit(
       common::visitors{
diff --git a/flang/test/Semantics/lcobound.f90 b/flang/test/Semantics/lcobound.f90
index ce2f001ce2ea72..f03f2cae03ec43 100644
--- a/flang/test/Semantics/lcobound.f90
+++ b/flang/test/Semantics/lcobound.f90
@@ -11,6 +11,9 @@ program lcobound_tests
   logical non_integer, logical_coarray[3,*]
   logical, parameter :: const_non_integer = .true.
   integer, allocatable :: lcobounds(:)
+  real bounded[2:3,4:5,*]
+
+  integer(kind=merge(kind(1),-1,all(lcobound(bounded)==[2,4,1]))) test_lcobound
 
   !___ standard-conforming statement with no optional arguments present ___
   lcobounds = lcobound(scalar_coarray)
@@ -50,28 +53,28 @@ program lcobound_tests
 
   !___ non-conforming statements ___
 
-  !ERROR: DIM=0 dimension is out of range for coarray with corank 1
+  !ERROR: DIM=0 dimension must be positive
   n = lcobound(scalar_coarray, dim=0)
 
-  !ERROR: DIM=0 dimension is out of range for coarray with corank 3
+  !ERROR: DIM=0 dimension must be positive
   n = lcobound(coarray_corank3, dim=0)
 
-  !ERROR: DIM=-1 dimension is out of range for coarray with corank 1
+  !ERROR: DIM=-1 dimension must be positive
   n = lcobound(scalar_coarray, dim=-1)
 
-  !ERROR: DIM=2 dimension is out of range for coarray with corank 1
+  !ERROR: DIM=2 dimension is out of range for corank-1 coarray
   n = lcobound(array_coarray, dim=2)
 
-  !ERROR: DIM=2 dimension is out of range for coarray with corank 1
+  !ERROR: DIM=2 dimension is out of range for corank-1 coarray
   n = lcobound(array_coarray, 2)
 
-  !ERROR: DIM=4 dimension is out of range for coarray with corank 3
+  !ERROR: DIM=4 dimension is out of range for corank-3 coarray
   n = lcobound(coarray_corank3, dim=4)
 
-  !ERROR: DIM=4 dimension is out of range for coarray with corank 3
+  !ERROR: DIM=4 dimension is out of range for corank-3 coarray
   n = lcobound(dim=4, coarray=coarray_corank3)
 
-  !ERROR: DIM=5 dimension is out of range for coarray with corank 3
+  !ERROR: DIM=5 dimension is out of range for corank-3 coarray
   n = lcobound(coarray_corank3, const_out_of_range_dim)
 
   !ERROR: No intrinsic or user-defined ASSIGNMENT(=) matches scalar INTEGER(4) and rank 1 array of INTEGER(4)
diff --git a/flang/test/Semantics/ucobound.f90 b/flang/test/Semantics/ucobound.f90
index f9da11a03a6b0f..d84c80cdd315c0 100644
--- a/flang/test/Semantics/ucobound.f90
+++ b/flang/test/Semantics/ucobound.f90
@@ -11,6 +11,9 @@ program ucobound_tests
   logical non_integer, logical_coarray[3,*]
   logical, parameter :: const_non_integer = .true.
   integer, allocatable :: ucobounds(:)
+  real bounded[2:3,4:5,*]
+
+  integer(kind=merge(kind(1),-1,ucobound(bounded,1)==3.and.ucobound(bounded,2)==5)) test_ucobound
 
   !___ standard-conforming statement with no optional arguments present ___
   ucobounds = ucobound(scalar_coarray)
@@ -50,28 +53,28 @@ program ucobound_tests
 
   !___ non-conforming statements ___
 
-  !ERROR: DIM=0 dimension is out of range for coarray with corank 1
+  !ERROR: DIM=0 dimension must be positive
   n = ucobound(scalar_coarray, dim=0)
 
-  !ERROR: DIM=0 dimension is out of range for coarray with corank 3
+  !ERROR: DIM=0 dimension must be positive
   n = ucobound(coarray_corank3, dim=0)
 
-  !ERROR: DIM=-1 dimension is out of range for coarray with corank 1
+  !ERROR: DIM=-1 dimension must be positive
   n = ucobound(scalar_coarray, dim=-1)
 
-  !ERROR: DIM=2 dimension is out of range for coarray with corank 1
+  !ERROR: DIM=2 dimension is out of range for corank-1 coarray
   n = ucobound(array_coarray, dim=2)
 
-  !ERROR: DIM=2 dimension is out of range for coarray with corank 1
+  !ERROR: DIM=2 dimension is out of range for corank-1 coarray
   n = ucobound(array_coarray, 2)
 
-  !ERROR: DIM=4 dimension is out of range for coarray with corank 3
+  !ERROR: DIM=4 dimension is out of range for corank-3 coarray
   n = ucobound(coarray_corank3, dim=4)
 
-  !ERROR: DIM=4 dimension is out of range for coarray with corank 3
+  !ERROR: DIM=4 dimension is out of range for corank-3 coarray
   n = ucobound(dim=4, coarray=coarray_corank3)
 
-  !ERROR: DIM=5 dimension is out of range for coarray with corank 3
+  !ERROR: DIM=5 dimension is out of range for corank-3 coarray
   n = ucobound(coarray_corank3, const_out_of_range_dim)
 
   !ERROR: No intrinsic or user-defined ASSIGNMENT(=) matches scalar INTEGER(4) and rank 1 array of INTEGER(4)