[Flang] Finalize IO operations after calling EndIO

[PeteSteinfeld]

Summary

The process of passing arguments to IO calls can cause allocations that
get referenced during EndIO calls. Calling "Finalize" causes these
allocations to be deallocated. This means that references to them in
the code in EndIO will be invalid. The fix is to delay the call to
"finalize" until after the call to EndIO.

This particularly causes problems with the IO items are strings that are
produced by calls to functions.

[clementval]

LGTM

flang/lib/Lower/IO.cpp

else
       genOutputItemList(converter, cookie, stmt.items, isFormatted,
                         csi.hasTransferConditionSpec(), ok,
                         /*inLoop=*/false);
   } else { // PRINT
     genOutputItemList(converter, cookie, std::get<1>(stmt.t), isFormatted,
                       csi.hasTransferConditionSpec(), ok,
                       /*inLoop=*/false);
   }
-  stmtCtx.finalize();
 
   builder.restoreInsertionPoint(insertPt);
   if constexpr (hasIOCtrl) {
     genIOReadSize(converter, loc, cookie, stmt.controls,
                   csi.hasErrorConditionSpec());
   }
   // Generate end statement call/s.
-  return genEndIO(converter, loc, cookie, csi, stmtCtx);
+  mlir::Value result = genEndIO(converter, loc, cookie, csi, stmtCtx);
+  stmtCtx.finalize();
+  return result;
 }
 
 void Fortran::lower::genPrintStatement(
     Fortran::lower::AbstractConverter &converter,
     const Fortran::parser::PrintStmt &stmt) {
   // PRINT does not take an io-control-spec. It only has a format specifier, so
   // it is a simplified case of WRITE.
   genDataTransferStmt</*isInput=*/false, /*ioCtrl=*/false>(converter, stmt);

flang/test/Lower/io-write.f90

+! RUN: bbc -emit-fir %s -o - | FileCheck %s
+
+! Test that IO item calls stackrestore in the right place 
+
+! CHECK-LABEL: func.func @_QQmain() {
+  character(3) string
+  write(string,getstring(6))
+! CHECK:  %[[Val_0:.*]] = fir.alloca i32 {adapt.valuebyref}
+! CHECK:  %[[Val_1:.*]] = fir.address_of(@_QFEstring) : !fir.ref<!fir.char<1,3>>
+! CHECK:  %[[Const_3:.*]] = arith.constant 3 : index
+! CHECK:  %[[Val_2:.*]] = fir.convert %[[Val_1]] : (!fir.ref<!fir.char<1,3>>) -> !fir.ref<i8>
+! CHECK:  %[[Val_3:.*]] = fir.convert %[[Const_3]] : (index) -> i64
+! CHECK:  %[[Const_6:.*]] = arith.constant 6 : i32
+! CHECK:  fir.store %[[Const_6]] to %[[Val_0]] : !fir.ref<i32>
+! CHECK:  %[[Val_4:.*]] = fir.load %[[Val_0]] : !fir.ref<i32>
+! CHECK:  %[[Val_5:.*]] = fir.convert %[[Val_4]] : (i32) -> i64
+! CHECK:  %[[Val_6:.*]] = fir.convert %[[Val_5]] : (i64) -> index
+! CHECK:  %[[Const_0:.*]] = arith.constant 0 : index
+! CHECK:  %[[Val_7:.*]] = arith.cmpi sgt, %[[Val_6]], %[[Const_0]] : index
+! CHECK:  %[[Val_8:.*]] = arith.select %[[Val_7]], %[[Val_6]], %[[Const_0]] : index
+! CHECK:  %[[Val_9:.*]] = fir.call @llvm.stacksave() : () -> !fir.ref<i8>
+! CHECK:  %[[Val_10:.*]] = fir.alloca !fir.char<1,?>(%[[Val_8]] : index) {bindc_name = ".result"}
+! CHECK:  %[[Val_11:.*]] = fir.call @_QFPgetstring(%[[Val_10]], %[[Val_8]], %[[Val_0]]) : (!fir.ref<!fir.char<1,?>>, index, !fir.ref<i32>) -> !fir.boxchar<1>
+! CHECK:  %[[Val_12:.*]] = fir.convert %[[Val_10]] : (!fir.ref<!fir.char<1,?>>) -> !fir.ref<i8>
+! CHECK:  %[[Val_13:.*]] = fir.convert %[[Val_8]] : (index) -> i64
+! CHECK:  %[[Val_14:.*]] = fir.zero_bits !fir.box<none>
+! CHECK:  %[[Const_0_i64:.*]] = arith.constant 0 : i64
+! CHECK:  %[[Val_15:.*]] = fir.convert %[[Const_0_i64]] : (i64) -> !fir.ref<!fir.llvm_ptr<i8>>
+! CHECK:  %[[Const_0_i64_0:.*]] = arith.constant 0 : i64
+! CHECK:  %[[Val_16:.*]] = fir.address_of(@_QQcl.{{.*}}) : !fir.ref<!fir.char<1,{{.*}}>>
+! CHECK:  %[[Val_17:.*]] = fir.convert %[[Val_16]] : (!fir.ref<!fir.char<1,{{.*}}>>) -> !fir.ref<i8>
+! CHECK:  %[[Val_18:.*]] = fir.call @_FortranAioBeginInternalFormattedOutput(%[[Val_2]], %[[Val_3]], %[[Val_12]], %[[Val_13]],
+! %[[Val_14]], %[[Val_15]], %[[Const_0_i64_0]], %17, %{{.*}}) : (!fir.ref<i8>, i64, !fir.ref<i8>, i64, !fir.box<none>, !fir.ref<!fir.llvm_ptr<i8>>, i64, !fir.ref<i8>, i32) -> !fir.ref<i8>
+! CHECK:  %[[Val_19:.*]] = fir.call @_FortranAioEndIoStatement(%18) : (!fir.ref<i8>) -> i32
+! CHECK:  fir.call @llvm.stackrestore(%[[Val_9]]) : (!fir.ref<i8>) -> ()
+  if (string/="hi") stop 'FAIL'
+contains
+  function getstring(n) result(r)
+    character(n) r
+    r = '("hi")'
+  end function
+end