[llvm] [AggressiveInstCombine] Memchr inline (PR #130525)

[Yingwei Zheng via llvm-commits]

================
@@ -1106,79 +1105,46 @@ void StrNCmpInliner::inlineCompare(Value *LHS, StringRef RHS, uint64_t N,
   }
 }
 
-/// Convert memchr with a small constant string into a switch
 static bool foldMemChr(CallInst *Call, DomTreeUpdater *DTU,
                        const DataLayout &DL) {
-  if (isa<Constant>(Call->getArgOperand(1)))
+  Value *Ptr = Call->getArgOperand(0);
+  Value *Val = Call->getArgOperand(1);
+  Value *Len = Call->getArgOperand(2);
+  
+  // If length is not a constant, we can't do the optimization
+  auto *LenC = dyn_cast<ConstantInt>(Len);
+  if (!LenC)
     return false;
-
-  StringRef Str;
-  Value *Base = Call->getArgOperand(0);
-  if (!getConstantStringInfo(Base, Str, /*TrimAtNul=*/false))
-    return false;
-
-  uint64_t N = Str.size();
-  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 > MemChrInlineThreshold)
-    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());
-  SmallVector<DominatorTree::UpdateType, 8> Updates;
-
-  BasicBlock *BBSuccess = BasicBlock::Create(
-      Call->getContext(), "memchr.success", BB->getParent(), BBNext);
-  IRB.SetInsertPoint(BBSuccess);
-  PHINode *IndexPHI = IRB.CreatePHI(IndexTy, N, "memchr.idx");
-  Value *FirstOccursLocation = IRB.CreateInBoundsPtrAdd(Base, IndexPHI);
-  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(), "memchr.case",
-                                            BB->getParent(), BBSuccess);
-    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});
+    
+  uint64_t Length = LenC->getZExtValue();
+  
+  // Check if this is a small memchr we should inline
+  if (Length <= MemChrInlineThreshold) {
+    IRBuilder<> IRB(Call);
+    
+    // Truncate the search value to i8
+    Value *ByteVal = IRB.CreateTrunc(Val, IRB.getInt8Ty());
+    
+    // Initialize result to null
+    Value *Result = ConstantPointerNull::get(cast<PointerType>(Call->getType()));
+    
+    // For each byte up to Length
+    for (unsigned i = 0; i < Length; i++) {
+      Value *CurPtr = i == 0 ? Ptr : 
+                      IRB.CreateGEP(IRB.getInt8Ty(), Ptr, 
+                                  ConstantInt::get(DL.getIndexType(Call->getType()), i));
+      Value *CurByte = IRB.CreateLoad(IRB.getInt8Ty(), CurPtr);
----------------
dtcxzyw wrote:

I think it is dangerous to put these load instructions in a basic block.
Consider the following case:
```
const char *arr = "a";
return memchr(arr, 'a', 3);
```
It is well-defined because it returns at the first occurrence.
After this transformation, we transform this call into three loads, which causes out-of-bound 
memory access.
See also https://en.cppreference.com/w/c/string/byte/memchr.

To avoid UB, we should transform this call into a if chain.

BTW, your implementation returns the *last* occurrence of ch :(


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