[InstCombine] Fix memrchr logic error that prevents folding

Correct a logic bug in the memrchr enhancement added in D123629 that makes it ineffective in a subset of cases. Reviewed By: nikic Differential Revision: https://reviews.llvm.org/D128856
llvm · Jun 30, 2022 · 3a743a5 · 3a743a5
1 parent ab2e1c0
commit 3a743a5
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 9 deletions.
diff --git a/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp b/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp
@@ -977,15 +977,16 @@ Value *LibCallSimplifier::optimizeMemRChr(CallInst *CI, IRBuilderBase &B) {
       // Fold memrchr(s, c, N) --> s + Pos for constant N > Pos.
       return B.CreateGEP(B.getInt8Ty(), SrcStr, B.getInt64(Pos));
 
-    if (Str.find(CharC->getZExtValue(), Pos) == StringRef::npos) {
-      // When there is just a single occurrence of C in S, fold
+    if (Str.find(Str[Pos]) == Pos) {
+      // When there is just a single occurrence of C in S, i.e., the one
+      // in Str[Pos], fold
       //   memrchr(s, c, N) --> N <= Pos ? null : s + Pos
       // for nonconstant N.
       Value *Cmp = B.CreateICmpULE(Size, ConstantInt::get(Size->getType(),
-							  Pos),
-				   "memrchr.cmp");
+                                                          Pos),
+                                   "memrchr.cmp");
       Value *SrcPlus = B.CreateGEP(B.getInt8Ty(), SrcStr, B.getInt64(Pos),
-				   "memrchr.ptr_plus");
+                                   "memrchr.ptr_plus");
       return B.CreateSelect(Cmp, NullPtr, SrcPlus, "memrchr.sel");
     }
   }

diff --git a/llvm/test/Transforms/InstCombine/memrchr-3.ll b/llvm/test/Transforms/InstCombine/memrchr-3.ll
@@ -196,8 +196,9 @@ define i8* @fold_memrchr_a12345_0_n(i64 %N) {
 
 define i8* @fold_memrchr_a12345_3_n(i64 %n) {
 ; CHECK-LABEL: @fold_memrchr_a12345_3_n(
-; CHECK-NEXT:    [[RET:%.*]] = call i8* @memrchr(i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 0), i32 3, i64 [[N:%.*]])
-; CHECK-NEXT:    ret i8* [[RET]]
+; CHECK-NEXT:    [[MEMRCHR_CMP:%.*]] = icmp ult i64 [[N:%.*]], 3
+; CHECK-NEXT:    [[MEMRCHR_SEL:%.*]] = select i1 [[MEMRCHR_CMP]], i8* null, i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 2)
+; CHECK-NEXT:    ret i8* [[MEMRCHR_SEL]]
 ;
 
   %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
@@ -210,8 +211,9 @@ define i8* @fold_memrchr_a12345_3_n(i64 %n) {
 
 define i8* @fold_memrchr_a12345_5_n(i64 %n) {
 ; CHECK-LABEL: @fold_memrchr_a12345_5_n(
-; CHECK-NEXT:    [[RET:%.*]] = call i8* @memrchr(i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 0), i32 5, i64 [[N:%.*]])
-; CHECK-NEXT:    ret i8* [[RET]]
+; CHECK-NEXT:    [[MEMRCHR_CMP:%.*]] = icmp ult i64 [[N:%.*]], 5
+; CHECK-NEXT:    [[MEMRCHR_SEL:%.*]] = select i1 [[MEMRCHR_CMP]], i8* null, i8* getelementptr inbounds ([5 x i8], [5 x i8]* @a12345, i64 0, i64 4)
+; CHECK-NEXT:    ret i8* [[MEMRCHR_SEL]]
 ;
 
   %ptr = getelementptr [5 x i8], [5 x i8]* @a12345, i32 0, i32 0
@@ -324,3 +326,17 @@ define i8* @call_memrchr_a123123_2_n(i64 %n) {
   %ret = call i8* @memrchr(i8* %ptr, i32 2, i64 %n)
   ret i8* %ret
 }
+
+
+; And again for 1 to exercise the other edge case.
+
+define i8* @call_memrchr_a123123_1_n(i64 %n) {
+; CHECK-LABEL: @call_memrchr_a123123_1_n(
+; CHECK-NEXT:    [[RET:%.*]] = call i8* @memrchr(i8* getelementptr inbounds ([6 x i8], [6 x i8]* @a123123, i64 0, i64 0), i32 1, i64 [[N:%.*]])
+; CHECK-NEXT:    ret i8* [[RET]]
+;
+
+  %ptr = getelementptr [6 x i8], [6 x i8]* @a123123, i32 0, i32 0
+  %ret = call i8* @memrchr(i8* %ptr, i32 1, i64 %n)
+  ret i8* %ret
+}