From 7774f5216717f80df5ac120676ed4ab440df86fb Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler@nvidia.com>
Date: Wed, 12 Mar 2025 09:21:33 -0700
Subject: [PATCH] [flang] Enforce control flow restrictions on CHANGE TEAM

Like DO CONCURRENT and CRITICAL constructs, control flow
into and out of a CHANGE TEAM construct is disallowed.
---
 flang/lib/Semantics/check-coarray.cpp  | 51 ++++++++++++++++++++++++--
 flang/lib/Semantics/check-coarray.h    | 21 +----------
 flang/test/Semantics/change_team02.f90 | 24 ++++++++++++
 3 files changed, 73 insertions(+), 23 deletions(-)
 create mode 100644 flang/test/Semantics/change_team02.f90

diff --git a/flang/lib/Semantics/check-coarray.cpp b/flang/lib/Semantics/check-coarray.cpp
index 833a08899308a..789f3449772ca 100644
--- a/flang/lib/Semantics/check-coarray.cpp
+++ b/flang/lib/Semantics/check-coarray.cpp
@@ -63,6 +63,39 @@ class CriticalBodyEnforce {
   parser::CharBlock criticalSourcePosition_;
 };
 
+class ChangeTeamBodyEnforce {
+public:
+  ChangeTeamBodyEnforce(
+      SemanticsContext &context, parser::CharBlock changeTeamSourcePosition)
+      : context_{context}, changeTeamSourcePosition_{changeTeamSourcePosition} {
+  }
+  std::set<parser::Label> labels() { return labels_; }
+  template <typename T> bool Pre(const T &) { return true; }
+  template <typename T> void Post(const T &) {}
+
+  template <typename T> bool Pre(const parser::Statement<T> &statement) {
+    currentStatementSourcePosition_ = statement.source;
+    if (statement.label.has_value()) {
+      labels_.insert(*statement.label);
+    }
+    return true;
+  }
+
+  void Post(const parser::ReturnStmt &) {
+    context_
+        .Say(currentStatementSourcePosition_,
+            "RETURN statement is not allowed in a CHANGE TEAM construct"_err_en_US)
+        .Attach(
+            changeTeamSourcePosition_, "Enclosing CHANGE TEAM construct"_en_US);
+  }
+
+private:
+  SemanticsContext &context_;
+  std::set<parser::Label> labels_;
+  parser::CharBlock currentStatementSourcePosition_;
+  parser::CharBlock changeTeamSourcePosition_;
+};
+
 template <typename T>
 static void CheckTeamType(SemanticsContext &context, const T &x) {
   if (const auto *expr{GetExpr(context, x)}) {
@@ -361,17 +394,29 @@ void CoarrayChecker::Leave(const parser::FormTeamStmt &x) {
 
 void CoarrayChecker::Enter(const parser::CriticalConstruct &x) {
   auto &criticalStmt{std::get<parser::Statement<parser::CriticalStmt>>(x.t)};
-
   const parser::Block &block{std::get<parser::Block>(x.t)};
   CriticalBodyEnforce criticalBodyEnforce{context_, criticalStmt.source};
   parser::Walk(block, criticalBodyEnforce);
-
-  // C1119
+  parser::Walk(std::get<parser::Statement<parser::EndCriticalStmt>>(x.t),
+      criticalBodyEnforce);
   LabelEnforce criticalLabelEnforce{
       context_, criticalBodyEnforce.labels(), criticalStmt.source, "CRITICAL"};
   parser::Walk(block, criticalLabelEnforce);
 }
 
+void CoarrayChecker::Enter(const parser::ChangeTeamConstruct &x) {
+  auto &changeTeamStmt{
+      std::get<parser::Statement<parser::ChangeTeamStmt>>(x.t)};
+  const parser::Block &block{std::get<parser::Block>(x.t)};
+  ChangeTeamBodyEnforce changeTeamBodyEnforce{context_, changeTeamStmt.source};
+  parser::Walk(block, changeTeamBodyEnforce);
+  parser::Walk(std::get<parser::Statement<parser::EndChangeTeamStmt>>(x.t),
+      changeTeamBodyEnforce);
+  LabelEnforce changeTeamLabelEnforce{context_, changeTeamBodyEnforce.labels(),
+      changeTeamStmt.source, "CHANGE TEAM"};
+  parser::Walk(block, changeTeamLabelEnforce);
+}
+
 // Check that coarray names and selector names are all distinct.
 void CoarrayChecker::CheckNamesAreDistinct(
     const std::list<parser::CoarrayAssociation> &list) {
diff --git a/flang/lib/Semantics/check-coarray.h b/flang/lib/Semantics/check-coarray.h
index a968585b48be7..f156959019383 100644
--- a/flang/lib/Semantics/check-coarray.h
+++ b/flang/lib/Semantics/check-coarray.h
@@ -12,26 +12,6 @@
 #include "flang/Semantics/semantics.h"
 #include <list>
 
-namespace Fortran::parser {
-class CharBlock;
-class MessageFixedText;
-struct ChangeTeamStmt;
-struct CriticalStmt;
-struct CoarrayAssociation;
-struct EndChangeTeamStmt;
-struct EventPostStmt;
-struct EventWaitStmt;
-struct FormTeamStmt;
-struct ImageSelector;
-struct NotifyWaitStmt;
-struct SyncAllStmt;
-struct SyncImagesStmt;
-struct SyncMemoryStmt;
-struct SyncTeamStmt;
-struct LockStmt;
-struct UnlockStmt;
-} // namespace Fortran::parser
-
 namespace Fortran::semantics {
 
 class CoarrayChecker : public virtual BaseChecker {
@@ -53,6 +33,7 @@ class CoarrayChecker : public virtual BaseChecker {
   void Leave(const parser::FormTeamStmt &);
 
   void Enter(const parser::CriticalConstruct &);
+  void Enter(const parser::ChangeTeamConstruct &);
 
 private:
   SemanticsContext &context_;
diff --git a/flang/test/Semantics/change_team02.f90 b/flang/test/Semantics/change_team02.f90
new file mode 100644
index 0000000000000..a4cba3639f27b
--- /dev/null
+++ b/flang/test/Semantics/change_team02.f90
@@ -0,0 +1,24 @@
+! RUN: %python %S/test_errors.py %s %flang_fc1
+subroutine test
+  use, intrinsic :: iso_fortran_env, only: team_type
+  type(team_type) team
+loop1: do j = 1, 1
+    goto 1 ! ok
+1 construct2: change team (team)
+      goto 2 ! ok
+      exit construct2 ! ok
+      !ERROR: EXIT must not leave a CHANGE TEAM statement
+      exit loop1
+      !ERROR: EXIT must not leave a CHANGE TEAM statement
+      exit
+      !ERROR: CYCLE must not leave a CHANGE TEAM statement
+      cycle
+      !ERROR: RETURN statement is not allowed in a CHANGE TEAM construct
+      return
+      !ERROR: Control flow escapes from CHANGE TEAM
+      goto 3
+      !ERROR: Control flow escapes from CHANGE TEAM
+      write(*,*,err=3)
+2   end team construct2
+3 end do loop1
+end