[llvm] [IR] Add BooleanMap matcher (PR #184463)

[Gergo Stomfai via llvm-commits]

https://github.com/stomfaig updated https://github.com/llvm/llvm-project/pull/184463

>From 2c7f2afb21b5d1fc294847fa83ed711956e0e219 Mon Sep 17 00:00:00 2001
From: stomfaig <stomfaig at gmail.com>
Date: Tue, 3 Mar 2026 23:35:46 +0000
Subject: [PATCH] add BooleanMap matcher

---
 llvm/include/llvm/IR/PatternMatch.h | 44 +++++++++++++++++++++++++++++
 llvm/unittests/IR/PatternMatch.cpp  | 39 +++++++++++++++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/llvm/include/llvm/IR/PatternMatch.h b/llvm/include/llvm/IR/PatternMatch.h
index ee99d627dde5d..8b257bacbb214 100644
--- a/llvm/include/llvm/IR/PatternMatch.h
+++ b/llvm/include/llvm/IR/PatternMatch.h
@@ -2333,6 +2333,50 @@ m_ZExtOrTruncOrSelf(const OpTy &Op) {
   return m_CombineOr(m_CombineOr(m_ZExt(Op), m_Trunc(Op)), Op);
 }
 
+struct BooleanMap_match {
+  Value *&Cond;
+  Constant *&TrueC;
+  Constant *&FalseC;
+
+  BooleanMap_match(Value *&C, Constant *&TC, Constant *&FC)
+      : Cond(C), TrueC(TC), FalseC(FC) {}
+
+  template <typename OpTy> bool match(OpTy *V) const {
+    // select(Cond, TrueC, FalseC) — captures both constants directly
+    if (PatternMatch::match(
+            V, m_Select(m_Value(Cond), m_Constant(TrueC), m_Constant(FalseC))))
+      return true;
+
+    Type *Ty = V->getType();
+    Value *CondV;
+
+    // zext(i1 Cond) is equivalent to select(Cond, 1, 0)
+    if (PatternMatch::match(V, m_ZExt(m_Value(CondV))) &&
+        CondV->getType()->isIntOrIntVectorTy(1)) {
+      Cond = CondV;
+      TrueC = ConstantInt::get(Ty, 1);
+      FalseC = ConstantInt::get(Ty, 0);
+      return true;
+    }
+
+    // sext(i1 Cond) is equivalent to select(Cond, -1, 0)
+    if (PatternMatch::match(V, m_SExt(m_Value(CondV))) &&
+        CondV->getType()->isIntOrIntVectorTy(1)) {
+      Cond = CondV;
+      TrueC = Constant::getAllOnesValue(Ty);
+      FalseC = ConstantInt::get(Ty, 0);
+      return true;
+    }
+
+    return false;
+  }
+};
+
+inline BooleanMap_match m_BooleanMap(Value *&C, Constant *&TrueC,
+                                     Constant *&FalseC) {
+  return BooleanMap_match(C, TrueC, FalseC);
+}
+
 template <typename OpTy>
 inline CastInst_match<OpTy, UIToFPInst> m_UIToFP(const OpTy &Op) {
   return CastInst_match<OpTy, UIToFPInst>(Op);
diff --git a/llvm/unittests/IR/PatternMatch.cpp b/llvm/unittests/IR/PatternMatch.cpp
index 220e35d37c0be..2f6c6327838dd 100644
--- a/llvm/unittests/IR/PatternMatch.cpp
+++ b/llvm/unittests/IR/PatternMatch.cpp
@@ -635,6 +635,45 @@ TEST_F(PatternMatchTest, ZExtSExtSelf) {
   EXPECT_TRUE(m_ZExtOrSExtOrSelf(m_One()).match(One64S));
 }
 
+TEST_F(PatternMatchTest, BooleanMap) {
+  Value *Alloca = IRB.CreateAlloca(IRB.getInt1Ty());
+  Value *SelectCond = IRB.CreateLoad(IRB.getInt1Ty(), Alloca);
+  Value *C1 = IRB.getInt1(1);
+  Value *C2 = IRB.getInt1(0);
+
+  Value *Cond;
+  Constant *T, *F;
+
+  Value *select = IRB.CreateSelect(SelectCond, C1, C2);
+  EXPECT_TRUE(m_BooleanMap(Cond, T, F).match(select));
+  EXPECT_EQ(Cond, SelectCond);
+  EXPECT_EQ(T, C1);
+  EXPECT_EQ(F, C2);
+
+  Value *zext = IRB.CreateZExt(C1, IntegerType::getInt64Ty(Ctx));
+  EXPECT_TRUE(m_BooleanMap(Cond, T, F).match(zext));
+  EXPECT_EQ(Cond, C1);
+  EXPECT_EQ(T, IRB.getInt64(1));
+  EXPECT_EQ(F, IRB.getInt64(0));
+
+  Value *sext = IRB.CreateSExt(C1, IntegerType::getInt64Ty(Ctx));
+  EXPECT_TRUE(m_BooleanMap(Cond, T, F).match(sext));
+  EXPECT_EQ(Cond, C1);
+  EXPECT_EQ(T, IRB.getInt64(-1));
+  EXPECT_EQ(F, IRB.getInt64(0));
+
+  // Negative: zext/sext of non-i1 should not match
+  Value *I8Val = IRB.getInt8(1);
+  Value *ZExtI8 = IRB.CreateZExt(I8Val, IntegerType::getInt64Ty(Ctx));
+  EXPECT_FALSE(m_BooleanMap(Cond, T, F).match(ZExtI8));
+  Value *SExtI8 = IRB.CreateSExt(I8Val, IntegerType::getInt64Ty(Ctx));
+  EXPECT_FALSE(m_BooleanMap(Cond, T, F).match(SExtI8));
+
+  // Negative: plain arithmetic instruction should not match
+  Value *Add = IRB.CreateAdd(IRB.getInt32(1), IRB.getInt32(2));
+  EXPECT_FALSE(m_BooleanMap(Cond, T, F).match(Add));
+}
+
 TEST_F(PatternMatchTest, BitCast) {
   Value *OneDouble = ConstantFP::get(IRB.getDoubleTy(), APFloat(1.0));
   Value *ScalableDouble = ConstantFP::get(