Insert moves after arguments

Bug: 79186787
Change-Id: I628b347a85a42cfd681054c454b08b8bcd014c5c
diff --git a/src/main/java/com/android/tools/r8/ir/code/BasicBlock.java b/src/main/java/com/android/tools/r8/ir/code/BasicBlock.java
index fb75931..d59df07 100644
--- a/src/main/java/com/android/tools/r8/ir/code/BasicBlock.java
+++ b/src/main/java/com/android/tools/r8/ir/code/BasicBlock.java
@@ -48,10 +48,7 @@
       assert instruction.getBlock() == this;
       assert !instruction.isArgument() || argumentsAllowed;
       assert !instruction.isDebugLocalRead() || !instruction.getDebugValues().isEmpty();
-      // TODO(b/79186787): Ensure DEX backend inserts Move *after* arguments.
-      if (!(instruction.isArgument()
-          || instruction.isMove()
-          || instruction.isDebugLocalsChange())) {
+      if (!instruction.isArgument()) {
         argumentsAllowed = false;
       }
     }
diff --git a/src/main/java/com/android/tools/r8/ir/regalloc/LinearScanRegisterAllocator.java b/src/main/java/com/android/tools/r8/ir/regalloc/LinearScanRegisterAllocator.java
index 769c281..de85da5 100644
--- a/src/main/java/com/android/tools/r8/ir/regalloc/LinearScanRegisterAllocator.java
+++ b/src/main/java/com/android/tools/r8/ir/regalloc/LinearScanRegisterAllocator.java
@@ -125,7 +125,9 @@
   // Mapping from basic blocks to the set of values live at entry to that basic block.
   private Map<BasicBlock, Set<Value>> liveAtEntrySets;
   // The value of the first argument, or null if the method has no arguments.
-  private Value firstArgumentValue;
+  protected Value firstArgumentValue;
+  // The value of the last argument, or null if the method has no arguments.
+  private Value lastArgumentValue;
 
   // The set of registers that are free for allocation.
   private TreeSet<Integer> freeRegisters = new TreeSet<>();
@@ -2511,7 +2513,8 @@
       Invoke invoke, PriorityQueue<Move> insertAtDefinition, InstructionListIterator insertAt) {
     Move move = insertAtDefinition.poll();
     // Rewind instruction iterator to the position where the first move needs to be inserted.
-    Instruction previousDefinition = move.src().definition;
+    Instruction previousDefinition =
+        move.src().isArgument() ? lastArgumentValue.definition : move.src().definition;
     while (insertAt.peekPrevious() != previousDefinition) {
       insertAt.previous();
     }
@@ -2519,7 +2522,8 @@
     insertAt.add(move);
     while (!insertAtDefinition.isEmpty()) {
       move = insertAtDefinition.poll();
-      Instruction currentDefinition = move.src().definition;
+      Instruction currentDefinition =
+          move.src().isArgument() ? lastArgumentValue.definition : move.src().definition;
       assert currentDefinition.getNumber() >= previousDefinition.getNumber();
       if (currentDefinition.getNumber() > previousDefinition.getNumber()) {
         // Move the instruction iterator forward to where this move needs to be inserted.
@@ -2572,6 +2576,7 @@
         last.getLiveIntervals().link(next.getLiveIntervals());
         last = next;
       }
+      lastArgumentValue = last;
     }
   }
 
diff --git a/src/main/java/com/android/tools/r8/ir/regalloc/SpillMoveSet.java b/src/main/java/com/android/tools/r8/ir/regalloc/SpillMoveSet.java
index 3401b17..1ea788d 100644
--- a/src/main/java/com/android/tools/r8/ir/regalloc/SpillMoveSet.java
+++ b/src/main/java/com/android/tools/r8/ir/regalloc/SpillMoveSet.java
@@ -10,6 +10,7 @@
 import com.android.tools.r8.ir.code.Instruction;
 import com.android.tools.r8.ir.code.InstructionListIterator;
 import com.android.tools.r8.ir.code.Position;
+import com.android.tools.r8.ir.code.Value;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -119,17 +120,33 @@
    */
   public int scheduleAndInsertMoves(int tempRegister) {
     for (BasicBlock block : code.blocks) {
-      InstructionListIterator it = block.listIterator();
-      while (it.hasNext()) {
-        Instruction instruction = it.next();
+      InstructionListIterator insertAt = block.listIterator();
+      if (block == code.blocks.getFirst()) {
+        // Move insertAt iterator to the first non-argument, such that moves for the arguments will
+        // be inserted after the last argument.
+        while (insertAt.hasNext() && insertAt.peekNext().isArgument()) {
+          insertAt.next();
+        }
+        // Generate moves for each argument.
+        for (Value argumentValue = allocator.firstArgumentValue;
+            argumentValue != null;
+            argumentValue = argumentValue.getNextConsecutive()) {
+          Instruction instruction = argumentValue.definition;
+          int number = instruction.getNumber();
+          if (needsMovesBeforeInstruction(number)) {
+            scheduleMovesBeforeInstruction(tempRegister, instruction, insertAt);
+          }
+        }
+      }
+
+      while (insertAt.hasNext()) {
+        Instruction instruction = insertAt.peekNext();
+        assert !instruction.isArgument();
         int number = instruction.getNumber();
         if (needsMovesBeforeInstruction(number)) {
-          // Move back so moves are inserted before the instruction.
-          it.previous();
-          scheduleMovesBeforeInstruction(tempRegister, number, it);
-          // Move past the instruction again.
-          it.next();
+          scheduleMovesBeforeInstruction(tempRegister, instruction, insertAt);
         }
+        insertAt.next();
       }
     }
     return usedTempRegisters;
@@ -218,14 +235,15 @@
   }
 
   private void scheduleMovesBeforeInstruction(
-      int tempRegister, int instruction, InstructionListIterator insertAt) {
+      int tempRegister, Instruction instruction, InstructionListIterator insertAt) {
+    int number = instruction.getNumber();
 
     Position position;
     if (insertAt.hasPrevious() && insertAt.peekPrevious().isMoveException()) {
       position = insertAt.peekPrevious().getPosition();
     } else {
       Instruction next = insertAt.peekNext();
-      assert next.getNumber() == instruction;
+      assert next.getNumber() == number || instruction.isArgument();
       position = next.getPosition();
       if (position.isNone() && next.isGoto()) {
         position = next.asGoto().getTarget().getPosition();
@@ -234,17 +252,17 @@
 
     // Spill and restore moves for the incoming edge.
     Set<SpillMove> inMoves =
-        instructionToInMoves.computeIfAbsent(instruction - 1, (k) -> new LinkedHashSet<>());
+        instructionToInMoves.computeIfAbsent(number - 1, (k) -> new LinkedHashSet<>());
     removeArgumentRestores(inMoves);
 
     // Spill and restore moves for the outgoing edge.
     Set<SpillMove> outMoves =
-        instructionToOutMoves.computeIfAbsent(instruction - 1, (k) -> new LinkedHashSet<>());
+        instructionToOutMoves.computeIfAbsent(number - 1, (k) -> new LinkedHashSet<>());
     removeArgumentRestores(outMoves);
 
     // Get the phi moves for this instruction and schedule them with the out going spill moves.
     Set<SpillMove> phiMoves =
-        instructionToPhiMoves.computeIfAbsent(instruction - 1, (k) -> new LinkedHashSet<>());
+        instructionToPhiMoves.computeIfAbsent(number - 1, (k) -> new LinkedHashSet<>());
 
     // Remove/rewrite moves that we can guarantee will not be needed.
     pruneParallelMoveSets(inMoves, outMoves, phiMoves);