refactored to more closley follow the IR PatternMatch implementation

Author: Aidan

llvm/include/llvm/CodeGen/SDPatternMatch.h

@@ -548,49 +548,39 @@ struct BinaryOpc_match {
 };
 
 /// Matching while capturing mask
-template <typename T0, typename T1> struct SDShuffle_match {
+template <typename T0, typename T1, typename T2> struct SDShuffle_match {
   T0 Op1;
   T1 Op2;
+  T2 Mask;
 
-  ArrayRef<int> &CapturedMask;
-
-  // capturing mask
-  SDShuffle_match(const T0 &Op1, const T1 &Op2, ArrayRef<int> &MaskRef)
-      : Op1(Op1), Op2(Op2), CapturedMask(MaskRef) {}
+  SDShuffle_match(const T0 &Op1, const T1 &Op2, const T2 &Mask)
+      : Op1(Op1), Op2(Op2), Mask(Mask) {}
 
   template <typename MatchContext>
   bool match(const MatchContext &Ctx, SDValue N) {
     if (auto *I = dyn_cast<ShuffleVectorSDNode>(N)) {
-      if (Op1.match(Ctx, I->getOperand(0)) &&
-          Op2.match(Ctx, I->getOperand(1))) {
-        CapturedMask = I->getMask();
-        return true;
-      }
+      return Op1.match(Ctx, I->getOperand(0)) &&
+             Op2.match(Ctx, I->getOperand(1)) && Mask.match(I->getMask());
     }
     return false;
   }
 };
 
-/// Matching against a specific match
-template <typename T0, typename T1> struct SDShuffle_maskMatch {
-  T0 Op1;
-  T1 Op2;
-  ArrayRef<int> SpecificMask;
-
-  SDShuffle_maskMatch(const T0 &Op1, const T1 &Op2, ArrayRef<int> Mask)
-      : Op1(Op1), Op2(Op2), SpecificMask(Mask) {}
-
-  template <typename MatchContext>
-  bool match(const MatchContext &Ctx, SDValue N) {
-    if (auto *I = dyn_cast<ShuffleVectorSDNode>(N)) {
-      return Op1.match(Ctx, I->getOperand(0)) &&
-             Op2.match(Ctx, I->getOperand(1)) &&
-             std::equal(SpecificMask.begin(), SpecificMask.end(),
-                        I->getMask().begin(), I->getMask().end());
-    }
-    return false;
+struct m_Mask {
+  ArrayRef<int> &MaskRef;
+  m_Mask(ArrayRef<int> &MaskRef) : MaskRef(MaskRef) {}
+  bool match(ArrayRef<int> Mask) {
+    MaskRef = Mask;
+    return true;
   }
 };
+
+struct m_SpecificMask {
+  ArrayRef<int> &MaskRef;
+  m_SpecificMask(ArrayRef<int> &MaskRef) : MaskRef(MaskRef) {}
+  bool match(ArrayRef<int> Mask) { return MaskRef == Mask; }
+};
+
 template <typename LHS_P, typename RHS_P, typename Pred_t,
           bool Commutable = false, bool ExcludeChain = false>
 struct MaxMin_match {
@@ -842,15 +832,14 @@ inline BinaryOpc_match<LHS, RHS> m_FRem(const LHS &L, const RHS &R) {
 }
 
 template <typename V1_t, typename V2_t>
-inline SDShuffle_match<V1_t, V2_t> m_Shuffle(const V1_t &v1, const V2_t &v2,
-                                             ArrayRef<int> &mask) {
-  return SDShuffle_match<V1_t, V2_t>(v1, v2, mask);
+inline BinaryOpc_match<V1_t, V2_t> m_Shuffle(const V1_t &v1, const V2_t &v2) {
+  return BinaryOpc_match<V1_t, V2_t>(ISD::VECTOR_SHUFFLE, v1, v2);
 }
 
-template <typename V1_t, typename V2_t>
-inline SDShuffle_maskMatch<V1_t, V2_t>
-m_ShuffleSpecificMask(const V1_t &v1, const V2_t &v2, ArrayRef<int> mask) {
-  return SDShuffle_maskMatch<V1_t, V2_t>(v1, v2, mask);
+template <typename V1_t, typename V2_t, typename Mask_t>
+inline SDShuffle_match<V1_t, V2_t, Mask_t>
+m_Shuffle(const V1_t &v1, const V2_t &v2, const Mask_t &mask) {
+  return SDShuffle_match<V1_t, V2_t, Mask_t>(v1, v2, mask);
 }
 
 template <typename LHS, typename RHS>

llvm/unittests/CodeGen/SelectionDAGPatternMatchTest.cpp

@@ -125,6 +125,7 @@ TEST_F(SelectionDAGPatternMatchTest, matchVecShuffle) {
   auto VInt32VT = EVT::getVectorVT(Context, Int32VT, 4);
   const std::array<int, 4> MaskData = {2, 0, 3, 1};
   const std::array<int, 4> OtherMaskData = {1, 2, 3, 4};
+  ArrayRef<int> Mask(MaskData);
   ArrayRef<int> CapturedMask;
 
   SDValue V0 = DAG->getCopyFromReg(DAG->getEntryNode(), DL, 1, VInt32VT);
@@ -133,13 +134,11 @@ TEST_F(SelectionDAGPatternMatchTest, matchVecShuffle) {
       DAG->getVectorShuffle(VInt32VT, DL, V0, V1, MaskData);
 
   using namespace SDPatternMatch;
+  EXPECT_TRUE(sd_match(VecShuffleWithMask, m_Shuffle(m_Value(), m_Value())));
   EXPECT_TRUE(sd_match(VecShuffleWithMask,
-                       m_Shuffle(m_Value(), m_Value(), CapturedMask)));
+                       m_Shuffle(m_Value(), m_Value(), m_Mask(CapturedMask))));
   EXPECT_TRUE(sd_match(VecShuffleWithMask,
-                       m_ShuffleSpecificMask(m_Value(), m_Value(), MaskData)));
-  EXPECT_FALSE(
-      sd_match(VecShuffleWithMask,
-               m_ShuffleSpecificMask(m_Value(), m_Value(), OtherMaskData)));
+                       m_Shuffle(m_Value(), m_Value(), m_SpecificMask(Mask))));
   EXPECT_TRUE(std::equal(MaskData.begin(), MaskData.end(), CapturedMask.begin(),
                          CapturedMask.end()));
   EXPECT_FALSE(std::equal(OtherMaskData.begin(), OtherMaskData.end(),