From 17019a04b2e4d0c2bd2bce1842304a99668e8479 Mon Sep 17 00:00:00 2001
From: Yichao Yu <yyc1992@gmail.com>
Date: Sat, 4 Nov 2017 18:22:13 -0400
Subject: [PATCH] Fix GC preserve optimization

Only remove the call if all of the arguments are eliminated.
---
 src/llvm-alloc-opt.cpp        | 26 +++++++++++++++++------
 test/llvmpasses/alloc-opt2.jl | 40 +++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/src/llvm-alloc-opt.cpp b/src/llvm-alloc-opt.cpp
index cb26f53451a2b..159e120889a6d 100644
--- a/src/llvm-alloc-opt.cpp
+++ b/src/llvm-alloc-opt.cpp
@@ -58,6 +58,24 @@ static bool isBundleOperand(CallInst *call, unsigned idx)
 #endif
 }
 
+static void removeGCPreserve(CallInst *call, Instruction *val)
+{
+    auto replace = ConstantPointerNull::get(cast<PointerType>(val->getType()));
+    call->replaceUsesOfWith(val, replace);
+    for (auto &arg: call->arg_operands()) {
+        if (!isa<Constant>(arg.get())) {
+            return;
+        }
+    }
+    while (!call->use_empty()) {
+        auto end = cast<Instruction>(*call->user_begin());
+        // gc_preserve_end returns void.
+        assert(end->use_empty());
+        end->eraseFromParent();
+    }
+    call->eraseFromParent();
+}
+
 /**
  * Promote `julia.gc_alloc_obj` which do not have escaping root to a alloca.
  * Uses that are not considered to escape the object (i.e. heap address) includes,
@@ -608,13 +626,7 @@ void AllocOpt::replaceUsesWith(Instruction *orig_inst, Instruction *new_inst,
             }
             // Also remove the preserve intrinsics so that it can be better optimized.
             if (gc_preserve_begin && gc_preserve_begin == call->getCalledFunction()) {
-                while (!call->use_empty()) {
-                    auto end = cast<Instruction>(*call->user_begin());
-                    // gc_preserve_end returns void.
-                    assert(end->use_empty());
-                    end->eraseFromParent();
-                }
-                call->eraseFromParent();
+                removeGCPreserve(call, orig_i);
                 return;
             }
             if (auto intrinsic = dyn_cast<IntrinsicInst>(call)) {
diff --git a/test/llvmpasses/alloc-opt2.jl b/test/llvmpasses/alloc-opt2.jl
index 18f9e780f4cf9..c23444fb0abb7 100644
--- a/test/llvmpasses/alloc-opt2.jl
+++ b/test/llvmpasses/alloc-opt2.jl
@@ -46,8 +46,48 @@ L3:
 """)
 # CHECK-LABEL: }
 
+# CHECK-LABEL: @preserve_branches2
+# CHECK: alloca i64
+# CHECK: call %jl_value_t*** @julia.ptls_states()
+# CHECK: L1:
+# CHECK-NEXT: call void @llvm.lifetime.start{{.*}}(i64 8,
+# CHECK-NEXT: @llvm.julia.gc_preserve_begin{{.*}}%jl_value_t addrspace(10)* %v2
+# CHECK-NEXT: @external_function()
+# CHECK-NEXT: br i1 %b2, label %L2, label %L3
+
+# CHECK: L2:
+# CHECK-NOT: call void @llvm.lifetime.end{{.*}}(i64 8,
+# CHECK: @external_function()
+# CHECK-NEXT: br label %L3
+
+# CHECK: L3:
+# CHECK-NEXT: call void @llvm.lifetime.end{{.*}}(i64 8,
+println("""
+define void @preserve_branches2(i8* %fptr, i1 %b, i1 %b2) {
+  %ptls = call %jl_value_t*** @julia.ptls_states()
+  %ptls_i8 = bitcast %jl_value_t*** %ptls to i8*
+  %v2 = call %jl_value_t addrspace(10)* @external_function2()
+  br i1 %b, label %L1, label %L3
+
+L1:
+  %v = call noalias %jl_value_t addrspace(10)* @julia.gc_alloc_obj(i8* %ptls_i8, $isz 8, %jl_value_t addrspace(10)* @tag)
+  %tok = call token (...) @llvm.julia.gc_preserve_begin(%jl_value_t addrspace(10)* %v, %jl_value_t addrspace(10)* %v2)
+  call void @external_function()
+  br i1 %b2, label %L2, label %L3
+
+L2:
+  call void @external_function()
+  br label %L3
+
+L3:
+  ret void
+}
+""")
+# CHECK-LABEL: }
+
 println("""
 declare void @external_function()
+declare %jl_value_t addrspace(10)* @external_function2()
 declare %jl_value_t*** @julia.ptls_states()
 declare noalias %jl_value_t addrspace(10)* @julia.gc_alloc_obj(i8*, $isz, %jl_value_t addrspace(10)*)
 declare i64 @julia.pointer_from_objref(%jl_value_t addrspace(11)*)