[flang-commits] [flang] [flang][OpenMP] Reassociate logical ATOMIC update expressions (PR #156961)

[Eugene Epshteyn via flang-commits]

================
@@ -167,37 +181,53 @@ struct ReassocRewriter : public evaluate::rewrite::Identity {
       }
       return common::visit(
           [&](auto &&s) {
-            using Expr = evaluate::Expr<T>;
-            using TypeS = llvm::remove_cvref_t<decltype(s)>;
-            // This visitor has to be semantically correct for all possible
-            // types of s even though at runtime s will only be one of the
-            // matched types.
-            // Limit the construction to the operation types that we tried
-            // to match (otherwise TypeS(op1, op2) would fail for non-binary
-            // operations).
-            if constexpr (common::HasMember<TypeS, MatchTypes>) {
-              Expr atom{*sub[atomIdx].ref};
-              Expr op1{*sub[(atomIdx + 1) % 3].ref};
-              Expr op2{*sub[(atomIdx + 2) % 3].ref};
-              return Expr(
-                  TypeS(atom, Expr(TypeS(std::move(op1), std::move(op2)))));
-            } else {
-              return Expr(TypeS(s));
-            }
+            // Build the new expression from the matched components.
+            return Reconstruct<T, MatchTypes>(s, *sub[atomIdx].ref,
+                *sub[(atomIdx + 1) % 3].ref, *sub[(atomIdx + 2) % 3].ref);
           },
           evaluate::match::deparen(x).u);
     }
     return Id::operator()(std::move(x), u);
   }
 
   template <typename T, typename U,
-      typename = std::enable_if_t<!is_numeric_v<T>>>
+      typename = std::enable_if_t<!is_numeric_v<T> && !is_logical_v<T>>>
   evaluate::Expr<T> operator()(
       evaluate::Expr<T> &&x, const U &u, NonIntegralTag = {}) {
     return Id::operator()(std::move(x), u);
   }
 
 private:
+  template <typename T, typename MatchTypes, typename S>
+  evaluate::Expr<T> Reconstruct(const S &op, evaluate::Expr<T> atom,
+      evaluate::Expr<T> op1, evaluate::Expr<T> op2) {
+    using TypeS = llvm::remove_cvref_t<decltype(op)>;
+    // This function has to be semantically correct for all possible types
+    // of S even though at runtime s will only be one of the matched types.
+    // Limit the construction to the operation types that we tried to match
+    // (otherwise TypeS(op1, op2) would fail for non-binary operations).
+    if constexpr (!common::HasMember<TypeS, MatchTypes>) {
+      return evaluate::Expr<T>(TypeS(op));
+    } else if constexpr (is_logical_v<T>) {
+      constexpr int K{T::kind};
+      if constexpr (std::is_same_v<TypeS, evaluate::LogicalOperation<K>>) {
+        // Logical operators take an extra argument in their constructor,
+        // so they need their own reconstruction code.
+        common::LogicalOperator opCode{op.logicalOperator};
+        return evaluate::Expr<T>(TypeS( //
+            opCode, std::move(atom),
+            evaluate::Expr<T>(TypeS( //
+                opCode, std::move(op1), std::move(op2)))));
+      }
+    } else {
+      // Generic reconstruction.
+      return evaluate::Expr<T>(TypeS( //
+          std::move(atom),
+          evaluate::Expr<T>(TypeS( //
----------------
eugeneepshteyn wrote:

Empty comments at the end of the line?

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