CHA improvement

booster-cha/src/main/kotlin/com/didiglobal/booster/cha/ClassHierarchy.kt

@@ -46,22 +46,42 @@ class ClassHierarchy<ClassFile : Any, ClassParser>(
  name: String?,
  filter: ClassFileParser<ClassFile>.(clazz: ClassFile) -> Boolean = { true }
  ): Set<ClassFile> {
- val fq = name ?: return emptySet()
- val stack = Stack<String>().apply {
-  push(fq)
-  }
+ return get(name)?.let {
+  getDerivedTypes(it, filter)
+ } ?: emptySet()
+ }
 
- val children = mutableSetOf<ClassFile>()
- while (stack.isNotEmpty()) {
- children += graph[ClassNode(stack.pop())].takeIf(Set<ClassNode>::isNotEmpty)?.mapNotNull {
- get(it.name)
- }?.filter {
- filter(this, it)
- }?.onEach {
- stack.push(getClassName(it))
- } ?: emptyList()
- }
- return children
+ fun getDerivedTypes(
+ clazz: ClassFile,
+ filter: ClassFileParser<ClassFile>.(clazz: ClassFile) -> Boolean = { true }
+ ): Set<ClassFile> {
+ val node = ClassNode(getClassName(clazz))
+ return graph.getSuccessors(node).filter {
+ get(it.name)?.let { v -> filter(this, v) } == true
+ }.mapNotNull {
+ get(it.name)
+ }.toSet()
+ }
+
+ fun getSuperTypes(
+ name: String?,
+ filter: ClassFileParser<ClassFile>.(clazz: ClassFile) -> Boolean = { true }
+ ): Set<ClassFile> {
+ return get(name)?.let {
+ getSuperTypes(it, filter)
+ } ?: emptySet()
+ }
+
+ fun getSuperTypes(
+ clazz: ClassFile,
+ filter: ClassFileParser<ClassFile>.(clazz: ClassFile) -> Boolean = { true }
+ ): Set<ClassFile> {
+ val node = ClassNode(getClassName(clazz))
+ return graph.getPredecessors(node).filter {
+ get(it.name)?.let { v -> filter(this, v) } == true
+ }.mapNotNull {
+ get(it.name)
+ }.toSet()
  }
 
  fun isInheritFrom(child: ClassFile, parent: ClassFile) = when {
@@ -117,24 +137,6 @@ class ClassHierarchy<ClassFile : Any, ClassParser>(
  } ?: false
  }
 
- fun getSuperTypes(clazz: ClassFile): Set<ClassFile> {
- val superName = getSuperName(clazz) ?: return emptySet()
- var parent = this[superName]
-
- if (superName == JAVA_LANG_OBJECT) {
- return parent?.let(::setOf) ?: throw ClassNotFoundException(superName)
- }
-
- val classes = mutableSetOf<ClassFile>()
-
- while (null != parent) {
- classes += parent
- parent = getSuperName(parent)?.let(this::get)
- }
-
- return classes
- }
-
 }
 
 typealias OnClassResolveFailed = (String) -> Unit

booster-graph/src/main/kotlin/com/didiglobal/booster/graph/Graph.kt

@@ -2,6 +2,7 @@ package com.didiglobal.booster.graph
 
 import java.io.PrintWriter
 import java.util.Collections
+import java.util.Stack
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.CopyOnWriteArraySet
 
@@ -34,9 +35,43 @@ class Graph<N : Node> private constructor(
 
  operator fun get(from: N): Set<N> = edges[from]?.let(Collections::unmodifiableSet) ?: emptySet()
 
+ fun getPredecessors(node: N): Set<N> {
+ val stack = Stack<N>().apply {
+ push(node)
+ }
+ val predecessors = mutableSetOf<N>()
+ while (stack.isNotEmpty()) {
+ val current = stack.pop()
+ predecessors += entries.filter {
+ it.to == current && it.from !in predecessors
+ }.map {
+ it.from
+ }.onEach {
+ stack.push(it)
+ }
+ }
+ return predecessors
+ }
+
+ fun getSuccessors(node: N): Set<N> {
+ val stack = Stack<N>().apply {
+ push(node)
+ }
+ val successors = mutableSetOf<N>()
+ while (stack.isNotEmpty()) {
+ val current = stack.pop()
+ successors += get(current).filter {
+ it !in successors
+ }.onEach {
+ stack.push(it)
+ }
+ }
+ return successors
+ }
+
  fun reversed(): Graph<N> {
  return edges.entries.fold(Builder<N>()) { builder, entry ->
- entry.value.forEach {  to ->
+ entry.value.forEach { to ->
  builder.addEdge(to, entry.key)
  }
  builder