[flang] Delay parse tree rewriting for I/O UNIT=func()

When an I/O statement's UNIT= specifier is a variable that is a
function reference, parse tree rewriting may determine the wrong type
of the result because generic resolution has not yet been performed.
So move this bit of parse tree rewriting into I/O semantic
checking so that the right handling (integer -> external file unit
number, character pointer -> internal I/O) applies.

Differential Revision: https://reviews.llvm.org/D135210

flang/lib/Semantics/check-io.cpp

@@ -542,17 +542,50 @@ void IoChecker::Enter(const parser::IoControlSpec::Size &var) {
 
 void IoChecker::Enter(const parser::IoUnit &spec) {
   if (const parser::Variable * var{std::get_if<parser::Variable>(&spec.u)}) {
-    if (stmt_ == IoStmtKind::Write) {
-      CheckForDefinableVariable(*var, "Internal file");
+    // Only now after generic resolution can it be known whether a function
+    // call appearing as UNIT=f() is an integer scalar external unit number
+    // or a character pointer for internal I/O.
+    const auto *expr{GetExpr(context_, *var)};
+    std::optional<evaluate::DynamicType> dyType;
+    if (expr) {
+      dyType = expr->GetType();
     }
-    if (const auto *expr{GetExpr(context_, *var)}) {
+    if (dyType && dyType->category() == TypeCategory::Integer) {
+      if (expr->Rank() != 0) {
+        context_.Say(parser::FindSourceLocation(*var),
+            "I/O unit number must be scalar"_err_en_US);
+      }
+      // In the case of an integer unit number variable, rewrite the parse
+      // tree as if the unit had been parsed as a FileUnitNumber in order
+      // to ease lowering.
+      auto &mutableSpec{const_cast<parser::IoUnit &>(spec)};
+      auto &mutableVar{std::get<parser::Variable>(mutableSpec.u)};
+      auto source{mutableVar.GetSource()};
+      auto typedExpr{std::move(mutableVar.typedExpr)};
+      auto newExpr{common::visit(
+          [](auto &&indirection) {
+            return parser::Expr{std::move(indirection)};
+          },
+          std::move(mutableVar.u))};
+      newExpr.source = source;
+      newExpr.typedExpr = std::move(typedExpr);
+      mutableSpec.u = parser::FileUnitNumber{
+          parser::ScalarIntExpr{parser::IntExpr{std::move(newExpr)}}};
+    } else if (!dyType || dyType->category() != TypeCategory::Character) {
+      SetSpecifier(IoSpecKind::Unit);
+      context_.Say(parser::FindSourceLocation(*var),
+          "I/O unit must be a character variable or a scalar integer expression"_err_en_US);
+    } else { // CHARACTER variable (internal I/O)
+      if (stmt_ == IoStmtKind::Write) {
+        CheckForDefinableVariable(*var, "Internal file");
+      }
       if (HasVectorSubscript(*expr)) {
         context_.Say(parser::FindSourceLocation(*var), // C1201
             "Internal file must not have a vector subscript"_err_en_US);
       }
+      SetSpecifier(IoSpecKind::Unit);
+      flags_.set(Flag::InternalUnit);
     }
-    SetSpecifier(IoSpecKind::Unit);
-    flags_.set(Flag::InternalUnit);
   } else if (std::get_if<parser::Star>(&spec.u)) {
     SetSpecifier(IoSpecKind::Unit);
     flags_.set(Flag::StarUnit);

flang/lib/Semantics/rewrite-parse-tree.cpp

@@ -41,7 +41,6 @@ public:
   void Post(parser::Name &);
   void Post(parser::SpecificationPart &);
   bool Pre(parser::ExecutionPart &);
-  void Post(parser::IoUnit &);
   void Post(parser::ReadStmt &);
   void Post(parser::WriteStmt &);
 
@@ -130,29 +129,6 @@ bool RewriteMutator::Pre(parser::ExecutionPart &x) {
   return true;
 }
 
-// Convert a syntactically ambiguous io-unit internal-file-variable to a
-// file-unit-number.
-void RewriteMutator::Post(parser::IoUnit &x) {
-  if (auto *var{std::get_if<parser::Variable>(&x.u)}) {
-    const parser::Name &last{parser::GetLastName(*var)};
-    DeclTypeSpec *type{last.symbol ? last.symbol->GetType() : nullptr};
-    if (!type || type->category() != DeclTypeSpec::Character) {
-      // If the Variable is not known to be character (any kind), transform
-      // the I/O unit in situ to a FileUnitNumber so that automatic expression
-      // constraint checking will be applied.
-      auto source{var->GetSource()};
-      auto expr{common::visit(
-          [](auto &&indirection) {
-            return parser::Expr{std::move(indirection)};
-          },
-          std::move(var->u))};
-      expr.source = source;
-      x.u = parser::FileUnitNumber{
-          parser::ScalarIntExpr{parser::IntExpr{std::move(expr)}}};
-    }
-  }
-}
-
 // When a namelist group name appears (without NML=) in a READ or WRITE
 // statement in such a way that it can be misparsed as a format expression,
 // rewrite the I/O statement's parse tree node as if the namelist group

flang/test/Semantics/io04.f90

@@ -87,7 +87,7 @@
   !ERROR: If UNIT=* appears, REC must not appear
   write(*, rec=13) 'Ok'
 
-  !ERROR: Must have INTEGER type, but is REAL(4)
+  !ERROR: I/O unit must be a character variable or a scalar integer expression
   write(unit, *) 'Ok'
 
   !ERROR: If ADVANCE appears, UNIT=internal-file must not appear

flang/test/Semantics/io13.f90

@@ -0,0 +1,53 @@
+! RUN: %python %S/test_errors.py %s %flang_fc1
+! Tests for UNIT=function()
+module m1
+  integer, target :: itarget
+  character(20), target :: ctarget
+  logical, target :: ltarget
+  interface gf
+    module procedure :: intf, pintf, pchf, logf, plogf
+  end interface
+ contains
+  integer function intf(n)
+    integer(1), intent(in) :: n
+    intf = n
+  end function
+  function pintf(n)
+    integer(2), intent(in) :: n
+    integer, pointer :: pintf
+    pintf => itarget
+    pintf = n
+  end function
+  function pchf(n)
+    integer(4), intent(in) :: n
+    character(:), pointer :: pchf
+    pchf => ctarget
+  end function
+  logical function logf(n)
+    integer(8), intent(in) :: n
+    logf = .true.
+  end function
+  function plogf(n)
+    integer(16), intent(in) :: n
+    logical, pointer :: plf
+    plf => ltarget
+  end function
+  subroutine test
+    write(intf(6_1),"('hi')")
+    write(pintf(6_2),"('hi')")
+    write(pchf(123_4),"('hi')")
+    write(gf(6_1),"('hi')")
+    write(gf(6_2),"('hi')")
+    write(gf(666_4),"('hi')")
+    !ERROR: I/O unit must be a character variable or a scalar integer expression
+    write(logf(666_8),"('hi')")
+    !ERROR: I/O unit must be a character variable or a scalar integer expression
+    write(plogf(666_16),"('hi')")
+    !ERROR: I/O unit must be a character variable or a scalar integer expression
+    write(gf(666_8),"('hi')")
+    !ERROR: I/O unit must be a character variable or a scalar integer expression
+    write(gf(666_16),"('hi')")
+    !ERROR: I/O unit must be a character variable or a scalar integer expression
+    write(null(),"('hi')")
+  end subroutine
+end module