1 files changed, 41 insertions, 18 deletions

diff --git a/compiler/optimizing/nodes.cc b/compiler/optimizing/nodes.cc
index 24a89bca4e..98c3096cae 100644
--- a/compiler/optimizing/nodes.cc
+++ b/compiler/optimizing/nodes.cc
@@ -606,8 +606,23 @@ static void UpdateInputsUsers(HInstruction* instruction) {
 void HBasicBlock::ReplaceAndRemoveInstructionWith(HInstruction* initial,
                                                   HInstruction* replacement) {
   DCHECK(initial->GetBlock() == this);
-  InsertInstructionBefore(replacement, initial);
-  initial->ReplaceWith(replacement);
+  if (initial->IsControlFlow()) {
+    // We can only replace a control flow instruction with another control flow instruction.
+    DCHECK(replacement->IsControlFlow());
+    DCHECK_EQ(replacement->GetId(), -1);
+    DCHECK_EQ(replacement->GetType(), Primitive::kPrimVoid);
+    DCHECK_EQ(initial->GetBlock(), this);
+    DCHECK_EQ(initial->GetType(), Primitive::kPrimVoid);
+    DCHECK(initial->GetUses().IsEmpty());
+    DCHECK(initial->GetEnvUses().IsEmpty());
+    replacement->SetBlock(this);
+    replacement->SetId(GetGraph()->GetNextInstructionId());
+    instructions_.InsertInstructionBefore(replacement, initial);
+    UpdateInputsUsers(replacement);
+  } else {
+    InsertInstructionBefore(replacement, initial);
+    initial->ReplaceWith(replacement);
+  }
   RemoveInstruction(initial);
 }
 
@@ -1576,7 +1591,6 @@ HInstruction* HGraph::InlineInto(HGraph* outer_graph, HInvoke* invoke) {
     // Replace the invoke with the return value of the inlined graph.
     if (last->IsReturn()) {
       return_value = last->InputAt(0);
-      invoke->ReplaceWith(return_value);
     } else {
       DCHECK(last->IsReturnVoid());
     }
@@ -1624,10 +1638,6 @@ HInstruction* HGraph::InlineInto(HGraph* outer_graph, HInvoke* invoke) {
       }
     }
 
-    if (return_value != nullptr) {
-      invoke->ReplaceWith(return_value);
-    }
-
     // Update the meta information surrounding blocks:
     // (1) the graph they are now in,
     // (2) the reverse post order of that graph,
@@ -1697,20 +1707,21 @@ HInstruction* HGraph::InlineInto(HGraph* outer_graph, HInvoke* invoke) {
   size_t parameter_index = 0;
   for (HInstructionIterator it(entry_block_->GetInstructions()); !it.Done(); it.Advance()) {
     HInstruction* current = it.Current();
+    HInstruction* replacement = nullptr;
     if (current->IsNullConstant()) {
-      current->ReplaceWith(outer_graph->GetNullConstant(current->GetDexPc()));
+      replacement = outer_graph->GetNullConstant(current->GetDexPc());
     } else if (current->IsIntConstant()) {
-      current->ReplaceWith(outer_graph->GetIntConstant(
-          current->AsIntConstant()->GetValue(), current->GetDexPc()));
+      replacement = outer_graph->GetIntConstant(
+          current->AsIntConstant()->GetValue(), current->GetDexPc());
     } else if (current->IsLongConstant()) {
-      current->ReplaceWith(outer_graph->GetLongConstant(
-          current->AsLongConstant()->GetValue(), current->GetDexPc()));
+      replacement = outer_graph->GetLongConstant(
+          current->AsLongConstant()->GetValue(), current->GetDexPc());
     } else if (current->IsFloatConstant()) {
-      current->ReplaceWith(outer_graph->GetFloatConstant(
-          current->AsFloatConstant()->GetValue(), current->GetDexPc()));
+      replacement = outer_graph->GetFloatConstant(
+          current->AsFloatConstant()->GetValue(), current->GetDexPc());
     } else if (current->IsDoubleConstant()) {
-      current->ReplaceWith(outer_graph->GetDoubleConstant(
-          current->AsDoubleConstant()->GetValue(), current->GetDexPc()));
+      replacement = outer_graph->GetDoubleConstant(
+          current->AsDoubleConstant()->GetValue(), current->GetDexPc());
     } else if (current->IsParameterValue()) {
       if (kIsDebugBuild
           && invoke->IsInvokeStaticOrDirect()
@@ -1720,13 +1731,25 @@ HInstruction* HGraph::InlineInto(HGraph* outer_graph, HInvoke* invoke) {
         size_t last_input_index = invoke->InputCount() - 1;
         DCHECK(parameter_index != last_input_index);
       }
-      current->ReplaceWith(invoke->InputAt(parameter_index++));
+      replacement = invoke->InputAt(parameter_index++);
     } else if (current->IsCurrentMethod()) {
-      current->ReplaceWith(outer_graph->GetCurrentMethod());
+      replacement = outer_graph->GetCurrentMethod();
     } else {
       DCHECK(current->IsGoto() || current->IsSuspendCheck());
       entry_block_->RemoveInstruction(current);
     }
+    if (replacement != nullptr) {
+      current->ReplaceWith(replacement);
+      // If the current is the return value then we need to update the latter.
+      if (current == return_value) {
+        DCHECK_EQ(entry_block_, return_value->GetBlock());
+        return_value = replacement;
+      }
+    }
+  }
+
+  if (return_value != nullptr) {
+    invoke->ReplaceWith(return_value);
   }
 
   // Finally remove the invoke from the caller.