4 files changed, 107 insertions, 8 deletions

diff --git a/compiler/optimizing/load_store_elimination.cc b/compiler/optimizing/load_store_elimination.cc
index 605fdae0f8..89ad85e0b4 100644
--- a/compiler/optimizing/load_store_elimination.cc
+++ b/compiler/optimizing/load_store_elimination.cc
@@ -199,6 +199,12 @@ class LSEVisitor : public HGraphDelegateVisitor {
     if (predecessors.size() == 0) {
       return;
     }
+    if (block->IsExitBlock()) {
+      // Exit block doesn't really merge values since the control flow ends in
+      // its predecessors. Each predecessor needs to make sure stores are kept
+      // if necessary.
+      return;
+    }
 
     ScopedArenaVector<HInstruction*>& heap_values = heap_values_for_[block->GetBlockId()];
     for (size_t i = 0; i < heap_values.size(); i++) {
@@ -233,15 +239,23 @@ class LSEVisitor : public HGraphDelegateVisitor {
         }
       }
 
-      if (merged_value == kUnknownHeapValue || ref_info->IsSingletonAndNonRemovable()) {
-        // There are conflicting heap values from different predecessors,
-        // or the heap value may be needed after method return or deoptimization.
-        // Keep the last store in each predecessor since future loads cannot be eliminated.
-        for (HBasicBlock* predecessor : predecessors) {
-          ScopedArenaVector<HInstruction*>& pred_values =
-              heap_values_for_[predecessor->GetBlockId()];
-          KeepIfIsStore(pred_values[i]);
+      if (ref_info->IsSingleton()) {
+        if (ref_info->IsSingletonAndNonRemovable() ||
+            (merged_value == kUnknownHeapValue &&
+             !block->IsSingleReturnOrReturnVoidAllowingPhis())) {
+          // The heap value may be needed after method return or deoptimization,
+          // or there are conflicting heap values from different predecessors and
+          // this block is not a single return,
+          // keep the last store in each predecessor since future loads may not
+          // be eliminated.
+          for (HBasicBlock* predecessor : predecessors) {
+            ScopedArenaVector<HInstruction*>& pred_values =
+                heap_values_for_[predecessor->GetBlockId()];
+            KeepIfIsStore(pred_values[i]);
+          }
         }
+      } else {
+        // Currenctly we don't eliminate stores to non-singletons.
       }
 
       if ((merged_value == nullptr) || !from_all_predecessors) {
@@ -549,6 +563,31 @@ class LSEVisitor : public HGraphDelegateVisitor {
     }
   }
 
+  // Keep necessary stores before exiting a method via return/throw.
+  void HandleExit(HBasicBlock* block) {
+    const ScopedArenaVector<HInstruction*>& heap_values =
+        heap_values_for_[block->GetBlockId()];
+    for (size_t i = 0; i < heap_values.size(); i++) {
+      HInstruction* heap_value = heap_values[i];
+      ReferenceInfo* ref_info = heap_location_collector_.GetHeapLocation(i)->GetReferenceInfo();
+      if (!ref_info->IsSingletonAndRemovable()) {
+        KeepIfIsStore(heap_value);
+      }
+    }
+  }
+
+  void VisitReturn(HReturn* instruction) OVERRIDE {
+    HandleExit(instruction->GetBlock());
+  }
+
+  void VisitReturnVoid(HReturnVoid* return_void) OVERRIDE {
+    HandleExit(return_void->GetBlock());
+  }
+
+  void VisitThrow(HThrow* throw_instruction) OVERRIDE {
+    HandleExit(throw_instruction->GetBlock());
+  }
+
   void HandleInvoke(HInstruction* instruction) {
     SideEffects side_effects = instruction->GetSideEffects();
     ScopedArenaVector<HInstruction*>& heap_values =
diff --git a/compiler/optimizing/nodes.cc b/compiler/optimizing/nodes.cc
index f4f6434678..fff61f5727 100644
--- a/compiler/optimizing/nodes.cc
+++ b/compiler/optimizing/nodes.cc
@@ -1810,6 +1810,11 @@ bool HBasicBlock::IsSingleReturn() const {
   return HasOnlyOneInstruction(*this) && GetLastInstruction()->IsReturn();
 }
 
+bool HBasicBlock::IsSingleReturnOrReturnVoidAllowingPhis() const {
+  return (GetFirstInstruction() == GetLastInstruction()) &&
+         (GetLastInstruction()->IsReturn() || GetLastInstruction()->IsReturnVoid());
+}
+
 bool HBasicBlock::IsSingleTryBoundary() const {
   return HasOnlyOneInstruction(*this) && GetLastInstruction()->IsTryBoundary();
 }
diff --git a/compiler/optimizing/nodes.h b/compiler/optimizing/nodes.h
index 29c78a1e34..6672901781 100644
--- a/compiler/optimizing/nodes.h
+++ b/compiler/optimizing/nodes.h
@@ -968,6 +968,7 @@ class HBasicBlock : public ArenaObject<kArenaAllocBasicBlock> {
 
   bool IsSingleGoto() const;
   bool IsSingleReturn() const;
+  bool IsSingleReturnOrReturnVoidAllowingPhis() const;
   bool IsSingleTryBoundary() const;
 
   // Returns true if this block emits nothing but a jump.
diff --git a/test/530-checker-lse/src/Main.java b/test/530-checker-lse/src/Main.java
index ca8108f058..c4cc3b0121 100644
--- a/test/530-checker-lse/src/Main.java
+++ b/test/530-checker-lse/src/Main.java
@@ -949,6 +949,56 @@ public class Main {
     return array[1] + array[i];
   }
 
+  /// CHECK-START: int Main.testExitMerge(boolean) load_store_elimination (before)
+  /// CHECK: NewInstance
+  /// CHECK: InstanceFieldSet
+  /// CHECK: InstanceFieldGet
+  /// CHECK: Return
+  /// CHECK: InstanceFieldSet
+  /// CHECK: Throw
+
+  /// CHECK-START: int Main.testExitMerge(boolean) load_store_elimination (after)
+  /// CHECK-NOT: NewInstance
+  /// CHECK-NOT: InstanceFieldSet
+  /// CHECK-NOT: InstanceFieldGet
+  /// CHECK: Return
+  /// CHECK-NOT: InstanceFieldSet
+  /// CHECK: Throw
+  private static int testExitMerge(boolean cond) {
+    TestClass obj = new TestClass();
+    if (cond) {
+      obj.i = 1;
+      return obj.i + 1;
+    } else {
+      obj.i = 2;
+      throw new Error();
+    }
+  }
+
+  /// CHECK-START: int Main.testExitMerge2(boolean) load_store_elimination (before)
+  /// CHECK: NewInstance
+  /// CHECK: InstanceFieldSet
+  /// CHECK: InstanceFieldGet
+  /// CHECK: InstanceFieldSet
+  /// CHECK: InstanceFieldGet
+
+  /// CHECK-START: int Main.testExitMerge2(boolean) load_store_elimination (after)
+  /// CHECK-NOT: NewInstance
+  /// CHECK-NOT: InstanceFieldSet
+  /// CHECK-NOT: InstanceFieldGet
+  private static int testExitMerge2(boolean cond) {
+    TestClass obj = new TestClass();
+    int res;
+    if (cond) {
+      obj.i = 1;
+      res = obj.i + 1;
+    } else {
+      obj.i = 2;
+      res = obj.j + 2;
+    }
+    return res;
+  }
+
   static void assertIntEquals(int result, int expected) {
     if (expected != result) {
       throw new Error("Expected: " + expected + ", found: " + result);
@@ -1039,6 +1089,10 @@ public class Main {
     assertIntEquals(testStoreStore().j, 43);
     assertIntEquals(testStoreStoreWithDeoptimize(new int[4]), 4);
 
+    assertIntEquals(testExitMerge(true), 2);
+    assertIntEquals(testExitMerge2(true), 2);
+    assertIntEquals(testExitMerge2(false), 2);
+
     int ret = testNoSideEffects(iarray);
     assertIntEquals(iarray[0], 101);
     assertIntEquals(iarray[1], 102);