[InstCombine] Simplify the pointer operand of store if writing to null is UB (#127979)

Proof: https://alive2.llvm.org/ce/z/mzVj-u
I will add some follow-up patches to avoid duplicate code, support more
memory instructions, and bypass gep instructions.

Author: dtcxzyw

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

@@ -1437,6 +1437,20 @@ Instruction *InstCombinerImpl::visitStoreInst(StoreInst &SI) {
   if (isa<UndefValue>(Val))
     return eraseInstFromFunction(SI);
 
+  // TODO: Add a helper to simplify the pointer operand for all memory
+  // instructions.
+  // store val, (select (cond, null, P)) -> store val, P
+  // store val, (select (cond, P, null)) -> store val, P
+  if (!NullPointerIsDefined(SI.getFunction(), SI.getPointerAddressSpace())) {
+    if (SelectInst *Sel = dyn_cast<SelectInst>(Ptr)) {
+      if (isa<ConstantPointerNull>(Sel->getOperand(1)))
+        return replaceOperand(SI, 1, Sel->getOperand(2));
+
+      if (isa<ConstantPointerNull>(Sel->getOperand(2)))
+        return replaceOperand(SI, 1, Sel->getOperand(1));
+    }
+  }
+
   return nullptr;
 }
 

llvm/test/Transforms/InstCombine/store.ll

@@ -345,6 +345,48 @@ define void @store_to_readonly_noalias(ptr readonly noalias %0) {
   ret void
 }
 
+define void @store_select_with_null(i1 %cond, ptr %p) {
+; CHECK-LABEL: @store_select_with_null(
+; CHECK-NEXT:    store i32 0, ptr [[SEL:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
+  %sel = select i1 %cond, ptr %p, ptr null
+  store i32 0, ptr %sel, align 4
+  ret void
+}
+
+define void @store_select_with_null_commuted(i1 %cond, ptr %p) {
+; CHECK-LABEL: @store_select_with_null_commuted(
+; CHECK-NEXT:    store i32 0, ptr [[SEL:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
+  %sel = select i1 %cond, ptr null, ptr %p
+  store i32 0, ptr %sel, align 4
+  ret void
+}
+
+define void @store_select_with_null_null_is_valid(i1 %cond, ptr %p) null_pointer_is_valid {
+; CHECK-LABEL: @store_select_with_null_null_is_valid(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], ptr [[P:%.*]], ptr null
+; CHECK-NEXT:    store i32 0, ptr [[SEL]], align 4
+; CHECK-NEXT:    ret void
+;
+  %sel = select i1 %cond, ptr %p, ptr null
+  store i32 0, ptr %sel, align 4
+  ret void
+}
+
+define void @store_select_with_unknown(i1 %cond, ptr %p, ptr %p2) {
+; CHECK-LABEL: @store_select_with_unknown(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], ptr [[P:%.*]], ptr [[P2:%.*]]
+; CHECK-NEXT:    store i32 0, ptr [[SEL]], align 4
+; CHECK-NEXT:    ret void
+;
+  %sel = select i1 %cond, ptr %p, ptr %p2
+  store i32 0, ptr %sel, align 4
+  ret void
+}
+
 !0 = !{!4, !4, i64 0}
 !1 = !{!"omnipotent char", !2}
 !2 = !{!"Simple C/C++ TBAA"}

llvm/test/Transforms/PhaseOrdering/load-store-sameval.ll

@@ -0,0 +1,44 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes='instcombine,early-cse<memssa>' -S %s | FileCheck %s
+
+; FIXME: We can remove the store instruction in the exit block
+define i32 @load_store_sameval(ptr %p, i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: define i32 @load_store_sameval(
+; CHECK-SAME: ptr [[P:%.*]], i1 [[COND1:%.*]], i1 [[COND2:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[COND1]], ptr null, ptr [[P]]
+; CHECK-NEXT:    [[PRE:%.*]] = load i32, ptr [[SPEC_SELECT]], align 4
+; CHECK-NEXT:    br label %[[BLOCK:.*]]
+; CHECK:       [[BLOCK]]:
+; CHECK-NEXT:    br label %[[BLOCK2:.*]]
+; CHECK:       [[BLOCK2]]:
+; CHECK-NEXT:    br i1 [[COND2]], label %[[BLOCK3:.*]], label %[[EXIT:.*]]
+; CHECK:       [[BLOCK3]]:
+; CHECK-NEXT:    [[LOAD:%.*]] = load double, ptr [[SPEC_SELECT]], align 8
+; CHECK-NEXT:    [[CMP:%.*]] = fcmp une double [[LOAD]], 0.000000e+00
+; CHECK-NEXT:    br i1 [[CMP]], label %[[BLOCK]], label %[[BLOCK2]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    store i32 [[PRE]], ptr [[P]], align 4
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  %spec.select = select i1 %cond1, ptr null, ptr %p
+  %pre = load i32, ptr %spec.select, align 4
+  br label %block
+
+block:
+  br label %block2
+
+block2:
+  br i1 %cond2, label %block3, label %exit
+
+block3:
+  %load = load double, ptr %spec.select, align 8
+  %cmp = fcmp une double %load, 0.000000e+00
+  br i1 %cmp, label %block, label %block2
+
+exit:
+  store i32 %pre, ptr %spec.select, align 4
+  ret i32 0
+}
+