[DAG] SDPatternMatch - Fix m_Reassociatable mismatching

llvm/include/llvm/CodeGen/SDPatternMatch.h

@@ -1315,19 +1315,12 @@ template <typename... PatternTs> struct ReassociatableOpc_match {
     if (Leaves.size() != NumPatterns)
       return false;
 
-    // Matches[I][J] == true iff sd_context_match(Leaves[I], Ctx,
-    // std::get<J>(Patterns)) == true
-    std::array<SmallBitVector, NumPatterns> Matches;
-    for (size_t I = 0; I != NumPatterns; I++) {
-      std::apply(
-          [&](auto &...P) {
-            (Matches[I].push_back(sd_context_match(Leaves[I], Ctx, P)), ...);
-          },
-          Patterns);
-    }
-
     SmallBitVector Used(NumPatterns);
-    return reassociatableMatchHelper(Matches, Used);
+    return std::apply(
+        [&](auto &...P) -> bool {
+          return reassociatableMatchHelper(Ctx, Leaves, Used, P...);
+        },
+        Patterns);
   }
 
   void collectLeaves(SDValue V, SmallVector<SDValue> &Leaves) {
@@ -1339,21 +1332,29 @@ template <typename... PatternTs> struct ReassociatableOpc_match {
     }
   }
 
+  // Searchs for a matching leaf for every sub-pattern.
+  template <typename MatchContext, typename PatternHd, typename... PatternTl>
   [[nodiscard]] inline bool
-  reassociatableMatchHelper(ArrayRef<SmallBitVector> Matches,
-                            SmallBitVector &Used, size_t Curr = 0) {
-    if (Curr == Matches.size())
-      return true;
-    for (size_t Match = 0, N = Matches[Curr].size(); Match < N; Match++) {
-      if (!Matches[Curr][Match] || Used[Match])
+  reassociatableMatchHelper(const MatchContext &Ctx, ArrayRef<SDValue> Leaves,
+                            SmallBitVector &Used, PatternHd &HeadPattern,
+                            PatternTl &...TailPatterns) {
+    for (size_t Match = 0, N = Used.size(); Match < N; Match++) {
+      if (Used[Match] || !(sd_context_match(Leaves[Match], Ctx, HeadPattern)))
         continue;
       Used[Match] = true;
-      if (reassociatableMatchHelper(Matches, Used, Curr + 1))
+      if (reassociatableMatchHelper(Ctx, Leaves, Used, TailPatterns...))
         return true;
       Used[Match] = false;
     }
     return false;
   }
+
+  template <typename MatchContext>
+  [[nodiscard]] inline bool reassociatableMatchHelper(const MatchContext &Ctx,
+                                                      ArrayRef<SDValue> Leaves,
+                                                      SmallBitVector &Used) {
+    return true;
+  }
 };
 
 template <typename... PatternTs>

llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp

@@ -842,6 +842,28 @@ TEST_F(SelectionDAGPatternMatchTest, matchReassociatableOp) {
   EXPECT_TRUE(sd_match(
       MUL, m_ReassociatableMul(m_Value(), m_Value(), m_Value(), m_Value())));
 
+  // (Op0 + Op1) + Op0 binds correctly, allowing commutation
+  SDValue ADD010 = DAG->getNode(ISD::ADD, DL, Int32VT, ADD01, Op0);
+  SDValue A, B;
+  EXPECT_TRUE(sd_match(
+      ADD010, m_ReassociatableAdd(m_Value(A), m_Value(B), m_Deferred(A))));
+  EXPECT_EQ(Op0, A);
+  EXPECT_EQ(Op1, B);
+  EXPECT_NE(A, B);
+
+  A.setNode(nullptr);
+  B.setNode(nullptr);
+  EXPECT_TRUE(sd_match(
+      ADD010, m_ReassociatableAdd(m_Value(A), m_Value(B), m_Deferred(B))));
+  EXPECT_EQ(Op0, B);
+  EXPECT_EQ(Op1, A);
+  EXPECT_NE(A, B);
+
+  A.setNode(nullptr);
+  B.setNode(nullptr);
+  EXPECT_FALSE(sd_match(
+      ADD010, m_ReassociatableAdd(m_Value(A), m_Deferred(A), m_Deferred(A))));
+
   // Op0 * (Op1 * (Op2 * Op3))
   SDValue MUL123 = DAG->getNode(ISD::MUL, DL, Int32VT, Op1, MUL23);
   SDValue MUL0123 = DAG->getNode(ISD::MUL, DL, Int32VT, Op0, MUL123);