Coalesce Return/ReturnVoid instructions

We do this to avoid generating the exit frame code several times.

We don't create new Return/ReturnVoid instructions after the
inliner, and CodeSinking looks like a logical place to add this
optimization.

Locally, speed-compiling in Pixel 5:
 * system server: -516.05KB (1.1%)
 * SystemUIGoogle: -216.50KB (0.8%)
 * AGSA: -3021.91KB (0.93%)

Bug: 203237752
Bug: 188926997
Bug: 194133637
Test: art/test/testrunner/testrunner.py --host --64 --optimizing -b
Change-Id: Ib149382462efac82a0bbc114cbdbc773e1a5228c

6 files changed, 270 insertions, 4 deletions

diff --git a/compiler/optimizing/code_sinking.cc b/compiler/optimizing/code_sinking.cc
index ac1cefe7c5..d759a16f48 100644
--- a/compiler/optimizing/code_sinking.cc
+++ b/compiler/optimizing/code_sinking.cc
@@ -29,11 +29,19 @@
 namespace art HIDDEN {
 
 bool CodeSinking::Run() {
-  HBasicBlock* exit = graph_->GetExitBlock();
-  if (exit == nullptr) {
+  if (graph_->GetExitBlock() == nullptr) {
     // Infinite loop, just bail.
     return false;
   }
+
+  UncommonBranchSinking();
+  ReturnSinking();
+  return true;
+}
+
+void CodeSinking::UncommonBranchSinking() {
+  HBasicBlock* exit = graph_->GetExitBlock();
+  DCHECK(exit != nullptr);
   // TODO(ngeoffray): we do not profile branches yet, so use throw instructions
   // as an indicator of an uncommon branch.
   for (HBasicBlock* exit_predecessor : exit->GetPredecessors()) {
@@ -60,7 +68,6 @@ bool CodeSinking::Run() {
       SinkCodeToUncommonBranch(exit_predecessor);
     }
   }
-  return true;
 }
 
 static bool IsInterestingInstruction(HInstruction* instruction) {
@@ -531,4 +538,79 @@ void CodeSinking::SinkCodeToUncommonBranch(HBasicBlock* end_block) {
   }
 }
 
+void CodeSinking::ReturnSinking() {
+  HBasicBlock* exit = graph_->GetExitBlock();
+  DCHECK(exit != nullptr);
+
+  int number_of_returns = 0;
+  bool saw_return = false;
+  for (HBasicBlock* pred : exit->GetPredecessors()) {
+    // TODO(solanes): We might have Return/ReturnVoid->TryBoundary->Exit. We can theoretically
+    // handle them and move them out of the TryBoundary. However, it is a border case and it adds
+    // codebase complexity.
+    if (pred->GetLastInstruction()->IsReturn() || pred->GetLastInstruction()->IsReturnVoid()) {
+      saw_return |= pred->GetLastInstruction()->IsReturn();
+      ++number_of_returns;
+    }
+  }
+
+  if (number_of_returns < 2) {
+    // Nothing to do.
+    return;
+  }
+
+  // `new_block` will coalesce the Return instructions into Phi+Return, or the ReturnVoid
+  // instructions into a ReturnVoid.
+  HBasicBlock* new_block = new (graph_->GetAllocator()) HBasicBlock(graph_, exit->GetDexPc());
+  if (saw_return) {
+    HPhi* new_phi = nullptr;
+    for (size_t i = 0; i < exit->GetPredecessors().size(); /*++i in loop*/) {
+      HBasicBlock* pred = exit->GetPredecessors()[i];
+      if (!pred->GetLastInstruction()->IsReturn()) {
+        ++i;
+        continue;
+      }
+
+      HReturn* ret = pred->GetLastInstruction()->AsReturn();
+      if (new_phi == nullptr) {
+        // Create the new_phi, if we haven't done so yet. We do it here since we need to know the
+        // type to assign to it.
+        new_phi = new (graph_->GetAllocator()) HPhi(graph_->GetAllocator(),
+                                                    kNoRegNumber,
+                                                    /*number_of_inputs=*/0,
+                                                    ret->InputAt(0)->GetType());
+        new_block->AddPhi(new_phi);
+      }
+      new_phi->AddInput(ret->InputAt(0));
+      pred->ReplaceAndRemoveInstructionWith(ret,
+                                            new (graph_->GetAllocator()) HGoto(ret->GetDexPc()));
+      pred->ReplaceSuccessor(exit, new_block);
+      // Since we are removing a predecessor, there's no need to increment `i`.
+    }
+    new_block->AddInstruction(new (graph_->GetAllocator()) HReturn(new_phi, exit->GetDexPc()));
+  } else {
+    for (size_t i = 0; i < exit->GetPredecessors().size(); /*++i in loop*/) {
+      HBasicBlock* pred = exit->GetPredecessors()[i];
+      if (!pred->GetLastInstruction()->IsReturnVoid()) {
+        ++i;
+        continue;
+      }
+
+      HReturnVoid* ret = pred->GetLastInstruction()->AsReturnVoid();
+      pred->ReplaceAndRemoveInstructionWith(ret,
+                                            new (graph_->GetAllocator()) HGoto(ret->GetDexPc()));
+      pred->ReplaceSuccessor(exit, new_block);
+      // Since we are removing a predecessor, there's no need to increment `i`.
+    }
+    new_block->AddInstruction(new (graph_->GetAllocator()) HReturnVoid(exit->GetDexPc()));
+  }
+
+  new_block->AddSuccessor(exit);
+  graph_->AddBlock(new_block);
+
+  // Recompute dominance since we added a new block.
+  graph_->ClearDominanceInformation();
+  graph_->ComputeDominanceInformation();
+}
+
 }  // namespace art
diff --git a/compiler/optimizing/code_sinking.h b/compiler/optimizing/code_sinking.h
index 14f62f2417..c743db40d9 100644
--- a/compiler/optimizing/code_sinking.h
+++ b/compiler/optimizing/code_sinking.h
@@ -39,10 +39,16 @@ class CodeSinking : public HOptimization {
   static constexpr const char* kCodeSinkingPassName = "code_sinking";
 
  private:
-  // Try to move code only used by `end_block` and all its post-dominated / dominated
+  // Tries to sink code to uncommon branches.
+  void UncommonBranchSinking();
+  // Tries to move code only used by `end_block` and all its post-dominated / dominated
   // blocks, to these blocks.
   void SinkCodeToUncommonBranch(HBasicBlock* end_block);
 
+  // Coalesces the Return/ReturnVoid instructions into one, if we have two or more. We do this to
+  // avoid generating the exit frame code several times.
+  void ReturnSinking();
+
   DISALLOW_COPY_AND_ASSIGN(CodeSinking);
 };
 
diff --git a/test/2262-checker-return-sinking/expected-stderr.txt b/test/2262-checker-return-sinking/expected-stderr.txt
new file mode 100644
index 0000000000..e69de29bb2
--- /dev/null
+++ b/test/2262-checker-return-sinking/expected-stderr.txt
diff --git a/test/2262-checker-return-sinking/expected-stdout.txt b/test/2262-checker-return-sinking/expected-stdout.txt
new file mode 100644
index 0000000000..e69de29bb2
--- /dev/null
+++ b/test/2262-checker-return-sinking/expected-stdout.txt
diff --git a/test/2262-checker-return-sinking/info.txt b/test/2262-checker-return-sinking/info.txt
new file mode 100644
index 0000000000..6cc9ae94a4
--- /dev/null
+++ b/test/2262-checker-return-sinking/info.txt
@@ -0,0 +1 @@
+Tests that we sink Return/ReturnVoid instructions and coalesce them.
diff --git a/test/2262-checker-return-sinking/src/Main.java b/test/2262-checker-return-sinking/src/Main.java
new file mode 100644
index 0000000000..469896484e
--- /dev/null
+++ b/test/2262-checker-return-sinking/src/Main.java
@@ -0,0 +1,177 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+public class Main {
+    public static void main(String[] args) {
+        // Switch case.
+        assertEquals("First case", $noinline$testSinkReturnSwitch(1));
+        assertEquals("Second case", $noinline$testSinkReturnSwitch(2));
+        assertEquals("default", $noinline$testSinkReturnSwitch(3));
+        $noinline$testSinkReturnVoidSwitch(1);
+
+        // If/else if/else.
+        assertEquals("First case", $noinline$testSinkReturnIfElse(1));
+        assertEquals("Second case", $noinline$testSinkReturnIfElse(2));
+        assertEquals("default", $noinline$testSinkReturnIfElse(3));
+        $noinline$testSinkReturnVoidIfElse(1);
+
+        // Non-trivial if cases.
+        assertEquals("First case", $noinline$testSinkReturnSeparatedReturns(1));
+        assertEquals("Second case", $noinline$testSinkReturnSeparatedReturns(2));
+        assertEquals("default", $noinline$testSinkReturnSeparatedReturns(3));
+        $noinline$testSinkReturnVoidSeparatedReturns(1);
+    }
+
+    /// CHECK-START: java.lang.String Main.$noinline$testSinkReturnSwitch(int) code_sinking (before)
+    /// CHECK:     Return
+    /// CHECK:     Return
+    /// CHECK:     Return
+    /// CHECK-NOT: Return
+
+    /// CHECK-START: java.lang.String Main.$noinline$testSinkReturnSwitch(int) code_sinking (after)
+    /// CHECK:     Return
+    /// CHECK-NOT: Return
+    private static String $noinline$testSinkReturnSwitch(int switch_id) {
+        switch (switch_id) {
+            case 1:
+                return "First case";
+            case 2:
+                return "Second case";
+            default:
+                return "default";
+        }
+    }
+
+    /// CHECK-START: void Main.$noinline$testSinkReturnVoidSwitch(int) code_sinking (before)
+    /// CHECK:     ReturnVoid
+    /// CHECK:     ReturnVoid
+    /// CHECK:     ReturnVoid
+    /// CHECK-NOT: ReturnVoid
+
+    /// CHECK-START: void Main.$noinline$testSinkReturnVoidSwitch(int) code_sinking (after)
+    /// CHECK:     ReturnVoid
+    /// CHECK-NOT: ReturnVoid
+    private static void $noinline$testSinkReturnVoidSwitch(int switch_id) {
+        switch (switch_id) {
+            case 1:
+                $noinline$emptyMethod();
+                return;
+            case 2:
+                $noinline$emptyMethod2();
+                return;
+            default:
+                $noinline$emptyMethod3();
+                return;
+        }
+    }
+
+    /// CHECK-START: java.lang.String Main.$noinline$testSinkReturnIfElse(int) code_sinking (before)
+    /// CHECK:     Return
+    /// CHECK:     Return
+    /// CHECK:     Return
+    /// CHECK-NOT: Return
+
+    /// CHECK-START: java.lang.String Main.$noinline$testSinkReturnIfElse(int) code_sinking (after)
+    /// CHECK:     Return
+    /// CHECK-NOT: Return
+    private static String $noinline$testSinkReturnIfElse(int id) {
+        if (id == 1) {
+            return "First case";
+        } else if (id == 2) {
+            return "Second case";
+        } else {
+            return "default";
+        }
+    }
+
+    /// CHECK-START: void Main.$noinline$testSinkReturnVoidIfElse(int) code_sinking (before)
+    /// CHECK:     ReturnVoid
+    /// CHECK:     ReturnVoid
+    /// CHECK:     ReturnVoid
+    /// CHECK-NOT: ReturnVoid
+
+    /// CHECK-START: void Main.$noinline$testSinkReturnVoidIfElse(int) code_sinking (after)
+    /// CHECK:     ReturnVoid
+    /// CHECK-NOT: ReturnVoid
+    private static void $noinline$testSinkReturnVoidIfElse(int id) {
+        if (id == 1) {
+            $noinline$emptyMethod();
+            return;
+        } else if (id == 2) {
+            $noinline$emptyMethod2();
+            return;
+        } else {
+            $noinline$emptyMethod3();
+            return;
+        }
+    }
+
+    /// CHECK-START: java.lang.String Main.$noinline$testSinkReturnSeparatedReturns(int) code_sinking (before)
+    /// CHECK:     Return
+    /// CHECK:     Return
+    /// CHECK:     Return
+    /// CHECK-NOT: Return
+
+    /// CHECK-START: java.lang.String Main.$noinline$testSinkReturnSeparatedReturns(int) code_sinking (after)
+    /// CHECK:     Return
+    /// CHECK-NOT: Return
+    private static String $noinline$testSinkReturnSeparatedReturns(int id) {
+        if (id == 1) {
+            return "First case";
+        }
+        $noinline$emptyMethod();
+
+        if (id == 2) {
+            return "Second case";
+        }
+
+        $noinline$emptyMethod2();
+        return "default";
+    }
+
+    /// CHECK-START: void Main.$noinline$testSinkReturnVoidSeparatedReturns(int) code_sinking (before)
+    /// CHECK:     ReturnVoid
+    /// CHECK:     ReturnVoid
+    /// CHECK:     ReturnVoid
+    /// CHECK-NOT: ReturnVoid
+
+    /// CHECK-START: void Main.$noinline$testSinkReturnVoidSeparatedReturns(int) code_sinking (after)
+    /// CHECK:     ReturnVoid
+    /// CHECK-NOT: ReturnVoid
+    private static void $noinline$testSinkReturnVoidSeparatedReturns(int id) {
+        if (id == 1) {
+            return;
+        }
+        $noinline$emptyMethod();
+
+        if (id == 2) {
+            return;
+        }
+
+        $noinline$emptyMethod2();
+        return;
+    }
+
+    private static void $noinline$emptyMethod() {}
+    private static void $noinline$emptyMethod2() {}
+    private static void $noinline$emptyMethod3() {}
+
+    private static void assertEquals(String expected, String result) {
+        if (expected != result) {
+            throw new Error("Expected: " + expected + ", found: " + result);
+        }
+    }
+}