[flang] Fix rewriting of misparsed statement functions

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

@@ -40,8 +40,11 @@ class RewriteMutator {
   template <typename T> void Post(T &) {}
 
   void Post(parser::Name &);
-  void Post(parser::SpecificationPart &);
-  bool Pre(parser::ExecutionPart &);
+  bool Pre(parser::MainProgram &);
+  bool Pre(parser::FunctionSubprogram &);
+  bool Pre(parser::SubroutineSubprogram &);
+  bool Pre(parser::SeparateModuleSubprogram &);
+  bool Pre(parser::BlockConstruct &);
   bool Pre(parser::ActionStmt &);
   void Post(parser::ReadStmt &);
   void Post(parser::WriteStmt &);
@@ -65,12 +68,11 @@ class RewriteMutator {
   bool Pre(parser::EndTypeStmt &) { return false; }
 
 private:
-  using stmtFuncType =
-      parser::Statement<common::Indirection<parser::StmtFunctionStmt>>;
+  void FixMisparsedStmtFuncs(parser::SpecificationPart &, parser::Block &);
+
   SemanticsContext &context_;
   bool errorOnUnresolvedName_{true};
   parser::Messages &messages_;
-  std::list<stmtFuncType> stmtFuncsToConvert_;
 };
 
 // Check that name has been resolved to a symbol
@@ -94,41 +96,67 @@ static bool ReturnsDataPointer(const Symbol &symbol) {
   return false;
 }
 
-// Find mis-parsed statement functions and move to stmtFuncsToConvert_ list.
-void RewriteMutator::Post(parser::SpecificationPart &x) {
-  auto &list{std::get<std::list<parser::DeclarationConstruct>>(x.t)};
+// Finds misparsed statement functions in a specification part, rewrites
+// them into array element assignment statements, and moves them into the
+// beginning of the corresponding (execution part's) block.
+void RewriteMutator::FixMisparsedStmtFuncs(
+    parser::SpecificationPart &specPart, parser::Block &block) {
+  auto &list{std::get<std::list<parser::DeclarationConstruct>>(specPart.t)};
+  auto origFirst{block.begin()}; // insert each elem before origFirst
   for (auto it{list.begin()}; it != list.end();) {
-    bool isAssignment{false};
-    if (auto *stmt{std::get_if<stmtFuncType>(&it->u)}) {
+    bool convert{false};
+    if (auto *stmt{std::get_if<
+            parser::Statement<common::Indirection<parser::StmtFunctionStmt>>>(
+            &it->u)}) {
       if (const Symbol *
           symbol{std::get<parser::Name>(stmt->statement.value().t).symbol}) {
         const Symbol &ultimate{symbol->GetUltimate()};
-        isAssignment =
+        convert =
             ultimate.has<ObjectEntityDetails>() || ReturnsDataPointer(ultimate);
-        if (isAssignment) {
-          stmtFuncsToConvert_.emplace_back(std::move(*stmt));
+        if (convert) {
+          auto newStmt{stmt->statement.value().ConvertToAssignment()};
+          newStmt.source = stmt->source;
+          block.insert(origFirst,
+              parser::ExecutionPartConstruct{
+                  parser::ExecutableConstruct{std::move(newStmt)}});
         }
       }
     }
-    if (isAssignment) {
+    if (convert) {
       it = list.erase(it);
     } else {
       ++it;
     }
   }
 }
 
-// Insert converted assignments at start of ExecutionPart.
-bool RewriteMutator::Pre(parser::ExecutionPart &x) {
-  auto origFirst{x.v.begin()}; // insert each elem before origFirst
-  for (stmtFuncType &sf : stmtFuncsToConvert_) {
-    auto stmt{sf.statement.value().ConvertToAssignment()};
-    stmt.source = sf.source;
-    x.v.insert(origFirst,
-        parser::ExecutionPartConstruct{
-            parser::ExecutableConstruct{std::move(stmt)}});
-  }
-  stmtFuncsToConvert_.clear();
+bool RewriteMutator::Pre(parser::MainProgram &program) {
+  FixMisparsedStmtFuncs(std::get<parser::SpecificationPart>(program.t),
+      std::get<parser::ExecutionPart>(program.t).v);
+  return true;
+}
+
+bool RewriteMutator::Pre(parser::FunctionSubprogram &func) {
+  FixMisparsedStmtFuncs(std::get<parser::SpecificationPart>(func.t),
+      std::get<parser::ExecutionPart>(func.t).v);
+  return true;
+}
+
+bool RewriteMutator::Pre(parser::SubroutineSubprogram &subr) {
+  FixMisparsedStmtFuncs(std::get<parser::SpecificationPart>(subr.t),
+      std::get<parser::ExecutionPart>(subr.t).v);
+  return true;
+}
+
+bool RewriteMutator::Pre(parser::SeparateModuleSubprogram &subp) {
+  FixMisparsedStmtFuncs(std::get<parser::SpecificationPart>(subp.t),
+      std::get<parser::ExecutionPart>(subp.t).v);
+  return true;
+}
+
+bool RewriteMutator::Pre(parser::BlockConstruct &block) {
+  FixMisparsedStmtFuncs(std::get<parser::BlockSpecificationPart>(block.t).v,
+      std::get<parser::Block>(block.t));
   return true;
 }
 

flang/test/Semantics/rewrite03.f90

@@ -0,0 +1,50 @@
+!RUN: %flang_fc1 -fdebug-unparse %s 2>&1 | FileCheck %s
+!Test rewriting of misparsed statement function definitions
+!into array element assignment statements.
+
+program main
+  real sf(1)
+  integer :: j = 1
+!CHECK: sf(int(j,kind=8))=1._4
+  sf(j) = 1.
+end
+
+function func
+  real sf(1)
+  integer :: j = 1
+!CHECK: sf(int(j,kind=8))=2._4
+  sf(j) = 2.
+  func = 0.
+end
+
+subroutine subr
+  real sf(1)
+  integer :: j = 1
+!CHECK: sf(int(j,kind=8))=3._4
+  sf(j) = 3.
+end
+
+module m
+  interface
+    module subroutine smp
+    end
+  end interface
+end
+submodule(m) sm
+ contains
+  module procedure smp
+    real sf(1)
+    integer :: j = 1
+!CHECK: sf(int(j,kind=8))=4._4
+    sf(j) = 4.
+  end
+end
+
+subroutine block
+  block
+    real sf(1)
+    integer :: j = 1
+!CHECK: sf(int(j,kind=8))=5._4
+    sf(j) = 5.
+  end block
+end