[flang-commits] [flang] [flang] Fold IS_IOSTAT_END() & IS_IOSTAT_EOR() (PR #70971)

[Peter Klausler via flang-commits]

https://github.com/klausler created https://github.com/llvm/llvm-project/pull/70971

These intrinsic functions are not particularly valuable -- one can just compare a value to IOSTAT_END or IOSTAT_EOR directly -- but they are in the standard and are allowed to appear in constant expressions, so here's code to fold them.

>From ffa3114840a3f404d76a1654f50b62dcec025f67 Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler at nvidia.com>
Date: Wed, 1 Nov 2023 11:30:09 -0700
Subject: [PATCH] [flang] Fold IS_IOSTAT_END() & IS_IOSTAT_EOR()

These intrinsic functions are not particularly valuable -- one can just
compare a value to IOSTAT_END or IOSTAT_EOR directly -- but they are
in the standard and are allowed to appear in constant expressions,
so here's code to fold them.
---
 flang/lib/Evaluate/fold-logical.cpp | 21 ++++++++++++++++++++-
 flang/test/Evaluate/fold-iostat.f90 | 12 ++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 flang/test/Evaluate/fold-iostat.f90

diff --git a/flang/lib/Evaluate/fold-logical.cpp b/flang/lib/Evaluate/fold-logical.cpp
index 9fc42adf805f468..bfedc32a33a8bad 100644
--- a/flang/lib/Evaluate/fold-logical.cpp
+++ b/flang/lib/Evaluate/fold-logical.cpp
@@ -9,6 +9,7 @@
 #include "fold-implementation.h"
 #include "fold-reduction.h"
 #include "flang/Evaluate/check-expression.h"
+#include "flang/Runtime/magic-numbers.h"
 
 namespace Fortran::evaluate {
 
@@ -194,6 +195,24 @@ Expr<Type<TypeCategory::Logical, KIND>> FoldIntrinsicFunction(
         }
       }
     }
+  } else if (name == "is_iostat_end") {
+    if (args[0] && args[0]->UnwrapExpr() &&
+        IsActuallyConstant(*args[0]->UnwrapExpr())) {
+      using Int64 = Type<TypeCategory::Integer, 8>;
+      return FoldElementalIntrinsic<T, Int64>(context, std::move(funcRef),
+          ScalarFunc<T, Int64>([](const Scalar<Int64> &x) {
+            return Scalar<T>{x.ToInt64() == FORTRAN_RUNTIME_IOSTAT_END};
+          }));
+    }
+  } else if (name == "is_iostat_eor") {
+    if (args[0] && args[0]->UnwrapExpr() &&
+        IsActuallyConstant(*args[0]->UnwrapExpr())) {
+      using Int64 = Type<TypeCategory::Integer, 8>;
+      return FoldElementalIntrinsic<T, Int64>(context, std::move(funcRef),
+          ScalarFunc<T, Int64>([](const Scalar<Int64> &x) {
+            return Scalar<T>{x.ToInt64() == FORTRAN_RUNTIME_IOSTAT_EOR};
+          }));
+    }
   } else if (name == "lge" || name == "lgt" || name == "lle" || name == "llt") {
     // Rewrite LGE/LGT/LLE/LLT into ASCII character relations
     auto *cx0{UnwrapExpr<Expr<SomeCharacter>>(args[0])};
@@ -348,7 +367,7 @@ Expr<Type<TypeCategory::Logical, KIND>> FoldIntrinsicFunction(
       name == "__builtin_ieee_support_underflow_control") {
     return Expr<T>{true};
   }
-  // TODO: is_iostat_end, is_iostat_eor, logical, matmul, parity
+  // TODO: logical, matmul, parity
   return Expr<T>{std::move(funcRef)};
 }
 
diff --git a/flang/test/Evaluate/fold-iostat.f90 b/flang/test/Evaluate/fold-iostat.f90
new file mode 100644
index 000000000000000..7b795d82fa0529a
--- /dev/null
+++ b/flang/test/Evaluate/fold-iostat.f90
@@ -0,0 +1,12 @@
+! RUN: %python %S/test_folding.py %s %flang_fc1
+module m
+  use iso_fortran_env
+  logical, parameter :: test_end1 = is_iostat_end(iostat_end)
+  logical, parameter :: test_end2 = .not. is_iostat_end(iostat_eor)
+  logical, parameter :: test_eor1 = is_iostat_eor(iostat_eor)
+  logical, parameter :: test_eor2 = .not. is_iostat_eor(iostat_end)
+  logical, parameter :: test_arr1 = &
+    all(is_iostat_end([iostat_end, iostat_eor]) .eqv. [.true., .false.])
+  logical, parameter :: test_arr2 = &
+    all(is_iostat_eor([iostat_end, iostat_eor]) .eqv. [.false., .true.])
+end