[flang-commits] [flang] [Flang][OpenMP] Add Semantic Checks for Atomic Capture Construct (PR #108516)

[Kiran Chandramohan via flang-commits]

================
@@ -1977,6 +1977,72 @@ void OmpStructureChecker::CheckAtomicUpdateStmt(
   ErrIfAllocatableVariable(var);
 }
 
+// TODO: Allow cond-update-stmt once compare clause is supported.
+void OmpStructureChecker::CheckAtomicCaptureConstruct(
+    const parser::OmpAtomicCapture &atomicCaptureConstruct) {
+  const Fortran::parser::AssignmentStmt &stmt1 =
+      std::get<Fortran::parser::OmpAtomicCapture::Stmt1>(
+          atomicCaptureConstruct.t)
+          .v.statement;
+  const auto &stmt1Var{std::get<Fortran::parser::Variable>(stmt1.t)};
+  const auto &stmt1Expr{std::get<Fortran::parser::Expr>(stmt1.t)};
+
+  const Fortran::parser::AssignmentStmt &stmt2 =
+      std::get<Fortran::parser::OmpAtomicCapture::Stmt2>(
+          atomicCaptureConstruct.t)
+          .v.statement;
+  const auto &stmt2Var{std::get<Fortran::parser::Variable>(stmt2.t)};
+  const auto &stmt2Expr{std::get<Fortran::parser::Expr>(stmt2.t)};
+
+  if (Fortran::semantics::checkForSingleVariableOnRHS(stmt1)) {
+    CheckAtomicCaptureStmt(stmt1);
+    if (Fortran::semantics::checkForSymbolMatch(stmt2)) {
+      // Atomic capture construct is of the form [capture-stmt, update-stmt]
+      CheckAtomicUpdateStmt(stmt2);
+    } else {
+      // Atomic capture construct is of the form [capture-stmt, write-stmt]
+      CheckAtomicWriteStmt(stmt2);
+    }
+    // Variable captured in stmt1 should be assigned in stmt2
+    const auto *e{GetExpr(context_, stmt1Expr)};
+    const auto *v{GetExpr(context_, stmt2Var)};
+    if (e && v) {
+      const Symbol &stmt2VarSymbol = evaluate::GetSymbolVector(*v).front();
+      const Symbol &stmt1ExprSymbol = evaluate::GetSymbolVector(*e).front();
+      if (stmt2VarSymbol != stmt1ExprSymbol) {
+        context_.Say(stmt1Expr.source,
+            "Captured variable %s "
+            "expected to be assigned in the second statement of "
+            "atomic capture construct"_err_en_US,
+            stmt1ExprSymbol.name().ToString());
+      }
+    }
+  } else if (Fortran::semantics::checkForSymbolMatch(stmt1) &&
+      Fortran::semantics::checkForSingleVariableOnRHS(stmt2)) {
+    // Atomic capture construct is of the form [update-stmt, capture-stmt]
+    CheckAtomicUpdateStmt(stmt1);
+    CheckAtomicCaptureStmt(stmt2);
+    // Variable updated in stmt1 should be captured in stmt2
+    const auto *e{GetExpr(context_, stmt2Expr)};
+    const auto *v{GetExpr(context_, stmt1Var)};
+    if (e && v) {
+      const Symbol &stmt1VarSymbol = evaluate::GetSymbolVector(*v).front();
+      const Symbol &stmt2ExprSymbol = evaluate::GetSymbolVector(*e).front();
+      if (stmt1VarSymbol != stmt2ExprSymbol)
+        context_.Say(stmt1Var.GetSource(),
+            "Updated variable %s "
+            "expected to be captured in the second statement of "
+            "atomic capture construct"_err_en_US,
+            stmt1Var.GetSource().ToString());
----------------
kiranchandramohan wrote:

`ToString()` is not needed.

https://github.com/llvm/llvm-project/pull/108516