1 files changed, 36 insertions, 8 deletions

diff --git a/clang/lib/AST/StmtOpenACC.cpp b/clang/lib/AST/StmtOpenACC.cpp
index 462a10d..39dfa19 100644
--- a/clang/lib/AST/StmtOpenACC.cpp
+++ b/clang/lib/AST/StmtOpenACC.cpp
@@ -326,16 +326,30 @@ OpenACCAtomicConstruct *OpenACCAtomicConstruct::Create(
 
 static std::pair<const Expr *, const Expr *> getBinaryOpArgs(const Expr *Op) {
   if (const auto *BO = dyn_cast<BinaryOperator>(Op)) {
-    assert(BO->getOpcode() == BO_Assign);
+    assert(BO->isAssignmentOp());
     return {BO->getLHS(), BO->getRHS()};
   }
 
   const auto *OO = cast<CXXOperatorCallExpr>(Op);
-  assert(OO->getOperator() == OO_Equal);
-
+  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
@@ -343,18 +357,17 @@ 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::Update:
   case OpenACCAtomicKind::Capture:
-    assert(false && "Only 'read'/'write' have 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).
     std::pair<const Expr *, const Expr *> BinaryArgs =
-        getBinaryOpArgs(cast<const Expr>(getAssociatedStmt()));
+        getBinaryOpArgs(AssocStmt);
     // We want the L-value for each side, so we ignore implicit casts.
     return {BinaryArgs.first->IgnoreImpCasts(),
             BinaryArgs.second->IgnoreImpCasts(), /*expr=*/nullptr};
@@ -364,13 +377,28 @@ OpenACCAtomicConstruct::getAssociatedStmtInfo() const {
     // 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(cast<const Expr>(getAssociatedStmt()));
+        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");