IoUnit allows arbitrary expression

akuhlens · akuhlens · commit 0cca990e5b7e · 2025-02-12T13:09:59.000-08:00
diff --git a/flang/include/flang/Parser/parse-tree.h b/flang/include/flang/Parser/parse-tree.h
@@ -2642,7 +2642,7 @@ WRAPPER_CLASS(FileUnitNumber, ScalarIntExpr);
 // symbols are known.
 struct IoUnit {
   UNION_CLASS_BOILERPLATE(IoUnit);
-  std::variant<Variable, FileUnitNumber, Star> u;
+  std::variant<Variable, common::Indirection<Expr>, Star> u;
 };
 
 // R1206 file-name-expr -> scalar-default-char-expr
diff --git a/flang/lib/Lower/IO.cpp b/flang/lib/Lower/IO.cpp
@@ -1841,7 +1841,9 @@ static mlir::Value genIOUnit(Fortran::lower::AbstractConverter &converter,
                              int defaultUnitNumber) {
   auto &builder = converter.getFirOpBuilder();
   if (iounit)
-    if (auto *e = std::get_if<Fortran::parser::FileUnitNumber>(&iounit->u))
+    if (auto *e =
+            std::get_if<Fortran::common::Indirection<Fortran::parser::Expr>>(
+                &iounit->u))
       return genIOUnitNumber(converter, loc, Fortran::semantics::GetExpr(*e),
                              ty, csi, stmtCtx);
   return builder.create<mlir::arith::ConstantOp>(
diff --git a/flang/lib/Parser/io-parsers.cpp b/flang/lib/Parser/io-parsers.cpp
@@ -23,8 +23,12 @@ namespace Fortran::parser {
 // R905 char-variable -> variable
 // "char-variable" is attempted first since it's not type constrained but
 // syntactically ambiguous with "file-unit-number", which is constrained.
+// Note, "file-unit-number" is replaced by "expr" to allow for better
+// error messages.
 TYPE_PARSER(construct<IoUnit>(variable / lookAhead(space / ",);\n"_ch)) ||
-    construct<IoUnit>(fileUnitNumber) || construct<IoUnit>(star))
+    construct<IoUnit>(
+        indirect(expr) / (lookAhead(space >> ",)"_ch) || atEndOfStmt)) ||
+    construct<IoUnit>(star))
 
 // R1202 file-unit-number -> scalar-int-expr
 TYPE_PARSER(construct<FileUnitNumber>(
diff --git a/flang/lib/Semantics/check-io.cpp b/flang/lib/Semantics/check-io.cpp
@@ -9,7 +9,9 @@
 #include "check-io.h"
 #include "definable.h"
 #include "flang/Common/format.h"
+#include "flang/Common/indirection.h"
 #include "flang/Evaluate/tools.h"
+#include "flang/Parser/characters.h"
 #include "flang/Parser/tools.h"
 #include "flang/Semantics/expression.h"
 #include "flang/Semantics/tools.h"
@@ -576,8 +578,9 @@ void IoChecker::Enter(const parser::IoUnit &spec) {
           std::move(mutableVar.u))};
       newExpr.source = source;
       newExpr.typedExpr = std::move(typedExpr);
-      mutableSpec.u = parser::FileUnitNumber{
-          parser::ScalarIntExpr{parser::IntExpr{std::move(newExpr)}}};
+      mutableSpec.u = common::Indirection<parser::Expr>{std::move(newExpr)};
+      SetSpecifier(IoSpecKind::Unit);
+      flags_.set(Flag::NumberUnit);
     } else if (!dyType || dyType->category() != TypeCategory::Character) {
       SetSpecifier(IoSpecKind::Unit);
       context_.Say(parser::FindSourceLocation(*var),
@@ -598,6 +601,26 @@ void IoChecker::Enter(const parser::IoUnit &spec) {
   } else if (std::get_if<parser::Star>(&spec.u)) {
     SetSpecifier(IoSpecKind::Unit);
     flags_.set(Flag::StarUnit);
+  } else if (const common::Indirection<parser::Expr> *pexpr{
+                 std::get_if<common::Indirection<parser::Expr>>(&spec.u)}) {
+    const auto *expr{GetExpr(context_, *pexpr)};
+    std::optional<evaluate::DynamicType> dyType;
+    if (expr) {
+      dyType = expr->GetType();
+    }
+    if (!expr || !dyType) {
+      context_.Say(parser::FindSourceLocation(*pexpr),
+          "I/O unit must be a character variable or scalar integer expression"_err_en_US);
+    } else if (dyType->category() != TypeCategory::Integer) {
+      context_.Say(parser::FindSourceLocation(*pexpr),
+          "I/O unit must be a character variable or a scalar integer expression, but is an expression of type %s"_err_en_US,
+          parser::ToUpperCaseLetters(dyType->AsFortran()));
+    } else if (expr->Rank() != 0) {
+      context_.Say(parser::FindSourceLocation(*pexpr),
+          "I/O unit number must be scalar"_err_en_US);
+    }
+    SetSpecifier(IoSpecKind::Unit);
+    flags_.set(Flag::NumberUnit);
   }
 }
 
diff --git a/flang/test/Semantics/io03.f90 b/flang/test/Semantics/io03.f90
@@ -171,6 +171,22 @@
   !ERROR: ID kind (2) is smaller than default INTEGER kind (4)
   read(10, id=id2, asynchronous='yes') jj
 
+  !ERROR: I/O unit must be a character variable or a scalar integer expression, but is an expression of type CHARACTER(1)
+  read((msg), *)
+  !ERROR: I/O unit must be a character variable or a scalar integer expression, but is an expression of type CHARACTER(KIND=1,LEN=8_8)
+  read("a string", *)
+  !ERROR: I/O unit must be a character variable or a scalar integer expression, but is an expression of type CHARACTER(1)
+  read(msg//msg, *)
+  !ERROR: I/O unit must be a character variable or a scalar integer expression, but is an expression of type LOGICAL(4)
+  read(.true., *)
+  !ERROR: I/O unit must be a character variable or a scalar integer expression, but is an expression of type REAL(4)
+  read(1.0, *)
+  read(internal_fileA, *)
+  !ERROR: I/O unit must be a character variable or a scalar integer expression, but is an expression of type CHARACTER(1)
+  read((internal_fileA), *)
+  !ERROR: I/O unit number must be scalar
+  read([1,2,3], *)
+
 9 continue
 end
 
diff --git a/flang/test/Semantics/io04.f90 b/flang/test/Semantics/io04.f90
@@ -138,6 +138,23 @@
 
   write(*, '(X)')
 
+  !ERROR: I/O unit must be a character variable or a scalar integer expression, but is an expression of type CHARACTER(1)
+  write((msg), *)
+  !ERROR: I/O unit must be a character variable or a scalar integer expression, but is an expression of type CHARACTER(KIND=1,LEN=8_8)
+  write("a string", *)
+  !ERROR: I/O unit must be a character variable or a scalar integer expression, but is an expression of type CHARACTER(1)
+  write(msg//msg, *)
+  !ERROR: I/O unit must be a character variable or a scalar integer expression, but is an expression of type LOGICAL(4)
+  write(.true., *)
+  !ERROR: I/O unit must be a character variable or a scalar integer expression, but is an expression of type REAL(4)
+  write(1.0, *)
+  write(internal_fileA, *)
+  !! Not sure why this isn't an error with this message: I/O unit must be a character variable or a scalar integer expression, but is an expression of type CHARACTER(1)
+  write((internal_fileA), *)
+  !ERROR: I/O unit number must be scalar
+  write([1,2,3], *)
+
+
   !ERROR: Output item must not be a procedure
   print*, procptr
   !ERROR: Output item must not be a procedure
diff --git a/flang/test/Semantics/unsigned-errors.f90 b/flang/test/Semantics/unsigned-errors.f90
@@ -64,7 +64,7 @@
 !ERROR: No intrinsic or user-defined ASSIGNMENT(=) matches operand types INTEGER(4) and UNSIGNED(4)
 j = 1u
 
-!ERROR: Must have INTEGER type, but is UNSIGNED(4)
+!ERROR: I/O unit must be a character variable or a scalar integer expression, but is an expression of type UNSIGNED(4)
 write(6u,*) 'hi'
 
 !ERROR: ARITHMETIC IF expression must not be an UNSIGNED expression
