1 files changed, 62 insertions, 14 deletions

diff --git a/clang/lib/AST/StmtOpenACC.cpp b/clang/lib/AST/StmtOpenACC.cpp
index 2b56c1e..39dfa19 100644
--- a/clang/lib/AST/StmtOpenACC.cpp
+++ b/clang/lib/AST/StmtOpenACC.cpp
@@ -324,6 +324,32 @@ OpenACCAtomicConstruct *OpenACCAtomicConstruct::Create(
   return Inst;
 }
 
+static std::pair<const Expr *, const Expr *> getBinaryOpArgs(const Expr *Op) {
+  if (const auto *BO = dyn_cast<BinaryOperator>(Op)) {
+    assert(BO->isAssignmentOp());
+    return {BO->getLHS(), BO->getRHS()};
+  }
+
+  const auto *OO = cast<CXXOperatorCallExpr>(Op);
+  assert(OO->isAssignmentOp());
+  return {OO->getArg(0), OO->getArg(1)};
+}
+
+static std::pair<bool, const Expr *> getUnaryOpArgs(const Expr *Op) {
+  if (const auto *UO = dyn_cast<UnaryOperator>(Op))
+    return {true, UO->getSubExpr()};
+
+  if (const auto *OpCall = dyn_cast<CXXOperatorCallExpr>(Op)) {
+    // Post-inc/dec have a second unused argument to differentiate it, so we
+    // accept -- or ++ as unary, or any operator call with only 1 arg.
+    if (OpCall->getNumArgs() == 1 || OpCall->getOperator() != OO_PlusPlus ||
+        OpCall->getOperator() != OO_MinusMinus)
+      return {true, OpCall->getArg(0)};
+  }
+
+  return {false, nullptr};
+}
+
 const OpenACCAtomicConstruct::StmtInfo
 OpenACCAtomicConstruct::getAssociatedStmtInfo() const {
   // This ends up being a vastly simplified version of SemaOpenACCAtomic, since
@@ -331,29 +357,51 @@ OpenACCAtomicConstruct::getAssociatedStmtInfo() const {
   // asserts to ensure we don't get off into the weeds.
   assert(getAssociatedStmt() && "invalid associated stmt?");
 
+  const Expr *AssocStmt = cast<const Expr>(getAssociatedStmt());
   switch (AtomicKind) {
-  case OpenACCAtomicKind::None:
-  case OpenACCAtomicKind::Write:
-  case OpenACCAtomicKind::Update:
   case OpenACCAtomicKind::Capture:
-    assert(false && "Only 'read' has been implemented here");
+    assert(false && "Only 'read'/'write'/'update' have been implemented here");
     return {};
   case OpenACCAtomicKind::Read: {
     // Read only supports the format 'v = x'; where both sides are a scalar
     // expression. This can come in 2 forms; BinaryOperator or
     // CXXOperatorCallExpr (rarely).
-    const Expr *AssignExpr = cast<const Expr>(getAssociatedStmt());
-    if (const auto *BO = dyn_cast<BinaryOperator>(AssignExpr)) {
-      assert(BO->getOpcode() == BO_Assign);
-      return {BO->getLHS()->IgnoreImpCasts(), BO->getRHS()->IgnoreImpCasts()};
-    }
-
-    const auto *OO = cast<CXXOperatorCallExpr>(AssignExpr);
-    assert(OO->getOperator() == OO_Equal);
-
-    return {OO->getArg(0)->IgnoreImpCasts(), OO->getArg(1)->IgnoreImpCasts()};
+    std::pair<const Expr *, const Expr *> BinaryArgs =
+        getBinaryOpArgs(AssocStmt);
+    // We want the L-value for each side, so we ignore implicit casts.
+    return {BinaryArgs.first->IgnoreImpCasts(),
+            BinaryArgs.second->IgnoreImpCasts(), /*expr=*/nullptr};
   }
+  case OpenACCAtomicKind::Write: {
+    // Write supports only the format 'x = expr', where the expression is scalar
+    // type, and 'x' is a scalar l value. As above, this can come in 2 forms;
+    // Binary Operator or CXXOperatorCallExpr.
+    std::pair<const Expr *, const Expr *> BinaryArgs =
+        getBinaryOpArgs(AssocStmt);
+    // We want the L-value for ONLY the X side, so we ignore implicit casts. For
+    // the right side (the expr), we emit it as an r-value so we need to
+    // maintain implicit casts.
+    return {/*v=*/nullptr, BinaryArgs.first->IgnoreImpCasts(),
+            BinaryArgs.second};
   }
+  case OpenACCAtomicKind::None:
+  case OpenACCAtomicKind::Update: {
+    std::pair<bool, const Expr *> UnaryArgs = getUnaryOpArgs(AssocStmt);
+    if (UnaryArgs.first)
+      return {/*v=*/nullptr, UnaryArgs.second->IgnoreImpCasts(),
+              /*expr=*/nullptr};
+
+    std::pair<const Expr *, const Expr *> BinaryArgs =
+        getBinaryOpArgs(AssocStmt);
+    // For binary args, we just store the RHS as an expression (in the
+    // expression slot), since the codegen just wants the whole thing for a
+    // recipe.
+    return {/*v=*/nullptr, BinaryArgs.first->IgnoreImpCasts(),
+            BinaryArgs.second};
+  }
+  }
+
+  llvm_unreachable("unknown OpenACC atomic kind");
 }
 
 OpenACCCacheConstruct *OpenACCCacheConstruct::CreateEmpty(const ASTContext &C,