[InstCombine] Fold nested selects

https://alive2.llvm.org/ce/z/GjCXkB
https://alive2.llvm.org/ce/z/Guz2tt

Fixes #59393

Author: LebedevRI

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

@@ -2669,6 +2669,80 @@ foldRoundUpIntegerWithPow2Alignment(SelectInst &SI,
   return R;
 }
 
+/// Look for patterns like
+///   %outer.cond = select i1 %inner.cond, i1 %alt.cond, i1 false
+///   %inner.sel = select i1 %inner.cond, i8 %inner.sel.t, i8 %inner.sel.f
+///   %outer.sel = select i1 %outer.cond, i8 %outer.sel.t, i8 %inner.sel
+/// and rewrite it as
+///   %inner.sel = select i1 %cond.alternative, i8 %sel.outer.t, i8 %sel.inner.t
+///   %sel.outer = select i1 %cond.inner, i8 %inner.sel, i8 %sel.inner.f
+static Instruction *foldNestedSelects(SelectInst &OuterSel,
+                                      InstCombiner::BuilderTy &Builder) {
+  // We must start with a `select`.
+  Value *OuterCond, *InnerSel, *OuterSelFalseVal;
+  match(&OuterSel, m_Select(m_Value(OuterCond), m_Value(InnerSel),
+                            m_Value(OuterSelFalseVal)));
+
+  // Canonicalize inversion of the outermost `select`'s condition.
+  if (match(OuterCond, m_Not(m_Value(OuterCond))))
+    std::swap(InnerSel, OuterSelFalseVal);
+
+  auto m_c_LogicalOp = [](auto L, auto R) {
+    return m_CombineOr(m_c_LogicalAnd(L, R), m_c_LogicalOr(L, R));
+  };
+
+  // The condition of the outermost select must be an `and`/`or`.
+  if (!match(OuterCond, m_c_LogicalOp(m_Value(), m_Value())))
+    return nullptr;
+
+  // To simplify logic, prefer the pattern variant with an `or`.
+  bool IsAndVariant = match(OuterCond, m_LogicalAnd());
+  if (match(OuterCond, m_LogicalAnd()))
+    std::swap(InnerSel, OuterSelFalseVal);
+
+  // Profitability check - avoid increasing instruction count.
+  if (none_of(ArrayRef<Value *>({OuterCond, InnerSel}),
+              [](Value *V) { return V->hasOneUse(); }))
+    return nullptr;
+
+  // The appropriate hand of the outermost `select` must be a select itself.
+  Value *InnerCond, *InnerSelTrueVal, *InnerSelFalseVal;
+  if (!match(InnerSel, m_Select(m_Value(InnerCond), m_Value(InnerSelTrueVal),
+                                m_Value(InnerSelFalseVal))))
+    return nullptr;
+
+  // Canonicalize inversion of the innermost `select`'s condition.
+  if (match(InnerCond, m_Not(m_Value(InnerCond))))
+    std::swap(InnerSelTrueVal, InnerSelFalseVal);
+
+  Value *AltCond = nullptr;
+  auto matchOuterCond = [OuterCond, m_c_LogicalOp, &AltCond](auto m_InnerCond) {
+    return match(OuterCond, m_c_LogicalOp(m_InnerCond, m_Value(AltCond)));
+  };
+
+  // Finally, match the condition that was driving the outermost `select`,
+  // it should be a logical operation between the condition that was driving
+  // the innermost `select` (after accounting for the possible inversions
+  // of the condition), and some other condition.
+  if (matchOuterCond(m_Specific(InnerCond))) {
+    // Done!
+  } else if (Value * NotInnerCond; matchOuterCond(m_CombineAnd(
+                 m_Not(m_Specific(InnerCond)), m_Value(NotInnerCond)))) {
+    // Done!
+    std::swap(InnerSelTrueVal, InnerSelFalseVal);
+    InnerCond = NotInnerCond;
+  } else // Not the pattern we were looking for.
+    return nullptr;
+
+  Value *SelInner = Builder.CreateSelect(
+      AltCond, IsAndVariant ? OuterSelFalseVal : InnerSelFalseVal,
+      IsAndVariant ? InnerSelTrueVal : OuterSelFalseVal);
+  SelInner->takeName(InnerSel);
+  return SelectInst::Create(InnerCond,
+                            IsAndVariant ? SelInner : InnerSelTrueVal,
+                            IsAndVariant ? InnerSelFalseVal : SelInner);
+}
+
 Instruction *InstCombinerImpl::foldSelectOfBools(SelectInst &SI) {
   Value *CondVal = SI.getCondition();
   Value *TrueVal = SI.getTrueValue();
@@ -3292,5 +3366,8 @@ Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
     }
   }
 
+  if (Instruction *I = foldNestedSelects(SI, Builder))
+    return I;
+
   return nullptr;
 }