From 76fb700884dae5b71a57f9530fad9d4a98ba0af0 Mon Sep 17 00:00:00 2001
From: Linus <23507341+Linus789@users.noreply.github.com>
Date: Tue, 30 May 2023 21:51:43 +0000
Subject: [PATCH] fix(spoof-wifi-connection): use updated instruction indices
 (#2199)

Co-authored-by: Linus789 <Linus789@users.noreply.github.com>
Co-authored-by: oSumAtrIX <johan.melkonyan1@web.de>
---
 .../AbstractTransformInstructionsPatch.kt     | 55 ++++++++++---------
 1 file changed, 30 insertions(+), 25 deletions(-)

diff --git a/src/main/kotlin/app/revanced/util/patch/AbstractTransformInstructionsPatch.kt b/src/main/kotlin/app/revanced/util/patch/AbstractTransformInstructionsPatch.kt
index 3a1b8f23..a8354c6b 100644
--- a/src/main/kotlin/app/revanced/util/patch/AbstractTransformInstructionsPatch.kt
+++ b/src/main/kotlin/app/revanced/util/patch/AbstractTransformInstructionsPatch.kt
@@ -21,39 +21,44 @@ internal abstract class AbstractTransformInstructionsPatch<T> : BytecodePatch()
 
     abstract fun transform(mutableMethod: MutableMethod, entry: T)
 
+    // Returns the patch indices as a Sequence, which will execute lazily.
+    private fun findPatchIndices(classDef: ClassDef, method: Method): Sequence<T>? {
+        return method.implementation?.instructions?.asSequence()?.withIndex()?.mapNotNull { (index, instruction) ->
+            filterMap(classDef, method, instruction, index)
+        }
+    }
+
     override fun execute(context: BytecodeContext): PatchResult {
-        // Find all instructions
+        // Find all methods to patch
         buildMap {
             context.classes.forEach { classDef ->
-                classDef.methods.let { methods ->
-                    buildMap methodList@{
-                        methods.forEach methods@{ method ->
-                            with(method.implementation?.instructions ?: return@methods) {
-                                ArrayDeque<T>().also { patchIndices ->
-                                    this.forEachIndexed { index, instruction ->
-                                        val result = filterMap(classDef, method, instruction, index)
-                                        if (result != null) {
-                                            patchIndices.add(result)
-                                        }
-                                    }
-                                }.also { if (it.isEmpty()) return@methods }.let { patches ->
-                                    put(method, patches)
-                                }
-                            }
+                val methods = buildList {
+                    classDef.methods.forEach { method ->
+                        // Since the Sequence executes lazily,
+                        // using any() results in only calling
+                        // filterMap until the first index has been found.
+                        val patchIndices = findPatchIndices(classDef, method)
+
+                        if (patchIndices?.any() == true) {
+                            add(method)
                         }
                     }
-                }.also { if (it.isEmpty()) return@forEach }.let { methodPatches ->
-                    put(classDef, methodPatches)
+                }
+
+                if (methods.isNotEmpty()) {
+                    put(classDef, methods)
                 }
             }
         }.forEach { (classDef, methods) ->
-            // And finally transform the instructions...
-            with(context.proxy(classDef).mutableClass) {
-                methods.forEach { (method, patches) ->
-                    val mutableMethod = findMutableMethodOf(method)
-                    while (!patches.isEmpty()) {
-                        transform(mutableMethod, patches.removeLast())
-                    }
+            // And finally transform the methods...
+            val mutableClass = context.proxy(classDef).mutableClass
+
+            methods.map(mutableClass::findMutableMethodOf).forEach methods@{ mutableMethod ->
+                val patchIndices = findPatchIndices(mutableClass, mutableMethod)?.toCollection(ArrayDeque())
+                    ?: return@methods
+
+                while (!patchIndices.isEmpty()) {
+                    transform(mutableMethod, patchIndices.removeLast())
                 }
             }
         }