fix: Merge all extensions before initializing lookup maps

ReVanced · Aug 6, 2024 · 8c4dd5b · 8c4dd5b
1 parent 736b3ee
commit 8c4dd5b
Show file tree

Hide file tree

Showing 3 changed files with 102 additions and 54 deletions.
diff --git a/src/main/kotlin/app/revanced/patcher/Patcher.kt b/src/main/kotlin/app/revanced/patcher/Patcher.kt
@@ -39,9 +39,6 @@ class Patcher(private val config: PatcherConfig) : Closeable {
             patch.addRecursively()
         }
 
-        fun Patch<*>.anyRecursively(predicate: (Patch<*>) -> Boolean): Boolean =
-            predicate(this) || dependencies.any { dependency -> dependency.anyRecursively(predicate) }
-
         context.allPatches.let { allPatches ->
             // Check, if what kind of resource mode is required.
             config.resourceMode = if (allPatches.any { patch -> patch.anyRecursively { it is ResourcePatch } }) {
@@ -99,6 +96,17 @@ class Patcher(private val config: PatcherConfig) : Closeable {
             context.resourceContext.decodeResources(config.resourceMode)
         }
 
+        logger.info("Merging extensions")
+
+        context.executablePatches.forEachRecursively { patch ->
+            if (patch is BytecodePatch && patch.extension != null) {
+                context.bytecodeContext.merge(patch.extension)
+            }
+        }
+
+        // Initialize lookup maps.
+        context.bytecodeContext.lookupMaps
+
         logger.info("Executing patches")
 
         val executedPatches = LinkedHashMap<Patch<*>, PatchResult>()
@@ -146,7 +154,7 @@ class Patcher(private val config: PatcherConfig) : Closeable {
         }
     }
 
-    override fun close() = context.bytecodeContext.lookupMaps.close()
+    override fun close() = context.bytecodeContext.close()
 
     /**
      * Compile and save patched APK files.

diff --git a/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt b/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt
@@ -31,7 +31,9 @@ import java.util.logging.Logger
  * @param config The [PatcherConfig] used to create this context.
  */
 @Suppress("MemberVisibilityCanBePrivate")
-class BytecodePatchContext internal constructor(private val config: PatcherConfig) : PatchContext<Set<PatcherResult.PatchedDexFile>> {
+class BytecodePatchContext internal constructor(private val config: PatcherConfig) :
+    PatchContext<Set<PatcherResult.PatchedDexFile>>,
+    Closeable {
     private val logger = Logger.getLogger(BytecodePatchContext::class.java.name)
 
     /**
@@ -57,6 +59,13 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi
      */
     internal val lookupMaps by lazy { LookupMaps(classes) }
 
+    /**
+     * A map for lookup by [merge].
+     */
+    internal val classesByType = mutableMapOf<String, ClassDef>().apply {
+        classes.forEach { classDef -> put(classDef.type, classDef) }
+    }
+
     /**
      * Merge an extension to [classes].
      *
@@ -66,11 +75,11 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi
         val extension = extensionInputStream.readAllBytes()
 
         RawDexIO.readRawDexFile(extension, 0, null).classes.forEach { classDef ->
-            val existingClass = lookupMaps.classesByType[classDef.type] ?: run {
+            val existingClass = classesByType[classDef.type] ?: run {
                 logger.fine("Adding class \"$classDef\"")
 
-                lookupMaps.classesByType[classDef.type] = classDef
                 classes += classDef
+                classesByType[classDef.type] = classDef
 
                 return@forEach
             }
@@ -254,6 +263,12 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi
             methodsByStrings.clear()
         }
     }
+
+    override fun close() {
+        lookupMaps.close()
+        classesByType.clear()
+        classes.clear()
+    }
 }
 
 /**

diff --git a/src/main/kotlin/app/revanced/patcher/patch/Patch.kt b/src/main/kotlin/app/revanced/patcher/patch/Patch.kt
@@ -85,6 +85,31 @@ sealed class Patch<C : PatchContext<*>>(
     override fun toString() = name ?: "Patch"
 }
 
+internal fun Patch<*>.anyRecursively(
+    visited: MutableSet<Patch<*>> = mutableSetOf(),
+    predicate: (Patch<*>) -> Boolean,
+): Boolean {
+    if (this in visited) return false
+
+    if (predicate(this)) return true
+
+    visited += this
+
+    return dependencies.any { it.anyRecursively(visited, predicate) }
+}
+
+internal fun Iterable<Patch<*>>.forEachRecursively(
+    visited: MutableSet<Patch<*>> = mutableSetOf(),
+    action: (Patch<*>) -> Unit,
+): Unit = forEach {
+    if (it in visited) return@forEach
+
+    visited += it
+    action(it)
+
+    it.dependencies.forEachRecursively(visited, action)
+}
+
 /**
  * A bytecode patch.
  *
@@ -127,7 +152,6 @@ class BytecodePatch internal constructor(
     finalizeBlock,
 ) {
     override fun execute(context: PatcherContext) = with(context.bytecodeContext) {
-        extension?.let(::merge)
         fingerprints.forEach { it.match(this) }
 
         execute(this)
@@ -332,6 +356,16 @@ sealed class PatchBuilder<C : PatchContext<*>>(
     internal abstract fun build(): Patch<C>
 }
 
+/**
+ * Builds a [Patch].
+ *
+ * @param B The [PatchBuilder] to build the patch with.
+ * @param block The block to build the patch.
+ *
+ * @return The built [Patch].
+ */
+private fun <B : PatchBuilder<*>> B.buildPatch(block: B.() -> Unit = {}) = apply(block).build()
+
 /**
  * A [BytecodePatchBuilder] builder.
  *
@@ -379,9 +413,10 @@ class BytecodePatchBuilder internal constructor(
      *
      * @param extension The name of the extension resource.
      */
+    @Suppress("NOTHING_TO_INLINE")
     inline fun extendWith(extension: String) = apply {
         this.extension = object {}.javaClass.classLoader.getResourceAsStream(extension)
-            ?: throw PatchException("Extension resource \"$extension\" not found")
+            ?: throw PatchException("Extension \"$extension\" not found")
     }
 
     override fun build() = BytecodePatch(
@@ -398,6 +433,24 @@ class BytecodePatchBuilder internal constructor(
     )
 }
 
+/**
+ * Create a new [BytecodePatch].
+ *
+ * @param name The name of the patch.
+ * If null, the patch is named "Patch" and will not be loaded by [PatchLoader].
+ * @param description The description of the patch.
+ * @param use Weather or not the patch should be used.
+ * @param block The block to build the patch.
+ *
+ * @return The created [BytecodePatch].
+ */
+fun bytecodePatch(
+    name: String? = null,
+    description: String? = null,
+    use: Boolean = true,
+    block: BytecodePatchBuilder.() -> Unit = {},
+) = BytecodePatchBuilder(name, description, use).buildPatch(block) as BytecodePatch
+
 /**
  * A [RawResourcePatch] builder.
  *
@@ -425,6 +478,23 @@ class RawResourcePatchBuilder internal constructor(
     )
 }
 
+/**
+ * Create a new [RawResourcePatch].
+ *
+ * @param name The name of the patch.
+ * If null, the patch is named "Patch" and will not be loaded by [PatchLoader].
+ * @param description The description of the patch.
+ * @param use Weather or not the patch should be used.
+ * @param block The block to build the patch.
+ * @return The created [RawResourcePatch].
+ */
+fun rawResourcePatch(
+    name: String? = null,
+    description: String? = null,
+    use: Boolean = true,
+    block: RawResourcePatchBuilder.() -> Unit = {},
+) = RawResourcePatchBuilder(name, description, use).buildPatch(block) as RawResourcePatch
+
 /**
  * A [ResourcePatch] builder.
  *
@@ -452,51 +522,6 @@ class ResourcePatchBuilder internal constructor(
     )
 }
 
-/**
- * Builds a [Patch].
- *
- * @param B The [PatchBuilder] to build the patch with.
- * @param block The block to build the patch.
- *
- * @return The built [Patch].
- */
-private fun <B : PatchBuilder<*>> B.buildPatch(block: B.() -> Unit = {}) = apply(block).build()
-
-/**
- * Create a new [BytecodePatch].
- *
- * @param name The name of the patch.
- * If null, the patch is named "Patch" and will not be loaded by [PatchLoader].
- * @param description The description of the patch.
- * @param use Weather or not the patch should be used.
- * @param block The block to build the patch.
- *
- * @return The created [BytecodePatch].
- */
-fun bytecodePatch(
-    name: String? = null,
-    description: String? = null,
-    use: Boolean = true,
-    block: BytecodePatchBuilder.() -> Unit = {},
-) = BytecodePatchBuilder(name, description, use).buildPatch(block) as BytecodePatch
-
-/**
- * Create a new [RawResourcePatch].
- *
- * @param name The name of the patch.
- * If null, the patch is named "Patch" and will not be loaded by [PatchLoader].
- * @param description The description of the patch.
- * @param use Weather or not the patch should be used.
- * @param block The block to build the patch.
- * @return The created [RawResourcePatch].
- */
-fun rawResourcePatch(
-    name: String? = null,
-    description: String? = null,
-    use: Boolean = true,
-    block: RawResourcePatchBuilder.() -> Unit = {},
-) = RawResourcePatchBuilder(name, description, use).buildPatch(block) as RawResourcePatch
-
 /**
  * Create a new [ResourcePatch].
  *