perf: resolve fingerprints using method maps (#185)

Co-authored-by: oSumAtrIX <johan.melkonyan1@web.de>

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

@@ -4,7 +4,8 @@ import app.revanced.patcher.data.Context
 import app.revanced.patcher.extensions.PatchExtensions.dependencies
 import app.revanced.patcher.extensions.PatchExtensions.patchName
 import app.revanced.patcher.extensions.PatchExtensions.requiresIntegrations
-import app.revanced.patcher.fingerprint.method.impl.MethodFingerprint.Companion.resolve
+import app.revanced.patcher.fingerprint.method.impl.MethodFingerprint
+import app.revanced.patcher.fingerprint.method.impl.MethodFingerprint.Companion.resolveUsingLookupMap
 import app.revanced.patcher.patch.*
 import app.revanced.patcher.util.VersionReader
 import brut.androlib.Androlib
@@ -322,10 +323,7 @@ class Patcher(private val options: PatcherOptions) {
                 context.resourceContext
             } else {
                 context.bytecodeContext.also { context ->
-                    (patchInstance as BytecodePatch).fingerprints?.resolve(
+                    (patchInstance as BytecodePatch).fingerprints?.resolveUsingLookupMap(context)
-                        context,
-                        context.classes.classes
-                    )
                 }
             }
 
@@ -345,14 +343,17 @@ class Patcher(private val options: PatcherOptions) {
         return sequence {
             if (mergeIntegrations) context.integrations.merge(logger)
 
+            logger.trace("Initialize lookup maps for method MethodFingerprint resolution")
+
+            MethodFingerprint.initializeFingerprintResolutionLookupMaps(context.bytecodeContext)
+
             // prevent from decoding the manifest twice if it is not needed
             if (resourceDecodingMode == ResourceDecodingMode.FULL) decodeResources(ResourceDecodingMode.FULL)
 
-            logger.trace("Executing all patches")
+            logger.info("Executing patches")
 
             val executedPatches = LinkedHashMap<String, ExecutedPatch>() // first is name
 
-
             context.patches.forEach { patch ->
                 val patchResult = executePatch(patch, executedPatches)
 

src/main/kotlin/app/revanced/patcher/fingerprint/method/impl/MethodFingerprint.kt

@@ -4,7 +4,9 @@ import app.revanced.patcher.data.BytecodeContext
 import app.revanced.patcher.extensions.MethodFingerprintExtensions.fuzzyPatternScanMethod
 import app.revanced.patcher.fingerprint.Fingerprint
 import app.revanced.patcher.fingerprint.method.annotation.FuzzyPatternScanMethod
+import app.revanced.patcher.patch.PatchResultError
 import app.revanced.patcher.util.proxy.ClassProxy
+import org.jf.dexlib2.AccessFlags
 import org.jf.dexlib2.Opcode
 import org.jf.dexlib2.iface.ClassDef
 import org.jf.dexlib2.iface.Method
@@ -12,19 +14,21 @@ import org.jf.dexlib2.iface.instruction.Instruction
 import org.jf.dexlib2.iface.instruction.ReferenceInstruction
 import org.jf.dexlib2.iface.reference.StringReference
 import org.jf.dexlib2.util.MethodUtil
+import java.util.*
 
 private typealias StringMatch = MethodFingerprintResult.MethodFingerprintScanResult.StringsScanResult.StringMatch
 private typealias StringsScanResult = MethodFingerprintResult.MethodFingerprintScanResult.StringsScanResult
+private typealias MethodClassPair = Pair<Method, ClassDef>
 
 /**
- * Represents the [MethodFingerprint] for a method.
+ * A fingerprint to resolve methods.
- * @param returnType The return type of the method.
+ *
- * @param accessFlags The access flags of the method.
+ * @param returnType The method's return type compared using [String.startsWith].
- * @param parameters The parameters of the method.
+ * @param accessFlags The method's exact access flags using values of [AccessFlags].
- * @param opcodes The list of opcodes of the method.
+ * @param parameters The parameters of the method. Partial matches allowed and follow the same rules as [returnType].
- * @param strings A list of strings which a method contains.
+ * @param opcodes An opcode pattern of the method's instructions. Wildcard or unknown opcodes can be specified by `null`.
+ * @param strings A list of the method's strings compared each using [String.contains].
  * @param customFingerprint A custom condition for this fingerprint.
- * A `null` opcode is equals to an unknown opcode.
  */
 abstract class MethodFingerprint(
     internal val returnType: String? = null,
@@ -40,6 +44,192 @@ abstract class MethodFingerprint(
     var result: MethodFingerprintResult? = null
 
     companion object {
+        /**
+         * A list of methods and the class they were found in.
+         */
+        private val methods = mutableListOf<MethodClassPair>()
+
+        /**
+         * Lookup map for methods keyed to the methods access flags, return type and parameter.
+         */
+        private val methodSignatureLookupMap = mutableMapOf<String, MutableList<MethodClassPair>>()
+
+        /**
+         * Lookup map for methods keyed to the strings contained in the method.
+         */
+        private val methodStringsLookupMap = mutableMapOf<String, MutableList<MethodClassPair>>()
+
+        /**
+         * Appends a string based on the parameter reference types of this method.
+         */
+        private fun StringBuilder.appendParameters(parameters: Iterable<CharSequence>) {
+            // Maximum parameters to use in the signature key.
+            // Some apps have methods with an incredible number of parameters (over 100 parameters have been seen).
+            // To keep the signature map from becoming needlessly bloated,
+            // group together in the same map entry all methods with the same access/return and 5 or more parameters.
+            // The value of 5 was chosen based on local performance testing and is not set in stone.
+            val maxSignatureParameters = 5
+            // Must append a unique value before the parameters to distinguish this key includes the parameters.
+            // If this is not appended, then methods with no parameters
+            // will collide with different keys that specify access/return but omit the parameters.
+            append("p:")
+            parameters.forEachIndexed { index, parameter ->
+                if (index >= maxSignatureParameters) return
+                append(parameter.first())
+            }
+        }
+
+        /**
+         * Initializes lookup maps for [MethodFingerprint] resolution
+         * using attributes of methods such as the method signature or strings.
+         *
+         * @param context The [BytecodeContext] containing the classes to initialize the lookup maps with.
+         */
+        internal fun initializeFingerprintResolutionLookupMaps(context: BytecodeContext) {
+            fun MutableMap<String, MutableList<MethodClassPair>>.add(
+                key: String,
+                methodClassPair: MethodClassPair
+            ) {
+                var methodClassPairs = this[key]
+
+                methodClassPairs ?: run {
+                    methodClassPairs = LinkedList<MethodClassPair>().also { this[key] = it }
+                }
+
+                methodClassPairs!!.add(methodClassPair)
+            }
+
+            if (methods.isNotEmpty()) throw PatchResultError("Map already initialized")
+
+            context.classes.classes.forEach { classDef ->
+                classDef.methods.forEach { method ->
+                    val methodClassPair = method to classDef
+
+                    // For fingerprints with no access or return type specified.
+                    methods += methodClassPair
+
+                    val accessFlagsReturnKey = method.accessFlags.toString() + method.returnType.first()
+
+                    // Add <access><returnType> as the key.
+                    methodSignatureLookupMap.add(accessFlagsReturnKey, methodClassPair)
+
+                    // Add <access><returnType>[parameters] as the key.
+                    methodSignatureLookupMap.add(
+                        buildString {
+                            append(accessFlagsReturnKey)
+                            appendParameters(method.parameterTypes)
+                        },
+                        methodClassPair
+                    )
+
+                    // Add strings contained in the method as the key.
+                    method.implementation?.instructions?.forEach instructions@{ instruction ->
+                        if (instruction.opcode != Opcode.CONST_STRING && instruction.opcode != Opcode.CONST_STRING_JUMBO)
+                            return@instructions
+
+                        val string = ((instruction as ReferenceInstruction).reference as StringReference).string
+
+                        methodStringsLookupMap.add(string, methodClassPair)
+                    }
+
+                    // In the future, the class type could be added to the lookup map.
+                    // This would require MethodFingerprint to be changed to include the class type.
+                }
+            }
+        }
+
+        /**
+         * Resolve a list of [MethodFingerprint] using the lookup map built by [initializeFingerprintResolutionLookupMaps].
+         *
+         * [MethodFingerprint] resolution is fast, but if many are present they can consume a noticeable
+         * amount of time because they are resolved in sequence.
+         *
+         * For apps with many fingerprints, resolving performance can be improved by:
+         * - Slowest: Specify [opcodes] and nothing else.
+         * - Fast: Specify [accessFlags], [returnType].
+         * - Faster: Specify [accessFlags], [returnType] and [parameters].
+         * - Fastest: Specify [strings], with at least one string being an exact (non-partial) match.
+         */
+        internal fun Iterable<MethodFingerprint>.resolveUsingLookupMap(context: BytecodeContext) {
+            if (methods.isEmpty()) throw PatchResultError("lookup map not initialized")
+
+            for (fingerprint in this) {
+                fingerprint.resolveUsingLookupMap(context)
+            }
+        }
+
+        /**
+         * Resolve a [MethodFingerprint] using the lookup map built by [initializeFingerprintResolutionLookupMaps].
+         *
+         * [MethodFingerprint] resolution is fast, but if many are present they can consume a noticeable
+         * amount of time because they are resolved in sequence.
+         *
+         * For apps with many fingerprints, resolving performance can be improved by:
+         * - Slowest: Specify [opcodes] and nothing else.
+         * - Fast: Specify [accessFlags], [returnType].
+         * - Faster: Specify [accessFlags], [returnType] and [parameters].
+         * - Fastest: Specify [strings], with at least one string being an exact (non-partial) match.
+         */
+        internal fun MethodFingerprint.resolveUsingLookupMap(context: BytecodeContext): Boolean {
+            /**
+             * Lookup [MethodClassPair]s that match the methods strings present in a [MethodFingerprint].
+             *
+             * @return A list of [MethodClassPair]s that match the methods strings present in a [MethodFingerprint].
+             */
+            fun MethodFingerprint.methodStringsLookup(): List<MethodClassPair>? {
+                strings?.forEach {
+                    val methods = methodStringsLookupMap[it]
+                    if (methods != null) return methods
+                }
+                return null
+            }
+
+            /**
+             * Lookup [MethodClassPair]s that match the method signature present in a [MethodFingerprint].
+             *
+             * @return A list of [MethodClassPair]s that match the method signature present in a [MethodFingerprint].
+             */
+            fun MethodFingerprint.methodSignatureLookup(): List<MethodClassPair> {
+                if (accessFlags == null) return methods
+
+                var returnTypeValue = returnType
+                if (returnTypeValue == null) {
+                    if (AccessFlags.CONSTRUCTOR.isSet(accessFlags)) {
+                        // Constructors always have void return type
+                        returnTypeValue = "V"
+                    } else {
+                        return methods
+                    }
+                }
+
+                val key = buildString {
+                    append(accessFlags)
+                    append(returnTypeValue.first())
+                    if (parameters != null) appendParameters(parameters)
+                }
+                return methodSignatureLookupMap[key]!!
+            }
+
+            /**
+             * Resolve a [MethodFingerprint] using a list of [MethodClassPair].
+             *
+             * @return True if the resolution was successful, false otherwise.
+             */
+            fun MethodFingerprint.resolveUsingMethodClassPair(classMethods: Iterable<MethodClassPair>): Boolean {
+                classMethods.forEach { classAndMethod ->
+                    if (resolve(context, classAndMethod.first, classAndMethod.second)) return true
+                }
+                return false
+            }
+
+            val methodsWithSameStrings = methodStringsLookup()
+            if (methodsWithSameStrings != null) if (resolveUsingMethodClassPair(methodsWithSameStrings)) return true
+
+            // No strings declared or none matched (partial matches are allowed).
+            // Use signature matching.
+            return resolveUsingMethodClassPair(methodSignatureLookup())
+        }
+
         /**
          * Resolve a list of [MethodFingerprint] against a list of [ClassDef].
          *
@@ -48,10 +238,10 @@ abstract class MethodFingerprint(
          * @return True if the resolution was successful, false otherwise.
          */
         fun Iterable<MethodFingerprint>.resolve(context: BytecodeContext, classes: Iterable<ClassDef>) {
-            for (fingerprint in this) // For each fingerprint
+            for (fingerprint in this) // For each fingerprint...
-                classes@ for (classDef in classes) // search through all classes for the fingerprint
+                classes@ for (classDef in classes) // ...search through all classes for the MethodFingerprint
                     if (fingerprint.resolve(context, classDef))
-                        break@classes // if the resolution succeeded, continue with the next fingerprint
+                        break@classes // ...if the resolution succeeded, continue with the next MethodFingerprint.
         }
 
         /**
@@ -144,6 +334,52 @@ abstract class MethodFingerprint(
             val patternScanResult = if (methodFingerprint.opcodes != null) {
                 method.implementation?.instructions ?: return false
 
+                fun Method.patternScan(
+                    fingerprint: MethodFingerprint
+                ): MethodFingerprintResult.MethodFingerprintScanResult.PatternScanResult? {
+                    val instructions = this.implementation!!.instructions
+                    val fingerprintFuzzyPatternScanThreshold = fingerprint.fuzzyPatternScanMethod?.threshold ?: 0
+
+                    val pattern = fingerprint.opcodes!!
+                    val instructionLength = instructions.count()
+                    val patternLength = pattern.count()
+
+                    for (index in 0 until instructionLength) {
+                        var patternIndex = 0
+                        var threshold = fingerprintFuzzyPatternScanThreshold
+
+                        while (index + patternIndex < instructionLength) {
+                            val originalOpcode = instructions.elementAt(index + patternIndex).opcode
+                            val patternOpcode = pattern.elementAt(patternIndex)
+
+                            if (patternOpcode != null && patternOpcode.ordinal != originalOpcode.ordinal) {
+                                // reaching maximum threshold (0) means,
+                                // the pattern does not match to the current instructions
+                                if (threshold-- == 0) break
+                            }
+
+                            if (patternIndex < patternLength - 1) {
+                                // if the entire pattern has not been scanned yet
+                                // continue the scan
+                                patternIndex++
+                                continue
+                            }
+                            // the pattern is valid, generate warnings if fuzzyPatternScanMethod is FuzzyPatternScanMethod
+                            val result =
+                                MethodFingerprintResult.MethodFingerprintScanResult.PatternScanResult(
+                                    index,
+                                    index + patternIndex
+                                )
+                            if (fingerprint.fuzzyPatternScanMethod !is FuzzyPatternScanMethod) return result
+                            result.warnings = result.createWarnings(pattern, instructions)
+
+                            return result
+                        }
+                    }
+
+                    return null
+                }
+
                 method.patternScan(methodFingerprint) ?: return false
             } else null
 
@@ -160,52 +396,6 @@ abstract class MethodFingerprint(
             return true
         }
 
-        private fun Method.patternScan(
-            fingerprint: MethodFingerprint
-        ): MethodFingerprintResult.MethodFingerprintScanResult.PatternScanResult? {
-            val instructions = this.implementation!!.instructions
-            val fingerprintFuzzyPatternScanThreshold = fingerprint.fuzzyPatternScanMethod?.threshold ?: 0
-
-            val pattern = fingerprint.opcodes!!
-            val instructionLength = instructions.count()
-            val patternLength = pattern.count()
-
-            for (index in 0 until instructionLength) {
-                var patternIndex = 0
-                var threshold = fingerprintFuzzyPatternScanThreshold
-
-                while (index + patternIndex < instructionLength) {
-                    val originalOpcode = instructions.elementAt(index + patternIndex).opcode
-                    val patternOpcode = pattern.elementAt(patternIndex)
-
-                    if (patternOpcode != null && patternOpcode.ordinal != originalOpcode.ordinal) {
-                        // reaching maximum threshold (0) means,
-                        // the pattern does not match to the current instructions
-                        if (threshold-- == 0) break
-                    }
-
-                    if (patternIndex < patternLength - 1) {
-                        // if the entire pattern has not been scanned yet
-                        // continue the scan
-                        patternIndex++
-                        continue
-                    }
-                    // the pattern is valid, generate warnings if fuzzyPatternScanMethod is FuzzyPatternScanMethod
-                    val result =
-                        MethodFingerprintResult.MethodFingerprintScanResult.PatternScanResult(
-                            index,
-                            index + patternIndex
-                        )
-                    if (fingerprint.fuzzyPatternScanMethod !is FuzzyPatternScanMethod) return result
-                    result.warnings = result.createWarnings(pattern, instructions)
-
-                    return result
-                }
-            }
-
-            return null
-        }
-
         private fun MethodFingerprintResult.MethodFingerprintScanResult.PatternScanResult.createWarnings(
             pattern: Iterable<Opcode?>, instructions: Iterable<Instruction>
         ) = buildList {