From f58a498849e6527cdd67943f79d794d9bc83bed9 Mon Sep 17 00:00:00 2001
From: oSumAtrIX <johan.melkonyan1@web.de>
Date: Tue, 19 Apr 2022 19:47:35 +0200
Subject: [PATCH] add: `applyProxies` method

Signed-off-by: oSumAtrIX <johan.melkonyan1@web.de>
---
 src/main/kotlin/app/revanced/patcher/Patcher.kt   | 12 ++----------
 .../kotlin/app/revanced/patcher/PatcherData.kt    |  6 +++---
 .../app/revanced/patcher/proxy/ClassProxy.kt      |  9 +++++++++
 .../revanced/patcher/util/ProxyBackedClassList.kt | 15 +++++++++++++++
 4 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/src/main/kotlin/app/revanced/patcher/Patcher.kt b/src/main/kotlin/app/revanced/patcher/Patcher.kt
index 3f3160b..34cadad 100644
--- a/src/main/kotlin/app/revanced/patcher/Patcher.kt
+++ b/src/main/kotlin/app/revanced/patcher/Patcher.kt
@@ -70,16 +70,8 @@ class Patcher(
     fun save(): Map<String, MemoryDataStore> {
         val newDexFile = object : DexFile {
             override fun getClasses(): Set<ClassDef> {
-                val classes = patcherData.classes
-                val internalClasses = classes.internalClasses
-                for (proxy in classes.proxies) {
-                    if (!proxy.proxyUsed) continue
-
-                    val index = internalClasses.indexOfFirst { it.type == proxy.immutableClass.type }
-                    internalClasses[index] = proxy.mutatedClass
-                }
-
-                return ListBackedSet(internalClasses)
+                patcherData.classes.applyProxies()
+                return ListBackedSet(patcherData.classes.internalClasses)
             }
 
             override fun getOpcodes(): Opcodes {
diff --git a/src/main/kotlin/app/revanced/patcher/PatcherData.kt b/src/main/kotlin/app/revanced/patcher/PatcherData.kt
index 97701f5..34bfe45 100644
--- a/src/main/kotlin/app/revanced/patcher/PatcherData.kt
+++ b/src/main/kotlin/app/revanced/patcher/PatcherData.kt
@@ -72,10 +72,10 @@ internal inline fun <T> Iterable<T>.findIndexed(predicate: (T) -> Boolean): Pair
     return null
 }
 
-fun PatcherData.proxy(classProxy: ClassDef): ClassProxy {
-    var proxy = this.classes.proxies.find { it.immutableClass.type == classProxy.type }
+fun PatcherData.proxy(classDef: ClassDef): ClassProxy {
+    var proxy = this.classes.proxies.find { it.immutableClass.type == classDef.type }
     if (proxy == null) {
-        proxy = ClassProxy(classProxy)
+        proxy = ClassProxy(classDef)
         this.classes.proxies.add(proxy)
     }
     return proxy
diff --git a/src/main/kotlin/app/revanced/patcher/proxy/ClassProxy.kt b/src/main/kotlin/app/revanced/patcher/proxy/ClassProxy.kt
index 9511985..0aedeef 100644
--- a/src/main/kotlin/app/revanced/patcher/proxy/ClassProxy.kt
+++ b/src/main/kotlin/app/revanced/patcher/proxy/ClassProxy.kt
@@ -16,6 +16,15 @@ class ClassProxy(
     internal var proxyUsed = false
     internal lateinit var mutatedClass: MutableClass
 
+    init {
+        // in the instance, that a [MutableClass] is being proxied,
+        // do not create an additional clone and reuse the [MutableClass] instance
+        if (immutableClass is MutableClass) {
+            mutatedClass = immutableClass
+            proxyUsed = true
+        }
+    }
+
     /**
      * Allocates and returns a mutable clone of the original class.
      * A patch should always use the original immutable class reference
diff --git a/src/main/kotlin/app/revanced/patcher/util/ProxyBackedClassList.kt b/src/main/kotlin/app/revanced/patcher/util/ProxyBackedClassList.kt
index 6d68388..07c426e 100644
--- a/src/main/kotlin/app/revanced/patcher/util/ProxyBackedClassList.kt
+++ b/src/main/kotlin/app/revanced/patcher/util/ProxyBackedClassList.kt
@@ -14,6 +14,21 @@ class ProxyBackedClassList(internal val internalClasses: MutableList<ClassDef>)
         proxies.add(classProxy)
     }
 
+    /**
+     * Apply all resolved classes into [internalClasses] and clear the [proxies] list.
+     */
+    fun applyProxies() {
+        proxies.forEachIndexed { i, proxy ->
+            if (!proxy.proxyUsed) return@forEachIndexed
+
+            val index = internalClasses.indexOfFirst { it.type == proxy.immutableClass.type }
+            internalClasses[index] = proxy.mutatedClass
+
+            proxies.removeAt(i) // FIXME: check if this could cause issues when multiple patches use the same proxy
+        }
+
+    }
+
     override val size get() = internalClasses.size
     override fun contains(element: ClassDef) = internalClasses.contains(element)
     override fun containsAll(elements: Collection<ClassDef>) = internalClasses.containsAll(elements)