refactor: Add helper method to find a free register for a specific insert index (#4693)

patches/src/main/kotlin/app/revanced/patches/shared/misc/spoof/SpoofVideoStreamsPatch.kt

@@ -10,6 +10,7 @@ import app.revanced.patcher.patch.BytecodePatchContext
 import app.revanced.patcher.patch.bytecodePatch
 import app.revanced.patcher.util.proxy.mutableTypes.MutableMethod.Companion.toMutable
 import app.revanced.patches.all.misc.resources.addResourcesPatch
+import app.revanced.util.findFreeRegister
 import app.revanced.util.findInstructionIndicesReversedOrThrow
 import app.revanced.util.getReference
 import app.revanced.util.indexOfFirstInstructionOrThrow
@@ -94,14 +95,14 @@ fun spoofVideoStreamsPatch(
                     getReference<MethodReference>()?.name == "newUrlRequestBuilder"
             }
             val urlRegister = getInstruction<FiveRegisterInstruction>(newRequestBuilderIndex).registerD
-            val freeRegister = getInstruction<OneRegisterInstruction>(newRequestBuilderIndex + 1).registerA
+            val freeRegister = findFreeRegister(newRequestBuilderIndex, urlRegister)
 
             addInstructions(
                 newRequestBuilderIndex,
                 """
-                move-object v$freeRegister, p1
+                    move-object v$freeRegister, p1
-                invoke-static { v$urlRegister, v$freeRegister }, $EXTENSION_CLASS_DESCRIPTOR->fetchStreams(Ljava/lang/String;Ljava/util/Map;)V
+                    invoke-static { v$urlRegister, v$freeRegister }, $EXTENSION_CLASS_DESCRIPTOR->fetchStreams(Ljava/lang/String;Ljava/util/Map;)V
-            """,
+                """
             )
         }
 

patches/src/main/kotlin/app/revanced/patches/youtube/layout/hide/general/HideLayoutComponentsPatch.kt

@@ -20,11 +20,10 @@ import app.revanced.patches.shared.misc.settings.preference.*
 import app.revanced.patches.youtube.misc.litho.filter.addLithoFilter
 import app.revanced.patches.youtube.misc.litho.filter.lithoFilterPatch
 import app.revanced.patches.youtube.misc.navigation.navigationBarHookPatch
-import app.revanced.patches.youtube.misc.playservice.is_19_47_or_greater
 import app.revanced.patches.youtube.misc.playservice.is_20_07_or_greater
-import app.revanced.patches.youtube.misc.playservice.versionCheckPatch
 import app.revanced.patches.youtube.misc.settings.PreferenceScreen
 import app.revanced.patches.youtube.misc.settings.settingsPatch
+import app.revanced.util.findFreeRegister
 import app.revanced.util.findInstructionIndicesReversedOrThrow
 import app.revanced.util.getReference
 import app.revanced.util.indexOfFirstInstructionOrThrow
@@ -123,7 +122,6 @@ val hideLayoutComponentsPatch = bytecodePatch(
         addResourcesPatch,
         hideLayoutComponentsResourcePatch,
         navigationBarHookPatch,
-        versionCheckPatch
     )
 
     compatibleWith(
@@ -254,17 +252,16 @@ val hideLayoutComponentsPatch = bytecodePatch(
         (if (is_20_07_or_greater) parseElementFromBufferFingerprint
         else parseElementFromBufferLegacyFingerprint).let {
             it.method.apply {
-                // Target code is a mess with a lot of register moves.
-                // There is no simple way to find a free register for all versions so this is hard coded.
-                val freeRegister = if (is_19_47_or_greater) 6 else 0
                 val byteArrayParameter = "p3"
                 val startIndex = it.patternMatch!!.startIndex
                 val conversionContextRegister = getInstruction<TwoRegisterInstruction>(startIndex).registerA
                 val returnEmptyComponentInstruction = instructions.last { it.opcode == Opcode.INVOKE_STATIC }
                 val returnEmptyComponentRegister = (returnEmptyComponentInstruction as FiveRegisterInstruction).registerC
+                val insertIndex = startIndex + 1
+                val freeRegister = findFreeRegister(insertIndex, conversionContextRegister, returnEmptyComponentRegister)
 
                 addInstructionsWithLabels(
-                    startIndex + 1,
+                    insertIndex,
                     """
                         invoke-static { v$conversionContextRegister, $byteArrayParameter }, $LAYOUT_COMPONENTS_FILTER_CLASS_DESCRIPTOR->filterMixPlaylists(Ljava/lang/Object;[B)Z
                         move-result v$freeRegister 

patches/src/main/kotlin/app/revanced/patches/youtube/layout/shortsplayer/OpenShortsInRegularPlayerPatch.kt

@@ -16,9 +16,9 @@ import app.revanced.patches.youtube.misc.playservice.versionCheckPatch
 import app.revanced.patches.youtube.misc.settings.PreferenceScreen
 import app.revanced.patches.youtube.misc.settings.settingsPatch
 import app.revanced.patches.youtube.shared.mainActivityOnCreateFingerprint
+import app.revanced.util.findFreeRegister
 import app.revanced.util.getReference
 import app.revanced.util.indexOfFirstInstructionOrThrow
-import com.android.tools.smali.dexlib2.iface.instruction.FiveRegisterInstruction
 import com.android.tools.smali.dexlib2.iface.instruction.OneRegisterInstruction
 import com.android.tools.smali.dexlib2.iface.reference.MethodReference
 
@@ -106,11 +106,12 @@ val openShortsInRegularPlayerPatch = bytecodePatch(
                     getReference<MethodReference>()?.returnType ==
                             "Lcom/google/android/libraries/youtube/player/model/PlaybackStartDescriptor;"
                 }
-                val freeRegister = getInstruction<FiveRegisterInstruction>(index).registerC
                 val playbackStartRegister = getInstruction<OneRegisterInstruction>(index + 1).registerA
+                val insertIndex = index + 2
+                val freeRegister = findFreeRegister(insertIndex, playbackStartRegister)
 
                 addInstructionsWithLabels(
-                    index + 2,
+                    insertIndex,
                     extensionInstructions(playbackStartRegister, freeRegister)
                 )
             }

patches/src/main/kotlin/app/revanced/patches/youtube/layout/startupshortsreset/DisableResumingShortsOnStartupPatch.kt

@@ -1,9 +1,7 @@
 package app.revanced.patches.youtube.layout.startupshortsreset
 
 import app.revanced.patcher.extensions.InstructionExtensions.addInstructions
-import app.revanced.patcher.extensions.InstructionExtensions.addInstructionsWithLabels
 import app.revanced.patcher.extensions.InstructionExtensions.getInstruction
-import app.revanced.patcher.extensions.InstructionExtensions.removeInstruction
 import app.revanced.patcher.patch.bytecodePatch
 import app.revanced.patches.all.misc.resources.addResources
 import app.revanced.patches.all.misc.resources.addResourcesPatch
@@ -12,11 +10,12 @@ import app.revanced.patches.youtube.misc.extension.sharedExtensionPatch
 import app.revanced.patches.youtube.misc.playservice.is_20_02_or_greater
 import app.revanced.patches.youtube.misc.settings.PreferenceScreen
 import app.revanced.patches.youtube.misc.settings.settingsPatch
+import app.revanced.util.addInstructionsAtControlFlowLabel
+import app.revanced.util.findFreeRegister
 import app.revanced.util.getReference
 import app.revanced.util.indexOfFirstInstructionOrThrow
 import app.revanced.util.indexOfFirstInstructionReversedOrThrow
 import com.android.tools.smali.dexlib2.Opcode
-import com.android.tools.smali.dexlib2.iface.instruction.FiveRegisterInstruction
 import com.android.tools.smali.dexlib2.iface.instruction.OneRegisterInstruction
 import com.android.tools.smali.dexlib2.iface.reference.MethodReference
 
@@ -80,23 +79,19 @@ val disableResumingShortsOnStartupPatch = bytecodePatch(
                             reference?.definingClass == "Lcom/google/common/util/concurrent/ListenableFuture;" &&
                             reference.name == "isDone"
                 }
-                val originalInstructionRegister =
+                val freeRegister = findFreeRegister(listenableInstructionIndex)
-                    getInstruction<FiveRegisterInstruction>(listenableInstructionIndex).registerC
-                val freeRegister =
-                    getInstruction<OneRegisterInstruction>(listenableInstructionIndex + 1).registerA
 
-                addInstructionsWithLabels(
+                addInstructionsAtControlFlowLabel(
-                    listenableInstructionIndex + 1,
+                    listenableInstructionIndex,
                     """
                         invoke-static {}, $EXTENSION_CLASS_DESCRIPTOR->disableResumingStartupShortsPlayer()Z
                         move-result v$freeRegister
                         if-eqz v$freeRegister, :show
                         return-void
                         :show
-                        invoke-interface {v$originalInstructionRegister}, Lcom/google/common/util/concurrent/ListenableFuture;->isDone()Z
+                        nop
                     """
                 )
-                removeInstruction(listenableInstructionIndex)
             }
         }
 

patches/src/main/kotlin/app/revanced/patches/youtube/misc/fix/playbackspeed/FIxPlaybackSpeedWhilePlayingPatch.kt

@@ -8,6 +8,7 @@ import app.revanced.patches.youtube.misc.extension.sharedExtensionPatch
 import app.revanced.patches.youtube.misc.playertype.playerTypeHookPatch
 import app.revanced.patches.youtube.misc.playservice.is_19_34_or_greater
 import app.revanced.patches.youtube.misc.playservice.versionCheckPatch
+import app.revanced.util.findFreeRegister
 import app.revanced.util.indexOfFirstInstructionOrThrow
 import com.android.tools.smali.dexlib2.Opcode
 import com.android.tools.smali.dexlib2.iface.instruction.TwoRegisterInstruction
@@ -42,13 +43,14 @@ val fixPlaybackSpeedWhilePlayingPatch = bytecodePatch{
         }
 
         playbackSpeedInFeedsFingerprint.method.apply {
-            val freeRegister = implementation!!.registerCount - parameters.size - 2
             val playbackSpeedIndex = indexOfGetPlaybackSpeedInstruction(this)
             val playbackSpeedRegister = getInstruction<TwoRegisterInstruction>(playbackSpeedIndex).registerA
             val returnIndex = indexOfFirstInstructionOrThrow(playbackSpeedIndex, Opcode.RETURN_VOID)
+            val insertIndex = playbackSpeedIndex + 1
+            val freeRegister = findFreeRegister(insertIndex, playbackSpeedRegister)
 
             addInstructionsWithLabels(
-                playbackSpeedIndex + 1,
+                insertIndex,
                 """
                     invoke-static { v$playbackSpeedRegister }, $EXTENSION_CLASS_DESCRIPTOR->playbackSpeedChanged(F)Z
                     move-result v$freeRegister

patches/src/main/kotlin/app/revanced/patches/youtube/misc/litho/filter/LithoFilterPatch.kt

@@ -8,7 +8,6 @@ import app.revanced.patcher.extensions.InstructionExtensions.addInstructionsWith
 import app.revanced.patcher.extensions.InstructionExtensions.getInstruction
 import app.revanced.patcher.extensions.InstructionExtensions.removeInstructions
 import app.revanced.patcher.extensions.InstructionExtensions.replaceInstruction
-import app.revanced.patcher.patch.PatchException
 import app.revanced.patcher.patch.bytecodePatch
 import app.revanced.patcher.util.smali.ExternalLabel
 import app.revanced.patches.youtube.misc.extension.sharedExtensionPatch
@@ -16,12 +15,14 @@ import app.revanced.patches.youtube.misc.playservice.is_19_18_or_greater
 import app.revanced.patches.youtube.misc.playservice.is_19_25_or_greater
 import app.revanced.patches.youtube.misc.playservice.is_20_05_or_greater
 import app.revanced.patches.youtube.misc.playservice.versionCheckPatch
+import app.revanced.util.findFreeRegister
 import app.revanced.util.getReference
 import app.revanced.util.indexOfFirstInstructionOrThrow
 import app.revanced.util.indexOfFirstInstructionReversedOrThrow
 import com.android.tools.smali.dexlib2.AccessFlags
 import com.android.tools.smali.dexlib2.Opcode
-import com.android.tools.smali.dexlib2.iface.instruction.*
+import com.android.tools.smali.dexlib2.iface.instruction.OneRegisterInstruction
+import com.android.tools.smali.dexlib2.iface.instruction.TwoRegisterInstruction
 import com.android.tools.smali.dexlib2.iface.reference.FieldReference
 import com.android.tools.smali.dexlib2.iface.reference.MethodReference
 
@@ -191,17 +192,7 @@ val lithoFilterPatch = bytecodePatch(
                 },
             ).registerA
 
-            // Find a free temporary register.
+            val freeRegister = findFreeRegister(insertHookIndex, identifierRegister, stringBuilderRegister)
-            val freeRegister = getInstruction<OneRegisterInstruction>(
-                // Immediately before is a StringBuilder append constant character.
-                indexOfFirstInstructionReversedOrThrow(insertHookIndex, Opcode.CONST_16),
-            ).registerA
-
-            // Verify the temp register will not clobber the method result register.
-            if (stringBuilderRegister == freeRegister) {
-                throw PatchException("Free register will clobber StringBuilder register")
-            }
-
             val invokeFilterInstructions = """
                 invoke-static { v$identifierRegister, v$stringBuilderRegister }, $EXTENSION_CLASS_DESCRIPTOR->filter(Ljava/lang/String;Ljava/lang/StringBuilder;)Z
                 move-result v$freeRegister

patches/src/main/kotlin/app/revanced/util/BytecodeUtils.kt

@@ -15,13 +15,135 @@ import app.revanced.patches.shared.misc.mapping.get
 import app.revanced.patches.shared.misc.mapping.resourceMappingPatch
 import app.revanced.patches.shared.misc.mapping.resourceMappings
 import com.android.tools.smali.dexlib2.Opcode
+import com.android.tools.smali.dexlib2.Opcode.*
 import com.android.tools.smali.dexlib2.iface.Method
+import com.android.tools.smali.dexlib2.iface.instruction.FiveRegisterInstruction
 import com.android.tools.smali.dexlib2.iface.instruction.Instruction
 import com.android.tools.smali.dexlib2.iface.instruction.OneRegisterInstruction
 import com.android.tools.smali.dexlib2.iface.instruction.ReferenceInstruction
+import com.android.tools.smali.dexlib2.iface.instruction.RegisterRangeInstruction
+import com.android.tools.smali.dexlib2.iface.instruction.ThreeRegisterInstruction
+import com.android.tools.smali.dexlib2.iface.instruction.TwoRegisterInstruction
 import com.android.tools.smali.dexlib2.iface.instruction.WideLiteralInstruction
 import com.android.tools.smali.dexlib2.iface.reference.Reference
 import com.android.tools.smali.dexlib2.util.MethodUtil
+import java.util.EnumSet
+
+/**
+ * Starting from and including the instruction at index [startIndex],
+ * finds the next register that is wrote to and not read from. If a return instruction
+ * is encountered, then the lowest unused register is returned.
+ *
+ * This method can return a non 4-bit register, and the calling code may need to temporarily
+ * swap register contents if a 4-bit register is required.
+ *
+ * @param startIndex Inclusive starting index.
+ * @param registersToExclude Registers to exclude, and consider as used. For most use cases,
+ *                           all registers used in injected code should be specified.
+ * @throws IllegalArgumentException If a branch or conditional statement is encountered
+ *                                  before a suitable register is found.
+ */
+internal fun Method.findFreeRegister(startIndex: Int, vararg registersToExclude: Int): Int {
+    if (implementation == null) {
+        throw IllegalArgumentException("Method has no implementation: $this")
+    }
+    if (startIndex < 0 || startIndex >= instructions.count()) {
+        throw IllegalArgumentException("startIndex out of bounds: $startIndex")
+    }
+
+    // All registers used by an instruction.
+    fun Instruction.getRegistersUsed() = when (this) {
+        is FiveRegisterInstruction -> listOf(registerC, registerD, registerE, registerF, registerG)
+        is ThreeRegisterInstruction -> listOf(registerA, registerB, registerC)
+        is TwoRegisterInstruction -> listOf(registerA, registerB)
+        is OneRegisterInstruction -> listOf(registerA)
+        is RegisterRangeInstruction -> (startRegister until (startRegister + registerCount)).toList()
+        else -> emptyList()
+    }
+
+    // Register that is written to by an instruction.
+    fun Instruction.getRegisterWritten() = when (this) {
+        is ThreeRegisterInstruction -> registerA
+        is TwoRegisterInstruction -> registerA
+        is OneRegisterInstruction -> registerA
+        else -> throw IllegalStateException("Not a write instruction: $this")
+    }
+
+    val writeOpcodes = EnumSet.of(
+        NEW_INSTANCE, NEW_ARRAY,
+        MOVE, MOVE_FROM16, MOVE_16, MOVE_WIDE, MOVE_WIDE_FROM16, MOVE_WIDE_16, MOVE_OBJECT,
+        MOVE_OBJECT_FROM16, MOVE_OBJECT_16, MOVE_RESULT, MOVE_RESULT_WIDE, MOVE_RESULT_OBJECT, MOVE_EXCEPTION,
+        IGET, IGET_WIDE, IGET_OBJECT, IGET_BOOLEAN, IGET_BYTE, IGET_CHAR, IGET_SHORT,
+        SGET, SGET_WIDE, SGET_OBJECT, SGET_BOOLEAN, SGET_BYTE, SGET_CHAR, SGET_SHORT,
+    )
+
+    val branchOpcodes = EnumSet.of(
+        GOTO, GOTO_16, GOTO_32,
+        IF_EQ, IF_NE, IF_LT, IF_GE, IF_GT, IF_LE,
+        IF_EQZ, IF_NEZ, IF_LTZ, IF_GEZ, IF_GTZ, IF_LEZ,
+    )
+
+    val returnOpcodes = EnumSet.of(
+        RETURN_VOID, RETURN, RETURN_WIDE, RETURN_OBJECT,
+    )
+
+    // Highest 4-bit register available, exclusive. Ideally return a free register less than this.
+    val maxRegister4Bits = 16
+    var bestFreeRegisterFound: Int? = null
+    val usedRegisters = registersToExclude.toMutableSet()
+
+    for (i in startIndex until instructions.count()) {
+        val instruction = getInstruction(i)
+
+        if (instruction.opcode in returnOpcodes) {
+            // Method returns. Use lowest register that hasn't been encountered.
+            val freeRegister = (0 until implementation!!.registerCount).find {
+                it !in usedRegisters
+            }
+            if (freeRegister != null) {
+                return freeRegister
+            }
+            if (bestFreeRegisterFound != null) {
+                return bestFreeRegisterFound;
+            }
+
+            // Somehow every method register was read from before any register was wrote to.
+            // In practice this never occurs.
+            throw IllegalArgumentException("Could not find a free register from startIndex: " +
+                    "$startIndex excluding: $registersToExclude")
+        }
+
+        if (instruction.opcode in branchOpcodes) {
+            if (bestFreeRegisterFound != null) {
+                return bestFreeRegisterFound;
+            }
+            // This method is simple and does not follow branching.
+            throw IllegalArgumentException("Encountered a branch statement before a free register could be found")
+        }
+
+        if (instruction.opcode in writeOpcodes) {
+            val freeRegister = instruction.getRegisterWritten()
+            if (freeRegister !in usedRegisters) {
+                if (freeRegister < maxRegister4Bits) {
+                    // Found an ideal register.
+                    return freeRegister
+                }
+
+                // Continue searching for a 4-bit register if available.
+                if (bestFreeRegisterFound == null || freeRegister < bestFreeRegisterFound) {
+                    bestFreeRegisterFound = freeRegister
+                }
+            }
+        }
+
+        usedRegisters.addAll(instruction.getRegistersUsed())
+    }
+
+    // Cannot be reached since a branch or return statement will
+    // be encountered before the end of the method.
+    throw IllegalStateException()
+}
+
 
 /**
  * Find the [MutableMethod] from a given [Method] in a [MutableClass].
@@ -395,7 +517,7 @@ fun Method.findInstructionIndicesReversedOrThrow(opcode: Opcode): List<Int> {
 
 internal fun MutableMethod.insertFeatureFlagBooleanOverride(literal: Long, extensionsMethod: String) {
     val literalIndex = indexOfFirstLiteralInstructionOrThrow(literal)
-    val index = indexOfFirstInstructionOrThrow(literalIndex, Opcode.MOVE_RESULT)
+    val index = indexOfFirstInstructionOrThrow(literalIndex, MOVE_RESULT)
     val register = getInstruction<OneRegisterInstruction>(index).registerA
 
     val operation = if (register < 16) {
@@ -423,7 +545,7 @@ fun BytecodePatchContext.forEachLiteralValueInstruction(
     classes.forEach { classDef ->
         classDef.methods.forEach { method ->
             method.implementation?.instructions?.forEachIndexed { index, instruction ->
-                if (instruction.opcode == Opcode.CONST &&
+                if (instruction.opcode == CONST &&
                     (instruction as WideLiteralInstruction).wideLiteral == literal
                 ) {
                     val mutableMethod = proxy(classDef).mutableClass.findMutableMethodOf(method)