From f69a2df9b9f4be965811be20a815d4e5b7a4e25d Mon Sep 17 00:00:00 2001
From: Juneyoung Lee <aqjune@gmail.com>
Date: Mon, 21 Jun 2021 16:50:54 +0900
Subject: [PATCH] [InstCombine] Fold icmp (select c,const,arg), null if icmp
 arg, null can be simplified

This patch folds icmp (select c,const,arg), null if icmp arg, null can be simplified.

Resolves llvm.org/pr48975.

Reviewed By: nikic, xbolva00

Differential Revision: https://reviews.llvm.org/D96663
---
 .../InstCombine/InstCombineCompares.cpp       | 22 +++++---
 .../InstCombine/assume-icmp-null-select.ll    | 51 +++++++++++++++++++
 2 files changed, 67 insertions(+), 6 deletions(-)
 create mode 100644 test/Transforms/InstCombine/assume-icmp-null-select.ll

diff --git a/lib/Transforms/InstCombine/InstCombineCompares.cpp b/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 189b49e584c..3e8526c9235 100644
--- a/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -3290,14 +3290,24 @@ Instruction *InstCombinerImpl::foldICmpInstWithConstantNotInt(ICmpInst &I) {
     // constant folded and the select turned into a bitwise or.
     Value *Op1 = nullptr, *Op2 = nullptr;
     ConstantInt *CI = nullptr;
-    if (Constant *C = dyn_cast<Constant>(LHSI->getOperand(1))) {
-      Op1 = ConstantExpr::getICmp(I.getPredicate(), C, RHSC);
+
+    auto SimplifyOp = [&](Value *V) {
+      Value *Op = nullptr;
+      if (Constant *C = dyn_cast<Constant>(V)) {
+        Op = ConstantExpr::getICmp(I.getPredicate(), C, RHSC);
+      } else if (RHSC->isNullValue()) {
+        // If null is being compared, check if it can be further simplified.
+        Op = SimplifyICmpInst(I.getPredicate(), V, RHSC, SQ);
+      }
+      return Op;
+    };
+    Op1 = SimplifyOp(LHSI->getOperand(1));
+    if (Op1)
       CI = dyn_cast<ConstantInt>(Op1);
-    }
-    if (Constant *C = dyn_cast<Constant>(LHSI->getOperand(2))) {
-      Op2 = ConstantExpr::getICmp(I.getPredicate(), C, RHSC);
+
+    Op2 = SimplifyOp(LHSI->getOperand(2));
+    if (Op2)
       CI = dyn_cast<ConstantInt>(Op2);
-    }
 
     // We only want to perform this transformation if it will not lead to
     // additional code. This is true if either both sides of the select
diff --git a/test/Transforms/InstCombine/assume-icmp-null-select.ll b/test/Transforms/InstCombine/assume-icmp-null-select.ll
new file mode 100644
index 00000000000..dd433f169ba
--- /dev/null
+++ b/test/Transforms/InstCombine/assume-icmp-null-select.ll
@@ -0,0 +1,51 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -instcombine -S | FileCheck %s
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define i8* @example(i8* dereferenceable(24) %x) {
+; CHECK-LABEL: @example(
+; CHECK-NEXT:    [[X2:%.*]] = bitcast i8* [[X:%.*]] to {}**
+; CHECK-NEXT:    [[Y:%.*]] = load {}*, {}** [[X2]], align 8
+; CHECK-NEXT:    [[Y_IS_NULL:%.*]] = icmp ne {}* [[Y]], null
+; CHECK-NEXT:    call void @llvm.assume(i1 [[Y_IS_NULL]])
+; CHECK-NEXT:    ret i8* [[X]]
+;
+  %x2 = bitcast i8* %x to {}**
+  %y = load {}*, {}** %x2, align 8
+  %y_is_null = icmp eq {}* %y, null
+
+  %x0 = getelementptr inbounds i8, i8* %x, i64 0
+  %res = select i1 %y_is_null, i8* null, i8* %x0
+
+  %nonnull = icmp ne i8* %res, null
+  call void @llvm.assume(i1 %nonnull)
+
+  ret i8* %res
+}
+
+; TODO: this should be folded to `ret i8* %x` as well.
+define i8* @example2(i8* %x) {
+; CHECK-LABEL: @example2(
+; CHECK-NEXT:    [[X2:%.*]] = bitcast i8* [[X:%.*]] to {}**
+; CHECK-NEXT:    [[Y:%.*]] = load {}*, {}** [[X2]], align 8
+; CHECK-NEXT:    [[Y_IS_NULL:%.*]] = icmp eq {}* [[Y]], null
+; CHECK-NEXT:    [[RES:%.*]] = select i1 [[Y_IS_NULL]], i8* null, i8* [[X]]
+; CHECK-NEXT:    [[NONNULL:%.*]] = icmp ne i8* [[RES]], null
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NONNULL]])
+; CHECK-NEXT:    ret i8* [[RES]]
+;
+  %x2 = bitcast i8* %x to {}**
+  %y = load {}*, {}** %x2, align 8
+  %y_is_null = icmp eq {}* %y, null
+
+  %x0 = getelementptr inbounds i8, i8* %x, i64 0
+  %res = select i1 %y_is_null, i8* null, i8* %x0
+
+  %nonnull = icmp ne i8* %res, null
+  call void @llvm.assume(i1 %nonnull)
+
+  ret i8* %res
+}
+
+declare void @llvm.assume(i1)