fix: Fixed writer & signature resolver, improved tests & speed, minor…

… refactoring

src/main/kotlin/app/revanced/patcher/Patcher.kt

@@ -2,7 +2,7 @@ package app.revanced.patcher
 
 import app.revanced.patcher.cache.Cache
 import app.revanced.patcher.patch.Patch
-import app.revanced.patcher.resolver.MethodResolver
+import app.revanced.patcher.resolver.SignatureResolver
 import app.revanced.patcher.signature.MethodSignature
 import lanchon.multidexlib2.BasicDexFileNamer
 import lanchon.multidexlib2.MultiDexIO
@@ -21,23 +21,20 @@ class Patcher(
     private val patches = mutableSetOf<Patch>()
 
     init {
-        // TODO: find a way to load all dex classes, the code below only loads the first .dex file
         val dexFile = MultiDexIO.readDexFile(true, input, BasicDexFileNamer(), Opcodes.getDefault(), null)
-        cache = Cache(dexFile.classes, MethodResolver(dexFile.classes, signatures).resolve())
+        cache = Cache(dexFile.classes, SignatureResolver(dexFile.classes, signatures).resolve())
     }
 
     fun save() {
         val newDexFile = object : DexFile {
-            override fun getClasses(): MutableSet<out ClassDef> {
-                // TODO: find a way to return a set with a custom iterator
-                // TODO: the iterator would return the proxied class matching the current index of the list
-                // TODO: instead of the original class
-                for (classProxy in cache.classProxy) {
-                    if (!classProxy.proxyUsed) continue
-                    // TODO: merge this class with cache.classes somehow in an iterator
-                    classProxy.mutatedClass
-                }
-                return cache.classes.toMutableSet()
+            override fun getClasses(): Set<ClassDef> {
+                // this is a slow workaround for now
+                val mutableClassList = cache.classes.toMutableList()
+                cache.classProxy
+                    .filter { it.proxyUsed }.forEach { proxy ->
+                        mutableClassList[proxy.originalIndex] = proxy.mutatedClass
+                    }
+                return mutableClassList.toSet()
             }
 
             override fun getOpcodes(): Opcodes {
@@ -46,8 +43,8 @@ class Patcher(
             }
         }
 
-        // TODO: not sure about maxDexPoolSize & we should use the multithreading overload for writeDexFile
-        MultiDexIO.writeDexFile(true, output, BasicDexFileNamer(), newDexFile, 10, null)
+        // TODO: we should use the multithreading capable overload for writeDexFile
+        MultiDexIO.writeDexFile(true, output, BasicDexFileNamer(), newDexFile, 50000, null)
     }
 
     fun addPatches(vararg patches: Patch) {
@@ -56,6 +53,7 @@ class Patcher(
 
     fun applyPatches(stopOnError: Boolean = false): Map<String, Result<Nothing?>> {
         return buildMap {
+            // TODO: after each patch execution we could clear left overs like proxied classes to safe memory
             for (patch in patches) {
                 val result: Result<Nothing?> = try {
                     val pr = patch.execute(cache)

src/main/kotlin/app/revanced/patcher/cache/Cache.kt

@@ -1,22 +1,25 @@
 package app.revanced.patcher.cache
 
 import app.revanced.patcher.proxy.ClassProxy
-import app.revanced.patcher.signature.MethodSignatureScanResult
+import app.revanced.patcher.signature.SignatureResolverResult
 import org.jf.dexlib2.iface.ClassDef
 
 class Cache(
     internal val classes: Set<ClassDef>,
     val resolvedMethods: MethodMap
 ) {
+    // TODO: currently we create ClassProxies at multiple places, which is why we could have merge conflicts
+    //  this can be solved by creating a dedicated method for creating class proxies,
+    //  if the class proxy already exists in the cached proxy list below
     internal val classProxy = mutableSetOf<ClassProxy>()
 
     fun findClass(predicate: (ClassDef) -> Boolean): ClassProxy? {
         // if we already proxied the class matching the predicate,
-        val proxiedClass = classProxy.singleOrNull{classProxy -> predicate(classProxy.immutableClass)}
+        val proxiedClass = classProxy.singleOrNull { classProxy -> predicate(classProxy.immutableClass) }
         // return that proxy
         if (proxiedClass != null) return proxiedClass
         // else search the original class list
-        val foundClass =  classes.singleOrNull(predicate) ?: return null
+        val foundClass = classes.singleOrNull(predicate) ?: return null
         // create a class proxy with the index of the class in the classes list
         // TODO: There might be a more elegant way to the comment above
         val classProxy = ClassProxy(foundClass, classes.indexOf(foundClass))
@@ -27,8 +30,8 @@ class Cache(
     }
 }
 
-class MethodMap : LinkedHashMap<String, MethodSignatureScanResult>() {
-    override fun get(key: String): MethodSignatureScanResult {
+class MethodMap : LinkedHashMap<String, SignatureResolverResult>() {
+    override fun get(key: String): SignatureResolverResult {
         return super.get(key) ?: throw MethodNotFoundException("Method $key was not found in the method cache")
     }
 }

src/main/kotlin/app/revanced/patcher/proxy/ClassProxy.kt

@@ -3,10 +3,9 @@ package app.revanced.patcher.proxy
 import app.revanced.patcher.proxy.mutableTypes.MutableClass
 import org.jf.dexlib2.iface.ClassDef
 
-
 class ClassProxy(
     val immutableClass: ClassDef,
-    val originalClassIndex: Int,
+    val originalIndex: Int,
 ) {
     internal var proxyUsed = false
     internal lateinit var mutatedClass: MutableClass

src/main/kotlin/app/revanced/patcher/resolver/MethodResolverScanResult.kt → src/main/kotlin/app/revanced/patcher/resolver/PatternScanData.kt

@@ -1,6 +1,6 @@
 package app.revanced.patcher.resolver
 
-internal data class MethodResolverScanResult(
+internal data class PatternScanData(
     val found: Boolean,
     val startIndex: Int? = 0,
     val endIndex: Int? = 0

src/main/kotlin/app/revanced/patcher/resolver/MethodResolver.kt → src/main/kotlin/app/revanced/patcher/resolver/SignatureResolver.kt

@@ -1,43 +1,51 @@
 package app.revanced.patcher.resolver
 
 import app.revanced.patcher.cache.MethodMap
-import app.revanced.patcher.signature.MethodSignatureScanResult
-import app.revanced.patcher.signature.PatternScanData
+import app.revanced.patcher.proxy.ClassProxy
 import app.revanced.patcher.signature.MethodSignature
+import app.revanced.patcher.signature.PatternScanResult
+import app.revanced.patcher.signature.SignatureResolverResult
 import org.jf.dexlib2.Opcode
 import org.jf.dexlib2.iface.ClassDef
 import org.jf.dexlib2.iface.Method
 
-// TODO: add logger
-internal class MethodResolver(private val classes: Set<ClassDef>, private val signatures: Array<MethodSignature>) {
+// TODO: add logger back
+internal class SignatureResolver(
+    private val classes: Set<ClassDef>,
+    private val methodSignatures: Array<MethodSignature>
+) {
     fun resolve(): MethodMap {
         val methodMap = MethodMap()
 
-        for (classDef in classes) {
-            for (method in classDef.methods) {
-                for (methodSignature in signatures) {
-                    if (methodMap.containsKey(methodSignature.name)) { // method already found for this sig
-                        continue
-                    }
+        for ((index, classDef) in classes.withIndex()) {
+            for (signature in methodSignatures) {
+                if (methodMap.containsKey(signature.name)) {
+                    continue
+                }
+
+                for (method in classDef.methods) {
+                    val (isMatch, patternScanData) = compareSignatureToMethod(signature, method)
 
-                    val (r, sr) = cmp(method, methodSignature)
-                    if (!r || sr == null) {
+                    if (!isMatch || patternScanData == null) {
                         continue
                     }
 
-                    methodMap[methodSignature.name] = MethodSignatureScanResult(
-                        method,
-                        PatternScanData(
-                            // sadly we cannot create contracts for a data class, so we must assert
-                            sr.startIndex!!,
-                            sr.endIndex!!
+                    // create class proxy, in case a patch needs mutability
+                    val classProxy = ClassProxy(classDef, index)
+                    methodMap[signature.name] = SignatureResolverResult(
+                        classProxy,
+                        method.name,
+                        PatternScanResult(
+                            patternScanData.startIndex!!,
+                            patternScanData.endIndex!!
                         )
                     )
                 }
             }
         }
 
-        for (signature in signatures) {
+        // TODO: remove?
+        for (signature in methodSignatures) {
             if (methodMap.containsKey(signature.name)) continue
         }
 
@@ -46,53 +54,57 @@ internal class MethodResolver(private val classes: Set<ClassDef>, private val si
 
     // These functions do not require the constructor values, so they can be static.
     companion object {
-        fun resolveMethod(classNode: ClassDef, signature: MethodSignature): MethodSignatureScanResult? {
-            for (method in classNode.methods) {
-                val (r, sr) = cmp(method, signature)
+        fun resolveFromProxy(classProxy: ClassProxy, signature: MethodSignature): SignatureResolverResult? {
+            for (method in classProxy.immutableClass.methods) {
+                val (r, sr) = compareSignatureToMethod(signature, method)
                 if (!r || sr == null) continue
-                return MethodSignatureScanResult(
-                    method,
-                    PatternScanData(0, 0) // opcode list is always ignored.
+                return SignatureResolverResult(
+                    classProxy,
+                    method.name,
+                    null
                 )
             }
             return null
         }
 
-        private fun cmp(method: Method, signature: MethodSignature): Pair<Boolean, MethodResolverScanResult?> {
+        private fun compareSignatureToMethod(
+            signature: MethodSignature,
+            method: Method
+        ): Pair<Boolean, PatternScanData?> {
             // TODO: compare as generic object if not primitive
             signature.returnType?.let { _ ->
                 if (signature.returnType != method.returnType) {
-                    return@cmp false to null
+                    return@compareSignatureToMethod false to null
                 }
             }
 
             signature.accessFlags?.let { _ ->
                 if (signature.accessFlags != method.accessFlags) {
-                    return@cmp false to null
+                    return@compareSignatureToMethod false to null
                 }
             }
 
             // TODO: compare as generic object if the parameter is not primitive
             signature.methodParameters?.let { _ ->
                 if (signature.methodParameters != method.parameters) {
-                    return@cmp false to null
+                    return@compareSignatureToMethod false to null
                 }
             }
 
             signature.opcodes?.let { _ ->
                 val result = method.implementation?.instructions?.scanFor(signature.opcodes)
-                return@cmp if (result != null && result.found) true to result else false to null
+                return@compareSignatureToMethod if (result != null && result.found) true to result else false to null
             }
 
-            return true to MethodResolverScanResult(true)
+            return true to PatternScanData(true)
         }
     }
 }
 
 private operator fun ClassDef.component1() = this
 private operator fun ClassDef.component2() = this.methods
 
-private fun <T> MutableIterable<T>.scanFor(pattern: Array<Opcode>): MethodResolverScanResult {
+private fun <T> MutableIterable<T>.scanFor(pattern: Array<Opcode>): PatternScanData {
     // TODO: create var for count?
     for (i in 0 until this.count()) {
         var occurrence = 0
@@ -101,12 +113,12 @@ private fun <T> MutableIterable<T>.scanFor(pattern: Array<Opcode>): MethodResolv
             if (!n.shouldSkip() && n != pattern[occurrence]) break
             if (++occurrence >= pattern.size) {
                 val current = i + occurrence
-                return MethodResolverScanResult(true, current - pattern.size, current)
+                return PatternScanData(true, current - pattern.size, current)
             }
         }
     }
 
-    return MethodResolverScanResult(false)
+    return PatternScanData(false)
 }
 
 // TODO: extend Opcode type, not T (requires a cast to Opcode)

src/main/kotlin/app/revanced/patcher/signature/SignatureResolverResult.kt

@@ -0,0 +1,20 @@
+package app.revanced.patcher.signature
+
+import app.revanced.patcher.proxy.ClassProxy
+import app.revanced.patcher.resolver.SignatureResolver
+
+data class SignatureResolverResult(
+    val definingClassProxy: ClassProxy,
+    val resolvedMethodName: String,
+    val scanData: PatternScanResult?
+) {
+    @Suppress("Unused") // TODO(Sculas): remove this when we have coverage for this method.
+    fun findParentMethod(signature: MethodSignature): SignatureResolverResult? {
+        return SignatureResolver.resolveFromProxy(definingClassProxy, signature)
+    }
+}
+
+data class PatternScanResult(
+    val startIndex: Int,
+    val endIndex: Int
+)