[llvm] [AggressiveInstCombine] Expand strchr/memchr with small constant strings (PR #98501)

[via llvm-commits]

================
@@ -1103,6 +1108,85 @@ void StrNCmpInliner::inlineCompare(Value *LHS, StringRef RHS, uint64_t N,
   }
 }
 
+/// Convert strchr/memchr with a small constant string into a switch
+static bool foldStrChr(CallInst *Call, LibFunc Func, DomTreeUpdater *DTU,
+                       const DataLayout &DL) {
+  assert((Func == LibFunc_strchr || Func == LibFunc_memchr) &&
+         "Unexpected LibFunc");
+  if (isa<Constant>(Call->getArgOperand(1)))
+    return false;
+
+  StringRef Str;
+  Value *Base = Call->getArgOperand(0);
+  if (!getConstantStringInfo(Base, Str, /*TrimAtNul=*/Func == LibFunc_strchr))
+    return false;
+
+  uint64_t N = Str.size();
+  if (Func == LibFunc_memchr) {
+    if (auto *ConstInt = dyn_cast<ConstantInt>(Call->getArgOperand(2))) {
+      uint64_t Val = ConstInt->getZExtValue();
+      // Ignore the case that n is larger than the size of string.
+      if (Val > N)
+        return false;
+      N = Val;
+    } else
+      return false;
+  }
+
+  if (N > StrChrInlineThreshold)
+    return false;
+
+  BasicBlock *BB = Call->getParent();
+  BasicBlock *BBNext = SplitBlock(BB, Call, DTU);
+  IRBuilder<> IRB(BB);
+  IntegerType *ByteTy = IRB.getInt8Ty();
+  BB->getTerminator()->eraseFromParent();
+  SwitchInst *SI = IRB.CreateSwitch(
+      IRB.CreateTrunc(Call->getArgOperand(1), ByteTy), BBNext, N);
+  Type *IndexTy = DL.getIndexType(Call->getType());
+
+  PHINode *PHI = PHINode::Create(Call->getType(), 2, "", BBNext->begin());
+  PHI->addIncoming(Constant::getNullValue(Call->getType()), BB);
+
+  SmallVector<DominatorTree::UpdateType, 8> Updates;
+
+  BasicBlock *BBSuccess = BasicBlock::Create(
+      Call->getContext(), "strchr.success", BB->getParent(), BBNext);
+  IRB.SetInsertPoint(BBSuccess);
+  PHINode *IndexPHI = IRB.CreatePHI(IndexTy, N);
+  Value *FirstOccursLocation = IRB.CreateInBoundsPtrAdd(Base, IndexPHI);
+  PHI->addIncoming(FirstOccursLocation, BBSuccess);
+  IRB.CreateBr(BBNext);
+  if (DTU)
+    Updates.push_back({DominatorTree::Insert, BBSuccess, BBNext});
+
+  SmallPtrSet<ConstantInt *, 4> Cases;
+  for (uint64_t I = 0; I < N; ++I) {
+    ConstantInt *CaseVal = ConstantInt::get(ByteTy, Str[I]);
+    if (!Cases.insert(CaseVal).second)
+      continue;
+
+    BasicBlock *BBCase = BasicBlock::Create(Call->getContext(), "strchr.case",
+                                            BB->getParent(), BBNext);
+    SI->addCase(CaseVal, BBCase);
+    IRB.SetInsertPoint(BBCase);
+    IndexPHI->addIncoming(ConstantInt::get(IndexTy, I), BBCase);
+    IRB.CreateBr(BBSuccess);
+    if (DTU) {
+      Updates.push_back({DominatorTree::Insert, BB, BBCase});
+      Updates.push_back({DominatorTree::Insert, BBCase, BBSuccess});
+    }
+  }
----------------
goldsteinn wrote:

Does the new code get properly optimized if the call is only used for equality?

Can you add a few tests / modify a few of your existing tests to test that case?

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