利用Trie结构将性能提升了10倍

Em3rs0n · Em3rs0n · commit 388e25f904e7 · 2018-11-24T04:30:38.000-05:00
借助新的并行计算工具构造了一个Trie结构，在保持APK Parser性能基本不变的情况下，将
自动适配表达式的性能提升了10倍。原先一个表达式求值需要耗费将近1秒，现在只需20ms不到。

旧的Benchmark记录

BM: Parse DexClasses takes 758 ms.
BM: Search android.support.v4.app takes 747 ms.
BM: Search com.tencent.mm.ui.contact takes 1239 ms.
BM: Search com.tencent.mm.plugin.sns.ui takes 1393 ms.
BM: Search com.tencent.mm takes 1635 ms.
BM: Search com.tencent.mm.ui.conversation takes 761 ms.
BM: Search com.tencent.mm.sdk.platformtools takes 860 ms.
BM: Search com.tencent.mm.sdk.platformtools takes 669 ms.
BM: Search com.tencent.mm.booter.notification takes 561 ms.
BM: Search com.tencent.mm.storage takes 871 ms.
BM: Search com.tencent.mm.booter.notification.queue takes 913 ms.
BM: Search com.tencent.mm.ui takes 518 ms.
BM: Search com.tencent.mm.modelsfs takes 626 ms.
BM: Search com.tencent.mm.ui.chatting takes 470 ms.
BM: Search over classes takes 4414 ms.

数据结构升级后的Benchmark记录

BM: Parse DexClasses takes 779 ms.
BM: Search android.support.v4.app takes 42 ms.
BM: Search com.tencent.mm takes 167 ms.
BM: Search com.tencent.mm.sdk.platformtools takes 5 ms.
BM: Search com.tencent.mm.booter.notification takes 1 ms.
BM: Search com.tencent.mm.booter.notification.queue takes 0 ms.
BM: Search com.tencent.mm.modelsfs takes 0 ms.
BM: Search com.tencent.mm.plugin.sns.ui takes 66 ms.
BM: Search com.tencent.mm.storage takes 1 ms.
BM: Search com.tencent.mm.ui takes 2 ms.
BM: Search com.tencent.mm.ui.chatting takes 24 ms.
BM: Search com.tencent.mm.ui.contact takes 19 ms.
BM: Search com.tencent.mm.ui.conversation takes 14 ms.
BM: Search over classes takes 558 ms.
diff --git a/src/main/kotlin/com/gh0u1l5/wechatmagician/spellbook/WechatGlobal.kt b/src/main/kotlin/com/gh0u1l5/wechatmagician/spellbook/WechatGlobal.kt
@@ -6,6 +6,7 @@ import com.gh0u1l5.wechatmagician.spellbook.SpellBook.getApplicationVersion
 import com.gh0u1l5.wechatmagician.spellbook.base.Version
 import com.gh0u1l5.wechatmagician.spellbook.base.WaitChannel
 import com.gh0u1l5.wechatmagician.spellbook.parser.ApkFile
+import com.gh0u1l5.wechatmagician.spellbook.parser.ClassTrie
 import com.gh0u1l5.wechatmagician.spellbook.util.BasicUtil.tryAsynchronously
 import de.robv.android.xposed.callbacks.XC_LoadPackage
 import java.lang.ref.WeakReference
@@ -75,7 +76,7 @@ object WechatGlobal {
      * 这些类名使用的是 JVM 标准中规定的类名格式, 例如 String 的类名会被表示为 "Ljava/lang/String;"
      * Refer: https://docs.oracle.com/javase/specs/jvms/se7/html/jvms-4.html#jvms-4.3
      */
-    @Volatile var wxClasses: Array<String>? = null
+    @Volatile var wxClasses: ClassTrie? = null
         get() {
             if (!wxUnitTestMode) {
                 initChannel.wait(INIT_TIMEOUT)
diff --git a/src/main/kotlin/com/gh0u1l5/wechatmagician/spellbook/parser/ApkFile.kt b/src/main/kotlin/com/gh0u1l5/wechatmagician/spellbook/parser/ApkFile.kt
@@ -1,5 +1,7 @@
 package com.gh0u1l5.wechatmagician.spellbook.parser
 
+import com.gh0u1l5.wechatmagician.spellbook.util.ParallelUtil.parallelForEach
+
 import java.io.Closeable
 import java.io.File
 import java.nio.ByteBuffer
@@ -23,13 +25,27 @@ class ApkFile(apkFile: File) : Closeable {
     override fun close() =
             zipFile.close()
 
-    val classTypes: Array<String> by lazy {
-        var ret = emptyArray<String>()
-        for (i in 1 until 1000) {
-            val path = if (i == 1) DEX_FILE else String.format(DEX_ADDITIONAL, i)
-            val entry = zipFile.getEntry(path) ?: break
-            val buffer = ByteBuffer.wrap(readEntry(entry))
-            ret += DexParser(buffer).parseClassTypes()
+    private fun getDexFilePath(idx: Int) =
+            if (idx == 1) DEX_FILE else String.format(DEX_ADDITIONAL, idx)
+
+    private fun isDexFileExist(idx: Int): Boolean {
+        val path = getDexFilePath(idx)
+        return zipFile.getEntry(path) != null
+    }
+
+    val classTypes: ClassTrie by lazy {
+        var last = 2
+        while (isDexFileExist(last)) last++
+
+        val ret = ClassTrie()
+        (1..last).parallelForEach { idx ->
+            val path = getDexFilePath(idx)
+            val entry = zipFile.getEntry(path)
+            val data = readEntry(entry)
+            val buffer = ByteBuffer.wrap(data)
+            DexParser(buffer).parseClassTypes().forEach { type ->
+                ret += type
+            }
         }
         return@lazy ret
     }
diff --git a/src/main/kotlin/com/gh0u1l5/wechatmagician/spellbook/parser/ClassTrie.kt b/src/main/kotlin/com/gh0u1l5/wechatmagician/spellbook/parser/ClassTrie.kt
@@ -0,0 +1,71 @@
+package com.gh0u1l5.wechatmagician.spellbook.parser
+
+import com.gh0u1l5.wechatmagician.spellbook.util.ParallelUtil.parallelForEach
+import java.util.concurrent.ConcurrentHashMap
+
+class ClassTrie {
+    companion object {
+        private fun convertJVMTypeToClassName(type: String) =
+                type.substring(1, type.length - 1).replace('/', '.')
+    }
+
+    private val head: TrieNode = TrieNode()
+
+    operator fun plusAssign(type: String) {
+        head.add(convertJVMTypeToClassName(type))
+    }
+
+    operator fun plusAssign(types: Array<String>) {
+        types.asList().parallelForEach { this += it }
+    }
+
+    fun search(packageName: String, depth: Int): List<String> {
+        return head.search(packageName, depth)
+    }
+
+    private class TrieNode {
+        val classes: MutableList<String> = ArrayList(50)
+
+        val children: MutableMap<String, TrieNode> = ConcurrentHashMap()
+
+        fun add(className: String) {
+            add(className, 0)
+        }
+
+        private fun add(className: String, pos: Int) {
+            val delimiterAt = className.indexOf('.', pos)
+            if (delimiterAt == -1) {
+                synchronized(this) {
+                    classes.add(className)
+                }
+                return
+            }
+            val pkg = className.substring(pos, delimiterAt)
+            if (pkg !in children) {
+                children[pkg] = TrieNode()
+            }
+            children[pkg]!!.add(className, delimiterAt + 1)
+        }
+
+        fun get(depth: Int = 0): List<String> {
+            if (depth == 0) {
+                return classes
+            }
+            return children.flatMap { it.value.get(depth - 1) }
+        }
+
+        fun search(packageName: String, depth: Int): List<String> {
+            return search(packageName, depth, 0)
+        }
+
+        private fun search(packageName: String, depth: Int, pos: Int): List<String> {
+            val delimiterAt = packageName.indexOf('.', pos)
+            if (delimiterAt == -1) {
+                val pkg = packageName.substring(pos)
+                return children[pkg]?.get(depth) ?: emptyList()
+            }
+            val pkg = packageName.substring(pos, delimiterAt)
+            return children[pkg]?.search(packageName, depth, delimiterAt + 1) ?: emptyList()
+        }
+    }
+}
diff --git a/src/main/kotlin/com/gh0u1l5/wechatmagician/spellbook/util/MirrorUtil.kt b/src/main/kotlin/com/gh0u1l5/wechatmagician/spellbook/util/MirrorUtil.kt
@@ -1,7 +1,6 @@
 package com.gh0u1l5.wechatmagician.spellbook.util
 
 import com.gh0u1l5.wechatmagician.spellbook.WechatGlobal
-import com.gh0u1l5.wechatmagician.spellbook.util.ParallelUtil.parallelMap
 
 /**
  * 封装了一批用于检查“自动适配表达式”的函数
@@ -57,7 +56,7 @@ object MirrorUtil {
      * WARN: 仅供单元测试使用
      */
     fun generateReportWithForceEval(instances: List<Any>): List<Pair<String, String>> {
-        return instances.parallelMap { instance ->
+        return instances.map { instance ->
             collectFields(instance).map {
                 val value = it.second
                 if (value is Lazy<*>) {
@@ -67,6 +66,6 @@ object MirrorUtil {
                 }
                 "${instance::class.java.canonicalName}.${it.first}" to it.second.toString()
             }
-        }.flatten()
+        }.flatten() // 为了 Benchmark 的准确性, 不对结果进行排序
     }
 }
diff --git a/src/main/kotlin/com/gh0u1l5/wechatmagician/spellbook/util/ReflectionUtil.kt b/src/main/kotlin/com/gh0u1l5/wechatmagician/spellbook/util/ReflectionUtil.kt
@@ -3,6 +3,7 @@ package com.gh0u1l5.wechatmagician.spellbook.util
 import com.gh0u1l5.wechatmagician.spellbook.WechatGlobal
 import com.gh0u1l5.wechatmagician.spellbook.base.Classes
 import com.gh0u1l5.wechatmagician.spellbook.parser.ApkFile
+import com.gh0u1l5.wechatmagician.spellbook.parser.ClassTrie
 import de.robv.android.xposed.XC_MethodHook
 import de.robv.android.xposed.XposedBridge.hookMethod
 import java.lang.reflect.Field
@@ -66,40 +67,21 @@ object ReflectionUtil {
      * 里面其他包拥有的类.
      *
      * @param loader 用于取出 [Class] 对象的加载器
-     * @param classes 所有已知的类名, 由于 Java 的 [ClassLoader] 对象不支持读取所有类名, 我们必须先通过其他手段
-     * 获取一个类名列表, 详情请参见 [ApkFile] 和 [WechatGlobal]
+     * @param trie 整个 APK 的包结构, 由于 Java 的 [ClassLoader] 对象不支持读取所有类名, 我们必须先通过其他手段
+     * 解析 APK 结构, 然后才能检索某个包内的所有类, 详情请参见 [ApkFile] 和 [WechatGlobal]
      * @param packageName 包名
      * @param depth 深度
      */
-    @JvmStatic fun findClassesFromPackage(loader: ClassLoader, classes: Array<String>, packageName: String, depth: Int = 0): Classes {
+    @JvmStatic fun findClassesFromPackage(loader: ClassLoader, trie: ClassTrie, packageName: String, depth: Int = 0): Classes {
         val key = "$depth-$packageName"
         val cached = classCache[key]
         if (cached != null) {
             return cached
         }
-
-        val packageLength = packageName.count { it == '.' } + 1
-        val packageDescriptor = "L${packageName.replace('.', '/')}"
-        val result = Classes(classes.filter { clazz ->
-            val currentPackageLength = clazz.count { it == '/' }
-            if (currentPackageLength < packageLength) {
-                return@filter false
-            }
-            // Check depth
-            val currentDepth = currentPackageLength - packageLength
-            if (depth != -1 && depth != currentDepth) {
-                return@filter false
-            }
-            // Check prefix
-            if (!clazz.startsWith(packageDescriptor)) {
-                return@filter false
-            }
-            return@filter true
-        }.mapNotNull {
-            findClassIfExists(it.substring(1, it.length - 1).replace('/', '.'), loader)
+        val classes = Classes(trie.search(packageName, depth).mapNotNull { name ->
+            findClassIfExists(name, loader)
         })
-
-        return result.also { classCache[key] = result }
+        return classes.also { classCache[key] = classes }
     }
 
     /**

Original file line number	Diff line number	Diff line change
`@@ -1,7 +1,6 @@`
`1`	`1`	`package com.gh0u1l5.wechatmagician.spellbook.util`
`2`	`2`
`3`	`3`	`import com.gh0u1l5.wechatmagician.spellbook.WechatGlobal`
`4`		`-import com.gh0u1l5.wechatmagician.spellbook.util.ParallelUtil.parallelMap`
`5`	`4`
`6`	`5`	`/**`
`7`	`6`	`* 封装了一批用于检查“自动适配表达式”的函数`
`@@ -57,7 +56,7 @@ object MirrorUtil {`
`57`	`56`	`* WARN: 仅供单元测试使用`
`58`	`57`	`*/`
`59`	`58`	`fun generateReportWithForceEval(instances: List<Any>): List<Pair<String, String>> {`
`60`		`- return instances.parallelMap { instance ->`
	`59`	`+ return instances.map { instance ->`
`61`	`60`	`collectFields(instance).map {`
`62`	`61`	`val value = it.second`
`63`	`62`	`if (value is Lazy<*>) {`
`@@ -67,6 +66,6 @@ object MirrorUtil {`
`67`	`66`	`}`
`68`	`67`	`"${instance::class.java.canonicalName}.${it.first}" to it.second.toString()`
`69`	`68`	`}`
`70`		`- }.flatten()`
	`69`	`+ }.flatten() // 为了 Benchmark 的准确性, 不对结果进行排序`
`71`	`70`	`}`
`72`	`71`	`}`