From e8e01143ec642be7a1f7e8936254071a86cd9630 Mon Sep 17 00:00:00 2001
From: Max Kazantsev <max.kazantsev@azul.com>
Date: Wed, 4 Jul 2018 08:01:26 +0000
Subject: [PATCH] [ImplicitNullChecks] Check for rewrite of register used in
 'test' instruction

The following code pattern:

       mov %rax, %rcx
       test %rax, %rax
       %rax = ....
       je  throw_npe
       mov(%rcx), %r9
       mov(%rax), %r10

gets transformed into the following incorrect code after implicit null check pass:
        mov %rax, %rcx
       %rax = ....
       faulting_load_op("movl (%rax), %r10", throw_npe)
       mov(%rcx), %r9

For implicit null check pass, if the register that is checked for null value (ie, the register used in the 'test' instruction) is written into before the condition jump, we should avoid doing the optimization.

Patch by Surya Kumari Jangala!

Differential Revision: https://reviews.llvm.org/D48627
Reviewed By: skatkov

llvm-svn: 336241
---
 llvm/lib/CodeGen/ImplicitNullChecks.cpp       | 28 ++++++++++-
 .../X86/implicit-null-chk-reg-rewrite.mir     | 49 +++++++++++++++++++
 2 files changed, 75 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/implicit-null-chk-reg-rewrite.mir

diff --git a/llvm/lib/CodeGen/ImplicitNullChecks.cpp b/llvm/lib/CodeGen/ImplicitNullChecks.cpp
index ab777d28cd4d..b3f6e1fc7fec 100644
--- a/llvm/lib/CodeGen/ImplicitNullChecks.cpp
+++ b/llvm/lib/CodeGen/ImplicitNullChecks.cpp
@@ -496,6 +496,32 @@ bool ImplicitNullChecks::analyzeBlockForNullChecks(
   if (NotNullSucc->pred_size() != 1)
     return false;
 
+  // To prevent the invalid transformation of the following code:
+  //
+  //   mov %rax, %rcx
+  //   test %rax, %rax
+  //   %rax = ...
+  //   je throw_npe
+  //   mov(%rcx), %r9
+  //   mov(%rax), %r10
+  //
+  // into:
+  //
+  //   mov %rax, %rcx
+  //   %rax = ....
+  //   faulting_load_op("movl (%rax), %r10", throw_npe)
+  //   mov(%rcx), %r9
+  //
+  // we must ensure that there are no instructions between the 'test' and
+  // conditional jump that modify %rax.
+  const unsigned PointerReg = MBP.LHS.getReg();
+
+  assert(MBP.ConditionDef->getParent() ==  &MBB && "Should be in basic block");
+
+  for (auto I = MBB.rbegin(); MBP.ConditionDef != &*I; ++I)
+    if (I->modifiesRegister(PointerReg, TRI))
+      return false;
+
   // Starting with a code fragment like:
   //
   //   test %rax, %rax
@@ -550,8 +576,6 @@ bool ImplicitNullChecks::analyzeBlockForNullChecks(
   // ptr could be some non-null invalid reference that never gets loaded from
   // because some_cond is always true.
 
-  const unsigned PointerReg = MBP.LHS.getReg();
-
   SmallVector<MachineInstr *, 8> InstsSeenSoFar;
 
   for (auto &MI : *NotNullSucc) {
diff --git a/llvm/test/CodeGen/X86/implicit-null-chk-reg-rewrite.mir b/llvm/test/CodeGen/X86/implicit-null-chk-reg-rewrite.mir
new file mode 100644
index 000000000000..78d0d1401c41
--- /dev/null
+++ b/llvm/test/CodeGen/X86/implicit-null-chk-reg-rewrite.mir
@@ -0,0 +1,49 @@
+# RUN: llc -mtriple=x86_64 -run-pass=implicit-null-checks %s -o - | FileCheck %s
+--- |
+
+  define i32 @reg-rewrite(i32* %x) {
+  entry:
+    br i1 undef, label %is_null, label %not_null, !make.implicit !0
+
+  is_null:
+    ret i32 42
+
+  not_null:
+    ret i32 100
+  }
+
+  !0 = !{}
+
+...
+---
+# Check that the TEST instruction is replaced with 
+# FAULTING_OP only if there are no instructions
+# between the TEST and conditional jump
+# that clobber the register used in TEST.
+name:            reg-rewrite
+
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$rdi' }
+
+body:             |
+  bb.0.entry:
+    liveins: $rdi
+
+    TEST64rr $rdi, $rdi, implicit-def $eflags
+    ; CHECK-LABEL: bb.0.entry
+    ; CHECK-NOT: FAULTING_OP
+    renamable $rdi = MOV64ri 5000
+    JE_1 %bb.2, implicit $eflags
+
+  bb.1.not_null:
+    liveins: $rdi, $rsi
+    
+    $rax = MOV64rm renamable $rdi, 1, $noreg, 4, $noreg
+    RETQ $eax
+
+  bb.2.is_null:
+    $eax = MOV32ri 200
+    RETQ $eax
+...