From b6eed9a82e0ce530d94a194c88615d6c272e1854 Mon Sep 17 00:00:00 2001
From: DianQK <dianqk@dianqk.net>
Date: Sun, 19 Feb 2023 16:42:33 +0800
Subject: [PATCH] [SimplifyCFG] Check if the return instruction causes
 undefined behavior

This should fix https://github.com/rust-lang/rust/issues/107681.

Return undefined to a noundef return value is undefined.

Example:

```
define noundef i32 @test_ret_noundef(i1 %cond) {
entry:
  br i1 %cond, label %bb1, label %bb2
bb1:
  br label %bb2
bb2:
  %r = phi i32 [ undef, %entry ], [ 1, %bb1 ]
  ret i32 %r
}
```

Differential Revision: https://reviews.llvm.org/D144319
---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     |  14 ++
 .../unreachable-eliminate-on-ret.ll           | 122 ++++++++++++++++++
 2 files changed, 136 insertions(+)
 create mode 100644 llvm/test/Transforms/SimplifyCFG/unreachable-eliminate-on-ret.ll

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 372a6698ce8ac..82ce8a6ce894a 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -7126,6 +7126,20 @@ static bool passingValueIsAlwaysUndefined(Value *V, Instruction *I, bool PtrValu
         return passingValueIsAlwaysUndefined(V, GEP, PtrValueMayBeModified);
       }
 
+    // Look through return.
+    if (ReturnInst *Ret = dyn_cast<ReturnInst>(Use)) {
+      bool HasNoUndefAttr =
+          Ret->getFunction()->hasRetAttribute(Attribute::NoUndef);
+      // Return undefined to a noundef return value is undefined.
+      if (isa<UndefValue>(C) && HasNoUndefAttr)
+        return true;
+      // Return null to a nonnull+noundef return value is undefined.
+      if (C->isNullValue() && HasNoUndefAttr &&
+          Ret->getFunction()->hasRetAttribute(Attribute::NonNull)) {
+        return !PtrValueMayBeModified;
+      }
+    }
+
     // Look through bitcasts.
     if (BitCastInst *BC = dyn_cast<BitCastInst>(Use))
       return passingValueIsAlwaysUndefined(V, BC, PtrValueMayBeModified);
diff --git a/llvm/test/Transforms/SimplifyCFG/unreachable-eliminate-on-ret.ll b/llvm/test/Transforms/SimplifyCFG/unreachable-eliminate-on-ret.ll
new file mode 100644
index 0000000000000..e250a5d02c563
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/unreachable-eliminate-on-ret.ll
@@ -0,0 +1,122 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S < %s | FileCheck %s
+
+define noundef i32 @test_ret_noundef(i1 %cond) {
+; CHECK-LABEL: @test_ret_noundef(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    ret i32 1
+;
+entry:
+  br i1 %cond, label %bb1, label %bb2
+
+bb1:
+  br label %bb2
+
+bb2:
+  %r = phi i32 [ undef, %entry ], [ 1, %bb1 ]
+  ret i32 %r
+}
+
+define i32 @test_ret(i1 %cond) {
+; CHECK-LABEL: @test_ret(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[COND:%.*]], i32 1, i32 undef
+; CHECK-NEXT:    ret i32 [[SPEC_SELECT]]
+;
+entry:
+  br i1 %cond, label %bb1, label %bb2
+
+bb1:
+  br label %bb2
+
+bb2:
+  %r = phi i32 [ undef, %entry ], [ 1, %bb1 ]
+  ret i32 %r
+}
+
+define nonnull noundef ptr @test_ret_ptr_nonnull_noundef(i1 %cond, ptr %x) {
+; CHECK-LABEL: @test_ret_ptr_nonnull_noundef(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND:%.*]])
+; CHECK-NEXT:    ret ptr [[X:%.*]]
+;
+entry:
+  br i1 %cond, label %bb1, label %bb2
+
+bb1:
+  br label %bb2
+
+bb2:
+  %r = phi ptr [ null, %entry ], [ %x, %bb1 ]
+  ret ptr %r
+}
+
+define nonnull noundef ptr @test_ret_ptr_nonnull_noundef_gep_nonzero(i1 %cond, ptr %x) {
+; CHECK-LABEL: @test_ret_ptr_nonnull_noundef_gep_nonzero(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[COND:%.*]], ptr [[X:%.*]], ptr null
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds ptr, ptr [[SPEC_SELECT]], i64 12
+; CHECK-NEXT:    ret ptr [[GEP]]
+;
+entry:
+  br i1 %cond, label %bb1, label %bb2
+
+bb1:
+  br label %bb2
+
+bb2:
+  %phi = phi ptr [ null, %entry ], [ %x, %bb1 ]
+  %gep = getelementptr inbounds ptr, ptr %phi, i64 12
+  ret ptr %gep
+}
+
+define nonnull ptr @test_ret_ptr_nonnull(i1 %cond, ptr %x) {
+; CHECK-LABEL: @test_ret_ptr_nonnull(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[COND:%.*]], ptr [[X:%.*]], ptr null
+; CHECK-NEXT:    ret ptr [[SPEC_SELECT]]
+;
+entry:
+  br i1 %cond, label %bb1, label %bb2
+
+bb1:
+  br label %bb2
+
+bb2:
+  %r = phi ptr [ null, %entry ], [ %x, %bb1 ]
+  ret ptr %r
+}
+
+define noundef ptr @test_ret_ptr_noundef(i1 %cond, ptr %x) {
+; CHECK-LABEL: @test_ret_ptr_noundef(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[COND:%.*]], ptr [[X:%.*]], ptr null
+; CHECK-NEXT:    ret ptr [[SPEC_SELECT]]
+;
+entry:
+  br i1 %cond, label %bb1, label %bb2
+
+bb1:
+  br label %bb2
+
+bb2:
+  %r = phi ptr [ null, %entry ], [ %x, %bb1 ]
+  ret ptr %r
+}
+
+define ptr @test_ret_ptr(i1 %cond, ptr %x) {
+; CHECK-LABEL: @test_ret_ptr(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[COND:%.*]], ptr [[X:%.*]], ptr null
+; CHECK-NEXT:    ret ptr [[SPEC_SELECT]]
+;
+entry:
+  br i1 %cond, label %bb1, label %bb2
+
+bb1:
+  br label %bb2
+
+bb2:
+  %r = phi ptr [ null, %entry ], [ %x, %bb1 ]
+  ret ptr %r
+}